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 io.grpc.Context.CheckReturnValue;
20 import java.io.Closeable;
21 import java.util.ArrayList;
22 import java.util.concurrent.Callable;
23 import java.util.concurrent.Executor;
24 import java.util.concurrent.ScheduledExecutorService;
25 import java.util.concurrent.ScheduledFuture;
26 import java.util.concurrent.TimeUnit;
27 import java.util.concurrent.TimeoutException;
28 import java.util.concurrent.atomic.AtomicReference;
29 import java.util.logging.Level;
30 import java.util.logging.Logger;
31 
32 /**
33  * A context propagation mechanism which can carry scoped-values across API boundaries and between
34  * threads. Examples of state propagated via context include:
35  * <ul>
36  *   <li>Security principals and credentials.</li>
37  *   <li>Local and distributed tracing information.</li>
38  * </ul>
39  *
40  * <p>A Context object can be {@link #attach attached} to the {@link Storage}, which effectively
41  * forms a <b>scope</b> for the context.  The scope is bound to the current thread.  Within a scope,
42  * its Context is accessible even across API boundaries, through {@link #current}.  The scope is
43  * later exited by {@link #detach detaching} the Context.
44  *
45  * <p>Context objects are immutable and inherit state from their parent. To add or overwrite the
46  * current state a new context object must be created and then attached, replacing the previously
47  * bound context. For example:
48  *
49  * <pre>
50  *   Context withCredential = Context.current().withValue(CRED_KEY, cred);
51  *   withCredential.run(new Runnable() {
52  *     public void run() {
53  *        readUserRecords(userId, CRED_KEY.get());
54  *     }
55  *   });
56  * </pre>
57  *
58  * <p>Contexts are also used to represent a scoped unit of work. When the unit of work is done the
59  * context can be cancelled. This cancellation will also cascade to all descendant contexts. You can
60  * add a {@link CancellationListener} to a context to be notified when it or one of its ancestors
61  * has been cancelled. Cancellation does not release the state stored by a context and it's
62  * perfectly valid to {@link #attach()} an already cancelled context to make it current. To cancel a
63  * context (and its descendants) you first create a {@link CancellableContext} and when you need to
64  * signal cancellation call {@link CancellableContext#cancel} or {@link
65  * CancellableContext#detachAndCancel}. For example:
66  * <pre>
67  *   CancellableContext withCancellation = Context.current().withCancellation();
68  *   try {
69  *     withCancellation.run(new Runnable() {
70  *       public void run() {
71  *         while (waitingForData() &amp;&amp; !Context.current().isCancelled()) {}
72  *       }
73  *     });
74  *     doSomeWork();
75  *   } catch (Throwable t) {
76  *      withCancellation.cancel(t);
77  *   }
78  * </pre>
79  *
80  * <p>Contexts can also be created with a timeout relative to the system nano clock which will
81  * cause it to automatically cancel at the desired time.
82  *
83  *
84  * <p>Notes and cautions on use:
85  * <ul>
86  *    <li>Every {@code attach()} should have a {@code detach()} in the same method. And every
87  * CancellableContext should be cancelled at some point. Breaking these rules may lead to memory
88  * leaks.
89  *    <li>While Context objects are immutable they do not place such a restriction on the state
90  * they store.</li>
91  *    <li>Context is not intended for passing optional parameters to an API and developers should
92  * take care to avoid excessive dependence on context when designing an API.</li>
93  *    <li>Do not mock this class.  Use {@link #ROOT} for a non-null instance.
94  * </ul>
95  */
96 /* @DoNotMock("Use ROOT for a non-null Context") // commented out to avoid dependencies  */
97 @CheckReturnValue
98 public class Context {
99 
100   private static final Logger log = Logger.getLogger(Context.class.getName());
101 
102   private static final PersistentHashArrayMappedTrie<Key<?>, Object> EMPTY_ENTRIES =
103       new PersistentHashArrayMappedTrie<Key<?>, Object>();
104 
105   // Long chains of contexts are suspicious and usually indicate a misuse of Context.
106   // The threshold is arbitrarily chosen.
107   // VisibleForTesting
108   static final int CONTEXT_DEPTH_WARN_THRESH = 1000;
109 
110   /**
111    * The logical root context which is the ultimate ancestor of all contexts. This context
112    * is not cancellable and so will not cascade cancellation or retain listeners.
113    *
114    * <p>Never assume this is the default context for new threads, because {@link Storage} may define
115    * a default context that is different from ROOT.
116    */
117   public static final Context ROOT = new Context(null, EMPTY_ENTRIES);
118 
119   // Lazy-loaded storage. Delaying storage initialization until after class initialization makes it
120   // much easier to avoid circular loading since there can still be references to Context as long as
121   // they don't depend on storage, like key() and currentContextExecutor(). It also makes it easier
122   // to handle exceptions.
123   private static final AtomicReference<Storage> storage = new AtomicReference<Storage>();
124 
125   // For testing
storage()126   static Storage storage() {
127     Storage tmp = storage.get();
128     if (tmp == null) {
129       tmp = createStorage();
130     }
131     return tmp;
132   }
133 
createStorage()134   private static Storage createStorage() {
135     // Note that this method may be run more than once
136     try {
137       Class<?> clazz = Class.forName("io.grpc.override.ContextStorageOverride");
138       // The override's constructor is prohibited from triggering any code that can loop back to
139       // Context
140       Storage newStorage = (Storage) clazz.getConstructor().newInstance();
141       storage.compareAndSet(null, newStorage);
142     } catch (ClassNotFoundException e) {
143       Storage newStorage = new ThreadLocalContextStorage();
144       // Must set storage before logging, since logging may call Context.current().
145       if (storage.compareAndSet(null, newStorage)) {
146         // Avoid logging if this thread lost the race, to avoid confusion
147         log.log(Level.FINE, "Storage override doesn't exist. Using default", e);
148       }
149     } catch (Exception e) {
150       throw new RuntimeException("Storage override failed to initialize", e);
151     }
152     // Re-retreive from storage since compareAndSet may have failed (returned false) in case of
153     // race.
154     return storage.get();
155   }
156 
157   /**
158    * Create a {@link Key} with the given debug name. Multiple different keys may have the same name;
159    * the name is intended for debugging purposes and does not impact behavior.
160    */
key(String name)161   public static <T> Key<T> key(String name) {
162     return new Key<T>(name);
163   }
164 
165   /**
166    * Create a {@link Key} with the given debug name and default value. Multiple different keys may
167    * have the same name; the name is intended for debugging purposes and does not impact behavior.
168    */
keyWithDefault(String name, T defaultValue)169   public static <T> Key<T> keyWithDefault(String name, T defaultValue) {
170     return new Key<T>(name, defaultValue);
171   }
172 
173   /**
174    * Return the context associated with the current scope, will never return {@code null}.
175    *
176    * <p>Will never return {@link CancellableContext} even if one is attached, instead a
177    * {@link Context} is returned with the same properties and lifetime. This is to avoid
178    * code stealing the ability to cancel arbitrarily.
179    */
current()180   public static Context current() {
181     Context current = storage().current();
182     if (current == null) {
183       return ROOT;
184     }
185     return current;
186   }
187 
188   private ArrayList<ExecutableListener> listeners;
189   private CancellationListener parentListener = new ParentListener();
190   final CancellableContext cancellableAncestor;
191   final PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries;
192   // The number parents between this context and the root context.
193   final int generation;
194 
195   /**
196    * Construct a context that cannot be cancelled and will not cascade cancellation from its parent.
197    */
Context(PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries, int generation)198   private Context(PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries, int generation) {
199     cancellableAncestor = null;
200     this.keyValueEntries = keyValueEntries;
201     this.generation = generation;
202     validateGeneration(generation);
203   }
204 
205   /**
206    * Construct a context that cannot be cancelled but will cascade cancellation from its parent if
207    * it is cancellable.
208    */
Context(Context parent, PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries)209   private Context(Context parent, PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries) {
210     cancellableAncestor = cancellableAncestor(parent);
211     this.keyValueEntries = keyValueEntries;
212     this.generation = parent == null ? 0 : parent.generation + 1;
213     validateGeneration(generation);
214   }
215 
216   /**
217    * Create a new context which is independently cancellable and also cascades cancellation from
218    * its parent. Callers <em>must</em> ensure that either {@link
219    * CancellableContext#cancel(Throwable)} or {@link CancellableContext#detachAndCancel(Context,
220    * Throwable)} are called at a later point, in order to allow this context to be garbage
221    * collected.
222    *
223    * <p>Sample usage:
224    * <pre>
225    *   Context.CancellableContext withCancellation = Context.current().withCancellation();
226    *   try {
227    *     withCancellation.run(new Runnable() {
228    *       public void run() {
229    *         Context current = Context.current();
230    *         while (!current.isCancelled()) {
231    *           keepWorking();
232    *         }
233    *       }
234    *     });
235    *   } finally {
236    *     withCancellation.cancel(null);
237    *   }
238    * </pre>
239    */
withCancellation()240   public CancellableContext withCancellation() {
241     return new CancellableContext(this);
242   }
243 
244   /**
245    * Create a new context which will cancel itself after the given {@code duration} from now.
246    * The returned context will cascade cancellation of its parent. Callers may explicitly cancel
247    * the returned context prior to the deadline just as for {@link #withCancellation()}. If the unit
248    * of work completes before the deadline, the context should be explicitly cancelled to allow
249    * it to be garbage collected.
250    *
251    * <p>Sample usage:
252    * <pre>
253    *   Context.CancellableContext withDeadline = Context.current()
254    *       .withDeadlineAfter(5, TimeUnit.SECONDS, scheduler);
255    *   try {
256    *     withDeadline.run(new Runnable() {
257    *       public void run() {
258    *         Context current = Context.current();
259    *         while (!current.isCancelled()) {
260    *           keepWorking();
261    *         }
262    *       }
263    *     });
264    *   } finally {
265    *     withDeadline.cancel(null);
266    *   }
267    * </pre>
268    */
withDeadlineAfter(long duration, TimeUnit unit, ScheduledExecutorService scheduler)269   public CancellableContext withDeadlineAfter(long duration, TimeUnit unit,
270                                               ScheduledExecutorService scheduler) {
271     return withDeadline(Deadline.after(duration, unit), scheduler);
272   }
273 
274   /**
275    * Create a new context which will cancel itself at the given {@link Deadline}.
276    * The returned context will cascade cancellation of its parent. Callers may explicitly cancel
277    * the returned context prior to the deadline just as for {@link #withCancellation()}. If the unit
278    * of work completes before the deadline, the context should be explicitly cancelled to allow
279    * it to be garbage collected.
280    *
281    * <p>Sample usage:
282    * <pre>
283    *   Context.CancellableContext withDeadline = Context.current()
284    *      .withDeadline(someReceivedDeadline, scheduler);
285    *   try {
286    *     withDeadline.run(new Runnable() {
287    *       public void run() {
288    *         Context current = Context.current();
289    *         while (!current.isCancelled() &amp;&amp; moreWorkToDo()) {
290    *           keepWorking();
291    *         }
292    *       }
293    *     });
294    *   } finally {
295    *     withDeadline.cancel(null);
296    *   }
297    * </pre>
298    */
withDeadline(Deadline deadline, ScheduledExecutorService scheduler)299   public CancellableContext withDeadline(Deadline deadline,
300       ScheduledExecutorService scheduler) {
301     checkNotNull(deadline, "deadline");
302     checkNotNull(scheduler, "scheduler");
303     return new CancellableContext(this, deadline, scheduler);
304   }
305 
306   /**
307    * Create a new context with the given key value set. The new context will cascade cancellation
308    * from its parent.
309    *
310    <pre>
311    *   Context withCredential = Context.current().withValue(CRED_KEY, cred);
312    *   withCredential.run(new Runnable() {
313    *     public void run() {
314    *        readUserRecords(userId, CRED_KEY.get());
315    *     }
316    *   });
317    * </pre>
318    *
319    */
withValue(Key<V> k1, V v1)320   public <V> Context withValue(Key<V> k1, V v1) {
321     PersistentHashArrayMappedTrie<Key<?>, Object> newKeyValueEntries = keyValueEntries.put(k1, v1);
322     return new Context(this, newKeyValueEntries);
323   }
324 
325   /**
326    * Create a new context with the given key value set. The new context will cascade cancellation
327    * from its parent.
328    */
withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2)329   public <V1, V2> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2) {
330     PersistentHashArrayMappedTrie<Key<?>, Object> newKeyValueEntries =
331         keyValueEntries.put(k1, v1).put(k2, v2);
332     return new Context(this, newKeyValueEntries);
333   }
334 
335   /**
336    * Create a new context with the given key value set. The new context will cascade cancellation
337    * from its parent.
338    */
withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2, Key<V3> k3, V3 v3)339   public <V1, V2, V3> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2, Key<V3> k3, V3 v3) {
340     PersistentHashArrayMappedTrie<Key<?>, Object> newKeyValueEntries =
341         keyValueEntries.put(k1, v1).put(k2, v2).put(k3, v3);
342     return new Context(this, newKeyValueEntries);
343   }
344 
345   /**
346    * Create a new context with the given key value set. The new context will cascade cancellation
347    * from its parent.
348    */
withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2, Key<V3> k3, V3 v3, Key<V4> k4, V4 v4)349   public <V1, V2, V3, V4> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2,
350       Key<V3> k3, V3 v3, Key<V4> k4, V4 v4) {
351     PersistentHashArrayMappedTrie<Key<?>, Object> newKeyValueEntries =
352         keyValueEntries.put(k1, v1).put(k2, v2).put(k3, v3).put(k4, v4);
353     return new Context(this, newKeyValueEntries);
354   }
355 
356   /**
357    * Create a new context which propagates the values of this context but does not cascade its
358    * cancellation.
359    */
fork()360   public Context fork() {
361     return new Context(keyValueEntries, generation + 1);
362   }
363 
canBeCancelled()364   boolean canBeCancelled() {
365     return cancellableAncestor != null;
366   }
367 
368   /**
369    * Attach this context, thus enter a new scope within which this context is {@link #current}.  The
370    * previously current context is returned. It is allowed to attach contexts where {@link
371    * #isCancelled()} is {@code true}.
372    *
373    * <p>Instead of using {@code attach()} and {@link #detach(Context)} most use-cases are better
374    * served by using the {@link #run(Runnable)} or {@link #call(java.util.concurrent.Callable)} to
375    * execute work immediately within a context's scope. If work needs to be done in other threads it
376    * is recommended to use the 'wrap' methods or to use a propagating executor.
377    *
378    * <p>All calls to {@code attach()} should have a corresponding {@link #detach(Context)} within
379    * the same method:
380    * <pre>{@code Context previous = someContext.attach();
381    * try {
382    *   // Do work
383    * } finally {
384    *   someContext.detach(previous);
385    * }}</pre>
386    */
attach()387   public Context attach() {
388     Context prev = storage().doAttach(this);
389     if (prev == null) {
390       return ROOT;
391     }
392     return prev;
393   }
394 
395   /**
396    * Reverse an {@code attach()}, restoring the previous context and exiting the current scope.
397    *
398    * <p>This context should be the same context that was previously {@link #attach attached}.  The
399    * provided replacement should be what was returned by the same {@link #attach attach()} call.  If
400    * an {@code attach()} and a {@code detach()} meet above requirements, they match.
401    *
402    * <p>It is expected that between any pair of matching {@code attach()} and {@code detach()}, all
403    * {@code attach()}es and {@code detach()}es are called in matching pairs.  If this method finds
404    * that this context is not {@link #current current}, either you or some code in-between are not
405    * detaching correctly, and a SEVERE message will be logged but the context to attach will still
406    * be bound.  <strong>Never</strong> use {@code Context.current().detach()}, as this will
407    * compromise this error-detecting mechanism.
408    */
detach(Context toAttach)409   public void detach(Context toAttach) {
410     checkNotNull(toAttach, "toAttach");
411     storage().detach(this, toAttach);
412   }
413 
414   // Visible for testing
isCurrent()415   boolean isCurrent() {
416     return current() == this;
417   }
418 
419   /**
420    * Is this context cancelled.
421    */
isCancelled()422   public boolean isCancelled() {
423     if (cancellableAncestor == null) {
424       return false;
425     } else {
426       return cancellableAncestor.isCancelled();
427     }
428   }
429 
430   /**
431    * If a context {@link #isCancelled()} then return the cause of the cancellation or
432    * {@code null} if context was cancelled without a cause. If the context is not yet cancelled
433    * will always return {@code null}.
434    *
435    * <p>The cancellation cause is provided for informational purposes only and implementations
436    * should generally assume that it has already been handled and logged properly.
437    */
cancellationCause()438   public Throwable cancellationCause() {
439     if (cancellableAncestor == null) {
440       return null;
441     } else {
442       return cancellableAncestor.cancellationCause();
443     }
444   }
445 
446   /**
447    * A context may have an associated {@link Deadline} at which it will be automatically cancelled.
448    * @return A {@link io.grpc.Deadline} or {@code null} if no deadline is set.
449    */
getDeadline()450   public Deadline getDeadline() {
451     if (cancellableAncestor == null) {
452       return null;
453     }
454     return cancellableAncestor.getDeadline();
455   }
456 
457   /**
458    * Add a listener that will be notified when the context becomes cancelled.
459    */
addListener(final CancellationListener cancellationListener, final Executor executor)460   public void addListener(final CancellationListener cancellationListener,
461                           final Executor executor) {
462     checkNotNull(cancellationListener, "cancellationListener");
463     checkNotNull(executor, "executor");
464     if (canBeCancelled()) {
465       ExecutableListener executableListener =
466           new ExecutableListener(executor, cancellationListener);
467       synchronized (this) {
468         if (isCancelled()) {
469           executableListener.deliver();
470         } else {
471           if (listeners == null) {
472             // Now that we have a listener we need to listen to our parent so
473             // we can cascade listener notification.
474             listeners = new ArrayList<>();
475             listeners.add(executableListener);
476             if (cancellableAncestor != null) {
477               cancellableAncestor.addListener(parentListener, DirectExecutor.INSTANCE);
478             }
479           } else {
480             listeners.add(executableListener);
481           }
482         }
483       }
484     }
485   }
486 
487   /**
488    * Remove a {@link CancellationListener}.
489    */
removeListener(CancellationListener cancellationListener)490   public void removeListener(CancellationListener cancellationListener) {
491     if (!canBeCancelled()) {
492       return;
493     }
494     synchronized (this) {
495       if (listeners != null) {
496         for (int i = listeners.size() - 1; i >= 0; i--) {
497           if (listeners.get(i).listener == cancellationListener) {
498             listeners.remove(i);
499             // Just remove the first matching listener, given that we allow duplicate
500             // adds we should allow for duplicates after remove.
501             break;
502           }
503         }
504         // We have no listeners so no need to listen to our parent
505         if (listeners.isEmpty()) {
506           if (cancellableAncestor != null) {
507             cancellableAncestor.removeListener(parentListener);
508           }
509           listeners = null;
510         }
511       }
512     }
513   }
514 
515   /**
516    * Notify all listeners that this context has been cancelled and immediately release
517    * any reference to them so that they may be garbage collected.
518    */
notifyAndClearListeners()519   void notifyAndClearListeners() {
520     if (!canBeCancelled()) {
521       return;
522     }
523     ArrayList<ExecutableListener> tmpListeners;
524     synchronized (this) {
525       if (listeners == null) {
526         return;
527       }
528       tmpListeners = listeners;
529       listeners = null;
530     }
531     // Deliver events to non-child context listeners before we notify child contexts. We do this
532     // to cancel higher level units of work before child units. This allows for a better error
533     // handling paradigm where the higher level unit of work knows it is cancelled and so can
534     // ignore errors that bubble up as a result of cancellation of lower level units.
535     for (int i = 0; i < tmpListeners.size(); i++) {
536       if (!(tmpListeners.get(i).listener instanceof ParentListener)) {
537         tmpListeners.get(i).deliver();
538       }
539     }
540     for (int i = 0; i < tmpListeners.size(); i++) {
541       if (tmpListeners.get(i).listener instanceof ParentListener) {
542         tmpListeners.get(i).deliver();
543       }
544     }
545     if (cancellableAncestor != null) {
546       cancellableAncestor.removeListener(parentListener);
547     }
548   }
549 
550   // Used in tests to ensure that listeners are defined and released when cancellation cascades.
551   // It's very important to ensure that we do not accidentally retain listeners.
listenerCount()552   int listenerCount() {
553     synchronized (this) {
554       return listeners == null ? 0 : listeners.size();
555     }
556   }
557 
558   /**
559    * Immediately run a {@link Runnable} with this context as the {@link #current} context.
560    * @param r {@link Runnable} to run.
561    */
run(Runnable r)562   public void run(Runnable r) {
563     Context previous = attach();
564     try {
565       r.run();
566     } finally {
567       detach(previous);
568     }
569   }
570 
571   /**
572    * Immediately call a {@link Callable} with this context as the {@link #current} context.
573    * @param c {@link Callable} to call.
574    * @return result of call.
575    */
576   @CanIgnoreReturnValue
call(Callable<V> c)577   public <V> V call(Callable<V> c) throws Exception {
578     Context previous = attach();
579     try {
580       return c.call();
581     } finally {
582       detach(previous);
583     }
584   }
585 
586   /**
587    * Wrap a {@link Runnable} so that it executes with this context as the {@link #current} context.
588    */
wrap(final Runnable r)589   public Runnable wrap(final Runnable r) {
590     return new Runnable() {
591       @Override
592       public void run() {
593         Context previous = attach();
594         try {
595           r.run();
596         } finally {
597           detach(previous);
598         }
599       }
600     };
601   }
602 
603   /**
604    * Wrap a {@link Callable} so that it executes with this context as the {@link #current} context.
605    */
606   public <C> Callable<C> wrap(final Callable<C> c) {
607     return new Callable<C>() {
608       @Override
609       public C call() throws Exception {
610         Context previous = attach();
611         try {
612           return c.call();
613         } finally {
614           detach(previous);
615         }
616       }
617     };
618   }
619 
620   /**
621    * Wrap an {@link Executor} so that it always executes with this context as the {@link #current}
622    * context. It is generally expected that {@link #currentContextExecutor(Executor)} would be
623    * used more commonly than this method.
624    *
625    * <p>One scenario in which this executor may be useful is when a single thread is sharding work
626    * to multiple threads.
627    *
628    * @see #currentContextExecutor(Executor)
629    */
630   public Executor fixedContextExecutor(final Executor e) {
631     class FixedContextExecutor implements Executor {
632       @Override
633       public void execute(Runnable r) {
634         e.execute(wrap(r));
635       }
636     }
637 
638     return new FixedContextExecutor();
639   }
640 
641   /**
642    * Create an executor that propagates the {@link #current} context when {@link Executor#execute}
643    * is called as the {@link #current} context of the {@code Runnable} scheduled. <em>Note that this
644    * is a static method.</em>
645    *
646    * @see #fixedContextExecutor(Executor)
647    */
648   public static Executor currentContextExecutor(final Executor e) {
649     class CurrentContextExecutor implements Executor {
650       @Override
651       public void execute(Runnable r) {
652         e.execute(Context.current().wrap(r));
653       }
654     }
655 
656     return new CurrentContextExecutor();
657   }
658 
659   /**
660    * Lookup the value for a key in the context inheritance chain.
661    */
662   private Object lookup(Key<?> key) {
663     return keyValueEntries.get(key);
664   }
665 
666   /**
667    * A context which inherits cancellation from its parent but which can also be independently
668    * cancelled and which will propagate cancellation to its descendants. To avoid leaking memory,
669    * every CancellableContext must have a defined lifetime, after which it is guaranteed to be
670    * cancelled.
671    *
672    * <p>This class must be cancelled by either calling {@link #close} or {@link #cancel}.
673    * {@link #close} is equivalent to calling {@code cancel(null)}. It is safe to call the methods
674    * more than once, but only the first call will have any effect. Because it's safe to call the
675    * methods multiple times, users are encouraged to always call {@link #close} at the end of
676    * the operation, and disregard whether {@link #cancel} was already called somewhere else.
677    *
678    * <p>Blocking code can use the try-with-resources idiom:
679    * <pre>
680    * try (CancellableContext c = Context.current()
681    *     .withDeadlineAfter(100, TimeUnit.MILLISECONDS, executor)) {
682    *   Context toRestore = c.attach();
683    *   try {
684    *     // do some blocking work
685    *   } finally {
686    *     c.detach(toRestore);
687    *   }
688    * }</pre>
689    *
690    * <p>Asynchronous code will have to manually track the end of the CancellableContext's lifetime,
691    * and cancel the context at the appropriate time.
692    */
693   public static final class CancellableContext extends Context implements Closeable {
694 
695     private final Deadline deadline;
696     private final Context uncancellableSurrogate;
697 
698     private boolean cancelled;
699     private Throwable cancellationCause;
700     private ScheduledFuture<?> pendingDeadline;
701 
702     /**
703      * Create a cancellable context that does not have a deadline.
704      */
705     private CancellableContext(Context parent) {
706       super(parent, parent.keyValueEntries);
707       deadline = parent.getDeadline();
708       // Create a surrogate that inherits from this to attach so that you cannot retrieve a
709       // cancellable context from Context.current()
710       uncancellableSurrogate = new Context(this, keyValueEntries);
711     }
712 
713     /**
714      * Create a cancellable context that has a deadline.
715      */
716     private CancellableContext(Context parent, Deadline deadline,
717         ScheduledExecutorService scheduler) {
718       super(parent, parent.keyValueEntries);
719       Deadline parentDeadline = parent.getDeadline();
720       if (parentDeadline != null && parentDeadline.compareTo(deadline) <= 0) {
721         // The new deadline won't have an effect, so ignore it
722         deadline = parentDeadline;
723       } else {
724         // The new deadline has an effect
725         if (!deadline.isExpired()) {
726           // The parent deadline was after the new deadline so we need to install a listener
727           // on the new earlier deadline to trigger expiration for this context.
728           pendingDeadline = deadline.runOnExpiration(new Runnable() {
729             @Override
730             public void run() {
731               try {
732                 cancel(new TimeoutException("context timed out"));
733               } catch (Throwable t) {
734                 log.log(Level.SEVERE, "Cancel threw an exception, which should not happen", t);
735               }
736             }
737           }, scheduler);
738         } else {
739           // Cancel immediately if the deadline is already expired.
740           cancel(new TimeoutException("context timed out"));
741         }
742       }
743       this.deadline = deadline;
744       uncancellableSurrogate = new Context(this, keyValueEntries);
745     }
746 
747 
748     @Override
749     public Context attach() {
750       return uncancellableSurrogate.attach();
751     }
752 
753     @Override
754     public void detach(Context toAttach) {
755       uncancellableSurrogate.detach(toAttach);
756     }
757 
758     /**
759      * Returns true if the Context is the current context.
760      *
761      * @deprecated This method violates some GRPC class encapsulation and should not be used.
762      *     If you must know whether a Context is the current context, check whether it is the same
763      *     object returned by {@link Context#current()}.
764      */
765     //TODO(spencerfang): The superclass's method is package-private, so this should really match.
766     @Override
767     @Deprecated
768     public boolean isCurrent() {
769       return uncancellableSurrogate.isCurrent();
770     }
771 
772     /**
773      * Cancel this context and optionally provide a cause (can be {@code null}) for the
774      * cancellation. This will trigger notification of listeners. It is safe to call this method
775      * multiple times. Only the first call will have any effect.
776      *
777      * <p>Calling {@code cancel(null)} is the same as calling {@link #close}.
778      *
779      * @return {@code true} if this context cancelled the context and notified listeners,
780      *    {@code false} if the context was already cancelled.
781      */
782     @CanIgnoreReturnValue
783     public boolean cancel(Throwable cause) {
784       boolean triggeredCancel = false;
785       synchronized (this) {
786         if (!cancelled) {
787           cancelled = true;
788           if (pendingDeadline != null) {
789             // If we have a scheduled cancellation pending attempt to cancel it.
790             pendingDeadline.cancel(false);
791             pendingDeadline = null;
792           }
793           this.cancellationCause = cause;
794           triggeredCancel = true;
795         }
796       }
797       if (triggeredCancel) {
798         notifyAndClearListeners();
799       }
800       return triggeredCancel;
801     }
802 
803     /**
804      * Cancel this context and detach it as the current context.
805      *
806      * @param toAttach context to make current.
807      * @param cause of cancellation, can be {@code null}.
808      */
809     public void detachAndCancel(Context toAttach, Throwable cause) {
810       try {
811         detach(toAttach);
812       } finally {
813         cancel(cause);
814       }
815     }
816 
817     @Override
818     public boolean isCancelled() {
819       synchronized (this) {
820         if (cancelled) {
821           return true;
822         }
823       }
824       // Detect cancellation of parent in the case where we have no listeners and
825       // record it.
826       if (super.isCancelled()) {
827         cancel(super.cancellationCause());
828         return true;
829       }
830       return false;
831     }
832 
833     @Override
834     public Throwable cancellationCause() {
835       if (isCancelled()) {
836         return cancellationCause;
837       }
838       return null;
839     }
840 
841     @Override
842     public Deadline getDeadline() {
843       return deadline;
844     }
845 
846     @Override
847     boolean canBeCancelled() {
848       return true;
849     }
850 
851     /**
852      * Cleans up this object by calling {@code cancel(null)}.
853      */
854     @Override
855     public void close() {
856       cancel(null);
857     }
858   }
859 
860   /**
861    * A listener notified on context cancellation.
862    */
863   public interface CancellationListener {
864     /**
865      * @param context the newly cancelled context.
866      */
867     public void cancelled(Context context);
868   }
869 
870   /**
871    * Key for indexing values stored in a context.
872    */
873   public static final class Key<T> {
874     private final String name;
875     private final T defaultValue;
876 
877     Key(String name) {
878       this(name, null);
879     }
880 
881     Key(String name, T defaultValue) {
882       this.name = checkNotNull(name, "name");
883       this.defaultValue = defaultValue;
884     }
885 
886     /**
887      * Get the value from the {@link #current()} context for this key.
888      */
889     @SuppressWarnings("unchecked")
890     public T get() {
891       return get(Context.current());
892     }
893 
894     /**
895      * Get the value from the specified context for this key.
896      */
897     @SuppressWarnings("unchecked")
898     public T get(Context context) {
899       T value = (T) context.lookup(this);
900       return value == null ? defaultValue : value;
901     }
902 
903     @Override
904     public String toString() {
905       return name;
906     }
907   }
908 
909   /**
910    * Defines the mechanisms for attaching and detaching the "current" context. The constructor for
911    * extending classes <em>must not</em> trigger any activity that can use Context, which includes
912    * logging, otherwise it can trigger an infinite initialization loop. Extending classes must not
913    * assume that only one instance will be created; Context guarantees it will only use one
914    * instance, but it may create multiple and then throw away all but one.
915    *
916    * <p>The default implementation will put the current context in a {@link ThreadLocal}.  If an
917    * alternative implementation named {@code io.grpc.override.ContextStorageOverride} exists in the
918    * classpath, it will be used instead of the default implementation.
919    *
920    * <p>This API is <a href="https://github.com/grpc/grpc-java/issues/2462">experimental</a> and
921    * subject to change.
922    */
923   public abstract static class Storage {
924     /**
925      * @deprecated This is an old API that is no longer used.
926      */
927     @Deprecated
928     public void attach(Context toAttach) {
929       throw new UnsupportedOperationException("Deprecated. Do not call.");
930     }
931 
932     /**
933      * Implements {@link io.grpc.Context#attach}.
934      *
935      * <p>Caution: {@link Context#attach()} interprets a return value of {@code null} to mean
936      * the same thing as {@link Context#ROOT}.
937      *
938      * <p>See also: {@link #current()}.
939 
940      * @param toAttach the context to be attached
941      * @return A {@link Context} that should be passed back into {@link #detach(Context, Context)}
942      *        as the {@code toRestore} parameter. {@code null} is a valid return value, but see
943      *        caution note.
944      */
945     public Context doAttach(Context toAttach) {
946       // This is a default implementation to help migrate existing Storage implementations that
947       // have an attach() method but no doAttach() method.
948       Context current = current();
949       attach(toAttach);
950       return current;
951     }
952 
953     /**
954      * Implements {@link io.grpc.Context#detach}
955      *
956      * @param toDetach the context to be detached. Should be, or be equivalent to, the current
957      *        context of the current scope
958      * @param toRestore the context to be the current.  Should be, or be equivalent to, the context
959      *        of the outer scope
960      */
961     public abstract void detach(Context toDetach, Context toRestore);
962 
963     /**
964      * Implements {@link io.grpc.Context#current}.
965      *
966      * <p>Caution: {@link Context} interprets a return value of {@code null} to mean the same
967      * thing as {@code Context{@link #ROOT}}.
968      *
969      * <p>See also {@link #doAttach(Context)}.
970      *
971      * @return The context of the current scope. {@code null} is a valid return value, but see
972      *        caution note.
973      */
974     public abstract Context current();
975   }
976 
977   /**
978    * Stores listener and executor pair.
979    */
980   private class ExecutableListener implements Runnable {
981     private final Executor executor;
982     private final CancellationListener listener;
983 
984     private ExecutableListener(Executor executor, CancellationListener listener) {
985       this.executor = executor;
986       this.listener = listener;
987     }
988 
989     private void deliver() {
990       try {
991         executor.execute(this);
992       } catch (Throwable t) {
993         log.log(Level.INFO, "Exception notifying context listener", t);
994       }
995     }
996 
997     @Override
998     public void run() {
999       listener.cancelled(Context.this);
1000     }
1001   }
1002 
1003   private class ParentListener implements CancellationListener {
1004     @Override
1005     public void cancelled(Context context) {
1006       if (Context.this instanceof CancellableContext) {
1007         // Record cancellation with its cancellationCause.
1008         ((CancellableContext) Context.this).cancel(context.cancellationCause());
1009       } else {
1010         notifyAndClearListeners();
1011       }
1012     }
1013   }
1014 
1015   @CanIgnoreReturnValue
1016   private static <T> T checkNotNull(T reference, Object errorMessage) {
1017     if (reference == null) {
1018       throw new NullPointerException(String.valueOf(errorMessage));
1019     }
1020     return reference;
1021   }
1022 
1023   private enum DirectExecutor implements Executor {
1024     INSTANCE;
1025 
1026     @Override
1027     public void execute(Runnable command) {
1028       command.run();
1029     }
1030 
1031     @Override
1032     public String toString() {
1033       return "Context.DirectExecutor";
1034     }
1035   }
1036 
1037   /**
1038    * Returns {@code parent} if it is a {@link CancellableContext}, otherwise returns the parent's
1039    * {@link #cancellableAncestor}.
1040    */
1041   static CancellableContext cancellableAncestor(Context parent) {
1042     if (parent == null) {
1043       return null;
1044     }
1045     if (parent instanceof CancellableContext) {
1046       return (CancellableContext) parent;
1047     }
1048     // The parent simply cascades cancellations.
1049     // Bypass the parent and reference the ancestor directly (may be null).
1050     return parent.cancellableAncestor;
1051   }
1052 
1053   /**
1054    * If the ancestry chain length is unreasonably long, then print an error to the log and record
1055    * the stack trace.
1056    */
1057   private static void validateGeneration(int generation) {
1058     if (generation == CONTEXT_DEPTH_WARN_THRESH) {
1059       log.log(
1060           Level.SEVERE,
1061           "Context ancestry chain length is abnormally long. "
1062               + "This suggests an error in application code. "
1063               + "Length exceeded: " + CONTEXT_DEPTH_WARN_THRESH,
1064           new Exception());
1065     }
1066   }
1067 
1068   // Not using the standard com.google.errorprone.annotations.CheckReturnValue because that will
1069   // introduce dependencies that some io.grpc.Context API consumers may not want.
1070   @interface CheckReturnValue {}
1071 
1072   @interface CanIgnoreReturnValue {}
1073 }
1074