Home | History | Annotate | Download | only in grpc
      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
     92   public void setUp() throws Exception {
     93     Context.ROOT.attach();
     94   }
     95 
     96   @After
     97   public void tearDown() throws Exception {
     98     scheduler.shutdown();
     99     assertEquals(Context.ROOT, Context.current());
    100   }
    101 
    102   @Test
    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
    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
    135   public void rootCanNeverHaveAListener() {
    136     Context root = Context.current();
    137     root.addListener(cancellationListener, MoreExecutors.directExecutor());
    138     assertEquals(0, root.listenerCount());
    139   }
    140 
    141   @Test
    142   public void rootIsNotCancelled() {
    143     assertFalse(Context.ROOT.isCancelled());
    144     assertNull(Context.ROOT.cancellationCause());
    145   }
    146 
    147   @Test
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    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
    579   public void rootHasNoDeadline() {
    580     assertNull(Context.ROOT.getDeadline());
    581   }
    582 
    583   @Test
    584   public void contextWithDeadlineHasDeadline() {
    585     Context.CancellableContext cancellableContext =
    586         Context.ROOT.withDeadlineAfter(1, TimeUnit.SECONDS, scheduler);
    587     assertNotNull(cancellableContext.getDeadline());
    588   }
    589 
    590   @Test
    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
    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
    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
    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
    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
    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
    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
    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
    715     public void execute(Runnable r) {
    716       runnables.add(r);
    717     }
    718   }
    719 
    720   @Test
    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
    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
    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
    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
    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
    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
    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
    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
    926   public void cancellableAncestorFork() {
    927     Context.CancellableContext cancellable = Context.current().withCancellation();
    928     Context fork = cancellable.fork();
    929     assertNull(fork.cancellableAncestor);
    930   }
    931 
    932   @Test
    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
    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
    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