api: Add a Supplier overload to Context

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

@@ -27,6 +27,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -554,6 +555,21 @@
     }
   }
 
+  /**
+    * Immediately supply a value from a {@link Supplier} with this context as the
+    * {@link #current} context.
+    * @param supplier {@link Supplier} to use to produce the value.
+    * @return result of supplier.
+    */
+  public <V> V supply(Supplier<V> supplier) {
+    Context previous = attach();
+    try {
+      return supplier.get();
+    } finally {
+      detach(previous);
+    }
+  }
+
   /**
    * Wrap a {@link Runnable} so that it executes with this context as the {@link #current} context.
    */

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

@@ -49,6 +49,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.logging.Handler;
 import java.util.logging.Level;
 import java.util.logging.LogRecord;
@@ -553,6 +554,46 @@ public Object call() {
     current.detach(Context.ROOT);
   }
 
+  @Test
+  public void testSupply() throws Exception {
+    Context base = Context.current().withValue(PET, "cat");
+    Context current = Context.current().withValue(PET, "fish");
+    current.attach();
+
+    final Object ret = new Object();
+    Supplier<Object> supplier = new Supplier<Object>() {
+      @Override
+      public Object get() {
+        runner.run();
+        return ret;
+      }
+    };
+
+    assertSame(ret, base.supply(supplier));
+    assertSame(base, observed);
+    assertSame(current, Context.current());
+
+    assertSame(ret, current.supply(supplier));
+    assertSame(current, observed);
+    assertSame(current, Context.current());
+
+    final TestError err = new TestError();
+    try {
+      base.supply(new Supplier<Object>() {
+        @Override
+        public Object get() {
+          throw err;
+        }
+      });
+      fail("Excepted exception");
+    } catch (TestError ex) {
+      assertSame(err, ex);
+    }
+    assertSame(current, Context.current());
+
+    current.detach(Context.ROOT);
+  }
+
   @Test
   public void currentContextExecutor() {
     QueuedExecutor queuedExecutor = new QueuedExecutor();