context: change storage API to return a restoreable context (#3292)

This API change allows storage implementations to put some state
information into the return value, giving it the ability to act
as a cookie. In environments where contexts can be replaced, the
current original context can be stashed there and be restored
when detach is called.

3 files changed, 154 insertions, 7 deletions

diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java
index b8c4fb2c5..8375fb6c6 100644
--- a/context/src/main/java/io/grpc/Context.java
+++ b/context/src/main/java/io/grpc/Context.java
@@ -385,9 +385,11 @@ public class Context {
    * }}</pre>
    */
   public Context attach() {
-    Context previous = current();
-    storage().attach(this);
-    return previous;
+    Context prev = storage().doAttach(this);
+    if (prev == null) {
+      return ROOT;
+    }
+    return prev;
   }
 
   /**
@@ -878,11 +880,33 @@ public class Context {
    */
   public abstract static class Storage {
     /**
+     * @deprecated This is an old API that is no longer used.
+     */
+    @Deprecated
+    public void attach(Context toAttach) {
+      throw new UnsupportedOperationException("Deprecated. Do not call.");
+    }
+
+    /**
      * Implements {@link io.grpc.Context#attach}.
      *
+     * <p>Caution: {@link Context#attach()} interprets a return value of {@code null} to mean
+     * the same thing as {@link Context#ROOT}.
+     *
+     * <p>See also: {@link #current()}.
+
      * @param toAttach the context to be attached
+     * @return A {@link Context} that should be passed back into {@link #detach(Context, Context)}
+     *        as the {@code toRestore} parameter. {@code null} is a valid return value, but see
+     *        caution note.
      */
-    public abstract void attach(Context toAttach);
+    public Context doAttach(Context toAttach) {
+      // This is a default implementation to help migrate existing Storage implementations that
+      // have an attach() method but no doAttach() method.
+      Context current = current();
+      attach(toAttach);
+      return current;
+    }
 
     /**
      * Implements {@link io.grpc.Context#detach}
@@ -895,7 +919,15 @@ public class Context {
     public abstract void detach(Context toDetach, Context toRestore);
 
     /**
-     * Implements {@link io.grpc.Context#current}.  Returns the context of the current scope.
+     * Implements {@link io.grpc.Context#current}.
+     *
+     * <p>Caution: {@link Context} interprets a return value of {@code null} to mean the same
+     * thing as {@code Context{@link #ROOT}}.
+     *
+     * <p>See also {@link #doAttach(Context)}.
+     *
+     * @return The context of the current scope. {@code null} is a valid return value, but see
+     *        caution note.
      */
     public abstract Context current();
   }
