diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java index c3d34117c2..c42257791f 100644 --- a/context/src/main/java/io/grpc/Context.java +++ b/context/src/main/java/io/grpc/Context.java @@ -99,8 +99,6 @@ public class Context { private static final PersistentHashArrayMappedTrie, Object> EMPTY_ENTRIES = new PersistentHashArrayMappedTrie, Object>(); - private static final Key DEADLINE_KEY = new Key("deadline"); - /** * The logical root context which is the ultimate ancestor of all contexts. This context * is not cancellable and so will not cascade cancellation or retain listeners. @@ -191,8 +189,7 @@ public class Context { */ private Context(PersistentHashArrayMappedTrie, Object> keyValueEntries) { cancellableAncestor = null; - // Not inheriting cancellation implies not inheriting a deadline too. - this.keyValueEntries = keyValueEntries.put(DEADLINE_KEY, null); + this.keyValueEntries = keyValueEntries; cascadesCancellation = false; canBeCancelled = false; } @@ -466,7 +463,10 @@ public class Context { * @return A {@link io.grpc.Deadline} or {@code null} if no deadline is set. */ public Deadline getDeadline() { - return DEADLINE_KEY.get(this); + if (cancellableAncestor == null) { + return null; + } + return cancellableAncestor.getDeadline(); } /** @@ -685,28 +685,19 @@ public class Context { */ public static final class CancellableContext extends Context { + private final Deadline deadline; + private final Context uncancellableSurrogate; + private boolean cancelled; private Throwable cancellationCause; - private final Context uncancellableSurrogate; private ScheduledFuture pendingDeadline; - /** - * If the parent deadline is before the given deadline there is no need to install the value - * or listen for its expiration as the parent context will already be listening for it. - */ - private static PersistentHashArrayMappedTrie, Object> deriveDeadline( - Context parent, Deadline deadline) { - Deadline parentDeadline = DEADLINE_KEY.get(parent); - return parentDeadline == null || deadline.isBefore(parentDeadline) - ? parent.keyValueEntries.put(DEADLINE_KEY, deadline) - : parent.keyValueEntries; - } - /** * Create a cancellable context that does not have a deadline. */ private CancellableContext(Context parent) { super(parent, parent.keyValueEntries, true); + deadline = parent.getDeadline(); // Create a surrogate that inherits from this to attach so that you cannot retrieve a // cancellable context from Context.current() uncancellableSurrogate = new Context(this, keyValueEntries); @@ -717,8 +708,13 @@ public class Context { */ private CancellableContext(Context parent, Deadline deadline, ScheduledExecutorService scheduler) { - super(parent, deriveDeadline(parent, deadline), true); - if (DEADLINE_KEY.get(this) == deadline) { + super(parent, parent.keyValueEntries, true); + Deadline parentDeadline = parent.getDeadline(); + if (parentDeadline != null && parentDeadline.compareTo(deadline) <= 0) { + // The new deadline won't have an effect, so ignore it + deadline = parentDeadline; + } else { + // The new deadline has an effect 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. @@ -737,6 +733,7 @@ public class Context { cancel(new TimeoutException("context timed out")); } } + this.deadline = deadline; uncancellableSurrogate = new Context(this, keyValueEntries); } @@ -829,6 +826,11 @@ public class Context { } return null; } + + @Override + public Deadline getDeadline() { + return deadline; + } } /**