diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java index a1329bf8b8..4aee3bea3b 100644 --- a/context/src/main/java/io/grpc/Context.java +++ b/context/src/main/java/io/grpc/Context.java @@ -17,6 +17,7 @@ package io.grpc; import io.grpc.Context.CheckReturnValue; +import io.grpc.PersistentHashArrayMappedTrie.Node; import java.io.Closeable; import java.util.ArrayList; import java.util.concurrent.Callable; @@ -99,9 +100,6 @@ public class Context { static final Logger log = Logger.getLogger(Context.class.getName()); - private static final PersistentHashArrayMappedTrie, Object> EMPTY_ENTRIES = - new PersistentHashArrayMappedTrie<>(); - // Long chains of contexts are suspicious and usually indicate a misuse of Context. // The threshold is arbitrarily chosen. // VisibleForTesting @@ -188,14 +186,14 @@ public class Context { private ArrayList listeners; private CancellationListener parentListener = new ParentListener(); final CancellableContext cancellableAncestor; - final PersistentHashArrayMappedTrie, Object> keyValueEntries; + final Node, Object> keyValueEntries; // The number parents between this context and the root context. final int generation; /** * Construct a context that cannot be cancelled and will not cascade cancellation from its parent. */ - private Context(PersistentHashArrayMappedTrie, Object> keyValueEntries, int generation) { + private Context(Node, Object> keyValueEntries, int generation) { this.cancellableAncestor = null; this.keyValueEntries = keyValueEntries; this.generation = generation; @@ -206,7 +204,7 @@ public class Context { * Construct a context that cannot be cancelled but will cascade cancellation from its parent if * it is cancellable. */ - private Context(Context parent, PersistentHashArrayMappedTrie, Object> keyValueEntries) { + private Context(Context parent, Node, Object> keyValueEntries) { this.cancellableAncestor = cancellableAncestor(parent); this.keyValueEntries = keyValueEntries; this.generation = parent.generation + 1; @@ -218,7 +216,7 @@ public class Context { */ private Context() { this.cancellableAncestor = null; - this.keyValueEntries = EMPTY_ENTRIES; + this.keyValueEntries = null; this.generation = 0; validateGeneration(generation); } @@ -351,7 +349,8 @@ public class Context { * are unrelated, have separate keys for them. */ public Context withValue(Key k1, V v1) { - PersistentHashArrayMappedTrie, Object> newKeyValueEntries = keyValueEntries.put(k1, v1); + Node, Object> newKeyValueEntries = + PersistentHashArrayMappedTrie.put(keyValueEntries, k1, v1); return new Context(this, newKeyValueEntries); } @@ -360,8 +359,9 @@ public class Context { * from its parent. */ public Context withValues(Key k1, V1 v1, Key k2, V2 v2) { - PersistentHashArrayMappedTrie, Object> newKeyValueEntries = - keyValueEntries.put(k1, v1).put(k2, v2); + Node, Object> newKeyValueEntries = + PersistentHashArrayMappedTrie.put(keyValueEntries, k1, v1); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k2, v2); return new Context(this, newKeyValueEntries); } @@ -370,8 +370,10 @@ public class Context { * from its parent. */ public Context withValues(Key k1, V1 v1, Key k2, V2 v2, Key k3, V3 v3) { - PersistentHashArrayMappedTrie, Object> newKeyValueEntries = - keyValueEntries.put(k1, v1).put(k2, v2).put(k3, v3); + Node, Object> newKeyValueEntries = + PersistentHashArrayMappedTrie.put(keyValueEntries, k1, v1); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k2, v2); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k3, v3); return new Context(this, newKeyValueEntries); } @@ -395,8 +397,11 @@ public class Context { */ public Context withValues(Key k1, V1 v1, Key k2, V2 v2, Key k3, V3 v3, Key k4, V4 v4) { - PersistentHashArrayMappedTrie, Object> newKeyValueEntries = - keyValueEntries.put(k1, v1).put(k2, v2).put(k3, v3).put(k4, v4); + Node, Object> newKeyValueEntries = + PersistentHashArrayMappedTrie.put(keyValueEntries, k1, v1); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k2, v2); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k3, v3); + newKeyValueEntries = PersistentHashArrayMappedTrie.put(newKeyValueEntries, k4, v4); return new Context(this, newKeyValueEntries); } @@ -703,13 +708,6 @@ public class Context { return new CurrentContextExecutor(); } - /** - * Lookup the value for a key in the context inheritance chain. - */ - Object lookup(Key key) { - return keyValueEntries.get(key); - } - /** * A context which inherits cancellation from its parent but which can also be independently * cancelled and which will propagate cancellation to its descendants. To avoid leaking memory, @@ -939,7 +937,7 @@ public class Context { */ @SuppressWarnings("unchecked") public T get(Context context) { - T value = (T) context.lookup(this); + T value = (T) PersistentHashArrayMappedTrie.get(context.keyValueEntries, this); return value == null ? defaultValue : value; } diff --git a/context/src/main/java/io/grpc/PersistentHashArrayMappedTrie.java b/context/src/main/java/io/grpc/PersistentHashArrayMappedTrie.java index 72ef37bf91..66ee869722 100644 --- a/context/src/main/java/io/grpc/PersistentHashArrayMappedTrie.java +++ b/context/src/main/java/io/grpc/PersistentHashArrayMappedTrie.java @@ -29,28 +29,14 @@ import java.util.Arrays; * Bagwell (2000). The rest of the implementation is ignorant of/ignores the * paper. */ -final class PersistentHashArrayMappedTrie { - private final Node root; +final class PersistentHashArrayMappedTrie { - PersistentHashArrayMappedTrie() { - this(null); - } - - private PersistentHashArrayMappedTrie(Node root) { - this.root = root; - } - - public int size() { - if (root == null) { - return 0; - } - return root.size(); - } + private PersistentHashArrayMappedTrie() {} /** * Returns the value with the specified key, or {@code null} if it does not exist. */ - public V get(K key) { + static V get(Node root, K key) { if (root == null) { return null; } @@ -60,12 +46,11 @@ final class PersistentHashArrayMappedTrie { /** * Returns a new trie where the key is set to the specified value. */ - public PersistentHashArrayMappedTrie put(K key, V value) { + static Node put(Node root, K key, V value) { if (root == null) { - return new PersistentHashArrayMappedTrie<>(new Leaf<>(key, value)); - } else { - return new PersistentHashArrayMappedTrie<>(root.put(key, value, key.hashCode(), 0)); + return new Leaf<>(key, value); } + return root.put(key, value, key.hashCode(), 0); } // Not actually annotated to avoid depending on guava