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; *
* CancellableContext withCancellation = Context.current().withCancellation();
* try {
- * executorService.execute(withCancellation.wrap(new Runnable() {
+ * executorService.execute(withCancellation.wrap(new Runnable() {
* public void run() {
* while (waitingForData() && !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> contextStack =
- new ThreadLocal>() {
- @Override
- protected ArrayDeque initialValue() {
- return new ArrayDeque();
- }
- };
-
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 localContext = new ThreadLocal() {
+ @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 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.
*
* Sample usage:
*
@@ -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 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() {
@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 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 logRef = new AtomicReference();
+ 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());