1 /*
2  * Copyright 2015 The gRPC Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package io.grpc;
18 
19 import static io.grpc.Context.cancellableAncestor;
20 import static org.hamcrest.core.IsInstanceOf.instanceOf;
21 import static org.junit.Assert.assertEquals;
22 import static org.junit.Assert.assertFalse;
23 import static org.junit.Assert.assertNotNull;
24 import static org.junit.Assert.assertNotSame;
25 import static org.junit.Assert.assertNull;
26 import static org.junit.Assert.assertSame;
27 import static org.junit.Assert.assertThat;
28 import static org.junit.Assert.assertTrue;
29 import static org.junit.Assert.fail;
30 
31 import com.google.common.util.concurrent.MoreExecutors;
32 import com.google.common.util.concurrent.SettableFuture;
33 import java.lang.reflect.Field;
34 import java.lang.reflect.Modifier;
35 import java.util.ArrayDeque;
36 import java.util.Queue;
37 import java.util.concurrent.Callable;
38 import java.util.concurrent.CountDownLatch;
39 import java.util.concurrent.Executor;
40 import java.util.concurrent.Executors;
41 import java.util.concurrent.Future;
42 import java.util.concurrent.ScheduledExecutorService;
43 import java.util.concurrent.ScheduledThreadPoolExecutor;
44 import java.util.concurrent.TimeUnit;
45 import java.util.concurrent.TimeoutException;
46 import java.util.concurrent.atomic.AtomicBoolean;
47 import java.util.concurrent.atomic.AtomicReference;
48 import java.util.logging.Handler;
49 import java.util.logging.Level;
50 import java.util.logging.LogRecord;
51 import java.util.logging.Logger;
52 import java.util.regex.Pattern;
53 import org.junit.After;
54 import org.junit.Before;
55 import org.junit.Test;
56 import org.junit.runner.RunWith;
57 import org.junit.runners.JUnit4;
58 
59 /**
60  * Tests for {@link Context}.
61  */
62 @RunWith(JUnit4.class)
63 @SuppressWarnings("CheckReturnValue") // false-positive in test for current ver errorprone plugin
64 public class ContextTest {
65 
66   private static final Context.Key<String> PET = Context.key("pet");
67   private static final Context.Key<String> FOOD = Context.keyWithDefault("food", "lasagna");
68   private static final Context.Key<String> COLOR = Context.key("color");
69   private static final Context.Key<Object> FAVORITE = Context.key("favorite");
70   private static final Context.Key<Integer> LUCKY = Context.key("lucky");
71 
72   private Context listenerNotifedContext;
73   private CountDownLatch deadlineLatch = new CountDownLatch(1);
74   private Context.CancellationListener cancellationListener = new Context.CancellationListener() {
75     @Override
76     public void cancelled(Context context) {
77       listenerNotifedContext = context;
78       deadlineLatch.countDown();
79     }
80   };
81 
82   private Context observed;
83   private Runnable runner = new Runnable() {
84     @Override
85     public void run() {
86       observed = Context.current();
87     }
88   };
89   private ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
90 
91   @Before
setUp()92   public void setUp() throws Exception {
93     Context.ROOT.attach();
94   }
95 
96   @After
tearDown()97   public void tearDown() throws Exception {
98     scheduler.shutdown();
99     assertEquals(Context.ROOT, Context.current());
100   }
101 
102   @Test
defaultContext()103   public void defaultContext() throws Exception {
104     final SettableFuture<Context> contextOfNewThread = SettableFuture.create();
105     Context contextOfThisThread = Context.ROOT.withValue(PET, "dog");
106     Context toRestore = contextOfThisThread.attach();
107     new Thread(new Runnable() {
108       @Override
109       public void run() {
110         contextOfNewThread.set(Context.current());
111       }
112       }).start();
113     assertNotNull(contextOfNewThread.get(5, TimeUnit.SECONDS));
114     assertNotSame(contextOfThisThread, contextOfNewThread.get());
115     assertSame(contextOfThisThread, Context.current());
116     contextOfThisThread.detach(toRestore);
117   }
118 
119   @Test
rootCanBeAttached()120   public void rootCanBeAttached() {
121     Context fork = Context.ROOT.fork();
122     Context toRestore1 = fork.attach();
123     Context toRestore2 = Context.ROOT.attach();
124     assertTrue(Context.ROOT.isCurrent());
125 
126     Context toRestore3 = fork.attach();
127     assertTrue(fork.isCurrent());
128 
129     fork.detach(toRestore3);
130     Context.ROOT.detach(toRestore2);
131     fork.detach(toRestore1);
132   }
133 
134   @Test
rootCanNeverHaveAListener()135   public void rootCanNeverHaveAListener() {
136     Context root = Context.current();
137     root.addListener(cancellationListener, MoreExecutors.directExecutor());
138     assertEquals(0, root.listenerCount());
139   }
140 
141   @Test
rootIsNotCancelled()142   public void rootIsNotCancelled() {
143     assertFalse(Context.ROOT.isCancelled());
144     assertNull(Context.ROOT.cancellationCause());
145   }
146 
147   @Test
attachedCancellableContextCannotBeCastFromCurrent()148   public void attachedCancellableContextCannotBeCastFromCurrent() {
149     Context initial = Context.current();
150     Context.CancellableContext base = initial.withCancellation();
151     base.attach();
152     assertFalse(Context.current() instanceof Context.CancellableContext);
153     assertNotSame(base, Context.current());
154     assertNotSame(initial, Context.current());
155     base.detachAndCancel(initial, null);
156     assertSame(initial, Context.current());
157   }
158 
159   @Test
attachingNonCurrentReturnsCurrent()160   public void attachingNonCurrentReturnsCurrent() {
161     Context initial = Context.current();
162     Context base = initial.withValue(PET, "dog");
163     assertSame(initial, base.attach());
164     assertSame(base, initial.attach());
165   }
166 
167   @Test
detachingNonCurrentLogsSevereMessage()168   public void detachingNonCurrentLogsSevereMessage() {
169     final AtomicReference<LogRecord> logRef = new AtomicReference<LogRecord>();
170     Handler handler = new Handler() {
171       @Override
172       public void publish(LogRecord record) {
173         logRef.set(record);
174       }
175 
176       @Override
177       public void flush() {
178       }
179 
180       @Override
181       public void close() throws SecurityException {
182       }
183     };
184     Logger logger = Logger.getLogger(Context.storage().getClass().getName());
185     try {
186       logger.addHandler(handler);
187       Context initial = Context.current();
188       Context base = initial.withValue(PET, "dog");
189       // Base is not attached
190       base.detach(initial);
191       assertSame(initial, Context.current());
192       assertNotNull(logRef.get());
193       assertEquals(Level.SEVERE, logRef.get().getLevel());
194     } finally {
195       logger.removeHandler(handler);
196     }
197   }
198 
199   @Test
valuesAndOverrides()200   public void valuesAndOverrides() {
201     Context base = Context.current().withValue(PET, "dog");
202     Context child = base.withValues(PET, null, FOOD, "cheese");
203 
204     base.attach();
205 
206     assertEquals("dog", PET.get());
207     assertEquals("lasagna", FOOD.get());
208     assertNull(COLOR.get());
209 
210     child.attach();
211 
212     assertNull(PET.get());
213     assertEquals("cheese", FOOD.get());
214     assertNull(COLOR.get());
215 
216     child.detach(base);
217 
218     // Should have values from base
219     assertEquals("dog", PET.get());
220     assertEquals("lasagna", FOOD.get());
221     assertNull(COLOR.get());
222 
223     base.detach(Context.ROOT);
224 
225     assertNull(PET.get());
226     assertEquals("lasagna", FOOD.get());
227     assertNull(COLOR.get());
228   }
229 
230   @Test
withValuesThree()231   public void withValuesThree() {
232     Object fav = new Object();
233     Context base = Context.current().withValues(PET, "dog", COLOR, "blue");
234     Context child = base.withValues(PET, "cat", FOOD, "cheese", FAVORITE, fav);
235 
236     Context toRestore = child.attach();
237 
238     assertEquals("cat", PET.get());
239     assertEquals("cheese", FOOD.get());
240     assertEquals("blue", COLOR.get());
241     assertEquals(fav, FAVORITE.get());
242 
243     child.detach(toRestore);
244   }
245 
246   @Test
withValuesFour()247   public void withValuesFour() {
248     Object fav = new Object();
249     Context base = Context.current().withValues(PET, "dog", COLOR, "blue");
250     Context child = base.withValues(PET, "cat", FOOD, "cheese", FAVORITE, fav, LUCKY, 7);
251 
252     Context toRestore = child.attach();
253 
254     assertEquals("cat", PET.get());
255     assertEquals("cheese", FOOD.get());
256     assertEquals("blue", COLOR.get());
257     assertEquals(fav, FAVORITE.get());
258     assertEquals(7, (int) LUCKY.get());
259 
260     child.detach(toRestore);
261   }
262 
263   @Test
cancelReturnsFalseIfAlreadyCancelled()264   public void cancelReturnsFalseIfAlreadyCancelled() {
265     Context.CancellableContext base = Context.current().withCancellation();
266     assertTrue(base.cancel(null));
267     assertTrue(base.isCancelled());
268     assertFalse(base.cancel(null));
269   }
270 
271   @Test
notifyListenersOnCancel()272   public void notifyListenersOnCancel() {
273     class SetContextCancellationListener implements Context.CancellationListener {
274       private final AtomicReference<Context> observed;
275 
276       public SetContextCancellationListener(AtomicReference<Context> observed) {
277         this.observed = observed;
278       }
279 
280       @Override
281       public void cancelled(Context context) {
282         observed.set(context);
283       }
284     }
285 
286     Context.CancellableContext base = Context.current().withCancellation();
287     final AtomicReference<Context> observed1 = new AtomicReference<Context>();
288     base.addListener(new SetContextCancellationListener(observed1), MoreExecutors.directExecutor());
289     final AtomicReference<Context> observed2 = new AtomicReference<Context>();
290     base.addListener(new SetContextCancellationListener(observed2), MoreExecutors.directExecutor());
291     assertNull(observed1.get());
292     assertNull(observed2.get());
293     base.cancel(null);
294     assertSame(base, observed1.get());
295     assertSame(base, observed2.get());
296 
297     final AtomicReference<Context> observed3 = new AtomicReference<Context>();
298     base.addListener(new SetContextCancellationListener(observed3), MoreExecutors.directExecutor());
299     assertSame(base, observed3.get());
300   }
301 
302   @Test
exceptionOfExecutorDoesntThrow()303   public void exceptionOfExecutorDoesntThrow() {
304     final AtomicReference<Throwable> loggedThrowable = new AtomicReference<Throwable>();
305     Handler logHandler = new Handler() {
306       @Override
307       public void publish(LogRecord record) {
308         Throwable thrown = record.getThrown();
309         if (thrown != null) {
310           if (loggedThrowable.get() == null) {
311             loggedThrowable.set(thrown);
312           } else {
313             loggedThrowable.set(new RuntimeException("Too many exceptions", thrown));
314           }
315         }
316       }
317 
318       @Override
319       public void close() {}
320 
321       @Override
322       public void flush() {}
323     };
324     Logger logger = Logger.getLogger(Context.class.getName());
325     logger.addHandler(logHandler);
326     try {
327       Context.CancellableContext base = Context.current().withCancellation();
328       final AtomicReference<Runnable> observed1 = new AtomicReference<Runnable>();
329       final Error err = new Error();
330       base.addListener(cancellationListener, new Executor() {
331         @Override
332         public void execute(Runnable runnable) {
333           observed1.set(runnable);
334           throw err;
335         }
336       });
337       assertNull(observed1.get());
338       assertNull(loggedThrowable.get());
339       base.cancel(null);
340       assertNotNull(observed1.get());
341       assertSame(err, loggedThrowable.get());
342 
343       final Error err2 = new Error();
344       loggedThrowable.set(null);
345       final AtomicReference<Runnable> observed2 = new AtomicReference<Runnable>();
346       base.addListener(cancellationListener, new Executor() {
347         @Override
348         public void execute(Runnable runnable) {
349           observed2.set(runnable);
350           throw err2;
351         }
352       });
353       assertNotNull(observed2.get());
354       assertSame(err2, loggedThrowable.get());
355     } finally {
356       logger.removeHandler(logHandler);
357     }
358   }
359 
360   @Test
cascadingCancellationNotifiesChild()361   public void cascadingCancellationNotifiesChild() {
362     // Root is not cancellable so we can't cascade from it
363     Context.CancellableContext base = Context.current().withCancellation();
364     assertEquals(0, base.listenerCount());
365     Context child = base.withValue(FOOD, "lasagna");
366     assertEquals(0, child.listenerCount());
367     child.addListener(cancellationListener, MoreExecutors.directExecutor());
368     assertEquals(1, child.listenerCount());
369     assertEquals(1, base.listenerCount()); // child is now listening to base
370     assertFalse(base.isCancelled());
371     assertFalse(child.isCancelled());
372     IllegalStateException cause = new IllegalStateException();
373     base.cancel(cause);
374     assertTrue(base.isCancelled());
375     assertSame(cause, base.cancellationCause());
376     assertSame(child, listenerNotifedContext);
377     assertTrue(child.isCancelled());
378     assertSame(cause, child.cancellationCause());
379     assertEquals(0, base.listenerCount());
380     assertEquals(0, child.listenerCount());
381   }
382 
383   @Test
cascadingCancellationWithoutListener()384   public void cascadingCancellationWithoutListener() {
385     Context.CancellableContext base = Context.current().withCancellation();
386     Context child = base.withCancellation();
387     Throwable t = new Throwable();
388     base.cancel(t);
389     assertTrue(child.isCancelled());
390     assertSame(t, child.cancellationCause());
391   }
392 
393   // Context#isCurrent() and Context.CancellableContext#isCurrent() are intended
394   // to be visible only for testing. The deprecation is meant for users.
395   @SuppressWarnings("deprecation")
396   @Test
cancellableContextIsAttached()397   public void cancellableContextIsAttached() {
398     Context.CancellableContext base = Context.current().withValue(FOOD, "fish").withCancellation();
399     assertFalse(base.isCurrent());
400     Context toRestore = base.attach();
401 
402     Context attached = Context.current();
403     assertSame("fish", FOOD.get());
404     assertFalse(attached.isCancelled());
405     assertNull(attached.cancellationCause());
406     assertTrue(attached.canBeCancelled());
407     assertTrue(attached.isCurrent());
408     assertTrue(base.isCurrent());
409 
410     attached.addListener(cancellationListener, MoreExecutors.directExecutor());
411     Throwable t = new Throwable();
412     base.cancel(t);
413     assertTrue(attached.isCancelled());
414     assertSame(t, attached.cancellationCause());
415     assertSame(attached, listenerNotifedContext);
416 
417     base.detach(toRestore);
418   }
419 
420   @Test
cancellableContextCascadesFromCancellableParent()421   public void cancellableContextCascadesFromCancellableParent() {
422     // Root is not cancellable so we can't cascade from it
423     Context.CancellableContext base = Context.current().withCancellation();
424     Context child = base.withCancellation();
425     child.addListener(cancellationListener, MoreExecutors.directExecutor());
426     assertFalse(base.isCancelled());
427     assertFalse(child.isCancelled());
428     IllegalStateException cause = new IllegalStateException();
429     base.cancel(cause);
430     assertTrue(base.isCancelled());
431     assertSame(cause, base.cancellationCause());
432     assertSame(child, listenerNotifedContext);
433     assertTrue(child.isCancelled());
434     assertSame(cause, child.cancellationCause());
435     assertEquals(0, base.listenerCount());
436     assertEquals(0, child.listenerCount());
437   }
438 
439   @Test
nonCascadingCancellationDoesNotNotifyForked()440   public void nonCascadingCancellationDoesNotNotifyForked() {
441     Context.CancellableContext base = Context.current().withCancellation();
442     Context fork = base.fork();
443     fork.addListener(cancellationListener, MoreExecutors.directExecutor());
444     assertEquals(0, base.listenerCount());
445     assertEquals(0, fork.listenerCount());
446     assertTrue(base.cancel(new Throwable()));
447     assertNull(listenerNotifedContext);
448     assertFalse(fork.isCancelled());
449     assertNull(fork.cancellationCause());
450   }
451 
452   @Test
testWrapRunnable()453   public void testWrapRunnable() throws Exception {
454     Context base = Context.current().withValue(PET, "cat");
455     Context current = Context.current().withValue(PET, "fish");
456     current.attach();
457 
458     base.wrap(runner).run();
459     assertSame(base, observed);
460     assertSame(current, Context.current());
461 
462     current.wrap(runner).run();
463     assertSame(current, observed);
464     assertSame(current, Context.current());
465 
466     final TestError err = new TestError();
467     try {
468       base.wrap(new Runnable() {
469         @Override
470         public void run() {
471           throw err;
472         }
473       }).run();
474       fail("Expected exception");
475     } catch (TestError ex) {
476       assertSame(err, ex);
477     }
478     assertSame(current, Context.current());
479 
480     current.detach(Context.ROOT);
481   }
482 
483   @Test
testWrapCallable()484   public void testWrapCallable() throws Exception {
485     Context base = Context.current().withValue(PET, "cat");
486     Context current = Context.current().withValue(PET, "fish");
487     current.attach();
488 
489     final Object ret = new Object();
490     Callable<Object> callable = new Callable<Object>() {
491       @Override
492       public Object call() {
493         runner.run();
494         return ret;
495       }
496     };
497 
498     assertSame(ret, base.wrap(callable).call());
499     assertSame(base, observed);
500     assertSame(current, Context.current());
501 
502     assertSame(ret, current.wrap(callable).call());
503     assertSame(current, observed);
504     assertSame(current, Context.current());
505 
506     final TestError err = new TestError();
507     try {
508       base.wrap(new Callable<Object>() {
509         @Override
510         public Object call() {
511           throw err;
512         }
513       }).call();
514       fail("Excepted exception");
515     } catch (TestError ex) {
516       assertSame(err, ex);
517     }
518     assertSame(current, Context.current());
519 
520     current.detach(Context.ROOT);
521   }
522 
523   @Test
currentContextExecutor()524   public void currentContextExecutor() throws Exception {
525     QueuedExecutor queuedExecutor = new QueuedExecutor();
526     Executor executor = Context.currentContextExecutor(queuedExecutor);
527     Context base = Context.current().withValue(PET, "cat");
528     Context previous = base.attach();
529     try {
530       executor.execute(runner);
531     } finally {
532       base.detach(previous);
533     }
534     assertEquals(1, queuedExecutor.runnables.size());
535     queuedExecutor.runnables.remove().run();
536     assertSame(base, observed);
537   }
538 
539   @Test
fixedContextExecutor()540   public void fixedContextExecutor() throws Exception {
541     Context base = Context.current().withValue(PET, "cat");
542     QueuedExecutor queuedExecutor = new QueuedExecutor();
543     base.fixedContextExecutor(queuedExecutor).execute(runner);
544     assertEquals(1, queuedExecutor.runnables.size());
545     queuedExecutor.runnables.remove().run();
546     assertSame(base, observed);
547   }
548 
549   @Test
typicalTryFinallyHandling()550   public void typicalTryFinallyHandling() throws Exception {
551     Context base = Context.current().withValue(COLOR, "blue");
552     Context previous = base.attach();
553     try {
554       assertTrue(base.isCurrent());
555       // Do something
556     } finally {
557       base.detach(previous);
558     }
559     assertFalse(base.isCurrent());
560   }
561 
562   @Test
typicalCancellableTryCatchFinallyHandling()563   public void typicalCancellableTryCatchFinallyHandling() throws Exception {
564     Context.CancellableContext base = Context.current().withCancellation();
565     Context previous = base.attach();
566     try {
567       // Do something
568       throw new IllegalStateException("Argh");
569     } catch (IllegalStateException ise) {
570       base.cancel(ise);
571     } finally {
572       base.detachAndCancel(previous, null);
573     }
574     assertTrue(base.isCancelled());
575     assertNotNull(base.cancellationCause());
576   }
577 
578   @Test
rootHasNoDeadline()579   public void rootHasNoDeadline() {
580     assertNull(Context.ROOT.getDeadline());
581   }
582 
583   @Test
contextWithDeadlineHasDeadline()584   public void contextWithDeadlineHasDeadline() {
585     Context.CancellableContext cancellableContext =
586         Context.ROOT.withDeadlineAfter(1, TimeUnit.SECONDS, scheduler);
587     assertNotNull(cancellableContext.getDeadline());
588   }
589 
590   @Test
earlierParentDeadlineTakesPrecedenceOverLaterChildDeadline()591   public void earlierParentDeadlineTakesPrecedenceOverLaterChildDeadline() throws Exception {
592     final Deadline sooner = Deadline.after(100, TimeUnit.MILLISECONDS);
593     final Deadline later = Deadline.after(1, TimeUnit.MINUTES);
594     Context.CancellableContext parent = Context.ROOT.withDeadline(sooner, scheduler);
595     Context.CancellableContext child = parent.withDeadline(later, scheduler);
596     assertSame(parent.getDeadline(), sooner);
597     assertSame(child.getDeadline(), sooner);
598     final CountDownLatch latch = new CountDownLatch(1);
599     final AtomicReference<Exception> error = new AtomicReference<Exception>();
600     child.addListener(new Context.CancellationListener() {
601       @Override
602       public void cancelled(Context context) {
603         try {
604           assertTrue(sooner.isExpired());
605           assertFalse(later.isExpired());
606         } catch (Exception e) {
607           error.set(e);
608         }
609         latch.countDown();
610       }
611     }, MoreExecutors.directExecutor());
612     latch.await(3, TimeUnit.SECONDS);
613     if (error.get() != null) {
614       throw error.get();
615     }
616   }
617 
618   @Test
earlierChldDeadlineTakesPrecedenceOverLaterParentDeadline()619   public void earlierChldDeadlineTakesPrecedenceOverLaterParentDeadline() {
620     Deadline sooner = Deadline.after(1, TimeUnit.HOURS);
621     Deadline later = Deadline.after(1, TimeUnit.DAYS);
622     Context.CancellableContext parent = Context.ROOT.withDeadline(later, scheduler);
623     Context.CancellableContext child = parent.withDeadline(sooner, scheduler);
624     assertSame(parent.getDeadline(), later);
625     assertSame(child.getDeadline(), sooner);
626   }
627 
628   @Test
forkingContextDoesNotCarryDeadline()629   public void forkingContextDoesNotCarryDeadline() {
630     Deadline deadline = Deadline.after(1, TimeUnit.HOURS);
631     Context.CancellableContext parent = Context.ROOT.withDeadline(deadline, scheduler);
632     Context fork = parent.fork();
633     assertNull(fork.getDeadline());
634   }
635 
636   @Test
cancellationDoesNotExpireDeadline()637   public void cancellationDoesNotExpireDeadline() {
638     Deadline deadline = Deadline.after(1, TimeUnit.HOURS);
639     Context.CancellableContext parent = Context.ROOT.withDeadline(deadline, scheduler);
640     parent.cancel(null);
641     assertFalse(deadline.isExpired());
642   }
643 
644   @Test
absoluteDeadlineTriggersAndPropagates()645   public void absoluteDeadlineTriggersAndPropagates() throws Exception {
646     Context base = Context.current().withDeadline(Deadline.after(1, TimeUnit.SECONDS), scheduler);
647     Context child = base.withValue(FOOD, "lasagna");
648     child.addListener(cancellationListener, MoreExecutors.directExecutor());
649     assertFalse(base.isCancelled());
650     assertFalse(child.isCancelled());
651     assertTrue(deadlineLatch.await(2, TimeUnit.SECONDS));
652     assertTrue(base.isCancelled());
653     assertTrue(base.cancellationCause() instanceof TimeoutException);
654     assertSame(child, listenerNotifedContext);
655     assertTrue(child.isCancelled());
656     assertSame(base.cancellationCause(), child.cancellationCause());
657   }
658 
659   @Test
relativeDeadlineTriggersAndPropagates()660   public void relativeDeadlineTriggersAndPropagates() throws Exception {
661     Context base = Context.current().withDeadline(Deadline.after(1, TimeUnit.SECONDS), scheduler);
662     Context child = base.withValue(FOOD, "lasagna");
663     child.addListener(cancellationListener, MoreExecutors.directExecutor());
664     assertFalse(base.isCancelled());
665     assertFalse(child.isCancelled());
666     assertTrue(deadlineLatch.await(2, TimeUnit.SECONDS));
667     assertTrue(base.isCancelled());
668     assertTrue(base.cancellationCause() instanceof TimeoutException);
669     assertSame(child, listenerNotifedContext);
670     assertTrue(child.isCancelled());
671     assertSame(base.cancellationCause(), child.cancellationCause());
672   }
673 
674   @Test
innerDeadlineCompletesBeforeOuter()675   public void innerDeadlineCompletesBeforeOuter() throws Exception {
676     Context base = Context.current().withDeadline(Deadline.after(2, TimeUnit.SECONDS), scheduler);
677     Context child = base.withDeadline(Deadline.after(1, TimeUnit.SECONDS), scheduler);
678     child.addListener(cancellationListener, MoreExecutors.directExecutor());
679     assertFalse(base.isCancelled());
680     assertFalse(child.isCancelled());
681     assertTrue(deadlineLatch.await(2, TimeUnit.SECONDS));
682     assertFalse(base.isCancelled());
683     assertSame(child, listenerNotifedContext);
684     assertTrue(child.isCancelled());
685     assertTrue(child.cancellationCause() instanceof TimeoutException);
686 
687     deadlineLatch = new CountDownLatch(1);
688     base.addListener(cancellationListener, MoreExecutors.directExecutor());
689     assertTrue(deadlineLatch.await(2, TimeUnit.SECONDS));
690     assertTrue(base.isCancelled());
691     assertTrue(base.cancellationCause() instanceof TimeoutException);
692     assertNotSame(base.cancellationCause(), child.cancellationCause());
693   }
694 
695   @Test
cancellationCancelsScheduledTask()696   public void cancellationCancelsScheduledTask() {
697     ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
698     try {
699       assertEquals(0, executor.getQueue().size());
700       Context.CancellableContext base
701           = Context.current().withDeadline(Deadline.after(1, TimeUnit.DAYS), executor);
702       assertEquals(1, executor.getQueue().size());
703       base.cancel(null);
704       executor.purge();
705       assertEquals(0, executor.getQueue().size());
706     } finally {
707       executor.shutdown();
708     }
709   }
710 
711   private static class QueuedExecutor implements Executor {
712     private final Queue<Runnable> runnables = new ArrayDeque<Runnable>();
713 
714     @Override
execute(Runnable r)715     public void execute(Runnable r) {
716       runnables.add(r);
717     }
718   }
719 
720   @Test
childContextListenerNotifiedAfterParentListener()721   public void childContextListenerNotifiedAfterParentListener() {
722     Context.CancellableContext parent = Context.current().withCancellation();
723     Context child = parent.withValue(COLOR, "red");
724     final AtomicBoolean childAfterParent = new AtomicBoolean();
725     final AtomicBoolean parentCalled = new AtomicBoolean();
726     child.addListener(new Context.CancellationListener() {
727       @Override
728       public void cancelled(Context context) {
729         if (parentCalled.get()) {
730           childAfterParent.set(true);
731         }
732       }
733     }, MoreExecutors.directExecutor());
734     parent.addListener(new Context.CancellationListener() {
735       @Override
736       public void cancelled(Context context) {
737         parentCalled.set(true);
738       }
739     }, MoreExecutors.directExecutor());
740     parent.cancel(null);
741     assertTrue(parentCalled.get());
742     assertTrue(childAfterParent.get());
743   }
744 
745   @Test
expiredDeadlineShouldCancelContextImmediately()746   public void expiredDeadlineShouldCancelContextImmediately() {
747     Context parent = Context.current();
748     assertFalse(parent.isCancelled());
749 
750     Context.CancellableContext context = parent.withDeadlineAfter(0, TimeUnit.SECONDS, scheduler);
751     assertTrue(context.isCancelled());
752     assertThat(context.cancellationCause(), instanceOf(TimeoutException.class));
753 
754     assertFalse(parent.isCancelled());
755     Deadline deadline = Deadline.after(-10, TimeUnit.SECONDS);
756     assertTrue(deadline.isExpired());
757     context = parent.withDeadline(deadline, scheduler);
758     assertTrue(context.isCancelled());
759     assertThat(context.cancellationCause(), instanceOf(TimeoutException.class));
760   }
761 
762   /**
763    * Tests initializing the {@link Context} class with a custom logger which uses Context's storage
764    * when logging.
765    */
766   @Test
initContextWithCustomClassLoaderWithCustomLogger()767   public void initContextWithCustomClassLoaderWithCustomLogger() throws Exception {
768     StaticTestingClassLoader classLoader =
769         new StaticTestingClassLoader(
770             getClass().getClassLoader(), Pattern.compile("io\\.grpc\\.[^.]+"));
771     Class<?> runnable =
772         classLoader.loadClass(LoadMeWithStaticTestingClassLoader.class.getName());
773 
774     ((Runnable) runnable.getDeclaredConstructor().newInstance()).run();
775   }
776 
777   /**
778    * Ensure that newly created threads can attach/detach a context.
779    * The current test thread already has a context manually attached in {@link #setUp()}.
780    */
781   @Test
newThreadAttachContext()782   public void newThreadAttachContext() throws Exception {
783     Context parent = Context.current().withValue(COLOR, "blue");
784     parent.call(new Callable<Object>() {
785       @Override
786       public Object call() throws Exception {
787         assertEquals("blue", COLOR.get());
788 
789         final Context child = Context.current().withValue(COLOR, "red");
790         Future<String> workerThreadVal = scheduler
791             .submit(new Callable<String>() {
792               @Override
793               public String call() {
794                 Context initial = Context.current();
795                 assertNotNull(initial);
796                 Context toRestore = child.attach();
797                 try {
798                   assertNotNull(toRestore);
799                   return COLOR.get();
800                 } finally {
801                   child.detach(toRestore);
802                   assertEquals(initial, Context.current());
803                 }
804               }
805             });
806         assertEquals("red", workerThreadVal.get());
807 
808         assertEquals("blue", COLOR.get());
809         return null;
810       }
811     });
812   }
813 
814   /**
815    * Similar to {@link #newThreadAttachContext()} but without giving the new thread a specific ctx.
816    */
817   @Test
newThreadWithoutContext()818   public void newThreadWithoutContext() throws Exception {
819     Context parent = Context.current().withValue(COLOR, "blue");
820     parent.call(new Callable<Object>() {
821       @Override
822       public Object call() throws Exception {
823         assertEquals("blue", COLOR.get());
824 
825         Future<String> workerThreadVal = scheduler
826             .submit(new Callable<String>() {
827               @Override
828               public String call() {
829                 assertNotNull(Context.current());
830                 return COLOR.get();
831               }
832             });
833         assertEquals(null, workerThreadVal.get());
834 
835         assertEquals("blue", COLOR.get());
836         return null;
837       }
838     });
839   }
840 
841   @Test
storageReturnsNullTest()842   public void storageReturnsNullTest() throws Exception {
843     Field storage = Context.class.getDeclaredField("storage");
844     assertTrue(Modifier.isFinal(storage.getModifiers()));
845     // use reflection to forcibly change the storage object to a test object
846     storage.setAccessible(true);
847     Object o = storage.get(null);
848     @SuppressWarnings("unchecked")
849     AtomicReference<Context.Storage> storageRef = (AtomicReference<Context.Storage>) o;
850     Context.Storage originalStorage = storageRef.get();
851     try {
852       storageRef.set(new Context.Storage() {
853         @Override
854         public Context doAttach(Context toAttach) {
855           return null;
856         }
857 
858         @Override
859         public void detach(Context toDetach, Context toRestore) {
860           // noop
861         }
862 
863         @Override
864         public Context current() {
865           return null;
866         }
867       });
868       // current() returning null gets transformed into ROOT
869       assertEquals(Context.ROOT, Context.current());
870 
871       // doAttach() returning null gets transformed into ROOT
872       Context blueContext = Context.current().withValue(COLOR, "blue");
873       Context toRestore = blueContext.attach();
874       assertEquals(Context.ROOT, toRestore);
875 
876       // final sanity check
877       blueContext.detach(toRestore);
878       assertEquals(Context.ROOT, Context.current());
879     } finally {
880       // undo the changes
881       storageRef.set(originalStorage);
882       storage.setAccessible(false);
883     }
884   }
885 
886   @Test
cancellableAncestorTest()887   public void cancellableAncestorTest() {
888     assertEquals(null, cancellableAncestor(null));
889 
890     Context c = Context.current();
891     assertFalse(c.canBeCancelled());
892     assertEquals(null, cancellableAncestor(c));
893 
894     Context.CancellableContext withCancellation = c.withCancellation();
895     assertEquals(withCancellation, cancellableAncestor(withCancellation));
896 
897     Context child = withCancellation.withValue(COLOR, "blue");
898     assertFalse(child instanceof Context.CancellableContext);
899     assertEquals(withCancellation, cancellableAncestor(child));
900 
901     Context grandChild = child.withValue(COLOR, "red");
902     assertFalse(grandChild instanceof Context.CancellableContext);
903     assertEquals(withCancellation, cancellableAncestor(grandChild));
904   }
905 
906   @Test
cancellableAncestorIntegrationTest()907   public void cancellableAncestorIntegrationTest() {
908     Context base = Context.current();
909 
910     Context blue = base.withValue(COLOR, "blue");
911     assertNull(blue.cancellableAncestor);
912     Context.CancellableContext cancellable = blue.withCancellation();
913     assertNull(cancellable.cancellableAncestor);
914     Context childOfCancel = cancellable.withValue(PET, "cat");
915     assertSame(cancellable, childOfCancel.cancellableAncestor);
916     Context grandChildOfCancel = childOfCancel.withValue(FOOD, "lasagna");
917     assertSame(cancellable, grandChildOfCancel.cancellableAncestor);
918 
919     Context.CancellableContext cancellable2 = childOfCancel.withCancellation();
920     assertSame(cancellable, cancellable2.cancellableAncestor);
921     Context childOfCancellable2 = cancellable2.withValue(PET, "dog");
922     assertSame(cancellable2, childOfCancellable2.cancellableAncestor);
923   }
924 
925   @Test
cancellableAncestorFork()926   public void cancellableAncestorFork() {
927     Context.CancellableContext cancellable = Context.current().withCancellation();
928     Context fork = cancellable.fork();
929     assertNull(fork.cancellableAncestor);
930   }
931 
932   @Test
cancellableContext_closeCancelsWithNullCause()933   public void cancellableContext_closeCancelsWithNullCause() throws Exception {
934     Context.CancellableContext cancellable = Context.current().withCancellation();
935     cancellable.close();
936     assertTrue(cancellable.isCancelled());
937     assertNull(cancellable.cancellationCause());
938   }
939 
940   @Test
errorWhenAncestryLengthLong()941   public void errorWhenAncestryLengthLong() {
942     final AtomicReference<LogRecord> logRef = new AtomicReference<LogRecord>();
943     Handler handler = new Handler() {
944       @Override
945       public void publish(LogRecord record) {
946         logRef.set(record);
947       }
948 
949       @Override
950       public void flush() {
951       }
952 
953       @Override
954       public void close() throws SecurityException {
955       }
956     };
957     Logger logger = Logger.getLogger(Context.class.getName());
958     try {
959       logger.addHandler(handler);
960       Context ctx = Context.current();
961       for (int i = 0; i < Context.CONTEXT_DEPTH_WARN_THRESH ; i++) {
962         assertNull(logRef.get());
963         ctx = ctx.fork();
964       }
965       ctx = ctx.fork();
966       assertNotNull(logRef.get());
967       assertNotNull(logRef.get().getThrown());
968       assertEquals(Level.SEVERE, logRef.get().getLevel());
969     } finally {
970       logger.removeHandler(handler);
971     }
972   }
973 
974   // UsedReflectively
975   public static final class LoadMeWithStaticTestingClassLoader implements Runnable {
976     @Override
run()977     public void run() {
978       Logger logger = Logger.getLogger(Context.class.getName());
979       logger.setLevel(Level.ALL);
980       Handler handler = new Handler() {
981         @Override
982         public void publish(LogRecord record) {
983           Context ctx = Context.current();
984           Context previous = ctx.attach();
985           ctx.detach(previous);
986         }
987 
988         @Override
989         public void flush() {
990         }
991 
992         @Override
993         public void close() throws SecurityException {
994         }
995       };
996       logger.addHandler(handler);
997 
998       try {
999         assertNotNull(Context.ROOT);
1000       } finally {
1001         logger.removeHandler(handler);
1002       }
1003     }
1004   }
1005 
1006   /** Allows more precise catch blocks than plain Error to avoid catching AssertionError. */
1007   private static final class TestError extends Error {}
1008 }
1009