context: Make CancellableContext implement Closeable (#3607)

2 files changed, 60 insertions, 12 deletions

diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java
index 910233fae..4734e8545 100644
--- a/context/src/main/java/io/grpc/Context.java
+++ b/context/src/main/java/io/grpc/Context.java
@@ -16,6 +16,7 @@
 
 package io.grpc;
 
+import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
@@ -657,8 +658,29 @@ public class Context {
    * cancelled and which will propagate cancellation to its descendants. To avoid leaking memory,
    * every CancellableContext must have a defined lifetime, after which it is guaranteed to be
    * cancelled.
+   *
+   * <p>This class must be cancelled by either calling {@link #close} or {@link #cancel}.
+   * {@link #close} is equivalent to calling {@code cancel(null)}. It is safe to call the methods
+   * more than once, but only the first call will have any effect. Because it's safe to call the
+   * methods multiple times, users are encouraged to always call {@link #close} at the end of
+   * the operation, and disregard whether {@link #cancel} was already called somewhere else.
+   *
+   * <p>Blocking code can use the try-with-resources idiom:
+   * <pre>
+   * try (CancellableContext c = Context.current()
+   *     .withDeadlineAfter(100, TimeUnit.MILLISECONDS, executor)) {
+   *   Context toRestore = c.attach();
+   *   try {
+   *     // do some blocking work
+   *   } finally {
+   *     c.detach(toRestore);
+   *   }
+   * }</pre>
+   *
+   * <p>Asynchronous code will have to manually track the end of the CancellableContext's lifetime,
+   * and cancel the context at the appropriate time.
    */
-  public static final class CancellableContext extends Context {
+  public static final class CancellableContext extends Context implements Closeable {
 
     private final Deadline deadline;
     private final Context uncancellableSurrogate;
@@ -739,7 +761,10 @@ public class Context {
 
     /**
      * Cancel this context and optionally provide a cause (can be {@code null}) for the
-     * cancellation. This will trigger notification of listeners.
+     * cancellation. This will trigger notification of listeners. It is safe to call this method
+     * multiple times. Only the first call will have any effect.
+     *
+     * <p>Calling {@code cancel(null)} is the same as calling {@link #close}.
      *
      * @return {@code true} if this context cancelled the context and notified listeners,
      *    {@code false} if the context was already cancelled.
@@ -811,6 +836,14 @@ public class Context {
     boolean canBeCancelled() {
       return true;
     }
+
+    /**
+     * Cleans up this object by calling {@code cancel(null)}.
+     */
+    @Override
+    public void close() {
+      cancel(null);
+    }
   }
 
   /**
diff --git a/context/src/test/java/io/grpc/ContextTest.java b/context/src/test/java/io/grpc/ContextTest.java
index 93191bdf2..4f8e082ab 100644
--- a/context/src/test/java/io/grpc/ContextTest.java
+++ b/context/src/test/java/io/grpc/ContextTest.java
@@ -95,13 +95,14 @@ public class ContextTest {
   @After
   public void tearDown() throws Exception {
     scheduler.shutdown();
+    assertEquals(Context.ROOT, Context.current());
   }
 
   @Test
   public void defaultContext() throws Exception {
     final SettableFuture<Context> contextOfNewThread = SettableFuture.create();
     Context contextOfThisThread = Context.ROOT.withValue(PET, "dog");
-    contextOfThisThread.attach();
+    Context toRestore = contextOfThisThread.attach();
     new Thread(new Runnable() {
       @Override
       public void run() {
@@ -111,16 +112,22 @@ public class ContextTest {
     assertNotNull(contextOfNewThread.get(5, TimeUnit.SECONDS));
     assertNotSame(contextOfThisThread, contextOfNewThread.get());
     assertSame(contextOfThisThread, Context.current());
+    contextOfThisThread.detach(toRestore);
   }
 
   @Test
   public void rootCanBeAttached() {
     Context fork = Context.ROOT.fork();
-    fork.attach();
-    Context.ROOT.attach();
+    Context toRestore1 = fork.attach();
+    Context toRestore2 = Context.ROOT.attach();
     assertTrue(Context.ROOT.isCurrent());
-    fork.attach();
+
+    Context toRestore3 = fork.attach();
     assertTrue(fork.isCurrent());
+
+    fork.detach(toRestore3);
+    Context.ROOT.detach(toRestore2);
+    fork.detach(toRestore1);
   }
 
   @Test
@@ -225,14 +232,14 @@ public class ContextTest {
     Context base = Context.current().withValues(PET, "dog", COLOR, "blue");
     Context child = base.withValues(PET, "cat", FOOD, "cheese", FAVORITE, fav);
 
-    child.attach();
+    Context toRestore = child.attach();
 
     assertEquals("cat", PET.get());
     assertEquals("cheese", FOOD.get());
     assertEquals("blue", COLOR.get());
     assertEquals(fav, FAVORITE.get());
 
-    base.attach();
+    child.detach(toRestore);
   }
 
   @Test
@@ -241,7 +248,7 @@ public class ContextTest {
     Context base = Context.current().withValues(PET, "dog", COLOR, "blue");
     Context child = base.withValues(PET, "cat", FOOD, "cheese", FAVORITE, fav, LUCKY, 7);
 
-    child.attach();
+    Context toRestore = child.attach();
 
     assertEquals("cat", PET.get());
     assertEquals("cheese", FOOD.get());
@@ -249,7 +256,7 @@ public class ContextTest {
     assertEquals(fav, FAVORITE.get());
     assertEquals(7, (int) LUCKY.get());
 
-    base.attach();
+    child.detach(toRestore);
   }
 
   @Test
@@ -389,7 +396,7 @@ public class ContextTest {
   public void cancellableContextIsAttached() {
     Context.CancellableContext base = Context.current().withValue(FOOD, "fish").withCancellation();
     assertFalse(base.isCurrent());
-    base.attach();
+    Context toRestore = base.attach();
 
     Context attached = Context.current();
     assertSame("fish", FOOD.get());
@@ -406,7 +413,7 @@ public class ContextTest {
     assertSame(t, attached.cancellationCause());
     assertSame(attached, listenerNotifedContext);
 
-    Context.ROOT.attach();
+    base.detach(toRestore);
   }
 
   @Test
@@ -922,6 +929,14 @@ public class ContextTest {
   }
 
   @Test
+  public void cancellableContext_closeCancelsWithNullCause() throws Exception {
+    Context.CancellableContext cancellable = Context.current().withCancellation();
+    cancellable.close();
+    assertTrue(cancellable.isCancelled());
+    assertNull(cancellable.cancellationCause());
+  }
+
+  @Test
   public void errorWhenAncestryLengthLong() {
     final AtomicReference<LogRecord> logRef = new AtomicReference<LogRecord>();
     Handler handler = new Handler() {