blob: 27a736750009b237574c941ad7e3d73638bbb807 [file] [log] [blame]
Louis Ryanf6121162015-06-24 13:54:13 -07001/*
2 * Copyright 2015, Google Inc. All rights reserved.
3 *
4 * Redistribution and use in source and binary forms, with or without
5 * modification, are permitted provided that the following conditions are
6 * met:
7 *
8 * * Redistributions of source code must retain the above copyright
9 * notice, this list of conditions and the following disclaimer.
10 * * Redistributions in binary form must reproduce the above
11 * copyright notice, this list of conditions and the following disclaimer
12 * in the documentation and/or other materials provided with the
13 * distribution.
14 *
15 * * Neither the name of Google Inc. nor the names of its
16 * contributors may be used to endorse or promote products derived from
17 * this software without specific prior written permission.
18 *
19 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30 */
31
32package io.grpc;
33
Louis Ryanf6121162015-06-24 13:54:13 -070034import java.util.ArrayList;
35import java.util.concurrent.Callable;
36import java.util.concurrent.Executor;
Louis Ryanf6121162015-06-24 13:54:13 -070037import java.util.concurrent.ScheduledExecutorService;
38import java.util.concurrent.ScheduledFuture;
39import java.util.concurrent.TimeUnit;
40import java.util.concurrent.TimeoutException;
41import java.util.logging.Level;
42import java.util.logging.Logger;
43
Louis Ryanf6121162015-06-24 13:54:13 -070044/**
Louis Ryan8dad7342016-02-18 16:30:06 -080045 * A context propagation mechanism which can carry scoped-values across API boundaries and between
46 * threads. Examples of state propagated via context include:
Louis Ryanf6121162015-06-24 13:54:13 -070047 * <ul>
Louis Ryanf6121162015-06-24 13:54:13 -070048 * <li>Security principals and credentials.</li>
Louis Ryan8dad7342016-02-18 16:30:06 -080049 * <li>Local and distributed tracing information.</li>
Louis Ryanf6121162015-06-24 13:54:13 -070050 * </ul>
51 *
52 * <p>Context objects make their state available by being attached to the executing thread using
53 * a {@link ThreadLocal}. The context object bound to a thread is considered {@link #current()}.
54 * Context objects are immutable and inherit state from their parent. To add or overwrite the
55 * current state a new context object must be created and then attached to the thread replacing the
56 * previously bound context. For example:
57 * <pre>
58 * Context withCredential = Context.current().withValue(CRED_KEY, cred);
59 * executorService.execute(withCredential.wrap(new Runnable() {
60 * public void run() {
61 * readUserRecords(userId, CRED_KEY.get());
62 * }
63 * }));
64
65 * </pre>
66 *
Louis Ryan8dad7342016-02-18 16:30:06 -080067 *
68 * <p>Contexts are also used to represent a scoped unit of work. When the unit of work is
69 * done the context can be cancelled. This cancellation will also cascade to all descendant
70 * contexts. You can add a {@link CancellationListener} to a context to be notified when it or
Louis Ryanf6121162015-06-24 13:54:13 -070071 * one of its ancestors has been cancelled. Cancellation does not release the state stored by
72 * a context and it's perfectly valid to {@link #attach()} an already cancelled context to a
73 * thread to make it current. To cancel a context (and its descendants) you first create a
74 * {@link CancellableContext} and when you need to signal cancellation call
75 * {@link CancellableContext#cancel} or {@link CancellableContext#detachAndCancel}. For example:
76 * <pre>
77 * CancellableContext withCancellation = Context.current().withCancellation();
78 * try {
Louis Ryan49b51352015-11-06 17:01:26 -080079 * executorService.execute(withCancellation.wrap(new Runnable() {
Louis Ryanf6121162015-06-24 13:54:13 -070080 * public void run() {
81 * while (waitingForData() &amp;&amp; !Context.current().isCancelled()) {}
82 * }
83 * });
84 * doSomeWork();
85 * } catch (Throwable t) {
86 * withCancellation.cancel(t);
87 * }
88 * </pre>
89 *
Louis Ryan8dad7342016-02-18 16:30:06 -080090 * <p>Contexts can also be created with a timeout relative to the system nano clock which will
91 * cause it to automatically cancel at the desired time.
92 *
Louis Ryanf6121162015-06-24 13:54:13 -070093 *
94 * <p>Notes and cautions on use:
95 * <ul>
96 * <li>While Context objects are immutable they do not place such a restriction on the state
97 * they store.</li>
98 * <li>Context is not intended for passing optional parameters to an API and developers should
99 * take care to avoid excessive dependence on context when designing an API.</li>
100 * <li>If Context is being used in an environment that needs to support class unloading it is the
101 * responsibility of the application to ensure that all contexts are properly cancelled.</li>
102 * </ul>
103 */
104public class Context {
105
Eric Andersonc8648dc2016-08-04 11:08:36 -0700106 private static final Logger log = Logger.getLogger(Context.class.getName());
Louis Ryanf6121162015-06-24 13:54:13 -0700107
Louis Ryanf6121162015-06-24 13:54:13 -0700108 private static final Object[][] EMPTY_ENTRIES = new Object[0][2];
109
Louis Ryan8dad7342016-02-18 16:30:06 -0800110 private static final Key<Deadline> DEADLINE_KEY = new Key<Deadline>("deadline");
111
Louis Ryanf6121162015-06-24 13:54:13 -0700112 /**
113 * The logical root context which is {@link #current()} if no other context is bound. This context
114 * is not cancellable and so will not cascade cancellation or retain listeners.
115 */
116 public static final Context ROOT = new Context(null);
117
118 /**
Louis Ryan49b51352015-11-06 17:01:26 -0800119 * Currently bound context.
120 */
121 private static final ThreadLocal<Context> localContext = new ThreadLocal<Context>() {
122 @Override
123 protected Context initialValue() {
124 return ROOT;
125 }
126 };
127
128 /**
Eric Anderson4a427b72016-01-23 09:24:00 -0800129 * Create a {@link Key} with the given debug name. Multiple different keys may have the same name;
130 * the name is intended for debugging purposes and does not impact behavior.
Louis Ryanf6121162015-06-24 13:54:13 -0700131 */
132 public static <T> Key<T> key(String name) {
133 return new Key<T>(name);
134 }
135
136 /**
Eric Anderson4a427b72016-01-23 09:24:00 -0800137 * Create a {@link Key} with the given debug name and default value. Multiple different keys may
138 * have the same name; the name is intended for debugging purposes and does not impact behavior.
Louis Ryanf6121162015-06-24 13:54:13 -0700139 */
140 public static <T> Key<T> keyWithDefault(String name, T defaultValue) {
141 return new Key<T>(name, defaultValue);
142 }
143
144 /**
145 * Return the context associated with the current thread, will never return {@code null} as
146 * the {@link #ROOT} context is implicitly associated with all threads.
147 *
148 * <p>Will never return {@link CancellableContext} even if one is attached, instead a
149 * {@link Context} is returned with the same properties and lifetime. This is to avoid
150 * code stealing the ability to cancel arbitrarily.
151 */
152 public static Context current() {
Louis Ryan49b51352015-11-06 17:01:26 -0800153 Context current = localContext.get();
154 if (current == null) {
Louis Ryanf6121162015-06-24 13:54:13 -0700155 return ROOT;
156 }
Louis Ryan49b51352015-11-06 17:01:26 -0800157 return current;
Louis Ryanf6121162015-06-24 13:54:13 -0700158 }
159
160 private final Context parent;
161 private final Object[][] keyValueEntries;
162 private final boolean cascadesCancellation;
163 private ArrayList<ExecutableListener> listeners;
Louis Ryan73e2a232015-11-13 17:18:12 -0800164 private CancellationListener parentListener = new ParentListener();
165 private final boolean canBeCancelled;
Louis Ryanf6121162015-06-24 13:54:13 -0700166
167 /**
168 * Construct a context that cannot be cancelled and will not cascade cancellation from its parent.
169 */
170 private Context(Context parent) {
171 this.parent = parent;
Louis Ryan8dad7342016-02-18 16:30:06 -0800172 // Not inheriting cancellation implies not inheriting a deadline too.
173 keyValueEntries = new Object[][]{{DEADLINE_KEY, null}};
Louis Ryanf6121162015-06-24 13:54:13 -0700174 cascadesCancellation = false;
Louis Ryan73e2a232015-11-13 17:18:12 -0800175 canBeCancelled = false;
Louis Ryanf6121162015-06-24 13:54:13 -0700176 }
177
178 /**
179 * Construct a context that cannot be cancelled but will cascade cancellation from its parent if
180 * it is cancellable.
181 */
182 private Context(Context parent, Object[][] keyValueEntries) {
183 this.parent = parent;
184 this.keyValueEntries = keyValueEntries;
185 cascadesCancellation = true;
Louis Ryan73e2a232015-11-13 17:18:12 -0800186 canBeCancelled = this.parent != null && this.parent.canBeCancelled;
187 }
188
189 /**
190 * Construct a context that can be cancelled and will cascade cancellation from its parent if
191 * it is cancellable.
192 */
193 private Context(Context parent, Object[][] keyValueEntries, boolean isCancellable) {
194 this.parent = parent;
195 this.keyValueEntries = keyValueEntries;
196 cascadesCancellation = true;
197 canBeCancelled = isCancellable;
Louis Ryanf6121162015-06-24 13:54:13 -0700198 }
199
200 /**
201 * Create a new context which is independently cancellable and also cascades cancellation from
202 * its parent. Callers should ensure that either {@link CancellableContext#cancel(Throwable)}
Louis Ryan49b51352015-11-06 17:01:26 -0800203 * or {@link CancellableContext#detachAndCancel(Context, Throwable)} are called to notify
204 * listeners and release the resources associated with them.
Louis Ryanf6121162015-06-24 13:54:13 -0700205 *
206 * <p>Sample usage:
207 * <pre>
208 * Context.CancellableContext withCancellation = Context.current().withCancellation();
209 * try {
210 * executorService.execute(withCancellation.wrap(new Runnable() {
211 * public void run() {
212 * Context current = Context.current();
213 * while (!current.isCancelled()) {
214 * keepWorking();
215 * }
216 * }
217 * });
218 * doSomethingRelatedWork();
219 * } catch (Throwable t) {
220 * withCancellation.cancel(t);
221 * }
222 * </pre>
223 */
224 public CancellableContext withCancellation() {
225 return new CancellableContext(this);
226 }
227
228 /**
Louis Ryanf6121162015-06-24 13:54:13 -0700229 * Create a new context which will cancel itself after the given {@code duration} from now.
230 * The returned context will cascade cancellation of its parent. Callers may explicitly cancel
231 * the returned context prior to the deadline just as for {@link #withCancellation()},
232 *
233 * <p>Sample usage:
234 * <pre>
235 * Context.CancellableContext withDeadline = Context.current().withDeadlineAfter(5,
Eric Anderson6e94cf32016-01-23 09:32:41 -0800236 * TimeUnit.SECONDS, scheduler);
Louis Ryanf6121162015-06-24 13:54:13 -0700237 * executorService.execute(withDeadline.wrap(new Runnable() {
238 * public void run() {
239 * Context current = Context.current();
240 * while (!current.isCancelled()) {
241 * keepWorking();
242 * }
243 * }
244 * });
245 * </pre>
246 */
Eric Anderson6e94cf32016-01-23 09:32:41 -0800247 public CancellableContext withDeadlineAfter(long duration, TimeUnit unit,
Louis Ryan8dad7342016-02-18 16:30:06 -0800248 ScheduledExecutorService scheduler) {
249 return withDeadline(Deadline.after(duration, unit), scheduler);
250 }
251
252 /**
253 * Create a new context which will cancel itself at the given {@link Deadline}.
254 * The returned context will cascade cancellation of its parent. Callers may explicitly cancel
255 * the returned context prior to the deadline just as for {@link #withCancellation()},
256 *
257 * <p>Sample usage:
258 * <pre>
259 * Context.CancellableContext withDeadline = Context.current()
260 * .withDeadline(someReceivedDeadline);
261 * executorService.execute(withDeadline.wrap(new Runnable() {
262 * public void run() {
263 * Context current = Context.current();
264 * while (!current.isCancelled()) {
265 * keepWorking();
266 * }
267 * }
268 * });
269 * </pre>
270 */
271 public CancellableContext withDeadline(Deadline deadline,
Eric Anderson6e94cf32016-01-23 09:32:41 -0800272 ScheduledExecutorService scheduler) {
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700273 checkNotNull(deadline, "deadline");
274 checkNotNull(scheduler, "scheduler");
Louis Ryan8dad7342016-02-18 16:30:06 -0800275 return new CancellableContext(this, deadline, scheduler);
Louis Ryanf6121162015-06-24 13:54:13 -0700276 }
277
278 /**
279 * Create a new context with the given key value set. The new context will cascade cancellation
280 * from its parent.
281 *
282 <pre>
283 * Context withCredential = Context.current().withValue(CRED_KEY, cred);
284 * executorService.execute(withCredential.wrap(new Runnable() {
285 * public void run() {
286 * readUserRecords(userId, CRED_KEY.get());
287 * }
288 * }));
289 * </pre>
290 *
291 */
292 public <V> Context withValue(Key<V> k1, V v1) {
293 return new Context(this, new Object[][]{{k1, v1}});
294 }
295
296 /**
297 * Create a new context with the given key value set. The new context will cascade cancellation
298 * from its parent.
299 */
300 public <V1, V2> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2) {
301 return new Context(this, new Object[][]{{k1, v1}, {k2, v2}});
302 }
303
304 /**
305 * Create a new context with the given key value set. The new context will cascade cancellation
306 * from its parent.
307 */
308 public <V1, V2, V3> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2, Key<V3> k3, V3 v3) {
309 return new Context(this, new Object[][]{{k1, v1}, {k2, v2}, {k3, v3}});
310 }
311
312 /**
Louis Ryan8dad7342016-02-18 16:30:06 -0800313 * Create a new context which propagates the values of this context but does not cascade its
Eric Andersonb5e6d422016-04-28 18:48:55 -0700314 * cancellation.
Louis Ryanf6121162015-06-24 13:54:13 -0700315 */
Eric Andersonb5e6d422016-04-28 18:48:55 -0700316 public Context fork() {
317 return new Context(this);
Louis Ryanf6121162015-06-24 13:54:13 -0700318 }
319
320 boolean canBeCancelled() {
321 // A context is cancellable if it cascades from its parent and its parent is
Louis Ryan73e2a232015-11-13 17:18:12 -0800322 // cancellable or is itself directly cancellable..
323 return canBeCancelled;
Louis Ryanf6121162015-06-24 13:54:13 -0700324 }
325
326 /**
327 * Attach this context to the thread and make it {@link #current}, the previously current context
Louis Ryan49b51352015-11-06 17:01:26 -0800328 * is returned. It is allowed to attach contexts where {@link #isCancelled()} is {@code true}.
Louis Ryan8dad7342016-02-18 16:30:06 -0800329 *
330 * <p>Instead of using {@link #attach()} & {@link #detach(Context)} most use-cases are better
331 * served by using the {@link #run(Runnable)} or {@link #call(java.util.concurrent.Callable)}
332 * to execute work immediately within a context. If work needs to be done in other threads
333 * it is recommended to use the 'wrap' methods or to use a propagating executor.
Louis Ryanf6121162015-06-24 13:54:13 -0700334 */
Louis Ryan49b51352015-11-06 17:01:26 -0800335 public Context attach() {
336 Context previous = current();
337 localContext.set(this);
338 return previous;
339 }
340
341 /**
342 * Detach the current context from the thread and attach the provided replacement. If this
343 * context is not {@link #current()} a SEVERE message will be logged but the context to attach
344 * will still be bound.
345 */
346 public void detach(Context toAttach) {
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700347 checkNotNull(toAttach, "toAttach");
Louis Ryan8dad7342016-02-18 16:30:06 -0800348 if (toAttach.attach() != this) {
Louis Ryan49b51352015-11-06 17:01:26 -0800349 // Log a severe message instead of throwing an exception as the context to attach is assumed
350 // to be the correct one and the unbalanced state represents a coding mistake in a lower
351 // layer in the stack that cannot be recovered from here.
Eric Andersonc8648dc2016-08-04 11:08:36 -0700352 log.log(Level.SEVERE, "Context was not attached when detaching",
Louis Ryan49b51352015-11-06 17:01:26 -0800353 new Throwable().fillInStackTrace());
354 }
Louis Ryanf6121162015-06-24 13:54:13 -0700355 }
356
357 // Visible for testing
358 boolean isCurrent() {
359 return current() == this;
360 }
361
362 /**
Louis Ryanf6121162015-06-24 13:54:13 -0700363 * Is this context cancelled.
364 */
365 public boolean isCancelled() {
366 if (parent == null || !cascadesCancellation) {
367 return false;
368 } else {
369 return parent.isCancelled();
370 }
371 }
372
373 /**
374 * If a context {@link #isCancelled()} then return the cause of the cancellation or
375 * {@code null} if context was cancelled without a cause. If the context is not yet cancelled
376 * will always return {@code null}.
377 *
Louis Ryan8dad7342016-02-18 16:30:06 -0800378 * <p>The cancellation cause is provided for informational purposes only and implementations
379 * should generally assume that it has already been handled and logged properly.
Louis Ryanf6121162015-06-24 13:54:13 -0700380 */
Louis Ryan8dad7342016-02-18 16:30:06 -0800381 public Throwable cancellationCause() {
Louis Ryanf6121162015-06-24 13:54:13 -0700382 if (parent == null || !cascadesCancellation) {
383 return null;
384 } else {
Louis Ryan8dad7342016-02-18 16:30:06 -0800385 return parent.cancellationCause();
Louis Ryanf6121162015-06-24 13:54:13 -0700386 }
387 }
388
389 /**
Louis Ryan8dad7342016-02-18 16:30:06 -0800390 * A context may have an associated {@link Deadline} at which it will be automatically cancelled.
391 * @return A {@link io.grpc.Deadline} or {@code null} if no deadline is set.
392 */
Louis Ryan8dad7342016-02-18 16:30:06 -0800393 public Deadline getDeadline() {
394 return DEADLINE_KEY.get(this);
395 }
396
397 /**
Louis Ryanf6121162015-06-24 13:54:13 -0700398 * Add a listener that will be notified when the context becomes cancelled.
399 */
400 public void addListener(final CancellationListener cancellationListener,
401 final Executor executor) {
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700402 checkNotNull(cancellationListener, "cancellationListener");
403 checkNotNull(executor, "executor");
Louis Ryan73e2a232015-11-13 17:18:12 -0800404 if (canBeCancelled) {
Louis Ryanf6121162015-06-24 13:54:13 -0700405 ExecutableListener executableListener =
406 new ExecutableListener(executor, cancellationListener);
407 synchronized (this) {
408 if (isCancelled()) {
409 executableListener.deliver();
410 } else {
411 if (listeners == null) {
412 // Now that we have a listener we need to listen to our parent so
413 // we can cascade listener notification.
414 listeners = new ArrayList<ExecutableListener>();
415 listeners.add(executableListener);
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700416 parent.addListener(parentListener, DirectExecutor.INSTANCE);
Louis Ryanf6121162015-06-24 13:54:13 -0700417 } else {
418 listeners.add(executableListener);
419 }
420 }
421 }
Louis Ryanf6121162015-06-24 13:54:13 -0700422 }
423 }
424
425 /**
426 * Remove a {@link CancellationListener}.
427 */
428 public void removeListener(CancellationListener cancellationListener) {
Louis Ryan73e2a232015-11-13 17:18:12 -0800429 if (!canBeCancelled) {
430 return;
431 }
Louis Ryanf6121162015-06-24 13:54:13 -0700432 synchronized (this) {
433 if (listeners != null) {
434 for (int i = listeners.size() - 1; i >= 0; i--) {
435 if (listeners.get(i).listener == cancellationListener) {
436 listeners.remove(i);
Louis Ryan73e2a232015-11-13 17:18:12 -0800437 // Just remove the first matching listener, given that we allow duplicate
438 // adds we should allow for duplicates after remove.
Louis Ryanf6121162015-06-24 13:54:13 -0700439 break;
440 }
441 }
442 // We have no listeners so no need to listen to our parent
443 if (listeners.isEmpty()) {
444 parent.removeListener(parentListener);
445 listeners = null;
446 }
447 }
448 }
449 }
450
451 /**
452 * Notify all listeners that this context has been cancelled and immediately release
453 * any reference to them so that they may be garbage collected.
454 */
455 void notifyAndClearListeners() {
Louis Ryan73e2a232015-11-13 17:18:12 -0800456 if (!canBeCancelled) {
457 return;
458 }
Louis Ryanf6121162015-06-24 13:54:13 -0700459 ArrayList<ExecutableListener> tmpListeners;
460 synchronized (this) {
461 if (listeners == null) {
462 return;
463 }
464 tmpListeners = listeners;
465 listeners = null;
466 }
Louis Ryan73e2a232015-11-13 17:18:12 -0800467 // Deliver events to non-child context listeners before we notify child contexts. We do this
468 // to cancel higher level units of work before child units. This allows for a better error
469 // handling paradigm where the higher level unit of work knows it is cancelled and so can
470 // ignore errors that bubble up as a result of cancellation of lower level units.
Louis Ryanf6121162015-06-24 13:54:13 -0700471 for (int i = 0; i < tmpListeners.size(); i++) {
Louis Ryan73e2a232015-11-13 17:18:12 -0800472 if (!(tmpListeners.get(i).listener instanceof ParentListener)) {
473 tmpListeners.get(i).deliver();
474 }
475 }
476 for (int i = 0; i < tmpListeners.size(); i++) {
477 if (tmpListeners.get(i).listener instanceof ParentListener) {
478 tmpListeners.get(i).deliver();
479 }
Louis Ryanf6121162015-06-24 13:54:13 -0700480 }
481 parent.removeListener(parentListener);
482 }
483
Louis Ryan8dad7342016-02-18 16:30:06 -0800484 // Used in tests to ensure that listeners are defined and released when cancellation cascades.
485 // It's very important to ensure that we do not accidentally retain listeners.
Louis Ryanf6121162015-06-24 13:54:13 -0700486 int listenerCount() {
487 synchronized (this) {
488 return listeners == null ? 0 : listeners.size();
489 }
490 }
491
492 /**
Louis Ryan8dad7342016-02-18 16:30:06 -0800493 * Immediately run a {@link Runnable} with this context as the {@link #current} context.
494 * @param r {@link Runnable} to run.
495 */
496 public void run(Runnable r) {
497 Context previous = attach();
498 try {
499 r.run();
500 } finally {
501 detach(previous);
502 }
503 }
504
505 /**
506 * Immediately call a {@link Callable} with this context as the {@link #current} context.
507 * @param c {@link Callable} to call.
508 * @return result of call.
509 */
510 public <V> V call(Callable<V> c) throws Exception {
511 Context previous = attach();
512 try {
513 return c.call();
514 } finally {
515 detach(previous);
516 }
517 }
518
519 /**
Louis Ryanf6121162015-06-24 13:54:13 -0700520 * Wrap a {@link Runnable} so that it executes with this context as the {@link #current} context.
521 */
522 public Runnable wrap(final Runnable r) {
523 return new Runnable() {
524 @Override
525 public void run() {
Louis Ryan49b51352015-11-06 17:01:26 -0800526 Context previous = attach();
Louis Ryanf6121162015-06-24 13:54:13 -0700527 try {
528 r.run();
529 } finally {
Louis Ryan49b51352015-11-06 17:01:26 -0800530 detach(previous);
Louis Ryanf6121162015-06-24 13:54:13 -0700531 }
532 }
533 };
534 }
535
536 /**
537 * Wrap a {@link Callable} so that it executes with this context as the {@link #current} context.
538 */
539 public <C> Callable<C> wrap(final Callable<C> c) {
540 return new Callable<C>() {
541 @Override
542 public C call() throws Exception {
Louis Ryan49b51352015-11-06 17:01:26 -0800543 Context previous = attach();
Louis Ryanf6121162015-06-24 13:54:13 -0700544 try {
545 return c.call();
546 } finally {
Louis Ryan49b51352015-11-06 17:01:26 -0800547 detach(previous);
Louis Ryanf6121162015-06-24 13:54:13 -0700548 }
549 }
550 };
551 }
552
553 /**
Louis Ryan8dad7342016-02-18 16:30:06 -0800554 * Wrap an {@link Executor} so that it always executes with this context as the {@link #current}
555 * context. It is generally expected that {@link #currentContextExecutor(Executor)} would be
556 * used more commonly than this method.
Eric Andersona5d9e342015-07-22 16:54:08 -0700557 *
Louis Ryan8dad7342016-02-18 16:30:06 -0800558 * <p>One scenario in which this executor may be useful is when a single thread is sharding work
559 * to multiple threads.
560 *
561 * @see #currentContextExecutor(Executor)
Eric Andersona5d9e342015-07-22 16:54:08 -0700562 */
Louis Ryan8dad7342016-02-18 16:30:06 -0800563 public Executor fixedContextExecutor(final Executor e) {
564 class FixedContextExecutor implements Executor {
Eric Andersona5d9e342015-07-22 16:54:08 -0700565 @Override
566 public void execute(Runnable r) {
567 e.execute(wrap(r));
568 }
569 }
570
Louis Ryan8dad7342016-02-18 16:30:06 -0800571 return new FixedContextExecutor();
Eric Andersona5d9e342015-07-22 16:54:08 -0700572 }
573
574 /**
575 * Create an executor that propagates the {@link #current} context when {@link Executor#execute}
Louis Ryan8dad7342016-02-18 16:30:06 -0800576 * is called as the {@link #current} context of the {@code Runnable} scheduled. <em>Note that this
Eric Andersona5d9e342015-07-22 16:54:08 -0700577 * is a static method.</em>
578 *
Louis Ryan8dad7342016-02-18 16:30:06 -0800579 * @see #fixedContextExecutor(Executor)
Eric Andersona5d9e342015-07-22 16:54:08 -0700580 */
Louis Ryan8dad7342016-02-18 16:30:06 -0800581 public static Executor currentContextExecutor(final Executor e) {
582 class CurrentContextExecutor implements Executor {
Eric Andersona5d9e342015-07-22 16:54:08 -0700583 @Override
584 public void execute(Runnable r) {
585 e.execute(Context.current().wrap(r));
586 }
587 }
588
Louis Ryan8dad7342016-02-18 16:30:06 -0800589 return new CurrentContextExecutor();
Eric Andersona5d9e342015-07-22 16:54:08 -0700590 }
591
592 /**
Louis Ryanf6121162015-06-24 13:54:13 -0700593 * Lookup the value for a key in the context inheritance chain.
594 */
Eric Andersoneba12fb2015-07-23 13:24:04 -0700595 private Object lookup(Key<?> key) {
Louis Ryanf6121162015-06-24 13:54:13 -0700596 for (int i = 0; i < keyValueEntries.length; i++) {
597 if (key.equals(keyValueEntries[i][0])) {
598 return keyValueEntries[i][1];
599 }
600 }
601 if (parent == null) {
602 return null;
603 }
604 return parent.lookup(key);
605 }
606
607 /**
608 * A context which inherits cancellation from its parent but which can also be independently
609 * cancelled and which will propagate cancellation to its descendants.
610 */
611 public static final class CancellableContext extends Context {
612
613 private boolean cancelled;
Louis Ryan8dad7342016-02-18 16:30:06 -0800614 private Throwable cancellationCause;
615 private final Context uncancellableSurrogate;
616 private ScheduledFuture<?> pendingDeadline;
617
618 /**
619 * If the parent deadline is before the given deadline there is no need to install the value
620 * or listen for its expiration as the parent context will already be listening for it.
621 */
622 private static Object[][] deriveDeadline(Context parent, Deadline deadline) {
623 Deadline parentDeadline = DEADLINE_KEY.get(parent);
624 return parentDeadline == null || deadline.isBefore(parentDeadline)
625 ? new Object[][]{{ DEADLINE_KEY, deadline}} :
626 EMPTY_ENTRIES;
627 }
Louis Ryanf6121162015-06-24 13:54:13 -0700628
629 /**
630 * Create a cancellable context that does not have a deadline.
631 */
632 private CancellableContext(Context parent) {
Louis Ryan73e2a232015-11-13 17:18:12 -0800633 super(parent, EMPTY_ENTRIES, true);
Louis Ryan8dad7342016-02-18 16:30:06 -0800634 // Create a surrogate that inherits from this to attach so that you cannot retrieve a
Louis Ryanf6121162015-06-24 13:54:13 -0700635 // cancellable context from Context.current()
Louis Ryan8dad7342016-02-18 16:30:06 -0800636 uncancellableSurrogate = new Context(this, EMPTY_ENTRIES);
Louis Ryanf6121162015-06-24 13:54:13 -0700637 }
638
639 /**
640 * Create a cancellable context that has a deadline.
641 */
Louis Ryan8dad7342016-02-18 16:30:06 -0800642 private CancellableContext(Context parent, Deadline deadline,
Eric Anderson6e94cf32016-01-23 09:32:41 -0800643 ScheduledExecutorService scheduler) {
Louis Ryan8dad7342016-02-18 16:30:06 -0800644 super(parent, deriveDeadline(parent, deadline), true);
645 if (DEADLINE_KEY.get(this) == deadline) {
buchgrfd8fd512016-03-23 21:08:44 +0100646 final TimeoutException cause = new TimeoutException("context timed out");
647 if (!deadline.isExpired()) {
648 // The parent deadline was after the new deadline so we need to install a listener
649 // on the new earlier deadline to trigger expiration for this context.
650 pendingDeadline = deadline.runOnExpiration(new Runnable() {
651 @Override
652 public void run() {
Eric Andersonc8648dc2016-08-04 11:08:36 -0700653 try {
654 cancel(cause);
655 } catch (Throwable t) {
656 log.log(Level.SEVERE, "Cancel threw an exception, which should not happen", t);
657 }
buchgrfd8fd512016-03-23 21:08:44 +0100658 }
659 }, scheduler);
660 } else {
661 // Cancel immediately if the deadline is already expired.
662 cancel(cause);
663 }
Louis Ryan8dad7342016-02-18 16:30:06 -0800664 }
665 uncancellableSurrogate = new Context(this, EMPTY_ENTRIES);
Louis Ryanf6121162015-06-24 13:54:13 -0700666 }
667
668
669 @Override
Louis Ryan49b51352015-11-06 17:01:26 -0800670 public Context attach() {
Louis Ryan8dad7342016-02-18 16:30:06 -0800671 return uncancellableSurrogate.attach();
Louis Ryanf6121162015-06-24 13:54:13 -0700672 }
673
674 @Override
Louis Ryan49b51352015-11-06 17:01:26 -0800675 public void detach(Context toAttach) {
Louis Ryan8dad7342016-02-18 16:30:06 -0800676 uncancellableSurrogate.detach(toAttach);
Louis Ryanf6121162015-06-24 13:54:13 -0700677 }
678
679 @Override
680 public boolean isCurrent() {
Louis Ryan8dad7342016-02-18 16:30:06 -0800681 return uncancellableSurrogate.isCurrent();
Louis Ryanf6121162015-06-24 13:54:13 -0700682 }
683
684 /**
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700685 * Cancel this context and optionally provide a cause (can be {@code null}) for the
686 * cancellation. This will trigger notification of listeners.
Louis Ryanf6121162015-06-24 13:54:13 -0700687 *
688 * @return {@code true} if this context cancelled the context and notified listeners,
689 * {@code false} if the context was already cancelled.
690 */
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700691 public boolean cancel(Throwable cause) {
Louis Ryanf6121162015-06-24 13:54:13 -0700692 boolean triggeredCancel = false;
693 synchronized (this) {
694 if (!cancelled) {
695 cancelled = true;
Louis Ryan8dad7342016-02-18 16:30:06 -0800696 if (pendingDeadline != null) {
Louis Ryanf6121162015-06-24 13:54:13 -0700697 // If we have a scheduled cancellation pending attempt to cancel it.
Louis Ryan8dad7342016-02-18 16:30:06 -0800698 pendingDeadline.cancel(false);
699 pendingDeadline = null;
Louis Ryanf6121162015-06-24 13:54:13 -0700700 }
Louis Ryan8dad7342016-02-18 16:30:06 -0800701 this.cancellationCause = cause;
Louis Ryanf6121162015-06-24 13:54:13 -0700702 triggeredCancel = true;
703 }
704 }
705 if (triggeredCancel) {
706 notifyAndClearListeners();
707 }
708 return triggeredCancel;
709 }
710
711 /**
Louis Ryan49b51352015-11-06 17:01:26 -0800712 * Cancel this context and detach it as the current context from the thread.
Louis Ryanf6121162015-06-24 13:54:13 -0700713 *
Louis Ryan49b51352015-11-06 17:01:26 -0800714 * @param toAttach context to make current.
715 * @param cause of cancellation, can be {@code null}.
Louis Ryanf6121162015-06-24 13:54:13 -0700716 */
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700717 public void detachAndCancel(Context toAttach, Throwable cause) {
Louis Ryanf6121162015-06-24 13:54:13 -0700718 try {
Louis Ryan49b51352015-11-06 17:01:26 -0800719 detach(toAttach);
Louis Ryanf6121162015-06-24 13:54:13 -0700720 } finally {
721 cancel(cause);
722 }
723 }
724
725 @Override
Louis Ryanf6121162015-06-24 13:54:13 -0700726 public boolean isCancelled() {
727 synchronized (this) {
728 if (cancelled) {
729 return true;
730 }
731 }
732 // Detect cancellation of parent in the case where we have no listeners and
733 // record it.
734 if (super.isCancelled()) {
Louis Ryan8dad7342016-02-18 16:30:06 -0800735 cancel(super.cancellationCause());
Louis Ryanf6121162015-06-24 13:54:13 -0700736 return true;
737 }
738 return false;
739 }
740
Louis Ryanf6121162015-06-24 13:54:13 -0700741 @Override
Louis Ryan8dad7342016-02-18 16:30:06 -0800742 public Throwable cancellationCause() {
Louis Ryanf6121162015-06-24 13:54:13 -0700743 if (isCancelled()) {
Louis Ryan8dad7342016-02-18 16:30:06 -0800744 return cancellationCause;
Louis Ryanf6121162015-06-24 13:54:13 -0700745 }
746 return null;
747 }
748 }
749
750 /**
751 * A listener notified on context cancellation.
752 */
753 public interface CancellationListener {
754 /**
755 * @param context the newly cancelled context.
756 */
757 public void cancelled(Context context);
758 }
759
760 /**
761 * Key for indexing values stored in a context.
762 */
763 public static class Key<T> {
764
765 private final String name;
766 private final T defaultValue;
767
768 Key(String name) {
769 this(name, null);
770 }
771
772 Key(String name, T defaultValue) {
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700773 this.name = checkNotNull(name, "name");
Louis Ryanf6121162015-06-24 13:54:13 -0700774 this.defaultValue = defaultValue;
775 }
776
777 /**
778 * Get the value from the {@link #current()} context for this key.
779 */
780 @SuppressWarnings("unchecked")
781 public T get() {
782 return get(Context.current());
783 }
784
785 /**
786 * Get the value from the specified context for this key.
787 */
788 @SuppressWarnings("unchecked")
789 public T get(Context context) {
790 T value = (T) context.lookup(this);
791 return value == null ? defaultValue : value;
792 }
793
794 @Override
Louis Ryanbe600862015-11-06 13:46:13 -0800795 public String toString() {
796 return name;
797 }
Louis Ryanf6121162015-06-24 13:54:13 -0700798 }
799
800 /**
801 * Stores listener & executor pair.
802 */
803 private class ExecutableListener implements Runnable {
804 private final Executor executor;
805 private final CancellationListener listener;
806
807 private ExecutableListener(Executor executor, CancellationListener listener) {
808 this.executor = executor;
809 this.listener = listener;
810 }
811
812 private void deliver() {
Eric Anderson63f14822015-10-26 17:03:30 -0700813 try {
814 executor.execute(this);
815 } catch (Throwable t) {
Eric Andersonc8648dc2016-08-04 11:08:36 -0700816 log.log(Level.INFO, "Exception notifying context listener", t);
Eric Anderson63f14822015-10-26 17:03:30 -0700817 }
Louis Ryanf6121162015-06-24 13:54:13 -0700818 }
819
820 @Override
821 public void run() {
822 listener.cancelled(Context.this);
823 }
824 }
Louis Ryan73e2a232015-11-13 17:18:12 -0800825
826 private class ParentListener implements CancellationListener {
827 @Override
828 public void cancelled(Context context) {
829 if (Context.this instanceof CancellableContext) {
Louis Ryan8dad7342016-02-18 16:30:06 -0800830 // Record cancellation with its cancellationCause.
831 ((CancellableContext) Context.this).cancel(context.cancellationCause());
Louis Ryan73e2a232015-11-13 17:18:12 -0800832 } else {
833 notifyAndClearListeners();
834 }
835 }
836 }
Kun Zhangc4f7f5c2016-09-01 16:00:43 -0700837
838 private static <T> T checkNotNull(T reference, Object errorMessage) {
839 if (reference == null) {
840 throw new NullPointerException(String.valueOf(errorMessage));
841 }
842 return reference;
843 }
844
845 private enum DirectExecutor implements Executor {
846 INSTANCE;
847
848 @Override
849 public void execute(Runnable command) {
850 command.run();
851 }
852
853 @Override
854 public String toString() {
855 return "Context.DirectExecutor";
856 }
857 }
Louis Ryanf6121162015-06-24 13:54:13 -0700858}