From 8a9afd618aec0e1bf5858232cac70c89a90fb96e Mon Sep 17 00:00:00 2001
From: Carl Mastrangelo <notcarl@google.com>
Date: Wed, 17 Jul 2019 01:12:21 -0700
Subject: [PATCH] context:  fix race between CancellableContext and Context

The `pendingDeadline` variable is modified from the ctor of CancellableContext, but it isn't final.  The cancellation can happen before the variable is assigned.  It's generally bad practice to leak the this reference from the ctor to other threads anyways.

This code refactors the deadline calculation and scheduling so that `pendingDeadline` is modified under the lock, and the `this` reference is not exposed.

Discovered by TSAN.
---
 context/src/main/java/io/grpc/Context.java | 68 ++++++++++++----------
 1 file changed, 37 insertions(+), 31 deletions(-)

diff --git a/context/src/main/java/io/grpc/Context.java b/context/src/main/java/io/grpc/Context.java
index 477adbe51c..4910744bc7 100644
--- a/context/src/main/java/io/grpc/Context.java
+++ b/context/src/main/java/io/grpc/Context.java
@@ -296,11 +296,21 @@ public class Context {
    *   }
    * </pre>
    */
-  public CancellableContext withDeadline(Deadline deadline,
-      ScheduledExecutorService scheduler) {
-    checkNotNull(deadline, "deadline");
+  public CancellableContext withDeadline(Deadline newDeadline, ScheduledExecutorService scheduler) {
+    checkNotNull(newDeadline, "deadline");
     checkNotNull(scheduler, "scheduler");
-    return new CancellableContext(this, deadline, scheduler);
+    Deadline existingDeadline = getDeadline();
+    boolean scheduleDeadlineCancellation = true;
+    if (existingDeadline != null && existingDeadline.compareTo(newDeadline) <= 0) {
+      // The new deadline won't have an effect, so ignore it
+      newDeadline = existingDeadline;
+      scheduleDeadlineCancellation = false;
+    }
+    CancellableContext newCtx = new CancellableContext(this, newDeadline);
+    if (scheduleDeadlineCancellation) {
+      newCtx.setUpDeadlineCancellation(newDeadline, scheduler);
+    }
+    return newCtx;
   }
 
   /**
@@ -713,37 +723,33 @@ public class Context {
     /**
      * Create a cancellable context that has a deadline.
      */
-    private CancellableContext(Context parent, Deadline deadline,
-        ScheduledExecutorService scheduler) {
+    private CancellableContext(Context parent, Deadline deadline) {
       super(parent, parent.keyValueEntries);
-      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.
-          pendingDeadline = deadline.runOnExpiration(new Runnable() {
-            @Override
-            public void run() {
-              try {
-                cancel(new TimeoutException("context timed out"));
-              } 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(new TimeoutException("context timed out"));
-        }
-      }
       this.deadline = deadline;
-      uncancellableSurrogate = new Context(this, keyValueEntries);
+      this.uncancellableSurrogate = new Context(this, keyValueEntries);
     }
 
+    private void setUpDeadlineCancellation(Deadline deadline, ScheduledExecutorService scheduler) {
+      if (!deadline.isExpired()) {
+        final class CancelOnExpiration implements Runnable {
+          @Override
+          public void run() {
+            try {
+              cancel(new TimeoutException("context timed out"));
+            } catch (Throwable t) {
+              log.log(Level.SEVERE, "Cancel threw an exception, which should not happen", t);
+            }
+          }
+        }
+
+        synchronized (this) {
+          pendingDeadline = deadline.runOnExpiration(new CancelOnExpiration(), scheduler);
+        }
+      } else {
+        // Cancel immediately if the deadline is already expired.
+        cancel(new TimeoutException("context timed out"));
+      }
+    }
 
     @Override
     public Context attach() {