context: compress cancellation ancestor chain (#3372)

Now that we have the copy of write keyvalue store (#3368), there
is no need to keep the full parent chain. We only need a
reference to the nearest cancellable ancestor. This optimization
should in theory make cancellations more efficient and also make
our data structs more GC friendly.

context/src/main/java/io/grpc/Context.java

@@ -179,20 +179,20 @@ public class Context {
     return current;
   }
 
-  private final Context parent;
   private final boolean cascadesCancellation;
   private ArrayList<ExecutableListener> listeners;
   private CancellationListener parentListener = new ParentListener();
   private final boolean canBeCancelled;
+  final CancellableContext cancellableAncestor;
   final PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries;
 
   /**
    * Construct a context that cannot be cancelled and will not cascade cancellation from its parent.
    */
-  private Context(Context parent) {
-    this.parent = parent;
+  private Context(PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries) {
+    cancellableAncestor = null;
     // Not inheriting cancellation implies not inheriting a deadline too.
-    keyValueEntries = parent.keyValueEntries.put(DEADLINE_KEY, null);
+    this.keyValueEntries = keyValueEntries.put(DEADLINE_KEY, null);
     cascadesCancellation = false;
     canBeCancelled = false;
   }
@@ -202,10 +202,10 @@ public class Context {
    * it is cancellable.
    */
   private Context(Context parent, PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries) {
-    this.parent = parent;
+    cancellableAncestor = cancellableAncestor(parent);
     this.keyValueEntries = keyValueEntries;
     cascadesCancellation = true;
-    canBeCancelled = this.parent != null && this.parent.canBeCancelled;
+    canBeCancelled = cancellableAncestor != null;
   }
 
   /**
@@ -216,7 +216,7 @@ public class Context {
       Context parent,
       PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries,
       boolean isCancellable) {
-    this.parent = parent;
+    cancellableAncestor = cancellableAncestor(parent);
     this.keyValueEntries = keyValueEntries;
     cascadesCancellation = true;
     canBeCancelled = isCancellable;
@@ -230,7 +230,7 @@ public class Context {
       PersistentHashArrayMappedTrie<Key<?>, Object> keyValueEntries,
       boolean cascadesCancellation,
       boolean isCancellable) {
-    this.parent = parent;
+    cancellableAncestor = cancellableAncestor(parent);
     this.keyValueEntries = keyValueEntries;
     this.cascadesCancellation = cascadesCancellation;
     canBeCancelled = isCancellable;
@@ -374,7 +374,7 @@ public class Context {
    * cancellation.
    */
   public Context fork() {
-    return new Context(this);
+    return new Context(keyValueEntries);
   }
 
   boolean canBeCancelled() {
@@ -438,10 +438,10 @@ public class Context {
    * Is this context cancelled.
    */
   public boolean isCancelled() {
-    if (parent == null || !cascadesCancellation) {
+    if (cancellableAncestor == null || !cascadesCancellation) {
       return false;
     } else {
-      return parent.isCancelled();
+      return cancellableAncestor.isCancelled();
     }
   }
 
@@ -454,10 +454,10 @@ public class Context {
    * should generally assume that it has already been handled and logged properly.
    */
   public Throwable cancellationCause() {
-    if (parent == null || !cascadesCancellation) {
+    if (cancellableAncestor == null || !cascadesCancellation) {
       return null;
     } else {
-      return parent.cancellationCause();
+      return cancellableAncestor.cancellationCause();
     }
   }
 
@@ -488,7 +488,9 @@ public class Context {
             // we can cascade listener notification.
             listeners = new ArrayList<ExecutableListener>();
             listeners.add(executableListener);
-            parent.addListener(parentListener, DirectExecutor.INSTANCE);
+            if (cancellableAncestor != null) {
+              cancellableAncestor.addListener(parentListener, DirectExecutor.INSTANCE);
+            }
           } else {
             listeners.add(executableListener);
           }
@@ -516,7 +518,9 @@ public class Context {
         }
         // We have no listeners so no need to listen to our parent
         if (listeners.isEmpty()) {
-          parent.removeListener(parentListener);
+          if (cancellableAncestor != null) {
+            cancellableAncestor.removeListener(parentListener);
+          }
           listeners = null;
         }
       }
@@ -553,7 +557,9 @@ public class Context {
         tmpListeners.get(i).deliver();
       }
     }
-    parent.removeListener(parentListener);
+    if (cancellableAncestor != null) {
+      cancellableAncestor.removeListener(parentListener);
+    }
   }
 
   // Used in tests to ensure that listeners are defined and released when cancellation cascades.
@@ -1000,4 +1006,20 @@ public class Context {
       return "Context.DirectExecutor";
     }
   }
+
+  /**
+   * Returns {@code parent} if it is a {@link CancellableContext}, otherwise returns the parent's
+   * {@link #cancellableAncestor}.
+   */
+  static CancellableContext cancellableAncestor(Context parent) {
+    if (parent == null || !parent.canBeCancelled()) {
+      return null;
+    }
+    if (parent instanceof CancellableContext) {
+      return (CancellableContext) parent;
+    }
+    // The parent simply cascades cancellations.
+    // Bypass the parent and reference the ancestor directly.
+    return parent.cancellableAncestor;
+  }
 }

context/src/test/java/io/grpc/ContextTest.java

@@ -16,6 +16,7 @@
 
 package io.grpc;
 
+import static io.grpc.Context.cancellableAncestor;
 import static org.hamcrest.core.IsInstanceOf.instanceOf;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -875,6 +876,52 @@ public class ContextTest {
     }
   }
 
+  @Test
+  public void cancellableAncestorTest() {
+    assertEquals(null, cancellableAncestor(null));
+
+    Context c = Context.current();
+    assertFalse(c.canBeCancelled());
+    assertEquals(null, cancellableAncestor(c));
+
+    Context.CancellableContext withCancellation = c.withCancellation();
+    assertEquals(withCancellation, cancellableAncestor(withCancellation));
+
+    Context child = withCancellation.withValue(COLOR, "blue");
+    assertFalse(child instanceof Context.CancellableContext);
+    assertEquals(withCancellation, cancellableAncestor(child));
+
+    Context grandChild = child.withValue(COLOR, "red");
+    assertFalse(grandChild instanceof Context.CancellableContext);
+    assertEquals(withCancellation, cancellableAncestor(grandChild));
+  }
+
+  @Test
+  public void cancellableAncestorIntegrationTest() {
+    Context base = Context.current();
+
+    Context blue = base.withValue(COLOR, "blue");
+    assertNull(blue.cancellableAncestor);
+    Context.CancellableContext cancellable = blue.withCancellation();
+    assertNull(cancellable.cancellableAncestor);
+    Context childOfCancel = cancellable.withValue(PET, "cat");
+    assertSame(cancellable, childOfCancel.cancellableAncestor);
+    Context grandChildOfCancel = childOfCancel.withValue(FOOD, "lasagna");
+    assertSame(cancellable, grandChildOfCancel.cancellableAncestor);
+
+    Context.CancellableContext cancellable2 = childOfCancel.withCancellation();
+    assertSame(cancellable, cancellable2.cancellableAncestor);
+    Context childOfCancellable2 = cancellable2.withValue(PET, "dog");
+    assertSame(cancellable2, childOfCancellable2.cancellableAncestor);
+  }
+
+  @Test
+  public void cancellableAncestorFork() {
+    Context.CancellableContext cancellable = Context.current().withCancellation();
+    Context fork = cancellable.fork();
+    assertNull(fork.cancellableAncestor);
+  }
+
   // UsedReflectively
   public static final class LoadMeWithStaticTestingClassLoader implements Runnable {
     @Override