private static final Logger log = Logger.getLogger(Context.class.getName());

private static final Object[][] EMPTY_ENTRIES = new Object[0][2];

private static final Key<Deadline> DEADLINE_KEY = new Key<Deadline>("deadline");

/**

public static final Context ROOT = new Context(null);

// One and only one of them is non-null

private static final Storage storage;

private static final Exception storageInitError;

static {

Storage newStorage = null;

Exception error = null;

try {

Class<?> clazz = Class.forName("io.grpc.override.ContextStorageOverride");

newStorage = (Storage) clazz.getConstructor().newInstance();

} catch (ClassNotFoundException e) {

if (log.isLoggable(Level.FINE)) {

// Avoid writing to logger because custom log handlers may try to use Context, which is

// problemantic (e.g., NullPointerException) because the Context class has not done loading

// at this point. The caveat is that in environments stderr may be disabled, thus this

// message would go nowhere.

System.err.println("io.grpc.Context: Storage override doesn't exist. Using default.");

e.printStackTrace();

}

newStorage = new ThreadLocalContextStorage();

} catch (Exception e) {

error = e;

}

storage = newStorage;

storageInitError = error;

}

// For testing

static Storage storage() {

if (storage == null) {

throw new RuntimeException("Storage override had failed to initialize", storageInitError);

}

return storage;

}

/**

public static <T> Key<T> key(String name) {

return new Key<T>(name);

}

/**

public static <T> Key<T> keyWithDefault(String name, T defaultValue) {

return new Key<T>(name, defaultValue);

}

/**

public static Context current() {

Context current = storage().current();

if (current == null) {

return ROOT;

}

return current;

}

private final Context parent;

private final Object[][] keyValueEntries;

private final boolean cascadesCancellation;

private ArrayList<ExecutableListener> listeners;

private CancellationListener parentListener = new ParentListener();

private final boolean canBeCancelled;

/**

private Context(Context parent) {

this.parent = parent;

// Not inheriting cancellation implies not inheriting a deadline too.

keyValueEntries = new Object[][]{{DEADLINE_KEY, null}};

cascadesCancellation = false;

canBeCancelled = false;

}

/**

private Context(Context parent, Object[][] keyValueEntries) {

this.parent = parent;

this.keyValueEntries = keyValueEntries;

cascadesCancellation = true;

canBeCancelled = this.parent != null && this.parent.canBeCancelled;

}

/**

private Context(Context parent, Object[][] keyValueEntries, boolean isCancellable) {

this.parent = parent;

this.keyValueEntries = keyValueEntries;

cascadesCancellation = true;

canBeCancelled = isCancellable;

}

/**

public CancellableContext withCancellation() {

return new CancellableContext(this);

}

/**

public CancellableContext withDeadlineAfter(long duration, TimeUnit unit,

ScheduledExecutorService scheduler) {

return withDeadline(Deadline.after(duration, unit), scheduler);

}

/**

public CancellableContext withDeadline(Deadline deadline,

ScheduledExecutorService scheduler) {

checkNotNull(deadline, "deadline");

checkNotNull(scheduler, "scheduler");

return new CancellableContext(this, deadline, scheduler);

}

/**

public <V> Context withValue(Key<V> k1, V v1) {

return new Context(this, new Object[][]{{k1, v1}});

}

/**

public <V1, V2> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2) {

return new Context(this, new Object[][]{{k1, v1}, {k2, v2}});

}

/**

public <V1, V2, V3> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2, Key<V3> k3, V3 v3) {

return new Context(this, new Object[][]{{k1, v1}, {k2, v2}, {k3, v3}});

}

/**

public <V1, V2, V3, V4> Context withValues(Key<V1> k1, V1 v1, Key<V2> k2, V2 v2,

Key<V3> k3, V3 v3, Key<V4> k4, V4 v4) {

return new Context(this, new Object[][]{{k1, v1}, {k2, v2}, {k3, v3}, {k4, v4}});

}

/**

public Context fork() {

return new Context(this);

}

boolean canBeCancelled() {

// A context is cancellable if it cascades from its parent and its parent is

// cancellable or is itself directly cancellable..

return canBeCancelled;

}

/**

public Context attach() {

Context previous = current();

storage().attach(this);

return previous;

}

/**

public void detach(Context toAttach) {

checkNotNull(toAttach, "toAttach");

storage().detach(this, toAttach);

}

// Visible for testing

boolean isCurrent() {

return current() == this;

}

/**

public boolean isCancelled() {

if (parent == null || !cascadesCancellation) {

return false;

} else {

return parent.isCancelled();

}

}

/**

public Throwable cancellationCause() {

if (parent == null || !cascadesCancellation) {

return null;

} else {

return parent.cancellationCause();

}

}

/**

public Deadline getDeadline() {

return DEADLINE_KEY.get(this);

}

/**

public void addListener(final CancellationListener cancellationListener,

final Executor executor) {

checkNotNull(cancellationListener, "cancellationListener");

checkNotNull(executor, "executor");

if (canBeCancelled) {

ExecutableListener executableListener =

new ExecutableListener(executor, cancellationListener);

synchronized (this) {

if (isCancelled()) {

executableListener.deliver();

} else {

if (listeners == null) {

// Now that we have a listener we need to listen to our parent so

// we can cascade listener notification.

listeners = new ArrayList<ExecutableListener>();

listeners.add(executableListener);

parent.addListener(parentListener, DirectExecutor.INSTANCE);

} else {

listeners.add(executableListener);

}

}

}

}

}

/**

public void removeListener(CancellationListener cancellationListener) {

if (!canBeCancelled) {

return;

}

synchronized (this) {

if (listeners != null) {

for (int i = listeners.size() - 1; i >= 0; i--) {

if (listeners.get(i).listener == cancellationListener) {

listeners.remove(i);

// Just remove the first matching listener, given that we allow duplicate

// adds we should allow for duplicates after remove.

break;

}

}

// We have no listeners so no need to listen to our parent

if (listeners.isEmpty()) {

parent.removeListener(parentListener);

listeners = null;

}

}

}

}

/**

void notifyAndClearListeners() {

if (!canBeCancelled) {

return;

}

ArrayList<ExecutableListener> tmpListeners;

synchronized (this) {

if (listeners == null) {

return;

}

tmpListeners = listeners;

listeners = null;

}

// Deliver events to non-child context listeners before we notify child contexts. We do this

// to cancel higher level units of work before child units. This allows for a better error

// handling paradigm where the higher level unit of work knows it is cancelled and so can

// ignore errors that bubble up as a result of cancellation of lower level units.

for (int i = 0; i < tmpListeners.size(); i++) {

if (!(tmpListeners.get(i).listener instanceof ParentListener)) {

tmpListeners.get(i).deliver();

}

}

for (int i = 0; i < tmpListeners.size(); i++) {

if (tmpListeners.get(i).listener instanceof ParentListener) {

tmpListeners.get(i).deliver();

}

}

parent.removeListener(parentListener);

}

// Used in tests to ensure that listeners are defined and released when cancellation cascades.

// It's very important to ensure that we do not accidentally retain listeners.

int listenerCount() {

synchronized (this) {

return listeners == null ? 0 : listeners.size();

}

}

/**

public void run(Runnable r) {

Context previous = attach();

try {

r.run();

} finally {

detach(previous);

}

}

/**

public <V> V call(Callable<V> c) throws Exception {

Context previous = attach();

try {

return c.call();

} finally {

detach(previous);

}

}

/**

public Runnable wrap(final Runnable r) {

return new Runnable() {

@Override

public void run() {

Context previous = attach();

try {

r.run();

} finally {

detach(previous);

}

}

};

}

/**

public <C> Callable<C> wrap(final Callable<C> c) {

return new Callable<C>() {

@Override

public C call() throws Exception {

Context previous = attach();

try {

return c.call();

} finally {

detach(previous);

}

}

};

}

/**

public Executor fixedContextExecutor(final Executor e) {

class FixedContextExecutor implements Executor {

@Override

public void execute(Runnable r) {

e.execute(wrap(r));

}

}

return new FixedContextExecutor();

}

/**

public static Executor currentContextExecutor(final Executor e) {

class CurrentContextExecutor implements Executor {

@Override

public void execute(Runnable r) {

e.execute(Context.current().wrap(r));

}

}

return new CurrentContextExecutor();

}

/**

private Object lookup(Key<?> key) {

for (int i = 0; i < keyValueEntries.length; i++) {

if (key.equals(keyValueEntries[i][0])) {

return keyValueEntries[i][1];

}

}

if (parent == null) {

return null;

}

return parent.lookup(key);

}

/**

public static final class CancellableContext extends Context {

private boolean cancelled;

private Throwable cancellationCause;

private final Context uncancellableSurrogate;

private ScheduledFuture<?> pendingDeadline;

/**

private static Object[][] deriveDeadline(Context parent, Deadline deadline) {

Deadline parentDeadline = DEADLINE_KEY.get(parent);

return parentDeadline == null || deadline.isBefore(parentDeadline)

? new Object[][]{{ DEADLINE_KEY, deadline}} :

EMPTY_ENTRIES;

}

/**

private CancellableContext(Context parent) {

super(parent, EMPTY_ENTRIES, true);

// Create a surrogate that inherits from this to attach so that you cannot retrieve a

// cancellable context from Context.current()

uncancellableSurrogate = new Context(this, EMPTY_ENTRIES);

}

/**

private CancellableContext(Context parent, Deadline deadline,

ScheduledExecutorService scheduler) {

super(parent, deriveDeadline(parent, deadline), true);

if (DEADLINE_KEY.get(this) == deadline) {

final TimeoutException cause = new TimeoutException("context timed out");

if (!deadline.isExpired()) {

// The parent deadline was after the new deadline so we need to install a listener

// on the new earlier deadline to trigger expiration for this context.

pendingDeadline = deadline.runOnExpiration(new Runnable() {

@Override

public void run() {

try {

cancel(cause);

} catch (Throwable t) {

log.log(Level.SEVERE, "Cancel threw an exception, which should not happen", t);

}

}

}, scheduler);

} else {

// Cancel immediately if the deadline is already expired.

cancel(cause);

}

}

uncancellableSurrogate = new Context(this, EMPTY_ENTRIES);

}

@Override

public Context attach() {

return uncancellableSurrogate.attach();

}

@Override

public void detach(Context toAttach) {

uncancellableSurrogate.detach(toAttach);

}

@Override

public boolean isCurrent() {

return uncancellableSurrogate.isCurrent();

}

/**

public boolean cancel(Throwable cause) {

boolean triggeredCancel = false;

synchronized (this) {

if (!cancelled) {

cancelled = true;

if (pendingDeadline != null) {

// If we have a scheduled cancellation pending attempt to cancel it.

pendingDeadline.cancel(false);

pendingDeadline = null;

}

this.cancellationCause = cause;

triggeredCancel = true;

}

}

if (triggeredCancel) {

notifyAndClearListeners();

}

return triggeredCancel;

}

/**

public void detachAndCancel(Context toAttach, Throwable cause) {

try {

detach(toAttach);

} finally {

cancel(cause);

}

}

@Override

public boolean isCancelled() {

synchronized (this) {

if (cancelled) {

return true;

}

}

// Detect cancellation of parent in the case where we have no listeners and

// record it.

if (super.isCancelled()) {

cancel(super.cancellationCause());

return true;

}

return false;

}

@Override

public Throwable cancellationCause() {

if (isCancelled()) {

return cancellationCause;

}

return null;

}

}

/**

public interface CancellationListener {

/**

public void cancelled(Context context);

}

/**

public static class Key<T> {

private final String name;

private final T defaultValue;

Key(String name) {

this(name, null);

}

Key(String name, T defaultValue) {

this.name = checkNotNull(name, "name");

this.defaultValue = defaultValue;

}

/**

@SuppressWarnings("unchecked")

public T get() {

return get(Context.current());

}

/**

@SuppressWarnings("unchecked")

public T get(Context context) {

T value = (T) context.lookup(this);

return value == null ? defaultValue : value;

}

@Override

public String toString() {

return name;

}

}

/**

public abstract static class Storage {

/**

public abstract void attach(Context toAttach);

/**

public abstract void detach(Context toDetach, Context toRestore);

/**

public abstract Context current();

}

/**

private class ExecutableListener implements Runnable {

private final Executor executor;

private final CancellationListener listener;

private ExecutableListener(Executor executor, CancellationListener listener) {

this.executor = executor;

this.listener = listener;

}

private void deliver() {

try {

executor.execute(this);

} catch (Throwable t) {

log.log(Level.INFO, "Exception notifying context listener", t);

}

}

@Override

public void run() {

listener.cancelled(Context.this);

}

}

private class ParentListener implements CancellationListener {

@Override

public void cancelled(Context context) {

if (Context.this instanceof CancellableContext) {

// Record cancellation with its cancellationCause.

((CancellableContext) Context.this).cancel(context.cancellationCause());

} else {

notifyAndClearListeners();

}

}

}

private static <T> T checkNotNull(T reference, Object errorMessage) {

if (reference == null) {

throw new NullPointerException(String.valueOf(errorMessage));

}

return reference;

}

private enum DirectExecutor implements Executor {

INSTANCE;

@Override

public void execute(Runnable command) {

command.run();

}

@Override

public String toString() {

return "Context.DirectExecutor";

}

}