feat: add keepalive feature to tear down streams in their absence

SamuelSchaub · SamuelSchaub · commit b89478d73ec5 · 2025-11-05T18:35:57.000Z
diff --git a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/StreamingSubscriberConnection.java b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/StreamingSubscriberConnection.java
@@ -63,6 +63,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -95,6 +96,7 @@ final class StreamingSubscriberConnection extends AbstractApiService implements
 
   private final SubscriberStub subscriberStub;
   private final int channelAffinity;
+  private final long protocolVersion;
   private final String subscription;
   private final SubscriptionName subscriptionNameObject;
   private final ScheduledExecutorService systemExecutor;
@@ -127,6 +129,17 @@ final class StreamingSubscriberConnection extends AbstractApiService implements
   private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false);
   private final SubscriberShutdownSettings subscriberShutdownSettings;
 
+  private final boolean enableKeepalive;
+  private static final long KEEP_ALIVE_SUPPORT_VERSION = 1;
+  private static final Duration CLIENT_PING_INTERVAL = Duration.ofSeconds(30);
+  private ScheduledFuture<?> pingSchedulerHandle;
+
+  private static final Duration SERVER_MONITOR_INTERVAL = Duration.ofSeconds(15);
+  private static final Duration SERVER_PING_TIMEOUT_DURATION = Duration.ofSeconds(15);
+  private final AtomicLong lastServerResponseTime;
+  private final AtomicLong lastClientPingTime;
+  private ScheduledFuture<?> serverMonitorHandle;
+
   private StreamingSubscriberConnection(Builder builder) {
     subscription = builder.subscription;
     subscriptionNameObject = SubscriptionName.parse(builder.subscription);
@@ -154,6 +167,7 @@ private StreamingSubscriberConnection(Builder builder) {
 
     subscriberStub = builder.subscriberStub;
     channelAffinity = builder.channelAffinity;
+    protocolVersion = builder.protocolVersion;
 
     MessageDispatcher.Builder messageDispatcherBuilder;
     if (builder.receiver != null) {
@@ -190,6 +204,9 @@ private StreamingSubscriberConnection(Builder builder) {
 
     flowControlSettings = builder.flowControlSettings;
     useLegacyFlowControl = builder.useLegacyFlowControl;
+    enableKeepalive = protocolVersion >= KEEP_ALIVE_SUPPORT_VERSION;
+    lastServerResponseTime = new AtomicLong(clock.nanoTime());
+    lastClientPingTime = new AtomicLong(-1L);
   }
 
   public StreamingSubscriberConnection setExactlyOnceDeliveryEnabled(
@@ -218,6 +235,12 @@ protected void doStop() {
     } finally {
       lock.unlock();
     }
+
+    if (enableKeepalive) {
+      stopClientPinger();
+      stopServerMonitor();
+    }
+
     runShutdown();
     notifyStopped();
   }
@@ -266,6 +289,10 @@ public void onStart(StreamController controller) {
 
     @Override
     public void onResponse(StreamingPullResponse response) {
+      if (enableKeepalive) {
+        lastServerResponseTime.set(clock.nanoTime());
+      }
+
       channelReconnectBackoffMillis.set(INITIAL_CHANNEL_RECONNECT_BACKOFF.toMillis());
 
       boolean exactlyOnceDeliveryEnabledResponse =
@@ -295,11 +322,19 @@ public void onResponse(StreamingPullResponse response) {
 
     @Override
     public void onError(Throwable t) {
+      if (enableKeepalive) {
+        stopClientPinger();
+        stopServerMonitor();
+      }
       errorFuture.setException(t);
     }
 
     @Override
     public void onComplete() {
+      if (enableKeepalive) {
+        stopClientPinger();
+        stopServerMonitor();
+      }
       logger.fine("Streaming pull terminated successfully!");
       errorFuture.set(null);
     }
@@ -336,6 +371,7 @@ private void initialize() {
                 this.useLegacyFlowControl
                     ? 0
                     : valueOrZero(flowControlSettings.getMaxOutstandingRequestBytes()))
+            .setProtocolVersion(protocolVersion)
             .build());
 
     /**
@@ -350,6 +386,11 @@ private void initialize() {
       lock.unlock();
     }
 
+    if (enableKeepalive) {
+      startClientPinger();
+      startServerMonitor();
+    }
+
     ApiFutures.addCallback(
         errorFuture,
         new ApiFutureCallback<Void>() {
@@ -366,6 +407,10 @@ public void onSuccess(@Nullable Void result) {
 
           @Override
           public void onFailure(Throwable cause) {
+            if (enableKeepalive) {
+              stopClientPinger();
+              stopServerMonitor();
+            }
             if (!isAlive()) {
               // we don't care about subscription failures when we're no longer running.
               logger.log(Level.FINE, "pull failure after service no longer running", cause);
@@ -410,6 +455,100 @@ private boolean isAlive() {
     return state == State.RUNNING || state == State.STARTING;
   }
 
+  private void startClientPinger() {
+    if (pingSchedulerHandle != null) {
+      pingSchedulerHandle.cancel(false);
+    }
+
+    pingSchedulerHandle =
+        systemExecutor.scheduleAtFixedRate(
+            () -> {
+              try {
+                lock.lock();
+                try {
+                  if (clientStream != null && isAlive()) {
+                    clientStream.send(StreamingPullRequest.newBuilder().build());
+                    lastClientPingTime.set(clock.nanoTime());
+                    logger.log(Level.FINEST, "Sent client keepalive ping");
+                  }
+                } finally {
+                  lock.unlock();
+                }
+              } catch (Exception e) {
+                logger.log(Level.FINE, "Error sending client keepalive ping", e);
+              }
+            },
+            CLIENT_PING_INTERVAL.getSeconds(),
+            CLIENT_PING_INTERVAL.getSeconds(),
+            TimeUnit.SECONDS);
+  }
+
+  private void stopClientPinger() {
+    if (pingSchedulerHandle != null) {
+      pingSchedulerHandle.cancel(false);
+      pingSchedulerHandle = null;
+    }
+  }
+
+  private void startServerMonitor() {
+    if (serverMonitorHandle != null) {
+      serverMonitorHandle.cancel(false);
+    }
+
+    serverMonitorHandle =
+        systemExecutor.scheduleAtFixedRate(
+            () -> {
+              try {
+                if (!isAlive()) {
+                  return;
+                }
+
+                long now = clock.nanoTime();
+                long lastResponse = lastServerResponseTime.get();
+                long lastPing = lastClientPingTime.get();
+
+                if (lastPing <= lastResponse) {
+                  return;
+                }
+
+                Duration elapsedSincePing = Duration.ofNanos(now - lastPing);
+                if (elapsedSincePing.compareTo(SERVER_PING_TIMEOUT_DURATION) <= 0) {
+                  return;
+                }
+
+                logger.log(
+                    Level.WARNING,
+                    "No response from server for {0} seconds. Closing stream.",
+                    elapsedSincePing.getSeconds());
+
+                lock.lock();
+                try {
+                  if (clientStream != null) {
+                    clientStream.closeSendWithError(
+                        Status.UNAVAILABLE
+                            .withDescription("Keepalive timeout with server")
+                            .asException());
+                  }
+                } finally {
+                  lock.unlock();
+                }
+                stopServerMonitor();
+              } catch (Exception e) {
+                logger.log(Level.FINE, "Error in server keepalive monitor", e);
+              }
+            },
+            SERVER_MONITOR_INTERVAL.getSeconds(),
+            SERVER_MONITOR_INTERVAL.getSeconds(),
+            TimeUnit.SECONDS);
+  }
+
+  private void stopServerMonitor() {
+    if (serverMonitorHandle != null) {
+      serverMonitorHandle.cancel(false);
+      serverMonitorHandle = null;
+    }
+  }
+
   public void setResponseOutstandingMessages(AckResponse ackResponse) {
     // We will close the futures with ackResponse - if there are multiple references to the same
     // future they will be handled appropriately
@@ -769,6 +908,7 @@ public static final class Builder {
     private Distribution ackLatencyDistribution;
     private SubscriberStub subscriberStub;
     private int channelAffinity;
+    private long protocolVersion;
     private FlowController flowController;
     private FlowControlSettings flowControlSettings;
     private boolean useLegacyFlowControl;
@@ -840,6 +980,11 @@ public Builder setChannelAffinity(int channelAffinity) {
       return this;
     }
 
+    public Builder setProtocolVersion(long protocolVersion) {
+      this.protocolVersion = protocolVersion;
+      return this;
+    }
+
     public Builder setFlowController(FlowController flowController) {
       this.flowController = flowController;
       return this;
diff --git a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Subscriber.java b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Subscriber.java
@@ -144,6 +144,7 @@ public class Subscriber extends AbstractApiService implements SubscriberInterfac
   private final boolean maxDurationPerAckExtensionDefaultUsed;
   private final java.time.Duration minDurationPerAckExtension;
   private final boolean minDurationPerAckExtensionDefaultUsed;
+  private final long protocolVersion;
 
   // The ExecutorProvider used to generate executors for processing messages.
   private final ExecutorProvider executorProvider;
@@ -182,6 +183,7 @@ private Subscriber(Builder builder) {
     maxDurationPerAckExtensionDefaultUsed = builder.maxDurationPerAckExtensionDefaultUsed;
     minDurationPerAckExtension = builder.minDurationPerAckExtension;
     minDurationPerAckExtensionDefaultUsed = builder.minDurationPerAckExtensionDefaultUsed;
+    protocolVersion = builder.protocolVersion;
 
     clock = builder.clock.isPresent() ? builder.clock.get() : CurrentMillisClock.getDefaultClock();
 
@@ -428,6 +430,7 @@ private void startStreamingConnections() {
                 .setEnableOpenTelemetryTracing(enableOpenTelemetryTracing)
                 .setTracer(tracer)
                 .setSubscriberShutdownSettings(subscriberShutdownSettings)
+                .setProtocolVersion(protocolVersion)
                 .build();
 
         streamingSubscriberConnections.add(streamingSubscriberConnection);
@@ -548,6 +551,8 @@ public static final class Builder {
     private boolean enableOpenTelemetryTracing = false;
     private OpenTelemetry openTelemetry = null;
 
+    private long protocolVersion = 0L;
+
     private SubscriberShutdownSettings subscriberShutdownSettings =
         SubscriberShutdownSettings.newBuilder().build();
 
@@ -771,6 +776,12 @@ Builder setClock(ApiClock clock) {
       return this;
     }
 
+    /** Gives the ability to override the protocol version */
+    public Builder setProtocolVersion(long protocolVersion) {
+      this.protocolVersion = protocolVersion;
+      return this;
+    }
+
     /**
      * OpenTelemetry will be enabled if setEnableOpenTelemetry is true and and instance of
      * OpenTelemetry has been provied. Warning: traces are subject to change. The name and
diff --git a/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/StreamingSubscriberConnectionTest.java b/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/StreamingSubscriberConnectionTest.java
