context: change storage API to return a restoreable context

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

@@ -385,9 +385,11 @@ boolean canBeCancelled() {
    * }}</pre>
    */
   public Context attach() {
-    Context previous = current();
-    storage().attach(this);
-    return previous;
+    Context prev = storage().doAttach(this);
+    if (prev == null) {
+      return ROOT;
+    }
+    return prev;
   }
 
   /**
@@ -877,12 +879,28 @@ public String toString() {
    * subject to change.
    */
   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}.
      *
      * @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.
      */
-    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}

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 @@ public void detach(Context toDetach, Context toRestore) {
       log.log(Level.SEVERE, "Context was not attached when detaching",
           new Throwable().fillInStackTrace());
     }
-    attach(toRestore);
+    doAttach(toRestore);
   }
 
   @Override

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

@@ -35,6 +35,7 @@
 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 +764,70 @@ public void initContextWithCustomClassLoaderWithCustomLogger() throws Exception
     ((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;
+      }
+    });
+  }
+
   // UsedReflectively
   public static final class LoadMeWithStaticTestingClassLoader implements Runnable {
     @Override