diff --git a/context/src/main/java/io/grpc/ThreadLocalContextStorage.java b/context/src/main/java/io/grpc/ThreadLocalContextStorage.java
index d0c482ca9..7c49d5ff7 100644
--- a/context/src/main/java/io/grpc/ThreadLocalContextStorage.java
+++ b/context/src/main/java/io/grpc/ThreadLocalContextStorage.java
@@ -31,8 +31,10 @@ final class ThreadLocalContextStorage extends Context.Storage {
   private static final ThreadLocal<Context> localContext = new ThreadLocal<Context>();
 
   @Override
-  public void attach(Context toAttach) {
+  public Context doAttach(Context toAttach) {
+    Context current = current();
     localContext.set(toAttach);
+    return current;
   }
 
   @Override
@@ -44,7 +46,7 @@ final class ThreadLocalContextStorage extends Context.Storage {
       log.log(Level.SEVERE, "Context was not attached when detaching",
           new Throwable().fillInStackTrace());
     }
-    attach(toRestore);
+    doAttach(toRestore);
   }
 
   @Override
diff --git a/context/src/test/java/io/grpc/ContextTest.java b/context/src/test/java/io/grpc/ContextTest.java
index c20d41afc..3c61550eb 100644
--- a/context/src/test/java/io/grpc/ContextTest.java
+++ b/context/src/test/java/io/grpc/ContextTest.java
@@ -29,12 +29,15 @@ import static org.junit.Assert.fail;
 
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.util.ArrayDeque;
 import java.util.Queue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -763,6 +766,116 @@ public class ContextTest {
     ((Runnable) runnable.getDeclaredConstructor().newInstance()).run();
   }
 
+  /**
+   * Ensure that newly created threads can attach/detach a context.
+   * The current test thread already has a context manually attached in {@link #setUp()}.
+   */
+  @Test
+  public void newThreadAttachContext() throws Exception {
+    Context parent = Context.current().withValue(COLOR, "blue");
+    parent.call(new Callable<Object>() {
+      @Override
+      public Object call() throws Exception {
+        assertEquals("blue", COLOR.get());
+
+        final Context child = Context.current().withValue(COLOR, "red");
+        Future<String> workerThreadVal = scheduler
+            .submit(new Callable<String>() {
+              @Override
+              public String call() {
+                Context initial = Context.current();
+                assertNotNull(initial);
+                Context toRestore = child.attach();
+                try {
+                  assertNotNull(toRestore);
+                  return COLOR.get();
+                } finally {
+                  child.detach(toRestore);
+                  assertEquals(initial, Context.current());
+                }
+              }
+            });
+        assertEquals("red", workerThreadVal.get());
+
+        assertEquals("blue", COLOR.get());
+        return null;
+      }
+    });
+  }
+
+  /**
+   * Similar to {@link #newThreadAttachContext()} but without giving the new thread a specific ctx.
+   */
+  @Test
+  public void newThreadWithoutContext() throws Exception {
+    Context parent = Context.current().withValue(COLOR, "blue");
+    parent.call(new Callable<Object>() {
+      @Override
+      public Object call() throws Exception {
+        assertEquals("blue", COLOR.get());
+
+        Future<String> workerThreadVal = scheduler
+            .submit(new Callable<String>() {
+              @Override
+              public String call() {
+                assertNotNull(Context.current());
+                return COLOR.get();
+              }
+            });
+        assertEquals(null, workerThreadVal.get());
+
+        assertEquals("blue", COLOR.get());
+        return null;
+      }
+    });
+  }
+
+  @Test
+  public void storageReturnsNullTest() throws Exception {
+    Class<?> contextClass = Class.forName("io.grpc.Context");
+    Field storage = contextClass.getDeclaredField("storage");
+    assertTrue(Modifier.isFinal(storage.getModifiers()));
+    // use reflection to forcibly change the storage object to a test object
+    storage.setAccessible(true);
+    Object o = storage.get(null);
+    @SuppressWarnings("unchecked")
+    AtomicReference<Context.Storage> storageRef = (AtomicReference<Context.Storage>) o;
+    Context.Storage originalStorage = storageRef.get();
+    try {
+      storageRef.set(new Context.Storage() {
+        @Override
+        public Context doAttach(Context toAttach) {
+          return null;
+        }
+
+        @Override
+        public void detach(Context toDetach, Context toRestore) {
+          // noop
+        }
+
+        @Override
+        public Context current() {
+          return null;
+        }
+      });
+      // current() returning null gets transformed into ROOT
+      assertEquals(Context.ROOT, Context.current());
+
+      // doAttach() returning null gets transformed into ROOT
+      Context blueContext = Context.current().withValue(COLOR, "blue");
+      Context toRestore = blueContext.attach();
+      assertEquals(Context.ROOT, toRestore);
+
+      // final sanity check
+      blueContext.detach(toRestore);
+      assertEquals(Context.ROOT, Context.current());
+    } finally {
+      // undo the changes
+      storageRef.set(originalStorage);
+      storage.setAccessible(false);
+    }
+  }
+
   // UsedReflectively
   public static final class LoadMeWithStaticTestingClassLoader implements Runnable {
     @Override