context: Lazy load storage

Using static initialization is possible, but quite complex to handle
logging and circular loading. Lazy loading prevents an entire class of
circular dependencies.

1 files changed, 36 insertions, 35 deletions

diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java
index 68228fd85..7501f46f6 100644
--- a/context/src/main/java/io/grpc/Context.java
+++ b/context/src/main/java/io/grpc/Context.java
@@ -19,11 +19,11 @@ package io.grpc;
 import java.util.ArrayList;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
-import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -109,37 +109,42 @@ public class Context {
    */
   public static final Context ROOT = new Context(null);
 
-  // One and only one of them is non-null
-  private static final Storage storage;
-  private static final LinkedBlockingQueue<ClassNotFoundException> storageOverrideNotFoundErrors =
-      new LinkedBlockingQueue<ClassNotFoundException>();
-  private static final Exception storageInitError;
+  // 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<Storage>();
 
-  static {
-    Storage newStorage = null;
-    Exception error = null;
+  // 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");
-      newStorage = (Storage) clazz.getConstructor().newInstance();
+      // 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) {
-      // Avoid writing to logger because custom log handlers may try to use Context, which is
-      // problemantic (e.g., NullPointerException) because the Context class has not done loading
-      // at this point.
-      storageOverrideNotFoundErrors.add(e);
-      newStorage = new ThreadLocalContextStorage();
+      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);
+      }
     } catch (Exception e) {
-      error = e;
-    }
-    storage = newStorage;
-    storageInitError = error;
-  }
-
-  // For testing
-  static Storage storage() {
-    if (storage == null) {
-      throw new RuntimeException("Storage override had failed to initialize", storageInitError);
+      throw new RuntimeException("Storage override failed to initialize", e);
     }
-    return storage;
+    // Re-retreive from storage since compareAndSet may have failed (returned false) in case of
+    // race.
+    return storage.get();
   }
 
   /**
@@ -213,13 +218,6 @@ public class Context {
     canBeCancelled = isCancellable;
   }
 
-  private static void maybeLogStorageOverrideNotFound() {
-    ClassNotFoundException e;
-    while ((e = storageOverrideNotFoundErrors.poll()) != null) {
-      log.log(Level.FINE, "Storage override doesn't exist. Using default.", e);
-    }
-  }
-
   /**
    * Create a new context which is independently cancellable and also cascades cancellation from
    * its parent. Callers <em>must</em> ensure that either {@link
@@ -387,7 +385,6 @@ public class Context {
    * }}</pre>
    */
   public Context attach() {
-    maybeLogStorageOverrideNotFound();
     Context previous = current();
     storage().attach(this);
     return previous;
@@ -867,7 +864,11 @@ public class Context {
   }
 
   /**
-   * Defines the mechanisms for attaching and detaching the "current" context.
+   * Defines the mechanisms for attaching and detaching the "current" context. The constructor for
+   * extending classes <em>must not</em> trigger any activity that can use Context, which includes
+   * logging, otherwise it can trigger an infinite initialization loop. Extending classes must not
+   * assume that only one instance will be created; Context guarantees it will only use one
+   * instance, but it may create multiple and then throw away all but one.
    *
    * <p>The default implementation will put the current context in a {@link ThreadLocal}.  If an
    * alternative implementation named {@code io.grpc.override.ContextStorageOverride} exists in the