From 49b513504c06652816aeebcb49f2f650337c1aa5 Mon Sep 17 00:00:00 2001
From: Louis Ryan <lryan@google.com>
Date: Fri, 6 Nov 2015 17:01:26 -0800
Subject: [PATCH] Remove notion of stack from thread-local binding, use simple
 swapping instead. This allows for binding state to be reset to known-good
 states precisely which in turn facilitates making it safe to have 'detach'
 not throw exceptions and instead log a severe error that attach/detach calls
 were not correctly balanced.

---
 core/src/main/java/io/grpc/Context.java     | 118 ++++++++--------
 core/src/test/java/io/grpc/ContextTest.java | 146 +++++++++-----------
 2 files changed, 125 insertions(+), 139 deletions(-)

diff --git a/core/src/main/java/io/grpc/Context.java b/core/src/main/java/io/grpc/Context.java
index d13bff22b3..2975dbc2fb 100644
--- a/core/src/main/java/io/grpc/Context.java
+++ b/core/src/main/java/io/grpc/Context.java
@@ -39,7 +39,6 @@ import io.grpc.internal.SharedResourceHolder;
 
 import java.io.Closeable;
 import java.io.IOException;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
@@ -88,7 +87,7 @@ import javax.annotation.Nullable;
  * <pre>
  *   CancellableContext withCancellation = Context.current().withCancellation();
  *   try {
-   *   executorService.execute(withCancellation.wrap(new Runnable() {
+ *     executorService.execute(withCancellation.wrap(new Runnable() {
  *       public void run() {
  *         while (waitingForData() &amp;&amp; !Context.current().isCancelled()) {}
  *       }
@@ -142,17 +141,6 @@ public class Context {
         }
       };
 
-  /**
-   * Stack of context objects which is used to record attach & detach history on a thread.
-   */
-  private static final ThreadLocal<ArrayDeque<Context>> contextStack =
-      new ThreadLocal<ArrayDeque<Context>>() {
-    @Override
-    protected ArrayDeque<Context> initialValue() {
-      return new ArrayDeque<Context>();
-    }
-  };
-
   private static final Object[][] EMPTY_ENTRIES = new Object[0][2];
 
   /**
@@ -161,6 +149,16 @@ public class Context {
    */
   public static final Context ROOT = new Context(null);
 
+  /**
+   * Currently bound context.
+   */
+  private static final ThreadLocal<Context> localContext = new ThreadLocal<Context>() {
+    @Override
+    protected Context initialValue() {
+      return ROOT;
+    }
+  };
+
   /**
    * Create a {@link Key} with the given name.
    */
@@ -184,11 +182,11 @@ public class Context {
    * code stealing the ability to cancel arbitrarily.
    */
   public static Context current() {
-    ArrayDeque<Context> stack = contextStack.get();
-    if (stack.isEmpty()) {
+    Context current = localContext.get();
+    if (current == null) {
       return ROOT;
     }
-    return stack.peekLast();
+    return current;
   }
 
   private final Context parent;
@@ -229,8 +227,8 @@ public class Context {
   /**
    * Create a new context which is independently cancellable and also cascades cancellation from
    * its parent. Callers should ensure that either {@link CancellableContext#cancel(Throwable)}
-   * or {@link CancellableContext#detachAndCancel(Throwable)} are called to notify listeners and
-   * release the resources associated with them.
+   * or {@link CancellableContext#detachAndCancel(Context, Throwable)} are called to notify
+   * listeners and release the resources associated with them.
    *
    * <p>Sample usage:
    * <pre>
@@ -343,11 +341,30 @@ public class Context {
 
   /**
    * Attach this context to the thread and make it {@link #current}, the previously current context
-   * will be restored when detach is called. It is allowed to attach contexts where
-   * {@link #isCancelled()} is {@code true}.
+   * is returned. It is allowed to attach contexts where {@link #isCancelled()} is {@code true}.
    */
-  public void attach() {
-    contextStack.get().addLast(this);
+  public Context attach() {
+    Context previous = current();
+    localContext.set(this);
+    return previous;
+  }
+
+  /**
+   * Detach the current context from the thread and attach the provided replacement. If this
+   * context is not {@link #current()} a SEVERE message will be logged but the context to attach
+   * will still be bound.
+   */
+  public void detach(Context toAttach) {
+    Preconditions.checkNotNull(toAttach);
+    Context previous = current();
+    if (previous != this) {
+      // Log a severe message instead of throwing an exception as the context to attach is assumed
+      // to be the correct one and the unbalanced state represents a coding mistake in a lower
+      // layer in the stack that cannot be recovered from here.
+      LOG.log(Level.SEVERE, "Context was not attached when detaching",
+          new Throwable().fillInStackTrace());
+    }
+    localContext.set(toAttach);
   }
 
   // Visible for testing
@@ -355,27 +372,6 @@ public class Context {
     return current() == this;
   }
 
-  /**
-   * Detach the current context from the thread and restore the context that was previously
-   * attached to the thread as the 'current' context.
-   *
-   * @throws java.lang.IllegalStateException if this context is not {@link #current()}.
-   */
-  public void detach() {
-    ArrayDeque<Context> stack = contextStack.get();
-    if (stack.isEmpty()) {
-      if (this == ROOT) {
-        throw new IllegalStateException("Cannot detach root");
-      } else {
-        throw new IllegalStateException("Cannot detach non-root context when root is current");
-      }
-    }
-    if (stack.peekLast() != this) {
-      throw new IllegalStateException("Cannot detach a context that is not current");
-    }
-    stack.removeLast();
-  }
-
   /**
    * Is this context cancelled.
    */
@@ -492,11 +488,11 @@ public class Context {
     return new Runnable() {
       @Override
       public void run() {
-        attach();
+        Context previous = attach();
         try {
           r.run();
         } finally {
-          detach();
+          detach(previous);
         }
       }
     };
@@ -509,11 +505,11 @@ public class Context {
     return new Callable<C>() {
       @Override
       public C call() throws Exception {
-        attach();
+        Context previous = attach();
         try {
           return c.call();
         } finally {
-          detach();
+          detach(previous);
         }
       }
     };
@@ -586,7 +582,7 @@ public class Context {
      */
     private CancellableContext(Context parent) {
       super(parent, EMPTY_ENTRIES);
-      // Create a dummy that inherits from this to attach and detach so that you cannot retrieve a
+      // Create a dummy that inherits from this to attach so that you cannot retrieve a
       // cancellable context from Context.current()
       dummy = new Context(this, EMPTY_ENTRIES);
     }
@@ -611,13 +607,13 @@ public class Context {
 
 
     @Override
-    public void attach() {
-      dummy.attach();
+    public Context attach() {
+      return dummy.attach();
     }
 
     @Override
-    public void detach() {
-      dummy.detach();
+    public void detach(Context toAttach) {
+      dummy.detach(toAttach);
     }
 
     @Override
@@ -627,17 +623,17 @@ public class Context {
 
     /**
      * Attach this context to the thread and return a {@link AutoCloseable} that can be
-     * used with try-with-resource statements to properly {@link #detach} and {@link #cancel}
-     * the context on completion.
+     * used with try-with-resource statements to properly attach the previously bound context
+     * when {@link AutoCloseable#close()} is called.
      *
      * @return a {@link java.io.Closeable} which can be used with try-with-resource blocks.
      */
     public Closeable attachAsCloseable() {
-      attach();
+      final Context previous = attach();
       return new Closeable() {
         @Override
         public void close() throws IOException {
-          detachAndCancel(null);
+          detachAndCancel(previous, null);
         }
       };
     }
@@ -670,14 +666,14 @@ public class Context {
     }
 
     /**
-     * Cancel this context and detach it from the current context from the thread and restore the
-     * context that was previously attached to the thread as the 'current' context.
+     * Cancel this context and detach it as the current context from the thread.
      *
-     * @throws java.lang.IllegalStateException if this context is not {@link #current()}.
+     * @param toAttach context to make current.
+     * @param cause of cancellation, can be {@code null}.
      */
-    public void detachAndCancel(@Nullable Throwable cause) {
+    public void detachAndCancel(Context toAttach, @Nullable Throwable cause) {
       try {
-        detach();
+        detach(toAttach);
       } finally {
         cancel(cause);
       }
diff --git a/core/src/test/java/io/grpc/ContextTest.java b/core/src/test/java/io/grpc/ContextTest.java
index 0787ab4a21..79be52a4ae 100644
--- a/core/src/test/java/io/grpc/ContextTest.java
+++ b/core/src/test/java/io/grpc/ContextTest.java
@@ -41,6 +41,7 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.common.util.concurrent.SettableFuture;
 
 import io.grpc.internal.SharedResourceHolder;
 
@@ -63,6 +64,7 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Handler;
+import java.util.logging.Level;
 import java.util.logging.LogRecord;
 import java.util.logging.Logger;
 
@@ -97,10 +99,7 @@ public class ContextTest {
 
   @Before
   public void setUp() throws Exception {
-    // Detach all contexts back to the root.
-    while (Context.current() != Context.ROOT) {
-      Context.current().detach();
-    }
+    Context.ROOT.attach();
   }
 
   @After
@@ -114,29 +113,25 @@ public class ContextTest {
   }
 
   @Test
-  public void rootIsAlwaysBound() {
-    Context root = Context.current();
-    try {
-      root.detach();
-    } catch (IllegalStateException ise) {
-      // Expected
-      assertTrue(Context.ROOT.isCurrent());
-      return;
-    }
-    fail("Attempting to detach root should fail");
+  public void rootIsAlwaysBound() throws Exception {
+    final SettableFuture<Boolean> rootIsBound = SettableFuture.create();
+    new Thread(new Runnable() {
+      @Override
+      public void run() {
+        rootIsBound.set(Context.current() == Context.ROOT);
+      }
+    }).start();
+    assertTrue(rootIsBound.get(5, TimeUnit.SECONDS));
   }
 
   @Test
   public void rootCanBeAttached() {
-    Context root = Context.current();
-    Context.CancellableContext fork = root.fork();
+    Context.CancellableContext fork = Context.ROOT.fork();
+    fork.attach();
+    Context.ROOT.attach();
+    assertTrue(Context.ROOT.isCurrent());
     fork.attach();
-    root.attach();
-    assertTrue(root.isCurrent());
-    root.detach();
     assertTrue(fork.isCurrent());
-    fork.detach();
-    assertTrue(root.isCurrent());
   }
 
   @Test
@@ -155,58 +150,53 @@ public class ContextTest {
   @Test
   public void attachedCancellableContextCannotBeCastFromCurrent() {
     Context initial = Context.current();
-    Context.CancellableContext base = Context.current().withCancellation();
+    Context.CancellableContext base = initial.withCancellation();
     base.attach();
     assertFalse(Context.current() instanceof Context.CancellableContext);
+    assertNotSame(base, Context.current());
     assertNotSame(initial, Context.current());
-    base.detachAndCancel(null);
-    assertTrue(initial.isCurrent());
+    base.detachAndCancel(initial, null);
+    assertSame(initial, Context.current());
   }
 
   @Test
-  public void detachingNonCurrentThrowsIllegalStateException() {
-    Context current = Context.current();
-    Context base = Context.current().withValue(PET, "dog");
-    try {
-      base.detach();
-      fail("Expected exception");
-    } catch (IllegalStateException ise) {
-      // expected
-    }
-    assertSame(current, Context.current());
-
-    base.attach();
-    try {
-      current.withValue(PET, "cat").detach();
-      fail("Expected exception");
-    } catch (IllegalStateException ise) {
-      // expected
-    }
-    assertSame(base, Context.current());
-    base.detach();
+  public void attachingNonCurrentReturnsCurrent() {
+    Context initial = Context.current();
+    Context base = initial.withValue(PET, "dog");
+    assertSame(initial, base.attach());
+    assertSame(base, initial.attach());
   }
 
   @Test
-  public void detachUnwindsAttach() {
-    Context base = Context.current().fork();
-    Context child = base.withValue(COLOR, "red");
-    Context grandchild = child.withValue(COLOR, "blue");
-    base.attach();
-    base.attach();
-    child.attach();
-    base.attach();
-    grandchild.attach();
-    assertTrue(grandchild.isCurrent());
-    grandchild.detach();
-    assertTrue(base.isCurrent());
-    base.detach();
-    assertTrue(child.isCurrent());
-    child.detach();
-    assertTrue(base.isCurrent());
-    base.detach();
-    assertTrue(base.isCurrent());
-    base.detach();
-    assertTrue(Context.ROOT.isCurrent());
+  public void detachingNonCurrentLogsSevereMessage() {
+    final AtomicReference<LogRecord> logRef = new AtomicReference<LogRecord>();
+    Handler handler = new Handler() {
+      @Override
+      public void publish(LogRecord record) {
+        logRef.set(record);
+      }
+
+      @Override
+      public void flush() {
+      }
+
+      @Override
+      public void close() throws SecurityException {
+      }
+    };
+    Logger logger = Logger.getLogger(Context.class.getName());
+    try {
+      logger.addHandler(handler);
+      Context initial = Context.current();
+      Context base = initial.withValue(PET, "dog");
+      // Base is not attached
+      base.detach(initial);
+      assertSame(initial, Context.current());
+      assertNotNull(logRef.get());
+      assertEquals(Level.SEVERE, logRef.get().getLevel());
+    } finally {
+      logger.removeHandler(handler);
+    }
   }
 
   @Test
@@ -226,14 +216,14 @@ public class ContextTest {
     assertEquals("cheese", FOOD.get());
     assertNull(COLOR.get());
 
-    child.detach();
+    child.detach(base);
 
     // Should have values from base
     assertEquals("dog", PET.get());
     assertEquals("lasagna", FOOD.get());
     assertNull(COLOR.get());
 
-    base.detach();
+    base.detach(Context.ROOT);
 
     assertNull(PET.get());
     assertEquals("lasagna", FOOD.get());
@@ -253,7 +243,7 @@ public class ContextTest {
     assertEquals("blue", COLOR.get());
     assertEquals(fav, FAVORITE.get());
 
-    child.detach();
+    base.attach();
   }
 
   @Test
@@ -407,7 +397,7 @@ public class ContextTest {
     assertSame(t, attached.cause());
     assertSame(attached, listenerNotifedContext);
 
-    base.detach();
+    Context.ROOT.attach();
   }
 
   @Test
@@ -469,7 +459,7 @@ public class ContextTest {
     }
     assertSame(current, Context.current());
 
-    current.detach();
+    current.detach(Context.ROOT);
   }
 
   @Test
@@ -509,7 +499,7 @@ public class ContextTest {
     }
     assertSame(current, Context.current());
 
-    current.detach();
+    current.detach(Context.ROOT);
   }
 
   @Test
@@ -517,11 +507,11 @@ public class ContextTest {
     QueuedExecutor queuedExecutor = new QueuedExecutor();
     Executor executor = Context.propagate(queuedExecutor);
     Context base = Context.current().withValue(PET, "cat");
-    base.attach();
+    Context previous = base.attach();
     try {
       executor.execute(runner);
     } finally {
-      base.detach();
+      base.detach(previous);
     }
     assertEquals(1, queuedExecutor.runnables.size());
     queuedExecutor.runnables.remove().run();
@@ -541,12 +531,12 @@ public class ContextTest {
   @Test
   public void typicalTryFinallyHandling() throws Exception {
     Context base = Context.current().withValue(COLOR, "blue");
-    base.attach();
+    Context previous = base.attach();
     try {
       assertTrue(base.isCurrent());
       // Do something
     } finally {
-      base.detach();
+      base.detach(previous);
     }
     assertFalse(base.isCurrent());
   }
@@ -554,14 +544,14 @@ public class ContextTest {
   @Test
   public void typicalCancellableTryCatchFinallyHandling() throws Exception {
     Context.CancellableContext base = Context.current().withCancellation();
-    base.attach();
+    Context previous = base.attach();
     try {
       // Do something
       throw new IllegalStateException("Argh");
     } catch (IllegalStateException ise) {
       base.cancel(ise);
     } finally {
-      base.detachAndCancel(null);
+      base.detachAndCancel(previous, null);
     }
     assertTrue(base.isCancelled());
     assertNotNull(base.cause());
@@ -572,11 +562,11 @@ public class ContextTest {
     Context current = Context.current();
     Context.CancellableContext base = Context.current().withValue(FOOD, "fish").withCancellation();
 
-    base.attach();
+    Context previous = base.attach();
     Context baseAttached = Context.current();
     // Should not be able to get back the CancellableContext
     assertNotSame(baseAttached, base);
-    base.detach();
+    base.detach(previous);
 
     Closeable c = base.attachAsCloseable();
     assertEquals("fish", FOOD.get());