fix: renew long-lived subscriptions

Signed-off-by: mattisonchao <mattisonchao@gmail.com>

Author: mattisonchao

client-api/src/main/java/io/oxia/client/api/OxiaClientBuilder.java

@@ -259,6 +259,44 @@ static OxiaClientBuilder create(String serviceAddress) {
      */
     OxiaClientBuilder connectionKeepAliveTime(Duration connectionKeepAlive);
 
+    /**
+     * Configure the maximum age of long-lived subscriptions.
+     *
+     * <p>Here, a subscription means a client operation that continuously receives shard assignments,
+     * notifications, or sequence updates. This setting does not expose the underlying transport and
+     * is not an inactivity timeout. The client transparently renews each subscription at a random age
+     * between half this value and this value, even when it is healthy.
+     *
+     * <p>Bounding the age lets the client recover when an intermediary loses its upstream connection
+     * while leaving the downstream connection apparently healthy. Randomization spreads renewals
+     * across clients, following the Kubernetes {@code client-go} Reflector pattern.
+     *
+     * <p>Renewal preserves logical progress: shard assignments restart from a complete snapshot,
+     * notifications continue after the last received offset, and sequence updates suppress the
+     * repeated current key returned when they restart.
+     *
+     * <p>Default is <code>10 minutes</code>, resulting in subscription ages between 5 and 10 minutes.
+     * Calling this method also re-enables the maximum age if it was previously disabled with {@link
+     * #disableSubscriptionMaxAge()}.
+     *
+     * @param subscriptionMaxAge the upper bound for a subscription's randomized age
+     * @return the builder instance
+     * @see <a
+     *     href="https://github.com/kubernetes/client-go/blob/master/tools/cache/reflector.go">Kubernetes
+     *     client-go Reflector</a>
+     */
+    OxiaClientBuilder subscriptionMaxAge(Duration subscriptionMaxAge);
+
+    /**
+     * Disable the maximum age for long-lived subscriptions.
+     *
+     * <p>With the maximum age disabled, a subscription can appear healthy indefinitely if an
+     * intermediary loses its upstream connection while keeping the downstream HTTP/2 connection open.
+     *
+     * @return the builder instance
+     */
+    OxiaClientBuilder disableSubscriptionMaxAge();
+
     /**
      * Configure the authentication plugin and its parameters.
      *

client/src/main/java/io/oxia/client/AsyncOxiaClientImpl.java

@@ -795,7 +795,8 @@ public Closeable getSequenceUpdates(
                 this.rpcProvider,
                 this.shardManager,
                 this.instrumentProvider,
-                x -> closed);
+                x -> closed,
+                this.scheduledExecutor);
     }
 
     @Override

client/src/main/java/io/oxia/client/ClientConfig.java

@@ -40,4 +40,49 @@ public record ClientConfig(
         @NonNull Duration connectionBackoffMaxDelay,
         Duration connectionKeepAliveTime,
         Duration connectionKeepAliveTimeout,
-        int maxConnectionPerNode) {}
+        int maxConnectionPerNode,
+        @Nullable Duration subscriptionMaxAge) {
+
+    public ClientConfig(
+            @NonNull String serviceAddress,
+            @NonNull Duration requestTimeout,
+            int maxRequestsPerBatch,
+            int maxBatchSize,
+            long maxPendingBytes,
+            int maxWriteBatchesInFlight,
+            int maxReadBatchesInFlight,
+            int batchingThreads,
+            @NonNull Duration sessionTimeout,
+            @NonNull String clientIdentifier,
+            OpenTelemetry openTelemetry,
+            @NonNull String namespace,
+            @Nullable Authentication authentication,
+            boolean enableTls,
+            @NonNull Duration connectionBackoffMinDelay,
+            @NonNull Duration connectionBackoffMaxDelay,
+            Duration connectionKeepAliveTime,
+            Duration connectionKeepAliveTimeout,
+            int maxConnectionPerNode) {
+        this(
+                serviceAddress,
+                requestTimeout,
+                maxRequestsPerBatch,
+                maxBatchSize,
+                maxPendingBytes,
+                maxWriteBatchesInFlight,
+                maxReadBatchesInFlight,
+                batchingThreads,
+                sessionTimeout,
+                clientIdentifier,
+                openTelemetry,
+                namespace,
+                authentication,
+                enableTls,
+                connectionBackoffMinDelay,
+                connectionBackoffMaxDelay,
+                connectionKeepAliveTime,
+                connectionKeepAliveTimeout,
+                maxConnectionPerNode,
+                OxiaClientBuilderImpl.DefaultSubscriptionMaxAge);
+    }
+}

client/src/main/java/io/oxia/client/OxiaClientBuilderImpl.java

@@ -57,6 +57,7 @@ public class OxiaClientBuilderImpl implements OxiaClientBuilder {
     public static final int DefaultBatchingThreads = 1;
     public static final Duration DefaultRequestTimeout = Duration.ofSeconds(30);
     public static final Duration DefaultSessionTimeout = Duration.ofSeconds(15);
+    public static final Duration DefaultSubscriptionMaxAge = Duration.ofMinutes(10);
     public static final String DefaultNamespace = "default";
     public static final boolean DefaultEnableTls = false;
     public static final int DefaultMaxConnectionPerNode = 1;
@@ -92,6 +93,8 @@ public class OxiaClientBuilderImpl implements OxiaClientBuilder {
     protected Duration connectionKeepAliveTime = Duration.ofSeconds(10);
     protected Duration connectionKeepAliveTimeout = Duration.ofSeconds(3);
 
+    @Nullable protected Duration subscriptionMaxAge = DefaultSubscriptionMaxAge;
+
     protected int maxConnectionsPerNode = DefaultMaxConnectionPerNode;
 
     @Nullable protected SharedResources sharedResources;
@@ -252,6 +255,29 @@ public OxiaClientBuilder connectionKeepAliveTime(Duration keepAliveTime) {
         return this;
     }
 
+    @Override
+    public OxiaClientBuilder subscriptionMaxAge(@NonNull Duration subscriptionMaxAge) {
+        final long maxAgeMillis;
+        try {
+            maxAgeMillis = subscriptionMaxAge.toMillis();
+        } catch (ArithmeticException e) {
+            throw new IllegalArgumentException(
+                    "subscriptionMaxAge is too large: " + subscriptionMaxAge, e);
+        }
+        if (maxAgeMillis < 2) {
+            throw new IllegalArgumentException(
+                    "subscriptionMaxAge must be at least 2 ms: " + subscriptionMaxAge);
+        }
+        this.subscriptionMaxAge = subscriptionMaxAge;
+        return this;
+    }
+
+    @Override
+    public OxiaClientBuilder disableSubscriptionMaxAge() {
+        this.subscriptionMaxAge = null;
+        return this;
+    }
+
     @Override
     public OxiaClientBuilder authentication(String authPluginClassName, String authParamsString)
             throws UnsupportedAuthenticationException {
@@ -392,7 +418,8 @@ public ClientConfig getClientConfig() {
                 connectionBackoffMaxDelay,
                 connectionKeepAliveTime,
                 connectionKeepAliveTimeout,
-                maxConnectionsPerNode);
+                maxConnectionsPerNode,
+                subscriptionMaxAge);
     }
 
     @Override

client/src/main/java/io/oxia/client/SequenceUpdates.java

@@ -15,7 +15,10 @@
  */
 package io.oxia.client;
 
+import static com.google.common.base.Throwables.getRootCause;
+
 import io.github.merlimat.slog.Logger;
+import io.grpc.Status;
 import io.opentelemetry.api.common.Attributes;
 import io.oxia.client.grpc.RpcProvider;
 import io.oxia.client.grpc.observer.CancelableStreamObserver;
@@ -27,6 +30,7 @@
 import io.oxia.proto.GetSequenceUpdatesResponse;
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import lombok.NonNull;
@@ -43,9 +47,11 @@ public class SequenceUpdates implements Closeable {
     private final ShardManager shardManager;
     private final Counter counterSequenceUpdatesReceived;
     private final Function<Void, Boolean> isClientClosed;
+    private final ScheduledExecutorService executor;
 
     private boolean closed = false;
     private CancelableStreamObserver<?> stream;
+    private String lastDeliveredSequenceKey;
 
     SequenceUpdates(
             @NonNull String key,
@@ -54,13 +60,15 @@ public class SequenceUpdates implements Closeable {
             @NonNull RpcProvider rpcProvider,
             @NonNull ShardManager shardManager,
             @NonNull InstrumentProvider instrumentProvider,
-            Function<Void, Boolean> isClientClosed) {
+            Function<Void, Boolean> isClientClosed,
+            @NonNull ScheduledExecutorService executor) {
         this.key = key;
         this.partitionKey = partitionKey;
         this.listener = listener;
         this.rpcProvider = rpcProvider;
         this.shardManager = shardManager;
         this.isClientClosed = isClientClosed;
+        this.executor = executor;
 
         this.counterSequenceUpdatesReceived =
                 instrumentProvider.newCounter(
@@ -73,7 +81,7 @@ public class SequenceUpdates implements Closeable {
     }
 
     private synchronized void createStream() {
-        if (closed) {
+        if (closed || isClientClosed.apply(null)) {
             return;
         }
 
@@ -115,20 +123,38 @@ public void close() throws IOException {
         }
     }
 
-    private void handleUpdate(@NonNull GetSequenceUpdatesResponse value) {
-        listener.accept(value.getHighestSequenceKey());
+    private synchronized void handleUpdate(@NonNull GetSequenceUpdatesResponse value) {
+        var highestSequenceKey = value.getHighestSequenceKey();
+        if (highestSequenceKey.equals(lastDeliveredSequenceKey)) {
+            // A renewed subscription starts with the server's current highest sequence key. Skip it
+            // when the previous subscription already delivered that value.
+            return;
+        }
+        lastDeliveredSequenceKey = highestSequenceKey;
+        listener.accept(highestSequenceKey);
         counterSequenceUpdatesReceived.increment();
     }
 
     private synchronized void handleError(@NonNull Throwable t) {
         if (closed || isClientClosed.apply(null)) {
             return;
         }
-        log.warn().exception(t).log("Failure while processing sequence updates");
-        createStream();
+        if (Status.fromThrowable(getRootCause(t)).getCode() == Status.Code.DEADLINE_EXCEEDED) {
+            log.debug("Sequence updates subscription reached its configured maximum age");
+        } else {
+            log.warn().exception(t).log("Failure while processing sequence updates");
+        }
+        scheduleRestart();
     }
 
     private synchronized void handleCompleted() {
-        createStream();
+        if (closed || isClientClosed.apply(null)) {
+            return;
+        }
+        scheduleRestart();
+    }
+
+    private void scheduleRestart() {
+        executor.execute(this::createStream);
     }
 }

client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java

@@ -37,6 +37,7 @@
 import io.oxia.proto.ListResponse;
 import io.oxia.proto.NotificationBatch;
 import io.oxia.proto.NotificationsRequest;
+import io.oxia.proto.OxiaClientGrpc;
 import io.oxia.proto.RangeScanRequest;
 import io.oxia.proto.RangeScanResponse;
 import io.oxia.proto.ReadRequest;
@@ -50,6 +51,7 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongFunction;
@@ -119,10 +121,10 @@ public void getShardAssignments(
                                 final var barrierObserver =
                                         ManagedObservers.toBarrierStreamObserver(guardedObserver, barrierFuture);
                                 try {
-                                    connectionManager
-                                            .getConnection(clientConfig.serviceAddress())
-                                            .stub()
-                                            .getShardAssignments(request, barrierObserver);
+                                    var stub =
+                                            withSubscriptionMaxAge(
+                                                    connectionManager.getConnection(clientConfig.serviceAddress()).stub());
+                                    stub.getShardAssignments(request, barrierObserver);
                                 } catch (Throwable error) {
                                     barrierFuture.completeExceptionally(OxiaStatusException.from(error));
                                 }
@@ -154,9 +156,10 @@ public void getNotifications(
                                 final var barrierObserver =
                                         ManagedObservers.toBarrierStreamObserver(guardedObserver, barrierFuture);
                                 try {
-                                    connectionManager
-                                            .getConnection(getLeader(request.getShard(), hint))
-                                            .stub()
+                                    withSubscriptionMaxAge(
+                                                    connectionManager
+                                                            .getConnection(getLeader(request.getShard(), hint))
+                                                            .stub())
                                             .getNotifications(request, barrierObserver);
                                 } catch (Throwable error) {
                                     barrierFuture.completeExceptionally(OxiaStatusException.from(error));
@@ -394,9 +397,10 @@ public void getSequenceUpdates(
                                 final var barrierObserver =
                                         ManagedObservers.toBarrierClientResponseObserver(observer, barrierFuture);
                                 try {
-                                    connectionManager
-                                            .getConnection(getLeader(request.getShard(), hint))
-                                            .stub()
+                                    withSubscriptionMaxAge(
+                                                    connectionManager
+                                                            .getConnection(getLeader(request.getShard(), hint))
+                                                            .stub())
                                             .getSequenceUpdates(request, barrierObserver);
                                 } catch (Throwable error) {
                                     barrierFuture.completeExceptionally(OxiaStatusException.from(error));
@@ -413,6 +417,20 @@ public void getSequenceUpdates(
         }
     }
 
+    private OxiaClientGrpc.OxiaClientStub withSubscriptionMaxAge(OxiaClientGrpc.OxiaClientStub stub) {
+        var maxAge = clientConfig.subscriptionMaxAge();
+        if (maxAge == null) {
+            return stub;
+        }
+
+        long maxAgeMillis = maxAge.toMillis();
+        long minAgeMillis = maxAgeMillis / 2;
+        // Match Kubernetes Reflector watches: renew each subscription at a random point in
+        // [maxAge/2, maxAge) to prevent hanging operations without synchronizing clients.
+        long ageMillis = ThreadLocalRandom.current().nextLong(minAgeMillis, maxAgeMillis);
+        return stub.withDeadlineAfter(ageMillis, TimeUnit.MILLISECONDS);
+    }
+
     @Override
     public void close() throws Exception {
         try {