From 4fb54659f991beb9f3741366e2a8e3bbffac958f Mon Sep 17 00:00:00 2001
From: Lawrence Qiu <lawrenceqiu@google.com>
Date: Fri, 12 Jan 2024 12:50:53 -0500
Subject: [PATCH] fix: Cancel the Timeout Task for HttpJson

---
 .../gax/httpjson/HttpJsonClientCallImpl.java  | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCallImpl.java b/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCallImpl.java
index b98a5319fc0..6b9073544fd 100644
--- a/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCallImpl.java
+++ b/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCallImpl.java
@@ -46,6 +46,7 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
@@ -121,6 +122,11 @@ final class HttpJsonClientCallImpl<RequestT, ResponseT>
   @GuardedBy("lock")
   private volatile boolean closed;
 
+  // Store the timeout future created by the deadline schedule executor. The future
+  // can be cancelled if a response has been received before the timeout.
+  @GuardedBy("lock")
+  private ScheduledFuture<?> timeoutFuture;
+
   HttpJsonClientCallImpl(
       ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor,
       String endpoint,
@@ -176,7 +182,11 @@ public void start(Listener<ResponseT> responseListener, HttpJsonMetadata request
       // The future timeout value is guaranteed to not be a negative value as the
       // RetryAlgorithm will not retry
       long timeoutMs = timeout.toMillis();
-      this.deadlineCancellationExecutor.schedule(this::timeout, timeoutMs, TimeUnit.MILLISECONDS);
+      // Assign the scheduled future so that it can be cancelled if the timeout task
+      // is not needed (response received prior to timeout)
+      timeoutFuture =
+          this.deadlineCancellationExecutor.schedule(
+              this::timeout, timeoutMs, TimeUnit.MILLISECONDS);
     }
   }
 
@@ -430,6 +440,14 @@ private void close(
         return;
       }
       closed = true;
+
+      // Cancel the timeout future if there is a timeout task created
+      if (timeoutFuture != null) {
+        // timeout() invokes close(), but cancelling a completed task should no-op
+        timeoutFuture.cancel(true);
+        timeoutFuture = null;
+      }
+
       // Best effort task cancellation (to not be confused with task's thread interruption).
       // If the task is in blocking I/O waiting for the server response, it will keep waiting for
       // the response from the server, but once response is received the task will exit silently.