diff --git a/client-api/src/main/java/io/oxia/client/api/OxiaClientBuilder.java b/client-api/src/main/java/io/oxia/client/api/OxiaClientBuilder.java index 79464fdc..0d287119 100644 --- a/client-api/src/main/java/io/oxia/client/api/OxiaClientBuilder.java +++ b/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. + * + *

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. + * + *

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. + * + *

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. + * + *

Default is 10 minutes, 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 Kubernetes + * client-go Reflector + */ + OxiaClientBuilder subscriptionMaxAge(Duration subscriptionMaxAge); + + /** + * Disable the maximum age for long-lived subscriptions. + * + *

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. * diff --git a/client/src/main/java/io/oxia/client/AsyncOxiaClientImpl.java b/client/src/main/java/io/oxia/client/AsyncOxiaClientImpl.java index 3477b771..d4dcc410 100644 --- a/client/src/main/java/io/oxia/client/AsyncOxiaClientImpl.java +++ b/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 diff --git a/client/src/main/java/io/oxia/client/ClientConfig.java b/client/src/main/java/io/oxia/client/ClientConfig.java index c96eff29..119bd63d 100644 --- a/client/src/main/java/io/oxia/client/ClientConfig.java +++ b/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); + } +} diff --git a/client/src/main/java/io/oxia/client/OxiaClientBuilderImpl.java b/client/src/main/java/io/oxia/client/OxiaClientBuilderImpl.java index cd10c28c..0e530118 100644 --- a/client/src/main/java/io/oxia/client/OxiaClientBuilderImpl.java +++ b/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 diff --git a/client/src/main/java/io/oxia/client/SequenceUpdates.java b/client/src/main/java/io/oxia/client/SequenceUpdates.java index 43d7a199..d2feaea9 100644 --- a/client/src/main/java/io/oxia/client/SequenceUpdates.java +++ b/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 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 isClientClosed) { + Function 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,8 +123,17 @@ 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 (lastDeliveredSequenceKey != null + && highestSequenceKey.compareTo(lastDeliveredSequenceKey) <= 0) { + // Sequence keys use fixed-width numeric suffixes, so lexical order matches sequence + // order. A renewal can replay an older snapshot; keep callbacks monotonic by skipping + // keys that were already delivered or superseded. + return; + } + lastDeliveredSequenceKey = highestSequenceKey; + listener.accept(highestSequenceKey); counterSequenceUpdatesReceived.increment(); } @@ -124,11 +141,22 @@ 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); } } diff --git a/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java b/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java index 4eca7073..7d37cf84 100644 --- a/client/src/main/java/io/oxia/client/grpc/GrpcRpcProvider.java +++ b/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 { diff --git a/client/src/main/java/io/oxia/client/notify/ShardNotificationReceiver.java b/client/src/main/java/io/oxia/client/notify/ShardNotificationReceiver.java index a0c4dd5d..52b7d622 100644 --- a/client/src/main/java/io/oxia/client/notify/ShardNotificationReceiver.java +++ b/client/src/main/java/io/oxia/client/notify/ShardNotificationReceiver.java @@ -15,10 +15,12 @@ */ package io.oxia.client.notify; +import static com.google.common.base.Throwables.getRootCause; import static io.oxia.client.api.Notification.KeyModified; import static lombok.AccessLevel.PACKAGE; import io.github.merlimat.slog.Logger; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import io.oxia.client.CompositeConsumer; import io.oxia.client.api.Notification; @@ -83,6 +85,7 @@ void start() { @Override public void onNext(NotificationBatch batch) { + backoff.reset(); if (offset.isPresent() && offset.getAsLong() >= batch.getOffset()) { // Ignore repeated notifications return; @@ -119,31 +122,41 @@ public void onError(Throwable t) { return; } + if (Status.fromThrowable(getRootCause(t)).getCode() == Status.Code.DEADLINE_EXCEEDED) { + backoff.reset(); + log.debug("Notifications subscription reached its configured maximum age"); + scheduleRestart(0); + return; + } + long retryDelayMillis = backoff.nextDelayMillis(); log.warn() .attr("retryInSeconds", retryDelayMillis / 1000.0) .exceptionMessage(t) .log("Error while receiving notifications"); + scheduleRestart(retryDelayMillis); + } + + @Override + public void onCompleted() { + if (!closed) { + scheduleRestart(0); + } + } + + private void scheduleRestart(long delayMillis) { notificationManager .getExecutor() .schedule( () -> { if (!closed) { - log.info("Retrying getting notifications"); start(); } }, - retryDelayMillis, + delayMillis, TimeUnit.MILLISECONDS); } - @Override - public void onCompleted() { - if (!closed) { - start(); - } - } - @RequiredArgsConstructor(access = PACKAGE) static class Factory { private final @NonNull RpcProvider rpcProvider; diff --git a/client/src/main/java/io/oxia/client/shard/ShardManager.java b/client/src/main/java/io/oxia/client/shard/ShardManager.java index 1c85a1ec..7f2fb1f8 100644 --- a/client/src/main/java/io/oxia/client/shard/ShardManager.java +++ b/client/src/main/java/io/oxia/client/shard/ShardManager.java @@ -26,6 +26,7 @@ import com.google.common.annotations.VisibleForTesting; import io.github.merlimat.slog.Logger; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import io.opentelemetry.api.common.Attributes; import io.oxia.client.CompositeConsumer; @@ -123,15 +124,27 @@ public void onError(Throwable error) { return; } } - log.warn().exceptionMessage(getRootCause(error)).log("Failed receiving shard assignments"); + final long retryDelayMillis; + if (Status.fromThrowable(getRootCause(error)).getCode() == Status.Code.DEADLINE_EXCEEDED + && initialAssignmentsFuture.isDone()) { + backoff.reset(); + log.debug("Shard assignments subscription reached its configured maximum age"); + retryDelayMillis = 0; + } else { + log.warn().exceptionMessage(getRootCause(error)).log("Failed receiving shard assignments"); + retryDelayMillis = backoff.nextDelayMillis(); + } + scheduleRestart(retryDelayMillis); + } + + private void scheduleRestart(long delayMillis) { asyncExecutor.schedule( () -> { if (!closed) { - log.info("Retry creating stream for shard assignments"); start(); } }, - backoff.nextDelayMillis(), + delayMillis, TimeUnit.MILLISECONDS); } @@ -142,15 +155,7 @@ public void onCompleted() { } log.warn("Stream closed while receiving shard assignments"); - asyncExecutor.schedule( - () -> { - if (!closed) { - log.info("Retry creating stream for shard assignments after stream closed"); - start(); - } - }, - backoff.nextDelayMillis(), - TimeUnit.MILLISECONDS); + scheduleRestart(backoff.nextDelayMillis()); } private void updateAssignments(io.oxia.proto.ShardAssignments shardAssignments) { diff --git a/client/src/test/java/io/oxia/client/OxiaClientBuilderTest.java b/client/src/test/java/io/oxia/client/OxiaClientBuilderTest.java index a84ca8d1..b9438651 100644 --- a/client/src/test/java/io/oxia/client/OxiaClientBuilderTest.java +++ b/client/src/test/java/io/oxia/client/OxiaClientBuilderTest.java @@ -42,6 +42,29 @@ void requestTimeout() { assertThatNoException().isThrownBy(() -> builder.requestTimeout(Duration.ofMillis(1))); } + @Test + void subscriptionMaxAge() { + assertThatThrownBy(() -> builder.subscriptionMaxAge(ZERO)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> builder.subscriptionMaxAge(Duration.ofMillis(1))) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> builder.subscriptionMaxAge(Duration.ofNanos(1))) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> builder.subscriptionMaxAge(Duration.ofMillis(-1))) + .isInstanceOf(IllegalArgumentException.class); + + var timeout = Duration.ofMinutes(1); + builder.subscriptionMaxAge(timeout); + var impl = (OxiaClientBuilderImpl) builder; + assertThat(impl.subscriptionMaxAge).isEqualTo(timeout); + + builder.disableSubscriptionMaxAge(); + assertThat(impl.subscriptionMaxAge).isNull(); + + builder.subscriptionMaxAge(timeout); + assertThat(impl.subscriptionMaxAge).isEqualTo(timeout); + } + @Test void batchLinger() { assertThatThrownBy(() -> builder.batchLinger(ZERO)) @@ -84,6 +107,7 @@ void loadConfigWithProperties() { Properties properties = new Properties(); properties.setProperty("serviceAddress", "address:5678"); properties.setProperty("requestTimeout", "1"); + properties.setProperty("subscriptionMaxAge", "5"); properties.setProperty("batchLinger", "2"); properties.setProperty("maxRequestsPerBatch", "3"); properties.setProperty("sessionTimeout", "4"); @@ -96,6 +120,7 @@ void loadConfigWithProperties() { OxiaClientBuilderImpl impl = (OxiaClientBuilderImpl) builder; assertThat(impl.serviceAddress).isEqualTo("address:5678"); assertThat(impl.requestTimeout).isEqualTo(Duration.ofMillis(1)); + assertThat(impl.subscriptionMaxAge).isEqualTo(Duration.ofMillis(5)); assertThat(impl.batchLinger).isEqualTo(Duration.ofMillis(2)); assertThat(impl.maxRequestsPerBatch).isEqualTo(3); assertThat(impl.sessionTimeout).isEqualTo(Duration.ofMillis(4)); diff --git a/client/src/test/java/io/oxia/client/SequenceUpdatesTest.java b/client/src/test/java/io/oxia/client/SequenceUpdatesTest.java new file mode 100644 index 00000000..0b487c32 --- /dev/null +++ b/client/src/test/java/io/oxia/client/SequenceUpdatesTest.java @@ -0,0 +1,86 @@ +/* + * Copyright © 2026 The Oxia Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.oxia.client; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.grpc.Status; +import io.oxia.client.grpc.RpcProvider; +import io.oxia.client.grpc.observer.CancelableStreamObserver; +import io.oxia.client.metrics.InstrumentProvider; +import io.oxia.client.shard.ShardManager; +import io.oxia.proto.GetSequenceUpdatesRequest; +import io.oxia.proto.GetSequenceUpdatesResponse; +import java.util.ArrayList; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +class SequenceUpdatesTest { + @Test + void ignoresCurrentSequenceKeyAfterSubscriptionRenewal() throws Exception { + var rpcProvider = mock(RpcProvider.class); + var shardManager = mock(ShardManager.class); + when(shardManager.getShardForKey(any())).thenReturn(0L); + var observerRef = new AtomicReference>(); + doAnswer( + invocation -> { + observerRef.set(invocation.getArgument(1)); + return null; + }) + .when(rpcProvider) + .getSequenceUpdates(any(GetSequenceUpdatesRequest.class), any()); + + var delivered = new ArrayList(); + var executor = Executors.newSingleThreadScheduledExecutor(); + try { + try (var updates = + new SequenceUpdates( + "key", + "partition", + delivered::add, + rpcProvider, + shardManager, + InstrumentProvider.NOOP, + ignored -> false, + executor)) { + var firstObserver = observerRef.get(); + var first = + new GetSequenceUpdatesResponse().setHighestSequenceKey("key-00000000000000000001"); + firstObserver.onNext(first); + + firstObserver.onError(Status.DEADLINE_EXCEEDED.asRuntimeException()); + await().untilAsserted(() -> assertThat(observerRef.get()).isNotSameAs(firstObserver)); + var renewedObserver = observerRef.get(); + renewedObserver.onNext( + new GetSequenceUpdatesResponse().setHighestSequenceKey("key-00000000000000000000")); + renewedObserver.onNext(first); + renewedObserver.onNext( + new GetSequenceUpdatesResponse().setHighestSequenceKey("key-00000000000000000002")); + + assertThat(delivered) + .containsExactly("key-00000000000000000001", "key-00000000000000000002"); + } + } finally { + executor.shutdownNow(); + } + } +} diff --git a/client/src/test/java/io/oxia/client/grpc/GrpcRpcProviderTest.java b/client/src/test/java/io/oxia/client/grpc/GrpcRpcProviderTest.java index 5612773c..55c97ff7 100644 --- a/client/src/test/java/io/oxia/client/grpc/GrpcRpcProviderTest.java +++ b/client/src/test/java/io/oxia/client/grpc/GrpcRpcProviderTest.java @@ -57,6 +57,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import lombok.NonNull; import org.junit.jupiter.api.Test; @@ -921,6 +922,237 @@ public void onCompleted() {} } } + @Test + void longLivedSubscriptionsUseRandomizedDeadline() throws Exception { + var assignmentsDeadlineMillis = new AtomicLong(-1); + var notificationsDeadlineMillis = new AtomicLong(-1); + var sequenceUpdatesDeadlineMillis = new AtomicLong(-1); + Server server = + ServerBuilder.forPort(0) + .directExecutor() + .addService( + new OxiaClientGrpc.OxiaClientImplBase() { + @Override + public void getShardAssignments( + ShardAssignmentsRequest request, + StreamObserver responseObserver) { + var deadline = Context.current().getDeadline(); + if (deadline != null) { + assignmentsDeadlineMillis.set(deadline.timeRemaining(TimeUnit.MILLISECONDS)); + } + responseObserver.onNext(new ShardAssignments()); + } + + @Override + public void getNotifications( + NotificationsRequest request, + StreamObserver responseObserver) { + var deadline = Context.current().getDeadline(); + if (deadline != null) { + notificationsDeadlineMillis.set( + deadline.timeRemaining(TimeUnit.MILLISECONDS)); + } + responseObserver.onNext(new NotificationBatch()); + } + + @Override + public void getSequenceUpdates( + GetSequenceUpdatesRequest request, + StreamObserver responseObserver) { + var deadline = Context.current().getDeadline(); + if (deadline != null) { + sequenceUpdatesDeadlineMillis.set( + deadline.timeRemaining(TimeUnit.MILLISECONDS)); + } + responseObserver.onNext(new GetSequenceUpdatesResponse()); + } + }) + .build() + .start(); + var address = "localhost:" + server.getPort(); + var executor = Executors.newSingleThreadScheduledExecutor(); + var config = + ((OxiaClientBuilderImpl) + OxiaClientBuilder.create(address) + .requestTimeout(Duration.ofSeconds(5)) + .subscriptionMaxAge(Duration.ofMillis(200))) + .getClientConfig(); + var received = new CountDownLatch(3); + var terminated = new CountDownLatch(3); + + try (var provider = new GrpcRpcProvider(config, executor, shardId -> address)) { + provider.getShardAssignments( + new ShardAssignmentsRequest(), + new StreamObserver<>() { + @Override + public void onNext(ShardAssignments value) { + received.countDown(); + } + + @Override + public void onError(Throwable t) { + terminated.countDown(); + } + + @Override + public void onCompleted() { + terminated.countDown(); + } + }); + provider.getNotifications( + new NotificationsRequest().setShard(1), + new StreamObserver<>() { + @Override + public void onNext(NotificationBatch value) { + received.countDown(); + } + + @Override + public void onError(Throwable t) { + terminated.countDown(); + } + + @Override + public void onCompleted() { + terminated.countDown(); + } + }); + provider.getSequenceUpdates( + new GetSequenceUpdatesRequest().setShard(1), + new CancelableStreamObserver<>() { + @Override + protected void handleNext(GetSequenceUpdatesResponse value) { + received.countDown(); + } + + @Override + protected void handleError(Throwable t) { + terminated.countDown(); + } + + @Override + protected void handleComplete() { + terminated.countDown(); + } + }); + + assertThat(received.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(terminated.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(assignmentsDeadlineMillis.get()).isBetween(1L, 200L); + assertThat(notificationsDeadlineMillis.get()).isBetween(1L, 200L); + assertThat(sequenceUpdatesDeadlineMillis.get()).isBetween(1L, 200L); + } finally { + executor.shutdownNow(); + server.shutdownNow(); + } + } + + @Test + void longLivedSubscriptionsCanDisableDeadline() throws Exception { + var assignmentsDeadline = new AtomicReference(); + var notificationsDeadline = new AtomicReference(); + var sequenceUpdatesDeadline = new AtomicReference(); + Server server = + ServerBuilder.forPort(0) + .directExecutor() + .addService( + new OxiaClientGrpc.OxiaClientImplBase() { + @Override + public void getShardAssignments( + ShardAssignmentsRequest request, + StreamObserver responseObserver) { + assignmentsDeadline.set(Context.current().getDeadline()); + responseObserver.onNext(new ShardAssignments()); + responseObserver.onCompleted(); + } + + @Override + public void getNotifications( + NotificationsRequest request, + StreamObserver responseObserver) { + notificationsDeadline.set(Context.current().getDeadline()); + responseObserver.onNext(new NotificationBatch()); + responseObserver.onCompleted(); + } + + @Override + public void getSequenceUpdates( + GetSequenceUpdatesRequest request, + StreamObserver responseObserver) { + sequenceUpdatesDeadline.set(Context.current().getDeadline()); + responseObserver.onNext(new GetSequenceUpdatesResponse()); + responseObserver.onCompleted(); + } + }) + .build() + .start(); + var address = "localhost:" + server.getPort(); + var executor = Executors.newSingleThreadScheduledExecutor(); + var config = + ((OxiaClientBuilderImpl) OxiaClientBuilder.create(address).disableSubscriptionMaxAge()) + .getClientConfig(); + var completed = new CountDownLatch(3); + + try (var provider = new GrpcRpcProvider(config, executor, shardId -> address)) { + provider.getShardAssignments( + new ShardAssignmentsRequest(), + new StreamObserver<>() { + @Override + public void onNext(ShardAssignments value) {} + + @Override + public void onError(Throwable t) { + completed.countDown(); + } + + @Override + public void onCompleted() { + completed.countDown(); + } + }); + provider.getNotifications( + new NotificationsRequest().setShard(1), + new StreamObserver<>() { + @Override + public void onNext(NotificationBatch value) {} + + @Override + public void onError(Throwable t) { + completed.countDown(); + } + + @Override + public void onCompleted() { + completed.countDown(); + } + }); + provider.getSequenceUpdates( + new GetSequenceUpdatesRequest().setShard(1), + new CancelableStreamObserver<>() { + @Override + protected void handleNext(GetSequenceUpdatesResponse value) {} + + @Override + protected void handleError(Throwable t) { + completed.countDown(); + } + + @Override + protected void handleComplete() { + completed.countDown(); + } + }); + + assertThat(completed.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(assignmentsDeadline.get()).isNull(); + assertThat(notificationsDeadline.get()).isNull(); + assertThat(sequenceUpdatesDeadline.get()).isNull(); + } finally { + executor.shutdownNow(); + server.shutdownNow(); + } + } + @Test void getNotificationsTimesOutSilentInitialAttempt() throws Exception { Server server = diff --git a/client/src/test/java/io/oxia/client/notify/ShardNotificationReceiverTest.java b/client/src/test/java/io/oxia/client/notify/ShardNotificationReceiverTest.java index c808c6ff..ad584324 100644 --- a/client/src/test/java/io/oxia/client/notify/ShardNotificationReceiverTest.java +++ b/client/src/test/java/io/oxia/client/notify/ShardNotificationReceiverTest.java @@ -196,8 +196,41 @@ public void recoveryFromError() { assertThat(requests).hasValue(2); } + @Test + void renewsImmediatelyAtConfiguredMaximumAge() { + @Cleanup("shutdownNow") + ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); + when(notificationManager.getExecutor()).thenReturn(executorService); + when(notificationManager.getCounterNotificationsReceived()).thenReturn(mock(Counter.class)); + when(notificationManager.getCounterNotificationsBatchesReceived()) + .thenReturn(mock(Counter.class)); + + assertThat( + responses.offer( + new NotificationWrapper( + null, Status.DEADLINE_EXCEEDED.asRuntimeException(), false))) + .isTrue(); + assertThat( + responses.offer( + new NotificationWrapper(newNotificationBatch("key1", created(1L)), null, false))) + .isTrue(); + try (var notificationReceiver = + new ShardNotificationReceiver( + rpcProvider, + shardId, + notificationCallback, + notificationManager, + OptionalLong.empty())) { + await().untilAsserted(() -> verify(notificationCallback).accept(new KeyCreated("key1", 1L))); + } + assertThat(requests).hasValue(2); + } + @Test public void recoveryFromEndOfStream() throws Exception { + @Cleanup("shutdownNow") + ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); + when(notificationManager.getExecutor()).thenReturn(executorService); when(notificationManager.getCounterNotificationsReceived()).thenReturn(mock(Counter.class)); when(notificationManager.getCounterNotificationsBatchesReceived()) .thenReturn(mock(Counter.class));