context: try out static final storage

context/src/main/java/io/grpc/Context.java

@@ -116,42 +116,42 @@ public class Context {
    */
   public static final Context ROOT = new Context(null, EMPTY_ENTRIES);
 
+  // Visible For testing
+  static Storage storage() {
+    return LazyStorage.storage;
+  }
+
   // Lazy-loaded storage. Delaying storage initialization until after class initialization makes it
   // much easier to avoid circular loading since there can still be references to Context as long as
   // they don't depend on storage, like key() and currentContextExecutor(). It also makes it easier
   // to handle exceptions.
-  private static final AtomicReference<Storage> storage = new AtomicReference<>();
+  private static final class LazyStorage {
+    static final Storage storage;
 
-  // For testing
-  static Storage storage() {
-    Storage tmp = storage.get();
-    if (tmp == null) {
-      tmp = createStorage();
-    }
-    return tmp;
-  }
-
-  private static Storage createStorage() {
-    // Note that this method may be run more than once
-    try {
-      Class<?> clazz = Class.forName("io.grpc.override.ContextStorageOverride");
-      // The override's constructor is prohibited from triggering any code that can loop back to
-      // Context
-      Storage newStorage = (Storage) clazz.getConstructor().newInstance();
-      storage.compareAndSet(null, newStorage);
-    } catch (ClassNotFoundException e) {
-      Storage newStorage = new ThreadLocalContextStorage();
-      // Must set storage before logging, since logging may call Context.current().
-      if (storage.compareAndSet(null, newStorage)) {
-        // Avoid logging if this thread lost the race, to avoid confusion
-        log.log(Level.FINE, "Storage override doesn't exist. Using default", e);
+    static {
+      AtomicReference<Throwable> deferredStorageFailure = new AtomicReference<>();
+      storage = createStorage(deferredStorageFailure);
+      Throwable failure = deferredStorageFailure.get();
+      // Logging must happen after storage has been set, as loggers may use Context.
+      if (failure != null) {
+        log.log(Level.FINE, "Storage override doesn't exist. Using default", failure);
+      }
+    }
+
+    private static Storage createStorage(
+        AtomicReference<? super ClassNotFoundException> deferredStorageFailure) {
+      try {
+        Class<?> clazz = Class.forName("io.grpc.override.ContextStorageOverride");
+        // The override's constructor is prohibited from triggering any code that can loop back to
+        // Context
+        return clazz.asSubclass(Storage.class).getConstructor().newInstance();
+      } catch (ClassNotFoundException e) {
+        deferredStorageFailure.set(e);
+        return new ThreadLocalContextStorage();
+      } catch (Exception e) {
+        throw new RuntimeException("Storage override failed to initialize", e);
       }
-    } catch (Exception e) {
-      throw new RuntimeException("Storage override failed to initialize", e);
     }
-    // Re-retreive from storage since compareAndSet may have failed (returned false) in case of
-    // race.
-    return storage.get();
   }
 
   /**

context/src/test/java/io/grpc/ContextTest.java

@@ -840,16 +840,19 @@ public class ContextTest {
 
   @Test
   public void storageReturnsNullTest() throws Exception {
-    Field storage = Context.class.getDeclaredField("storage");
+    Class<?> lazyStorageClass = Class.forName("io.grpc.Context$LazyStorage");
+    Field storage = lazyStorageClass.getDeclaredField("storage");
     assertTrue(Modifier.isFinal(storage.getModifiers()));
     // use reflection to forcibly change the storage object to a test object
     storage.setAccessible(true);
+    Field modifiersField = Field.class.getDeclaredField("modifiers");
+    modifiersField.setAccessible(true);
+    int storageModifiers = modifiersField.getInt(storage);
+    modifiersField.set(storage, storageModifiers & ~Modifier.FINAL);
     Object o = storage.get(null);
-    @SuppressWarnings("unchecked")
-    AtomicReference<Context.Storage> storageRef = (AtomicReference<Context.Storage>) o;
-    Context.Storage originalStorage = storageRef.get();
+    Context.Storage originalStorage = (Context.Storage) o;
     try {
-      storageRef.set(new Context.Storage() {
+      storage.set(null, new Context.Storage() {
         @Override
         public Context doAttach(Context toAttach) {
           return null;
@@ -878,8 +881,10 @@ public class ContextTest {
       assertEquals(Context.ROOT, Context.current());
     } finally {
       // undo the changes
-      storageRef.set(originalStorage);
+      storage.set(null, originalStorage);
       storage.setAccessible(false);
+      modifiersField.set(storage, storageModifiers | Modifier.FINAL);
+      modifiersField.setAccessible(false);
     }
   }