From 772b3a529ce719c5a16eb16f99e160573e33cc8e Mon Sep 17 00:00:00 2001 From: contrueCT Date: Sat, 15 Aug 2026 17:25:54 +0800 Subject: [PATCH 1/6] fix(pd): keep KvClient watches alive after reconnect failures --- .../hugegraph/pd/client/AbstractClient.java | 9 +- .../apache/hugegraph/pd/client/KvClient.java | 216 ++++++++---- .../hugegraph/pd/client/KvClientTest.java | 325 ++++++++++++++++++ 3 files changed, 478 insertions(+), 72 deletions(-) diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java index 570bacc6bf..15b5865feb 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java @@ -137,7 +137,7 @@ private String resetStub() { Exception ex = null; for (int i = 0; i < proxy.getHostCount(); i++) { String host = proxy.nextHost(); - close(); + closeConnections(); channel = ManagedChannelBuilder.forTarget(host).usePlaintext().build(); PDBlockingStub blockingStub = @@ -149,7 +149,7 @@ private String resetStub() { Metapb.Member leader = members.getLeader(); leaderHost = leader.getGrpcUrl(); if (!host.equals(leaderHost)) { - close(); + closeConnections(); channel = ManagedChannelBuilder.forTarget(leaderHost).usePlaintext().build(); } proxy.setBlockingStub(setBlockingParams(createBlockingStub(), config)); @@ -268,13 +268,16 @@ protected void streamingCall(MethodDescriptor method, @Override public void close() { + closeConnections(); + } + + private void closeConnections() { closeChannel(channel); if (stubs != null) { for (AbstractBlockingStub stub : stubs.values()) { closeChannel((ManagedChannel) stub.getChannel()); } } - } private void closeChannel(ManagedChannel channel) { diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index 6197c891ad..48a36486f8 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -24,10 +24,13 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; -import java.util.function.BiConsumer; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import org.apache.hugegraph.pd.common.PDException; @@ -56,13 +59,25 @@ @Slf4j public class KvClient extends AbstractClient implements Closeable { - private AtomicLong clientId = new AtomicLong(0); - private Semaphore semaphore = new Semaphore(1); - private AtomicBoolean closed = new AtomicBoolean(false); - private Set observers = ConcurrentHashMap.newKeySet(); + private static final long RECONNECT_DELAY_MS = 1000L; + + private final AtomicLong clientId = new AtomicLong(0); + private final Semaphore semaphore = new Semaphore(1); + private final AtomicBoolean closed = new AtomicBoolean(false); + private final Set subscriptions = ConcurrentHashMap.newKeySet(); + private final ScheduledExecutorService reconnectExecutor; public KvClient(PDConfig pdConfig) { + this(pdConfig, Executors.newSingleThreadScheduledExecutor(runnable -> { + Thread thread = new Thread(runnable, "pd-kv-watch-reconnect"); + thread.setDaemon(true); + return thread; + })); + } + + KvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { super(pdConfig); + this.reconnectExecutor = reconnectExecutor; } @Override @@ -140,35 +155,26 @@ private void onEvent(WatchResponse value, Consumer consumer) { } } - private StreamObserver getObserver(String key, Consumer consumer, - BiConsumer listenWrapper, - long client) { - StreamObserver observer = getObserver(key, consumer, listenWrapper); - observers.add(observer); - return observer; - } - - private StreamObserver getObserver(String key, Consumer consumer, - BiConsumer listenWrapper) { + private StreamObserver getObserver(WatchSubscription subscription) { return new StreamObserver() { @Override public void onNext(WatchResponse value) { + if (subscription.observer.get() != this) { + return; + } switch (value.getState()) { case Starting: boolean b = clientId.compareAndSet(0, value.getClientId()); if (b) { - // observers.put(value.getClientId(), this); log.info("set watch client id to :{}", value.getClientId()); } release(); break; case Started: - onEvent(value, consumer); + onEvent(value, subscription.consumer); break; case Leader_Changed: - clientId.set(0); - release(); - listenWrapper.accept(key, consumer); + requestReconnect(subscription, this); break; case Alive: // only for check client is alive, do nothing @@ -180,49 +186,122 @@ public void onNext(WatchResponse value) { @Override public void onError(Throwable t) { - release(); - if (!closed.get()) { - clientId.set(0); - listenWrapper.accept(key, consumer); - } + requestReconnect(subscription, this); } @Override public void onCompleted() { - + requestReconnect(subscription, this); } }; } public void listen(String key, Consumer consumer) throws PDException { - long value = clientId.get(); - StreamObserver observer = getObserver(key, consumer, listenWrapper, value); - acquire(); + listen(key, consumer, false); + } + + public void listenPrefix(String prefix, Consumer consumer) throws PDException { + listen(prefix, consumer, true); + } + + private void listen(String key, Consumer consumer, boolean prefix) throws PDException { + WatchSubscription subscription = new WatchSubscription(key, consumer, prefix); + subscriptions.add(subscription); try { - WatchRequest k = - WatchRequest.newBuilder().setClientId(clientId.get()).setKey(key).build(); - streamingCall(KvServiceGrpc.getWatchMethod(), k, observer, 1); - } catch (Exception e) { - release(); - throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, e); + if (!startWatch(subscription)) { + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "KvClient is closed"); + } + } catch (PDException e) { + subscription.observer.set(null); + subscriptions.remove(subscription); + throw e; } } - public void listenPrefix(String prefix, Consumer consumer) throws PDException { - long value = clientId.get(); - StreamObserver observer = - getObserver(prefix, consumer, prefixListenWrapper, value); + private boolean startWatch(WatchSubscription subscription) throws PDException { + if (closed.get()) { + return false; + } + + StreamObserver observer = getObserver(subscription); + subscription.observer.set(observer); + if (closed.get()) { + subscription.observer.compareAndSet(observer, null); + return false; + } + acquire(); + if (closed.get()) { + subscription.observer.compareAndSet(observer, null); + release(); + return false; + } + + WatchRequest request = WatchRequest.newBuilder() + .setClientId(clientId.get()) + .setKey(subscription.key) + .build(); try { - WatchRequest k = - WatchRequest.newBuilder().setClientId(clientId.get()).setKey(prefix).build(); - streamingCall(KvServiceGrpc.getWatchPrefixMethod(), k, observer, 1); + if (subscription.prefix) { + streamingCall(KvServiceGrpc.getWatchPrefixMethod(), request, observer, 1); + } else { + streamingCall(KvServiceGrpc.getWatchMethod(), request, observer, 1); + } + return true; } catch (Exception e) { release(); + if (e instanceof PDException) { + throw (PDException) e; + } throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, e); } } + private void requestReconnect(WatchSubscription subscription, + StreamObserver sourceObserver) { + if (closed.get() || + !subscription.observer.compareAndSet(sourceObserver, null)) { + return; + } + clientId.set(0L); + release(); + scheduleReconnect(subscription); + } + + private void scheduleReconnect(WatchSubscription subscription) { + if (closed.get() || !subscription.reconnectScheduled.compareAndSet(false, true)) { + return; + } + try { + reconnectExecutor.schedule(() -> reconnect(subscription), RECONNECT_DELAY_MS, + TimeUnit.MILLISECONDS); + } catch (RuntimeException e) { + subscription.reconnectScheduled.set(false); + if (!closed.get()) { + log.warn("Failed to schedule watch reconnect for key {}", subscription.key, e); + } + } + } + + private void reconnect(WatchSubscription subscription) { + subscription.reconnectScheduled.set(false); + if (closed.get()) { + return; + } + try { + startWatch(subscription); + } catch (PDException e) { + log.warn("Failed to reconnect watch for key {}", subscription.key, e); + StreamObserver observer = subscription.observer.get(); + if (observer != null) { + requestReconnect(subscription, observer); + } else { + scheduleReconnect(subscription); + } + } + } + private void acquire() { if (clientId.get() == 0L) { try { @@ -340,41 +419,40 @@ public LockResponse keepAlive(String key) throws PDException { @Override public void close() { - for (StreamObserver o : observers) { + if (!closed.compareAndSet(false, true)) { + return; + } + reconnectExecutor.shutdownNow(); + release(); + for (WatchSubscription subscription : subscriptions) { try { - if (o != null) { - o.onCompleted(); + StreamObserver observer = + subscription.observer.getAndSet(null); + if (observer != null) { + observer.onCompleted(); } } catch (Exception e) { - + log.warn("Failed to close watch for key {}", subscription.key, e); } } - observers.clear(); - closed.set(true); + subscriptions.clear(); super.close(); } - BiConsumer listenWrapper = (key, consumer) -> { - try { - listen(key, consumer); - } catch (PDException e) { - try { - log.warn("start listen with warning:", e); - Thread.sleep(1000); - } catch (InterruptedException ex) { - } - } - }; + private final class WatchSubscription { - BiConsumer prefixListenWrapper = (key, consumer) -> { - try { - listenPrefix(key, consumer); - } catch (PDException e) { - try { - log.warn("start listenPrefix with warning:", e); - Thread.sleep(1000); - } catch (InterruptedException ex) { - } + private final String key; + private final Consumer consumer; + private final boolean prefix; + private final AtomicReference> observer; + private final AtomicBoolean reconnectScheduled; + + private WatchSubscription(String key, Consumer consumer, boolean prefix) { + this.key = key; + this.consumer = consumer; + this.prefix = prefix; + this.observer = new AtomicReference<>(); + this.reconnectScheduled = new AtomicBoolean(false); } - }; + } } diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index ae44fb234e..409ad7290a 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -18,23 +18,48 @@ package org.apache.hugegraph.pd.client; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import java.lang.reflect.Field; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +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.AtomicReference; import java.util.function.Consumer; import org.apache.commons.lang3.StringUtils; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.grpc.Metapb; +import org.apache.hugegraph.pd.grpc.PDGrpc; +import org.apache.hugegraph.pd.grpc.Pdpb; import org.apache.hugegraph.pd.grpc.kv.KResponse; +import org.apache.hugegraph.pd.grpc.kv.KvServiceGrpc; import org.apache.hugegraph.pd.grpc.kv.ScanPrefixResponse; import org.apache.hugegraph.pd.grpc.kv.WatchEvent; import org.apache.hugegraph.pd.grpc.kv.WatchKv; +import org.apache.hugegraph.pd.grpc.kv.WatchRequest; import org.apache.hugegraph.pd.grpc.kv.WatchResponse; import org.apache.hugegraph.pd.grpc.kv.WatchState; import org.apache.hugegraph.pd.grpc.kv.WatchType; +import org.junit.After; import org.junit.Before; import org.junit.Test; +import io.grpc.MethodDescriptor; +import io.grpc.Server; +import io.grpc.ServerBuilder; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.AbstractStub; +import io.grpc.stub.StreamObserver; public class KvClientTest extends BaseClientTest { @@ -45,6 +70,11 @@ public void setUp() { client = new KvClient<>(getPdConfig()); } + @After + public void tearDown() { + client.close(); + } + @Test public void testCreateStub() { // Setup @@ -67,6 +97,189 @@ public void testCreateBlockingStub() { } } + @Test + public void testTransportInitializationDoesNotCloseClient() throws Exception { + AtomicReference grpcAddress = new AtomicReference<>(); + Server server = ServerBuilder.forPort(0) + .addService(new PDGrpc.PDImplBase() { + @Override + public void getMembers( + Pdpb.GetMembersRequest request, + StreamObserver observer) { + Metapb.Member leader = + Metapb.Member.newBuilder() + .setGrpcUrl(grpcAddress.get()) + .build(); + observer.onNext( + Pdpb.GetMembersResponse.newBuilder() + .setLeader(leader) + .build()); + observer.onCompleted(); + } + }) + .build() + .start(); + grpcAddress.set("127.0.0.1:" + server.getPort()); + InitializableKvClient testClient = + new InitializableKvClient(PDConfig.of(grpcAddress.get()) + .setAuthority(user, pwd)); + + try { + testClient.initializeTransport(); + + assertThat(isClosed(testClient)).isFalse(); + testClient.close(); + assertThat(isClosed(testClient)).isTrue(); + } finally { + testClient.close(); + server.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + + @Test + public void testReconnectRetriesAfterFirstFailure() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + Consumer consumer = mock(Consumer.class); + context.client.listen("key", consumer); + context.client.failNextCalls(1); + + context.client.call(0).observer.onError(new RuntimeException("disconnected")); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(2); + assertThat(context.reconnectTasks).hasSize(1); + + context.runNextReconnect(); + + assertThat(context.client.calls).hasSize(3); + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.client.call(2).methodName) + .isEqualTo(KvServiceGrpc.getWatchMethod().getFullMethodName()); + assertThat(context.client.call(2).request.getKey()).isEqualTo("key"); + + WatchResponse started = WatchResponse.newBuilder() + .setState(WatchState.Starting) + .setClientId(1L) + .build(); + WatchEvent event = WatchEvent.newBuilder().setType(WatchType.Put).build(); + WatchResponse futureEvent = WatchResponse.newBuilder() + .setState(WatchState.Started) + .setClientId(1L) + .addEvents(event) + .build(); + context.client.call(2).observer.onNext(started); + context.client.call(2).observer.onNext(futureEvent); + verify(consumer).accept(futureEvent); + } + } + + @Test + public void testReconnectRetriesAfterConsecutiveFailures() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + context.client.failNextCalls(2); + + context.client.call(0).observer.onError(new RuntimeException("disconnected")); + context.runNextReconnect(); + context.runNextReconnect(); + context.runNextReconnect(); + + assertThat(context.client.calls).hasSize(4); + assertThat(context.reconnectTasks).isEmpty(); + } + } + + @Test + public void testLeaderChangedSchedulesReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + StreamObserver observer = context.client.call(0).observer; + observer.onNext(WatchResponse.newBuilder() + .setState(WatchState.Starting) + .setClientId(7L) + .build()); + + observer.onNext( + WatchResponse.newBuilder().setState(WatchState.Leader_Changed).build()); + + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.reconnectDelaysMs.get(0)).isGreaterThan(0L); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(2); + assertThat(context.client.call(1).request.getClientId()).isZero(); + } + } + + @Test + public void testCompletedSchedulesReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + + context.client.call(0).observer.onCompleted(); + + assertThat(context.reconnectTasks).hasSize(1); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(2); + } + } + + @Test + public void testObserverSchedulesOnlyOneReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + StreamObserver observer = context.client.call(0).observer; + + observer.onError(new RuntimeException("disconnected")); + observer.onCompleted(); + + assertThat(context.reconnectTasks).hasSize(1); + } + } + + @Test + public void testStaleObserverDoesNotScheduleReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + StreamObserver staleObserver = context.client.call(0).observer; + staleObserver.onError(new RuntimeException("disconnected")); + context.runNextReconnect(); + + staleObserver.onCompleted(); + + assertThat(context.client.calls).hasSize(2); + assertThat(context.reconnectTasks).isEmpty(); + } + } + + @Test + public void testPrefixReconnectPreservesPrefixMethod() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listenPrefix("prefix", response -> { }); + context.client.call(0).observer.onCompleted(); + context.runNextReconnect(); + + assertThat(context.client.calls).hasSize(2); + assertThat(context.client.call(1).methodName) + .isEqualTo(KvServiceGrpc.getWatchPrefixMethod().getFullMethodName()); + assertThat(context.client.call(1).request.getKey()).isEqualTo("prefix"); + } + } + + @Test + public void testCloseStopsScheduledReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + context.client.call(0).observer.onError(new RuntimeException("disconnected")); + Runnable reconnect = context.takeNextReconnect(); + + context.client.close(); + reconnect.run(); + + assertThat(context.client.calls).hasSize(1); + assertThat(context.reconnectTasks).isEmpty(); + verify(context.reconnectExecutor).shutdownNow(); + } + } + String key = "key"; String value = "value"; @@ -114,4 +327,116 @@ public void testPutAndGet() throws Exception { } } + + private static boolean isClosed(KvClient client) throws Exception { + Field field = KvClient.class.getDeclaredField("closed"); + field.setAccessible(true); + return ((AtomicBoolean) field.get(client)).get(); + } + + private WatchTestContext newWatchTestContext() { + ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); + Deque reconnectTasks = new ArrayDeque<>(); + List reconnectDelaysMs = new ArrayList<>(); + doAnswer(invocation -> { + reconnectTasks.addLast(invocation.getArgument(0)); + long delay = invocation.getArgument(1); + TimeUnit unit = invocation.getArgument(2); + reconnectDelaysMs.add(unit.toMillis(delay)); + return mock(ScheduledFuture.class); + }).when(reconnectExecutor).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + TestKvClient testClient = new TestKvClient(getPdConfig(), reconnectExecutor); + return new WatchTestContext(testClient, reconnectExecutor, + reconnectTasks, reconnectDelaysMs); + } + + private static class InitializableKvClient extends KvClient { + + InitializableKvClient(PDConfig pdConfig) { + super(pdConfig); + } + + void initializeTransport() throws PDException { + getStub(); + } + } + + private static class WatchTestContext implements AutoCloseable { + + private final TestKvClient client; + private final ScheduledExecutorService reconnectExecutor; + private final Deque reconnectTasks; + private final List reconnectDelaysMs; + + WatchTestContext(TestKvClient client, + ScheduledExecutorService reconnectExecutor, + Deque reconnectTasks, + List reconnectDelaysMs) { + this.client = client; + this.reconnectExecutor = reconnectExecutor; + this.reconnectTasks = reconnectTasks; + this.reconnectDelaysMs = reconnectDelaysMs; + } + + Runnable takeNextReconnect() { + assertThat(reconnectTasks).isNotEmpty(); + return reconnectTasks.removeFirst(); + } + + void runNextReconnect() { + takeNextReconnect().run(); + } + + @Override + public void close() { + client.close(); + } + } + + private static class TestKvClient extends KvClient { + + private final List calls = new ArrayList<>(); + private int remainingFailures; + + TestKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { + super(pdConfig, reconnectExecutor); + } + + void failNextCalls(int failures) { + this.remainingFailures = failures; + } + + WatchCall call(int index) { + return this.calls.get(index); + } + + @Override + protected void streamingCall(MethodDescriptor method, + ReqT request, + StreamObserver responseObserver, + int retry) throws PDException { + this.calls.add(new WatchCall(method.getFullMethodName(), + (WatchRequest) request, + (StreamObserver) responseObserver)); + if (this.remainingFailures > 0) { + this.remainingFailures--; + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "PD is still unreachable"); + } + } + } + + private static class WatchCall { + + private final String methodName; + private final WatchRequest request; + private final StreamObserver observer; + + WatchCall(String methodName, WatchRequest request, + StreamObserver observer) { + this.methodName = methodName; + this.request = request; + this.observer = observer; + } + } } From 3c2f4dfe7a57aaa85992113a0134a99bb200090d Mon Sep 17 00:00:00 2001 From: contrueCT Date: Sun, 16 Aug 2026 23:47:30 +0800 Subject: [PATCH 2/6] fix(pd): address kv watch review feedback --- .../hugegraph/pd/client/AbstractClient.java | 3 + .../apache/hugegraph/pd/client/KvClient.java | 105 ++++-- .../hugegraph/pd/client/KvClientTest.java | 344 +++++++++++++++++- 3 files changed, 417 insertions(+), 35 deletions(-) diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java index 15b5865feb..fda7e98ca9 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java @@ -261,8 +261,11 @@ protected void streamingCall(MethodDescriptor method, proxy.setStub(null); } streamingCall(method, request, responseObserver, ++retry); + return; } } + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "RPC streaming call failed", e); } } diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index 48a36486f8..5d8b153516 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -51,6 +51,7 @@ import org.apache.hugegraph.pd.grpc.kv.WatchResponse; import org.apache.hugegraph.pd.grpc.kv.WatchType; +import io.grpc.Status; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.AbstractStub; import io.grpc.stub.StreamObserver; @@ -60,9 +61,22 @@ public class KvClient extends AbstractClient implements Closeable { private static final long RECONNECT_DELAY_MS = 1000L; - - private final AtomicLong clientId = new AtomicLong(0); - private final Semaphore semaphore = new Semaphore(1); + private static final Set NON_RETRYABLE_WATCH_ERRORS = + Set.of(Status.Code.CANCELLED, + Status.Code.INVALID_ARGUMENT, + Status.Code.NOT_FOUND, + Status.Code.ALREADY_EXISTS, + Status.Code.PERMISSION_DENIED, + Status.Code.FAILED_PRECONDITION, + Status.Code.OUT_OF_RANGE, + Status.Code.UNIMPLEMENTED, + Status.Code.DATA_LOSS, + Status.Code.UNAUTHENTICATED); + + private final AtomicLong lockClientId = new AtomicLong(0); + private final AtomicLong watchClientId = new AtomicLong(0); + private final Semaphore lockSemaphore = new Semaphore(1); + private final Semaphore watchSemaphore = new Semaphore(1); private final AtomicBoolean closed = new AtomicBoolean(false); private final Set subscriptions = ConcurrentHashMap.newKeySet(); private final ScheduledExecutorService reconnectExecutor; @@ -142,7 +156,7 @@ public TTLResponse putTTL(String key, String value, long ttl) throws PDException private void onEvent(WatchResponse value, Consumer consumer) { log.info("receive message for {},event Count:{}", value, value.getEventsCount()); - clientId.compareAndSet(0L, value.getClientId()); + watchClientId.compareAndSet(0L, value.getClientId()); if (value.getEventsCount() != 0) { try { consumer.accept((T) value); @@ -164,11 +178,11 @@ public void onNext(WatchResponse value) { } switch (value.getState()) { case Starting: - boolean b = clientId.compareAndSet(0, value.getClientId()); + boolean b = watchClientId.compareAndSet(0, value.getClientId()); if (b) { log.info("set watch client id to :{}", value.getClientId()); } - release(); + release(watchSemaphore); break; case Started: onEvent(value, subscription.consumer); @@ -186,7 +200,11 @@ public void onNext(WatchResponse value) { @Override public void onError(Throwable t) { - requestReconnect(subscription, this); + if (isRetryableWatchError(t)) { + requestReconnect(subscription, this); + } else { + stopWatch(subscription, this, t); + } } @Override @@ -231,15 +249,15 @@ private boolean startWatch(WatchSubscription subscription) throws PDException { return false; } - acquire(); + acquire(watchClientId, watchSemaphore); if (closed.get()) { subscription.observer.compareAndSet(observer, null); - release(); + release(watchSemaphore); return false; } WatchRequest request = WatchRequest.newBuilder() - .setClientId(clientId.get()) + .setClientId(watchClientId.get()) .setKey(subscription.key) .build(); try { @@ -250,7 +268,7 @@ private boolean startWatch(WatchSubscription subscription) throws PDException { } return true; } catch (Exception e) { - release(); + release(watchSemaphore); if (e instanceof PDException) { throw (PDException) e; } @@ -264,11 +282,28 @@ private void requestReconnect(WatchSubscription subscription, !subscription.observer.compareAndSet(sourceObserver, null)) { return; } - clientId.set(0L); - release(); + watchClientId.set(0L); + release(watchSemaphore); scheduleReconnect(subscription); } + private static boolean isRetryableWatchError(Throwable throwable) { + Status.Code code = Status.fromThrowable(throwable).getCode(); + return !NON_RETRYABLE_WATCH_ERRORS.contains(code); + } + + private void stopWatch(WatchSubscription subscription, + StreamObserver sourceObserver, + Throwable throwable) { + if (!subscription.observer.compareAndSet(sourceObserver, null)) { + return; + } + release(watchSemaphore); + subscriptions.remove(subscription); + log.error("Watch for key {} stopped after a non-retryable error: {}", + subscription.key, Status.fromThrowable(throwable), throwable); + } + private void scheduleReconnect(WatchSubscription subscription) { if (closed.get() || !subscription.reconnectScheduled.compareAndSet(false, true)) { return; @@ -302,7 +337,7 @@ private void reconnect(WatchSubscription subscription) { } } - private void acquire() { + private void acquire(AtomicLong clientId, Semaphore semaphore) { if (clientId.get() == 0L) { try { semaphore.acquire(); @@ -316,7 +351,7 @@ private void acquire() { } } - private void release() { + private void release(Semaphore semaphore) { try { if (semaphore.availablePermits() == 0) { semaphore.release(); @@ -355,65 +390,68 @@ public Map getWatchMap(T response) { } public LockResponse lock(String key, long ttl) throws PDException { - acquire(); + acquire(lockClientId, lockSemaphore); LockResponse response; try { LockRequest k = - LockRequest.newBuilder().setKey(key).setClientId(clientId.get()).setTtl(ttl) + LockRequest.newBuilder().setKey(key).setClientId(lockClientId.get()).setTtl(ttl) .build(); response = blockingUnaryCall(KvServiceGrpc.getLockMethod(), k); handleErrors(response.getHeader()); - clientId.compareAndSet(0, response.getClientId()); + lockClientId.compareAndSet(0, response.getClientId()); } catch (Exception e) { throw e; } finally { - release(); + release(lockSemaphore); } return response; } public LockResponse lockWithoutReentrant(String key, long ttl) throws PDException { - acquire(); + acquire(lockClientId, lockSemaphore); LockResponse response; try { LockRequest k = - LockRequest.newBuilder().setKey(key).setClientId(clientId.get()).setTtl(ttl) + LockRequest.newBuilder().setKey(key).setClientId(lockClientId.get()).setTtl(ttl) .build(); response = blockingUnaryCall(KvServiceGrpc.getLockWithoutReentrantMethod(), k); handleErrors(response.getHeader()); - clientId.compareAndSet(0, response.getClientId()); + lockClientId.compareAndSet(0, response.getClientId()); } catch (Exception e) { throw e; } finally { - release(); + release(lockSemaphore); } return response; } public LockResponse isLocked(String key) throws PDException { - LockRequest k = LockRequest.newBuilder().setKey(key).setClientId(clientId.get()).build(); + LockRequest k = + LockRequest.newBuilder().setKey(key).setClientId(lockClientId.get()).build(); LockResponse response = blockingUnaryCall(KvServiceGrpc.getIsLockedMethod(), k); handleErrors(response.getHeader()); return response; } public LockResponse unlock(String key) throws PDException { - assert clientId.get() != 0; - LockRequest k = LockRequest.newBuilder().setKey(key).setClientId(clientId.get()).build(); + assert lockClientId.get() != 0; + LockRequest k = + LockRequest.newBuilder().setKey(key).setClientId(lockClientId.get()).build(); LockResponse response = blockingUnaryCall(KvServiceGrpc.getUnlockMethod(), k); handleErrors(response.getHeader()); - clientId.compareAndSet(0L, response.getClientId()); - assert clientId.get() == response.getClientId(); + lockClientId.compareAndSet(0L, response.getClientId()); + assert lockClientId.get() == response.getClientId(); return response; } public LockResponse keepAlive(String key) throws PDException { - assert clientId.get() != 0; - LockRequest k = LockRequest.newBuilder().setKey(key).setClientId(clientId.get()).build(); + assert lockClientId.get() != 0; + LockRequest k = + LockRequest.newBuilder().setKey(key).setClientId(lockClientId.get()).build(); LockResponse response = blockingUnaryCall(KvServiceGrpc.getKeepAliveMethod(), k); handleErrors(response.getHeader()); - clientId.compareAndSet(0L, response.getClientId()); - assert clientId.get() == response.getClientId(); + lockClientId.compareAndSet(0L, response.getClientId()); + assert lockClientId.get() == response.getClientId(); return response; } @@ -423,7 +461,8 @@ public void close() { return; } reconnectExecutor.shutdownNow(); - release(); + release(lockSemaphore); + release(watchSemaphore); for (WatchSubscription subscription : subscriptions) { try { StreamObserver observer = diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index 409ad7290a..ad14467aca 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -18,10 +18,12 @@ package org.apache.hugegraph.pd.client; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import java.lang.reflect.Field; @@ -43,6 +45,8 @@ import org.apache.hugegraph.pd.grpc.Pdpb; import org.apache.hugegraph.pd.grpc.kv.KResponse; import org.apache.hugegraph.pd.grpc.kv.KvServiceGrpc; +import org.apache.hugegraph.pd.grpc.kv.LockRequest; +import org.apache.hugegraph.pd.grpc.kv.LockResponse; import org.apache.hugegraph.pd.grpc.kv.ScanPrefixResponse; import org.apache.hugegraph.pd.grpc.kv.WatchEvent; import org.apache.hugegraph.pd.grpc.kv.WatchKv; @@ -54,9 +58,15 @@ import org.junit.Before; import org.junit.Test; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Server; import io.grpc.ServerBuilder; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.AbstractStub; import io.grpc.stub.StreamObserver; @@ -136,6 +146,44 @@ public void getMembers( } } + @Test(timeout = 2000L) + public void testStreamingFailurePropagatesToListenCaller() { + ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); + StatusRuntimeException failure = + Status.UNAVAILABLE.withDescription("stream failed").asRuntimeException(); + PDConfig config = PDConfig.of("first:8686,second:8686").setAuthority(user, pwd); + try (SequentialStreamingKvClient testClient = + new SequentialStreamingKvClient(config, reconnectExecutor, + new FailingChannel(failure), + new FailingChannel(failure), + new FailingChannel(failure), + new FailingChannel(failure))) { + assertThatThrownBy(() -> testClient.listen("key", response -> { })) + .isInstanceOf(PDException.class) + .hasCause(failure); + assertThatThrownBy(() -> testClient.listen("key", response -> { })) + .isInstanceOf(PDException.class) + .hasCause(failure); + } + } + + @Test(timeout = 2000L) + public void testStreamingFailureRetriesNextPeer() throws Exception { + ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); + StatusRuntimeException failure = + Status.UNAVAILABLE.withDescription("first peer failed").asRuntimeException(); + AtomicBoolean secondPeerCalled = new AtomicBoolean(false); + PDConfig config = PDConfig.of("first:8686,second:8686").setAuthority(user, pwd); + try (SequentialStreamingKvClient testClient = + new SequentialStreamingKvClient(config, reconnectExecutor, + new FailingChannel(failure), + new SuccessfulChannel(secondPeerCalled))) { + testClient.listen("key", response -> { }); + + assertThat(secondPeerCalled).isTrue(); + } + } + @Test public void testReconnectRetriesAfterFirstFailure() throws Exception { try (WatchTestContext context = newWatchTestContext()) { @@ -152,6 +200,7 @@ public void testReconnectRetriesAfterFirstFailure() throws Exception { assertThat(context.client.calls).hasSize(3); assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L); assertThat(context.client.call(2).methodName) .isEqualTo(KvServiceGrpc.getWatchMethod().getFullMethodName()); assertThat(context.client.call(2).request.getKey()).isEqualTo("key"); @@ -185,6 +234,73 @@ public void testReconnectRetriesAfterConsecutiveFailures() throws Exception { assertThat(context.client.calls).hasSize(4); assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L, 1000L); + } + } + + @Test + public void testWatchReconnectPreservesLockClientId() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.respondWithLockClientId(42L); + context.client.lock("lock", 1000L); + context.client.listen("key", response -> { }); + + assertThat(context.client.call(0).request.getClientId()).isZero(); + StreamObserver observer = context.client.call(0).observer; + observer.onNext(WatchResponse.newBuilder() + .setState(WatchState.Starting) + .setClientId(7L) + .build()); + observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.runNextReconnect(); + + context.client.keepAlive("lock"); + assertThat(context.client.lockCall(0).methodName) + .isEqualTo(KvServiceGrpc.getLockMethod().getFullMethodName()); + assertThat(context.client.lockCall(0).request.getClientId()).isZero(); + assertThat(context.client.lockCall(1).methodName) + .isEqualTo(KvServiceGrpc.getKeepAliveMethod().getFullMethodName()); + assertThat(context.client.lockCall(1).request.getClientId()).isEqualTo(42L); + } + } + + @Test + public void testPermanentErrorStopsReconnect() throws Exception { + for (Status status : List.of(Status.CANCELLED, + Status.INVALID_ARGUMENT, + Status.NOT_FOUND, + Status.ALREADY_EXISTS, + Status.PERMISSION_DENIED, + Status.FAILED_PRECONDITION, + Status.OUT_OF_RANGE, + Status.UNIMPLEMENTED, + Status.DATA_LOSS, + Status.UNAUTHENTICATED)) { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + StreamObserver observer = context.client.call(0).observer; + + observer.onError(status.withDescription("permanent watch error") + .asRuntimeException()); + observer.onCompleted(); + + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.client.calls).hasSize(1); + } + } + } + + @Test + public void testRetryableErrorSchedulesReconnect() throws Exception { + for (Status status : List.of(Status.UNAVAILABLE, Status.UNKNOWN)) { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + + context.client.call(0).observer.onError(status.asRuntimeException()); + + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.reconnectDelaysMs).containsExactly(1000L); + } } } @@ -202,7 +318,7 @@ public void testLeaderChangedSchedulesReconnect() throws Exception { WatchResponse.newBuilder().setState(WatchState.Leader_Changed).build()); assertThat(context.reconnectTasks).hasSize(1); - assertThat(context.reconnectDelaysMs.get(0)).isGreaterThan(0L); + assertThat(context.reconnectDelaysMs).containsExactly(1000L); context.runNextReconnect(); assertThat(context.client.calls).hasSize(2); assertThat(context.client.call(1).request.getClientId()).isZero(); @@ -217,6 +333,7 @@ public void testCompletedSchedulesReconnect() throws Exception { context.client.call(0).observer.onCompleted(); assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.reconnectDelaysMs).containsExactly(1000L); context.runNextReconnect(); assertThat(context.client.calls).hasSize(2); } @@ -232,21 +349,77 @@ public void testObserverSchedulesOnlyOneReconnect() throws Exception { observer.onCompleted(); assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.reconnectDelaysMs).containsExactly(1000L); + } + } + + @Test + public void testSubscriptionsReconnectIndependently() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + Consumer firstConsumer = mock(Consumer.class); + Consumer secondConsumer = mock(Consumer.class); + context.client.listen("first", firstConsumer); + context.client.call(0).observer.onNext(startingResponse(11L)); + context.client.listenPrefix("second", secondConsumer); + + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.client.call(1).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + + assertThat(context.reconnectTasks).hasSize(2); + assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L); + + context.runNextReconnect(); + assertThat(context.client.call(2).methodName) + .isEqualTo(KvServiceGrpc.getWatchMethod().getFullMethodName()); + assertThat(context.client.call(2).request.getKey()).isEqualTo("first"); + context.client.call(2).observer.onNext(startingResponse(12L)); + WatchResponse firstEvent = eventResponse(12L); + context.client.call(2).observer.onNext(firstEvent); + + context.runNextReconnect(); + assertThat(context.client.call(3).methodName) + .isEqualTo(KvServiceGrpc.getWatchPrefixMethod().getFullMethodName()); + assertThat(context.client.call(3).request.getKey()).isEqualTo("second"); + context.client.call(3).observer.onNext(startingResponse(12L)); + WatchResponse secondEvent = eventResponse(13L); + context.client.call(3).observer.onNext(secondEvent); + + verify(firstConsumer).accept(firstEvent); + verify(secondConsumer).accept(secondEvent); + verify(firstConsumer, never()).accept(secondEvent); + verify(secondConsumer, never()).accept(firstEvent); } } @Test public void testStaleObserverDoesNotScheduleReconnect() throws Exception { try (WatchTestContext context = newWatchTestContext()) { - context.client.listen("key", response -> { }); + Consumer consumer = mock(Consumer.class); + context.client.listen("key", consumer); StreamObserver staleObserver = context.client.call(0).observer; staleObserver.onError(new RuntimeException("disconnected")); context.runNextReconnect(); staleObserver.onCompleted(); + WatchEvent event = WatchEvent.newBuilder().setType(WatchType.Put).build(); + WatchResponse response = WatchResponse.newBuilder() + .setState(WatchState.Started) + .setClientId(8L) + .addEvents(event) + .build(); + staleObserver.onNext(response); assertThat(context.client.calls).hasSize(2); assertThat(context.reconnectTasks).isEmpty(); + verify(consumer, never()).accept(any()); + + StreamObserver currentObserver = context.client.call(1).observer; + currentObserver.onNext(WatchResponse.newBuilder() + .setState(WatchState.Starting) + .setClientId(8L) + .build()); + currentObserver.onNext(response); + verify(consumer).accept(response); } } @@ -334,6 +507,21 @@ private static boolean isClosed(KvClient client) throws Exception { return ((AtomicBoolean) field.get(client)).get(); } + private static WatchResponse startingResponse(long clientId) { + return WatchResponse.newBuilder() + .setState(WatchState.Starting) + .setClientId(clientId) + .build(); + } + + private static WatchResponse eventResponse(long clientId) { + return WatchResponse.newBuilder() + .setState(WatchState.Started) + .setClientId(clientId) + .addEvents(WatchEvent.newBuilder().setType(WatchType.Put).build()) + .build(); + } + private WatchTestContext newWatchTestContext() { ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); Deque reconnectTasks = new ArrayDeque<>(); @@ -361,6 +549,125 @@ void initializeTransport() throws PDException { } } + private static class SequentialStreamingKvClient extends KvClient { + + private final Deque channels = new ArrayDeque<>(); + + SequentialStreamingKvClient(PDConfig pdConfig, + ScheduledExecutorService reconnectExecutor, + Channel... channels) { + super(pdConfig, reconnectExecutor); + for (Channel channel : channels) { + this.channels.addLast(channel); + } + } + + @Override + protected AbstractStub getStub() { + return KvServiceGrpc.newStub(this.channels.removeFirst()); + } + } + + private static class FailingChannel extends Channel { + + private final StatusRuntimeException failure; + + FailingChannel(StatusRuntimeException failure) { + this.failure = failure; + } + + @Override + public String authority() { + return "test"; + } + + @Override + public ClientCall newCall( + MethodDescriptor method, CallOptions callOptions) { + return new FailingClientCall<>(this.failure); + } + } + + private static class FailingClientCall extends ClientCall { + + private final StatusRuntimeException failure; + + FailingClientCall(StatusRuntimeException failure) { + this.failure = failure; + } + + @Override + public void start(Listener responseListener, Metadata headers) { + throw this.failure; + } + + @Override + public void request(int numMessages) { + } + + @Override + public void cancel(String message, Throwable cause) { + } + + @Override + public void halfClose() { + } + + @Override + public void sendMessage(ReqT message) { + } + } + + private static class SuccessfulChannel extends Channel { + + private final AtomicBoolean called; + + SuccessfulChannel(AtomicBoolean called) { + this.called = called; + } + + @Override + public String authority() { + return "test"; + } + + @Override + public ClientCall newCall( + MethodDescriptor method, CallOptions callOptions) { + return new SuccessfulClientCall<>(this.called); + } + } + + private static class SuccessfulClientCall extends ClientCall { + + private final AtomicBoolean called; + + SuccessfulClientCall(AtomicBoolean called) { + this.called = called; + } + + @Override + public void start(Listener responseListener, Metadata headers) { + this.called.set(true); + } + + @Override + public void request(int numMessages) { + } + + @Override + public void cancel(String message, Throwable cause) { + } + + @Override + public void halfClose() { + } + + @Override + public void sendMessage(ReqT message) { + } + } + private static class WatchTestContext implements AutoCloseable { private final TestKvClient client; @@ -396,7 +703,9 @@ public void close() { private static class TestKvClient extends KvClient { private final List calls = new ArrayList<>(); + private final List lockCalls = new ArrayList<>(); private int remainingFailures; + private long lockClientId; TestKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { super(pdConfig, reconnectExecutor); @@ -410,6 +719,26 @@ WatchCall call(int index) { return this.calls.get(index); } + void respondWithLockClientId(long clientId) { + this.lockClientId = clientId; + } + + LockCall lockCall(int index) { + return this.lockCalls.get(index); + } + + @Override + protected RespT blockingUnaryCall( + MethodDescriptor method, ReqT request) { + LockRequest lockRequest = (LockRequest) request; + this.lockCalls.add(new LockCall(method.getFullMethodName(), lockRequest)); + return (RespT) LockResponse.newBuilder() + .setHeader(AbstractClient.okHeader) + .setClientId(this.lockClientId) + .setSucceed(true) + .build(); + } + @Override protected void streamingCall(MethodDescriptor method, ReqT request, @@ -439,4 +768,15 @@ private static class WatchCall { this.observer = observer; } } + + private static class LockCall { + + private final String methodName; + private final LockRequest request; + + LockCall(String methodName, LockRequest request) { + this.methodName = methodName; + this.request = request; + } + } } From 4f22b40c15eca329765d24acfb820f9b98b054ec Mon Sep 17 00:00:00 2001 From: contrueCT Date: Mon, 31 Aug 2026 01:32:49 +0800 Subject: [PATCH 3/6] fix(pd): keep watch reconnect executor responsive --- .../apache/hugegraph/pd/client/KvClient.java | 30 ++++++++++++++- .../hugegraph/pd/client/KvClientTest.java | 37 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index 2221cff60f..a005b25573 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -238,6 +238,11 @@ private void listen(String key, Consumer consumer, boolean prefix) throws PDE } private boolean startWatch(WatchSubscription subscription) throws PDException { + return startWatch(subscription, true); + } + + private boolean startWatch(WatchSubscription subscription, + boolean waitForPermit) throws PDException { if (closed.get()) { return false; } @@ -249,7 +254,12 @@ private boolean startWatch(WatchSubscription subscription) throws PDException { return false; } - acquire(watchClientId, watchSemaphore); + if (waitForPermit) { + acquire(watchClientId, watchSemaphore); + } else if (!tryAcquire(watchClientId, watchSemaphore)) { + subscription.observer.compareAndSet(observer, null); + return false; + } if (closed.get()) { subscription.observer.compareAndSet(observer, null); release(watchSemaphore); @@ -325,7 +335,9 @@ private void reconnect(WatchSubscription subscription) { return; } try { - startWatch(subscription); + if (!startWatch(subscription, false)) { + scheduleReconnect(subscription); + } } catch (PDException e) { log.warn("Failed to reconnect watch for key {}", subscription.key, e); StreamObserver observer = subscription.observer.get(); @@ -337,6 +349,20 @@ private void reconnect(WatchSubscription subscription) { } } + private boolean tryAcquire(AtomicLong clientId, Semaphore semaphore) { + if (clientId.get() != 0L) { + return true; + } + if (!semaphore.tryAcquire()) { + return false; + } + if (clientId.get() != 0L) { + semaphore.release(); + } + log.info("wait for client starting...."); + return true; + } + private void acquire(AtomicLong clientId, Semaphore semaphore) { if (clientId.get() == 0L) { try { diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index ad14467aca..5d4a62fcef 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -31,6 +31,9 @@ import java.util.ArrayList; import java.util.Deque; import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; @@ -391,6 +394,40 @@ public void testSubscriptionsReconnectIndependently() throws Exception { } } + @Test(timeout = 5000L) + public void testBlockedReconnectReschedulesOtherSubscription() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("first", response -> { }); + context.client.call(0).observer.onNext(startingResponse(11L)); + context.client.listenPrefix("second", response -> { }); + + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.client.call(1).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.runNextReconnect(); + + assertThat(context.client.calls).hasSize(3); + ExecutorService reconnectRunner = Executors.newSingleThreadExecutor(); + try { + Future secondReconnect = reconnectRunner.submit(context::runNextReconnect); + secondReconnect.get(1L, TimeUnit.SECONDS); + } finally { + reconnectRunner.shutdownNow(); + assertThat(reconnectRunner.awaitTermination(1L, TimeUnit.SECONDS)).isTrue(); + } + + assertThat(context.client.calls).hasSize(3); + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L, 1000L); + + context.client.call(2).observer.onNext(startingResponse(12L)); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(4); + assertThat(context.client.call(3).methodName) + .isEqualTo(KvServiceGrpc.getWatchPrefixMethod().getFullMethodName()); + assertThat(context.client.call(3).request.getKey()).isEqualTo("second"); + } + } + @Test public void testStaleObserverDoesNotScheduleReconnect() throws Exception { try (WatchTestContext context = newWatchTestContext()) { From 4202f266c217bc1d7c99027bda1a3539326567e9 Mon Sep 17 00:00:00 2001 From: contrueCT Date: Mon, 31 Aug 2026 14:19:14 +0800 Subject: [PATCH 4/6] fix(pd): harden kv watch reconnect lifecycle --- .../hugegraph/pd/client/AbstractClient.java | 20 +- .../apache/hugegraph/pd/client/KvClient.java | 309 ++++++--- .../hugegraph/pd/client/KvClientTest.java | 606 +++++++++++++++--- .../org/apache/hugegraph/SchemaDriver.java | 19 +- .../apache/hugegraph/SchemaDriverTest.java | 76 +++ 5 files changed, 854 insertions(+), 176 deletions(-) create mode 100644 hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java index fda7e98ca9..24d26ac3df 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java @@ -26,7 +26,6 @@ import java.util.function.Predicate; import java.util.stream.Stream; -import org.apache.commons.lang3.StringUtils; import org.apache.hugegraph.pd.client.interceptor.Authentication; import org.apache.hugegraph.pd.common.KVPair; import org.apache.hugegraph.pd.common.PDException; @@ -128,12 +127,15 @@ protected AbstractStub getStub() throws PDException { return setAsyncParams(proxy.getStub(), config); } + protected synchronized void invalidateAsyncStub() { + proxy.setStub(null); + } + protected abstract AbstractStub createStub(); protected abstract AbstractBlockingStub createBlockingStub(); private String resetStub() { - String leaderHost = ""; Exception ex = null; for (int i = 0; i < proxy.getHostCount(); i++) { String host = proxy.nextHost(); @@ -147,7 +149,7 @@ private String resetStub() { .setHeader(header).build(); GetMembersResponse members = blockingStub.getMembers(request); Metapb.Member leader = members.getLeader(); - leaderHost = leader.getGrpcUrl(); + String leaderHost = leader.getGrpcUrl(); if (!host.equals(leaderHost)) { closeConnections(); channel = ManagedChannelBuilder.forTarget(leaderHost).usePlaintext().build(); @@ -155,10 +157,9 @@ private String resetStub() { proxy.setBlockingStub(setBlockingParams(createBlockingStub(), config)); proxy.setStub(setAsyncParams(createStub(), config)); log.info("AbstractClient connect to host = {} success", leaderHost); - break; + return leaderHost; } catch (StatusRuntimeException se) { ex = se; - continue; } catch (Exception e) { ex = e; String msg = @@ -166,11 +167,14 @@ private String resetStub() { e.getMessage()); log.error(msg, e); } + proxy.setBlockingStub(null); + proxy.setStub(null); } - if (StringUtils.isEmpty(leaderHost) && ex != null) { + closeConnections(); + if (ex != null) { log.error(String.format("connect to %s with error: ", config.getServerHost()), ex); } - return leaderHost; + return ""; } protected RespT blockingUnaryCall( @@ -258,7 +262,7 @@ protected void streamingCall(MethodDescriptor method, if (e instanceof StatusRuntimeException) { if (retry < proxy.getHostCount()) { synchronized (this) { - proxy.setStub(null); + invalidateAsyncStub(); } streamingCall(method, request, responseObserver, ++retry); return; diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index a005b25573..023d3684e3 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -22,10 +22,12 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -61,9 +63,9 @@ public class KvClient extends AbstractClient implements Closeable { private static final long RECONNECT_DELAY_MS = 1000L; + private static final long WATCH_START_TIMEOUT_MS = 5000L; private static final Set NON_RETRYABLE_WATCH_ERRORS = - Set.of(Status.Code.CANCELLED, - Status.Code.INVALID_ARGUMENT, + Set.of(Status.Code.INVALID_ARGUMENT, Status.Code.NOT_FOUND, Status.Code.ALREADY_EXISTS, Status.Code.PERMISSION_DENIED, @@ -74,12 +76,11 @@ public class KvClient extends AbstractClient implements Status.Code.UNAUTHENTICATED); private final AtomicLong lockClientId = new AtomicLong(0); - private final AtomicLong watchClientId = new AtomicLong(0); private final Semaphore lockSemaphore = new Semaphore(1); - private final Semaphore watchSemaphore = new Semaphore(1); private final AtomicBoolean closed = new AtomicBoolean(false); private final Set subscriptions = ConcurrentHashMap.newKeySet(); private final ScheduledExecutorService reconnectExecutor; + private long transportGeneration; public KvClient(PDConfig pdConfig) { this(pdConfig, Executors.newSingleThreadScheduledExecutor(runnable -> { @@ -156,7 +157,6 @@ public TTLResponse putTTL(String key, String value, long ttl) throws PDException private void onEvent(WatchResponse value, Consumer consumer) { log.debug("receive message for {},event Count:{}", value, value.getEventsCount()); - watchClientId.compareAndSet(0L, value.getClientId()); if (value.getEventsCount() != 0) { try { consumer.accept((T) value); @@ -173,22 +173,21 @@ private StreamObserver getObserver(WatchSubscription subscription return new StreamObserver() { @Override public void onNext(WatchResponse value) { - if (subscription.observer.get() != this) { + if (!acceptFirstFrame(subscription, this)) { return; } switch (value.getState()) { case Starting: - boolean b = watchClientId.compareAndSet(0, value.getClientId()); + boolean b = subscription.clientId.compareAndSet(0, value.getClientId()); if (b) { log.info("set watch client id to :{}", value.getClientId()); } - release(watchSemaphore); break; case Started: onEvent(value, subscription.consumer); break; case Leader_Changed: - requestReconnect(subscription, this); + requestReconnect(subscription, this, true); break; case Alive: // only for check client is alive, do nothing @@ -201,7 +200,7 @@ public void onNext(WatchResponse value) { @Override public void onError(Throwable t) { if (isRetryableWatchError(t)) { - requestReconnect(subscription, this); + requestReconnect(subscription, this, shouldRotateWatchTransport(t)); } else { stopWatch(subscription, this, t); } @@ -209,21 +208,37 @@ public void onError(Throwable t) { @Override public void onCompleted() { - requestReconnect(subscription, this); + requestReconnect(subscription, this, false); } }; } public void listen(String key, Consumer consumer) throws PDException { - listen(key, consumer, false); + listen(key, consumer, throwable -> { }, false); + } + + public void listen(String key, Consumer consumer, + Consumer errorConsumer) throws PDException { + listen(key, consumer, errorConsumer, false); } public void listenPrefix(String prefix, Consumer consumer) throws PDException { - listen(prefix, consumer, true); + listen(prefix, consumer, throwable -> { }, true); + } + + public void listenPrefix(String prefix, Consumer consumer, + Consumer errorConsumer) throws PDException { + listen(prefix, consumer, errorConsumer, true); } - private void listen(String key, Consumer consumer, boolean prefix) throws PDException { - WatchSubscription subscription = new WatchSubscription(key, consumer, prefix); + private void listen(String key, Consumer consumer, + Consumer errorConsumer, + boolean prefix) throws PDException { + Objects.requireNonNull(key, "key"); + Objects.requireNonNull(consumer, "consumer"); + Objects.requireNonNull(errorConsumer, "errorConsumer"); + WatchSubscription subscription = + new WatchSubscription(key, consumer, errorConsumer, prefix); subscriptions.add(subscription); try { if (!startWatch(subscription)) { @@ -231,54 +246,55 @@ private void listen(String key, Consumer consumer, boolean prefix) throws PDE "KvClient is closed"); } } catch (PDException e) { - subscription.observer.set(null); + cleanupFailedStart(subscription); + subscriptions.remove(subscription); + throw e; + } catch (RuntimeException e) { + cleanupFailedStart(subscription); subscriptions.remove(subscription); throw e; } } - private boolean startWatch(WatchSubscription subscription) throws PDException { - return startWatch(subscription, true); - } - - private boolean startWatch(WatchSubscription subscription, - boolean waitForPermit) throws PDException { - if (closed.get()) { - return false; - } - - StreamObserver observer = getObserver(subscription); - subscription.observer.set(observer); - if (closed.get()) { - subscription.observer.compareAndSet(observer, null); - return false; + private void cleanupFailedStart(WatchSubscription subscription) { + synchronized (subscription) { + subscription.observer.set(null); + cancelStartTimeout(subscription); } + } - if (waitForPermit) { - acquire(watchClientId, watchSemaphore); - } else if (!tryAcquire(watchClientId, watchSemaphore)) { - subscription.observer.compareAndSet(observer, null); - return false; - } - if (closed.get()) { - subscription.observer.compareAndSet(observer, null); - release(watchSemaphore); + private boolean startWatch(WatchSubscription subscription) throws PDException { + if (closed.get() || !subscriptions.contains(subscription)) { return false; } WatchRequest request = WatchRequest.newBuilder() - .setClientId(watchClientId.get()) + .setClientId(subscription.clientId.get()) .setKey(subscription.key) .build(); + StreamObserver observer = getObserver(subscription); try { - if (subscription.prefix) { - streamingCall(KvServiceGrpc.getWatchPrefixMethod(), request, observer, 1); - } else { - streamingCall(KvServiceGrpc.getWatchMethod(), request, observer, 1); + synchronized (this) { + synchronized (subscription) { + if (closed.get() || !subscriptions.contains(subscription)) { + return false; + } + if (subscription.observer.get() != null) { + return true; + } + subscription.firstFrameReceived = false; + subscription.attemptGeneration = this.transportGeneration; + subscription.observer.set(observer); + } + if (subscription.prefix) { + streamingCall(KvServiceGrpc.getWatchPrefixMethod(), request, observer, 1); + } else { + streamingCall(KvServiceGrpc.getWatchMethod(), request, observer, 1); + } } + scheduleStartTimeout(subscription, observer); return true; } catch (Exception e) { - release(watchSemaphore); if (e instanceof PDException) { throw (PDException) e; } @@ -286,36 +302,137 @@ private boolean startWatch(WatchSubscription subscription, } } + private boolean acceptFirstFrame(WatchSubscription subscription, + StreamObserver sourceObserver) { + synchronized (subscription) { + if (subscription.observer.get() != sourceObserver) { + return false; + } + subscription.firstFrameReceived = true; + cancelStartTimeout(subscription); + return true; + } + } + + private void scheduleStartTimeout(WatchSubscription subscription, + StreamObserver sourceObserver) + throws PDException { + ScheduledFuture timeout; + try { + timeout = reconnectExecutor.schedule( + () -> onStartTimeout(subscription, sourceObserver), + WATCH_START_TIMEOUT_MS, TimeUnit.MILLISECONDS); + } catch (RuntimeException e) { + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "Failed to schedule watch start timeout", e); + } + synchronized (subscription) { + if (closed.get() || subscription.observer.get() != sourceObserver || + subscription.firstFrameReceived) { + timeout.cancel(false); + return; + } + ScheduledFuture previous = subscription.startTimeout.getAndSet(timeout); + if (previous != null) { + previous.cancel(false); + } + } + } + + private void onStartTimeout(WatchSubscription subscription, + StreamObserver sourceObserver) { + synchronized (subscription) { + if (subscription.observer.get() != sourceObserver || + subscription.firstFrameReceived) { + return; + } + } + log.warn("Watch for key {} did not receive its first response in {} ms", + subscription.key, WATCH_START_TIMEOUT_MS); + requestReconnect(subscription, sourceObserver, true); + } + private void requestReconnect(WatchSubscription subscription, - StreamObserver sourceObserver) { - if (closed.get() || - !subscription.observer.compareAndSet(sourceObserver, null)) { - return; + StreamObserver sourceObserver, + boolean rotateTransport) { + long attemptGeneration; + synchronized (subscription) { + if (closed.get() || subscription.observer.get() != sourceObserver) { + return; + } + subscription.observer.set(null); + cancelStartTimeout(subscription); + subscription.clientId.set(0L); + attemptGeneration = subscription.attemptGeneration; + } + if (rotateTransport) { + invalidateAttemptStub(attemptGeneration); } - watchClientId.set(0L); - release(watchSemaphore); scheduleReconnect(subscription); } + private synchronized void invalidateAttemptStub(long attemptGeneration) { + if (attemptGeneration != this.transportGeneration) { + return; + } + this.transportGeneration++; + invalidateAsyncStub(); + } + + private void cancelStartTimeout(WatchSubscription subscription) { + ScheduledFuture timeout = subscription.startTimeout.getAndSet(null); + if (timeout != null) { + timeout.cancel(false); + } + } + private static boolean isRetryableWatchError(Throwable throwable) { Status.Code code = Status.fromThrowable(throwable).getCode(); return !NON_RETRYABLE_WATCH_ERRORS.contains(code); } + private static boolean shouldRotateWatchTransport(Throwable throwable) { + return Status.fromThrowable(throwable).getCode() == Status.Code.UNAVAILABLE; + } + private void stopWatch(WatchSubscription subscription, StreamObserver sourceObserver, Throwable throwable) { - if (!subscription.observer.compareAndSet(sourceObserver, null)) { - return; + synchronized (subscription) { + if (subscription.observer.get() != sourceObserver) { + return; + } + subscription.observer.set(null); + cancelStartTimeout(subscription); } - release(watchSemaphore); subscriptions.remove(subscription); log.error("Watch for key {} stopped after a non-retryable error: {}", subscription.key, Status.fromThrowable(throwable), throwable); + notifyWatchStopped(subscription, unwrapWatchError(throwable)); + } + + private static Throwable unwrapWatchError(Throwable throwable) { + Throwable result = throwable; + while (result instanceof PDException && result.getCause() != null) { + result = result.getCause(); + } + return result; + } + + private void notifyWatchStopped(WatchSubscription subscription, Throwable throwable) { + if (!subscription.terminated.compareAndSet(false, true)) { + return; + } + try { + subscription.errorConsumer.accept(throwable); + } catch (RuntimeException e) { + log.warn("Failed to report stopped watch for key {}", subscription.key, e); + } } private void scheduleReconnect(WatchSubscription subscription) { - if (closed.get() || !subscription.reconnectScheduled.compareAndSet(false, true)) { + if (closed.get() || !subscriptions.contains(subscription) || + !subscription.reconnectScheduled.compareAndSet(false, true)) { return; } try { @@ -325,44 +442,42 @@ private void scheduleReconnect(WatchSubscription subscription) { subscription.reconnectScheduled.set(false); if (!closed.get()) { log.warn("Failed to schedule watch reconnect for key {}", subscription.key, e); + subscriptions.remove(subscription); + notifyWatchStopped(subscription, e); } } } private void reconnect(WatchSubscription subscription) { subscription.reconnectScheduled.set(false); - if (closed.get()) { + if (closed.get() || !subscriptions.contains(subscription)) { return; } try { - if (!startWatch(subscription, false)) { + if (!startWatch(subscription)) { scheduleReconnect(subscription); } - } catch (PDException e) { - log.warn("Failed to reconnect watch for key {}", subscription.key, e); + } catch (RuntimeException | PDException e) { StreamObserver observer = subscription.observer.get(); - if (observer != null) { - requestReconnect(subscription, observer); + if (!isRetryableWatchError(e)) { + log.error("Watch reconnect for key {} failed permanently", + subscription.key, e); + if (observer != null) { + stopWatch(subscription, observer, e); + } else { + subscriptions.remove(subscription); + notifyWatchStopped(subscription, unwrapWatchError(e)); + } + } else if (observer != null) { + log.warn("Failed to reconnect watch for key {}", subscription.key, e); + requestReconnect(subscription, observer, shouldRotateWatchTransport(e)); } else { + log.warn("Failed to reconnect watch for key {}", subscription.key, e); scheduleReconnect(subscription); } } } - private boolean tryAcquire(AtomicLong clientId, Semaphore semaphore) { - if (clientId.get() != 0L) { - return true; - } - if (!semaphore.tryAcquire()) { - return false; - } - if (clientId.get() != 0L) { - semaphore.release(); - } - log.info("wait for client starting...."); - return true; - } - private void acquire(AtomicLong clientId, Semaphore semaphore) { if (clientId.get() == 0L) { try { @@ -488,36 +603,52 @@ public void close() { } reconnectExecutor.shutdownNow(); release(lockSemaphore); - release(watchSemaphore); - for (WatchSubscription subscription : subscriptions) { - try { - StreamObserver observer = - subscription.observer.getAndSet(null); - if (observer != null) { - observer.onCompleted(); + synchronized (this) { + for (WatchSubscription subscription : subscriptions) { + try { + StreamObserver observer; + synchronized (subscription) { + observer = subscription.observer.getAndSet(null); + cancelStartTimeout(subscription); + } + if (observer != null) { + observer.onCompleted(); + } + } catch (Exception e) { + log.warn("Failed to close watch for key {}", subscription.key, e); } - } catch (Exception e) { - log.warn("Failed to close watch for key {}", subscription.key, e); } + subscriptions.clear(); + super.close(); } - subscriptions.clear(); - super.close(); } private final class WatchSubscription { private final String key; private final Consumer consumer; + private final Consumer errorConsumer; private final boolean prefix; private final AtomicReference> observer; private final AtomicBoolean reconnectScheduled; - - private WatchSubscription(String key, Consumer consumer, boolean prefix) { + private final AtomicBoolean terminated; + private final AtomicLong clientId; + private final AtomicReference> startTimeout; + private boolean firstFrameReceived; + private long attemptGeneration; + + private WatchSubscription(String key, Consumer consumer, + Consumer errorConsumer, + boolean prefix) { this.key = key; this.consumer = consumer; + this.errorConsumer = errorConsumer; this.prefix = prefix; this.observer = new AtomicReference<>(); this.reconnectScheduled = new AtomicBoolean(false); + this.terminated = new AtomicBoolean(false); + this.clientId = new AtomicLong(0L); + this.startTimeout = new AtomicReference<>(); } } } diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index 5d4a62fcef..da0cd24604 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -25,19 +25,19 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import java.lang.reflect.Field; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Deque; import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; +import java.util.concurrent.CountDownLatch; 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.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -58,12 +58,14 @@ import org.apache.hugegraph.pd.grpc.kv.WatchState; import org.apache.hugegraph.pd.grpc.kv.WatchType; import org.junit.After; +import org.junit.Assert; import org.junit.Before; import org.junit.Test; import io.grpc.CallOptions; import io.grpc.Channel; import io.grpc.ClientCall; +import io.grpc.ManagedChannel; import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Server; @@ -149,6 +151,40 @@ public void getMembers( } } + @Test + public void testStubCreationFailureIsReportedAsPdException() throws Exception { + AtomicReference grpcAddress = new AtomicReference<>(); + Server server = ServerBuilder.forPort(0) + .addService(new PDGrpc.PDImplBase() { + @Override + public void getMembers( + Pdpb.GetMembersRequest request, + StreamObserver observer) { + Metapb.Member leader = + Metapb.Member.newBuilder() + .setGrpcUrl(grpcAddress.get()) + .build(); + observer.onNext( + Pdpb.GetMembersResponse.newBuilder() + .setLeader(leader) + .build()); + observer.onCompleted(); + } + }) + .build() + .start(); + grpcAddress.set("127.0.0.1:" + server.getPort()); + try (FailingStubCreationKvClient testClient = + new FailingStubCreationKvClient( + PDConfig.of(grpcAddress.get()).setAuthority(user, pwd))) { + assertThatThrownBy(testClient::initializeTransport) + .isInstanceOf(PDException.class) + .hasMessageContaining("PD unreachable"); + } finally { + server.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + @Test(timeout = 2000L) public void testStreamingFailurePropagatesToListenCaller() { ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); @@ -172,18 +208,41 @@ public void testStreamingFailurePropagatesToListenCaller() { @Test(timeout = 2000L) public void testStreamingFailureRetriesNextPeer() throws Exception { + AtomicReference firstAddress = new AtomicReference<>(); + AtomicReference secondAddress = new AtomicReference<>(); + MutableLeaderService firstPdService = new MutableLeaderService(firstAddress); + MutableLeaderService secondPdService = new MutableLeaderService(secondAddress); + RecordingWatchService secondWatchService = new RecordingWatchService(12L); + Server firstServer = ServerBuilder.forPort(0) + .addService(firstPdService) + .build() + .start(); + Server secondServer = ServerBuilder.forPort(0) + .addService(secondPdService) + .addService(secondWatchService) + .build() + .start(); + firstAddress.set("127.0.0.1:" + firstServer.getPort()); + secondAddress.set("127.0.0.1:" + secondServer.getPort()); ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); + when(reconnectExecutor.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))) + .thenReturn(mock(ScheduledFuture.class)); StatusRuntimeException failure = Status.UNAVAILABLE.withDescription("first peer failed").asRuntimeException(); - AtomicBoolean secondPeerCalled = new AtomicBoolean(false); - PDConfig config = PDConfig.of("first:8686,second:8686").setAuthority(user, pwd); - try (SequentialStreamingKvClient testClient = - new SequentialStreamingKvClient(config, reconnectExecutor, - new FailingChannel(failure), - new SuccessfulChannel(secondPeerCalled))) { + PDConfig config = PDConfig.of(firstAddress.get() + "," + secondAddress.get()) + .setAuthority(user, pwd); + try (FailFirstStreamingKvClient testClient = + new FailFirstStreamingKvClient(config, reconnectExecutor, failure)) { testClient.listen("key", response -> { }); - assertThat(secondPeerCalled).isTrue(); + Assert.assertTrue(secondWatchService.started.await(2L, TimeUnit.SECONDS)); + assertThat(firstPdService.calls).hasValue(1); + assertThat(secondPdService.calls).hasValue(1); + assertThat(secondWatchService.calls).hasValue(1); + assertThat(testClient.stubCreations).hasValue(2); + } finally { + firstServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + secondServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); } } @@ -269,8 +328,7 @@ public void testWatchReconnectPreservesLockClientId() throws Exception { @Test public void testPermanentErrorStopsReconnect() throws Exception { - for (Status status : List.of(Status.CANCELLED, - Status.INVALID_ARGUMENT, + for (Status status : List.of(Status.INVALID_ARGUMENT, Status.NOT_FOUND, Status.ALREADY_EXISTS, Status.PERMISSION_DENIED, @@ -293,6 +351,18 @@ public void testPermanentErrorStopsReconnect() throws Exception { } } + @Test + public void testChannelResetCancellationSchedulesReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + + context.client.call(0).observer.onError(Status.CANCELLED.asRuntimeException()); + + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.client.stubInvalidations).isZero(); + } + } + @Test public void testRetryableErrorSchedulesReconnect() throws Exception { for (Status status : List.of(Status.UNAVAILABLE, Status.UNKNOWN)) { @@ -303,6 +373,11 @@ public void testRetryableErrorSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); + if (status.getCode() == Status.Code.UNAVAILABLE) { + assertThat(context.client.stubInvalidations).isEqualTo(1); + } else { + assertThat(context.client.stubInvalidations).isZero(); + } } } } @@ -322,12 +397,196 @@ public void testLeaderChangedSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); + assertThat(context.client.stubInvalidations).isEqualTo(1); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(2); + assertThat(context.client.call(1).request.getClientId()).isZero(); + } + } + + @Test + public void testLeaderChangedReconnectsToNewPeer() throws Exception { + AtomicReference leaderAddress = new AtomicReference<>(); + RecordingWatchService firstWatchService = new RecordingWatchService(11L); + RecordingWatchService secondWatchService = new RecordingWatchService(12L); + Server firstServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(leaderAddress)) + .addService(firstWatchService) + .build() + .start(); + Server secondServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(leaderAddress)) + .addService(secondWatchService) + .build() + .start(); + String firstAddress = "127.0.0.1:" + firstServer.getPort(); + String secondAddress = "127.0.0.1:" + secondServer.getPort(); + leaderAddress.set(firstAddress); + ScheduledExecutorService executor = mock(ScheduledExecutorService.class); + Deque reconnectTasks = new ArrayDeque<>(); + CountDownLatch reconnectScheduled = new CountDownLatch(1); + doAnswer(invocation -> { + long delay = invocation.getArgument(1); + if (invocation.getArgument(2).toMillis(delay) == 1000L) { + reconnectTasks.addLast(invocation.getArgument(0)); + reconnectScheduled.countDown(); + } + return mock(ScheduledFuture.class); + }).when(executor).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + PDConfig config = PDConfig.of(firstAddress + "," + secondAddress) + .setAuthority(user, pwd); + + try (KvClient testClient = new KvClient<>(config, executor)) { + testClient.listen("key", response -> { }); + Assert.assertTrue(firstWatchService.started.await(2L, TimeUnit.SECONDS)); + + leaderAddress.set(secondAddress); + firstWatchService.observer.get().onNext( + WatchResponse.newBuilder().setState(WatchState.Leader_Changed).build()); + Assert.assertTrue(reconnectScheduled.await(2L, TimeUnit.SECONDS)); + assertThat(reconnectTasks).hasSize(1); + reconnectTasks.removeFirst().run(); + + Assert.assertTrue(secondWatchService.started.await(2L, TimeUnit.SECONDS)); + assertThat(firstWatchService.calls).hasValue(1); + assertThat(secondWatchService.calls).hasValue(1); + } finally { + firstServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + secondServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + + @Test + public void testMissingFirstFrameSchedulesReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + + assertThat(context.watchTimeoutTasks).hasSize(1); + assertThat(context.watchTimeoutDelaysMs).containsExactly(5000L); + context.runNextWatchTimeout(); + + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.client.stubInvalidations).isEqualTo(1); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(2); + assertThat(context.client.call(1).request.getClientId()).isZero(); + } + } + + @Test + public void testFirstFramePreventsWatchTimeoutReconnect() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + context.client.call(0).observer.onNext(startingResponse(7L)); + + context.runNextWatchTimeout(); + + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.client.stubInvalidations).isZero(); + } + } + + @Test + public void testStaleWatchTimeoutDoesNotReplaceCurrentObserver() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("key", response -> { }); + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); context.runNextReconnect(); + + context.runNextWatchTimeout(); + assertThat(context.client.calls).hasSize(2); + assertThat(context.reconnectTasks).isEmpty(); + context.client.call(1).observer.onNext(startingResponse(9L)); assertThat(context.client.call(1).request.getClientId()).isZero(); } } + @Test + public void testPermanentReconnectFailureStopsSubscription() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + AtomicReference terminalError = new AtomicReference<>(); + context.client.listen("key", response -> { }, terminalError::set); + StatusRuntimeException failure = Status.UNAUTHENTICATED.asRuntimeException(); + context.client.failNextCallWith( + new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "permanent reconnect failure", failure)); + + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.runNextReconnect(); + + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.client.subscriptionCount()).isZero(); + assertThat(terminalError.get()).isSameAs(failure); + } + } + + @Test + public void testPermanentAsyncErrorIsReportedOnce() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + List terminalErrors = new ArrayList<>(); + context.client.listenPrefix("prefix", response -> { }, terminalErrors::add); + StatusRuntimeException failure = Status.PERMISSION_DENIED.asRuntimeException(); + StreamObserver observer = context.client.call(0).observer; + + observer.onError(failure); + observer.onCompleted(); + + assertThat(terminalErrors).containsExactly(failure); + assertThat(context.client.subscriptionCount()).isZero(); + } + } + + @Test + public void testInvalidWatchDoesNotLeakSubscription() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + assertThatThrownBy(() -> context.client.listen(null, response -> { })) + .isInstanceOf(NullPointerException.class); + + assertThat(context.client.subscriptionCount()).isZero(); + context.client.listen("valid", response -> { }); + assertThat(context.client.calls).hasSize(1); + } + } + + @Test + public void testStreamingStartFailureDoesNotLeakSubscription() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.failNextCalls(1); + + assertThatThrownBy(() -> context.client.listen("failed", response -> { })) + .isInstanceOf(PDException.class); + assertThat(context.client.subscriptionCount()).isZero(); + + context.client.listen("valid", response -> { }); + assertThat(context.client.subscriptionCount()).isEqualTo(1); + assertThat(context.client.calls).hasSize(2); + assertThat(context.client.call(1).request.getKey()).isEqualTo("valid"); + } + } + + @Test + public void testUncheckedReconnectFailureIsRetried() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + Consumer activeConsumer = mock(Consumer.class); + context.client.listen("key", response -> { }); + context.client.listen("active", activeConsumer); + context.client.failNextCallWith(new IllegalStateException("reconnect setup failed")); + + context.client.call(0).observer.onError(Status.UNKNOWN.asRuntimeException()); + context.runNextReconnect(); + + assertThat(context.reconnectTasks).hasSize(1); + assertThat(context.client.stubInvalidations).isZero(); + WatchResponse activeEvent = eventResponse(21L); + context.client.call(1).observer.onNext(activeEvent); + verify(activeConsumer).accept(activeEvent); + context.runNextReconnect(); + assertThat(context.client.calls).hasSize(4); + assertThat(context.reconnectTasks).isEmpty(); + } + } + @Test public void testCompletedSchedulesReconnect() throws Exception { try (WatchTestContext context = newWatchTestContext()) { @@ -337,6 +596,7 @@ public void testCompletedSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); + assertThat(context.client.stubInvalidations).isZero(); context.runNextReconnect(); assertThat(context.client.calls).hasSize(2); } @@ -364,6 +624,7 @@ public void testSubscriptionsReconnectIndependently() throws Exception { context.client.listen("first", firstConsumer); context.client.call(0).observer.onNext(startingResponse(11L)); context.client.listenPrefix("second", secondConsumer); + assertThat(context.client.call(1).request.getClientId()).isZero(); context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); context.client.call(1).observer.onError(Status.UNAVAILABLE.asRuntimeException()); @@ -383,6 +644,7 @@ public void testSubscriptionsReconnectIndependently() throws Exception { assertThat(context.client.call(3).methodName) .isEqualTo(KvServiceGrpc.getWatchPrefixMethod().getFullMethodName()); assertThat(context.client.call(3).request.getKey()).isEqualTo("second"); + assertThat(context.client.call(3).request.getClientId()).isZero(); context.client.call(3).observer.onNext(startingResponse(12L)); WatchResponse secondEvent = eventResponse(13L); context.client.call(3).observer.onNext(secondEvent); @@ -394,8 +656,8 @@ public void testSubscriptionsReconnectIndependently() throws Exception { } } - @Test(timeout = 5000L) - public void testBlockedReconnectReschedulesOtherSubscription() throws Exception { + @Test + public void testReconnectsDoNotSharePermit() throws Exception { try (WatchTestContext context = newWatchTestContext()) { context.client.listen("first", response -> { }); context.client.call(0).observer.onNext(startingResponse(11L)); @@ -404,27 +666,47 @@ public void testBlockedReconnectReschedulesOtherSubscription() throws Exception context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); context.client.call(1).observer.onError(Status.UNAVAILABLE.asRuntimeException()); context.runNextReconnect(); + context.runNextReconnect(); - assertThat(context.client.calls).hasSize(3); - ExecutorService reconnectRunner = Executors.newSingleThreadExecutor(); - try { - Future secondReconnect = reconnectRunner.submit(context::runNextReconnect); - secondReconnect.get(1L, TimeUnit.SECONDS); - } finally { - reconnectRunner.shutdownNow(); - assertThat(reconnectRunner.awaitTermination(1L, TimeUnit.SECONDS)).isTrue(); - } + assertThat(context.client.calls).hasSize(4); + assertThat(context.client.call(2).request.getClientId()).isZero(); + assertThat(context.client.call(3).request.getClientId()).isZero(); + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L); + } + } + + @Test + public void testActiveSubscriptionDoesNotRestoreFailedSubscriptionClientId() + throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("first", response -> { }); + context.client.call(0).observer.onNext(startingResponse(11L)); + context.client.listen("second", response -> { }); + + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.client.call(1).observer.onNext(eventResponse(11L)); + context.runNextReconnect(); assertThat(context.client.calls).hasSize(3); - assertThat(context.reconnectTasks).hasSize(1); - assertThat(context.reconnectDelaysMs).containsExactly(1000L, 1000L, 1000L); + assertThat(context.client.call(2).request.getClientId()).isZero(); + } + } - context.client.call(2).observer.onNext(startingResponse(12L)); + @Test + public void testStaleChannelErrorDoesNotInvalidateFreshStub() throws Exception { + try (WatchTestContext context = newWatchTestContext()) { + context.client.listen("first", response -> { }); + context.client.listen("second", response -> { }); + StreamObserver oldSecondObserver = context.client.call(1).observer; + + context.client.call(0).observer.onError(Status.UNAVAILABLE.asRuntimeException()); + context.runNextReconnect(); + oldSecondObserver.onError(Status.CANCELLED.asRuntimeException()); + + assertThat(context.client.stubInvalidations).isEqualTo(1); context.runNextReconnect(); assertThat(context.client.calls).hasSize(4); - assertThat(context.client.call(3).methodName) - .isEqualTo(KvServiceGrpc.getWatchPrefixMethod().getFullMethodName()); - assertThat(context.client.call(3).request.getKey()).isEqualTo("second"); } } @@ -490,6 +772,54 @@ public void testCloseStopsScheduledReconnect() throws Exception { } } + @Test(timeout = 3000L) + public void testCloseWaitsForRunningReconnectAndClosesItsChannel() throws Exception { + ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); + Deque reconnectTasks = new ArrayDeque<>(); + doAnswer(invocation -> { + long delay = invocation.getArgument(1); + TimeUnit unit = invocation.getArgument(2); + if (unit.toMillis(delay) == 1000L) { + reconnectTasks.addLast(invocation.getArgument(0)); + } + return mock(ScheduledFuture.class); + }).when(reconnectExecutor).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + ManagedChannel reconnectChannel = mock(ManagedChannel.class); + when(reconnectChannel.shutdownNow()).thenReturn(reconnectChannel); + when(reconnectChannel.awaitTermination(anyLong(), any(TimeUnit.class))).thenReturn(true); + CloseRaceKvClient testClient = + new CloseRaceKvClient(getPdConfig(), reconnectExecutor, reconnectChannel); + Thread reconnectThread = null; + Thread closeThread = null; + try { + testClient.listen("key", response -> { }); + testClient.observer.get().onError(Status.UNKNOWN.asRuntimeException()); + Runnable reconnect = reconnectTasks.removeFirst(); + reconnectThread = new Thread(reconnect, "test-watch-reconnect"); + reconnectThread.start(); + Assert.assertTrue(testClient.reconnectStarted.await(1L, TimeUnit.SECONDS)); + + closeThread = new Thread(testClient::close, "test-kv-client-close"); + closeThread.start(); + Assert.assertTrue(testClient.closeEntered.await(1L, TimeUnit.SECONDS)); + assertThat(testClient.closeReturned.await(200L, TimeUnit.MILLISECONDS)).isFalse(); + } finally { + testClient.allowReconnectStart.countDown(); + if (reconnectThread != null) { + reconnectThread.join(1000L); + } + if (closeThread != null) { + closeThread.join(1000L); + } + } + + assertThat(reconnectThread).isNotNull(); + assertThat(closeThread).isNotNull(); + assertThat(reconnectThread.isAlive()).isFalse(); + assertThat(closeThread.isAlive()).isFalse(); + verify(reconnectChannel).shutdownNow(); + } + String key = "key"; String value = "value"; @@ -562,17 +892,26 @@ private static WatchResponse eventResponse(long clientId) { private WatchTestContext newWatchTestContext() { ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); Deque reconnectTasks = new ArrayDeque<>(); + Deque watchTimeoutTasks = new ArrayDeque<>(); List reconnectDelaysMs = new ArrayList<>(); + List watchTimeoutDelaysMs = new ArrayList<>(); doAnswer(invocation -> { - reconnectTasks.addLast(invocation.getArgument(0)); long delay = invocation.getArgument(1); TimeUnit unit = invocation.getArgument(2); - reconnectDelaysMs.add(unit.toMillis(delay)); + long delayMs = unit.toMillis(delay); + if (delayMs == 1000L) { + reconnectTasks.addLast(invocation.getArgument(0)); + reconnectDelaysMs.add(delayMs); + } else { + watchTimeoutTasks.addLast(invocation.getArgument(0)); + watchTimeoutDelaysMs.add(delayMs); + } return mock(ScheduledFuture.class); }).when(reconnectExecutor).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); TestKvClient testClient = new TestKvClient(getPdConfig(), reconnectExecutor); return new WatchTestContext(testClient, reconnectExecutor, - reconnectTasks, reconnectDelaysMs); + reconnectTasks, watchTimeoutTasks, + reconnectDelaysMs, watchTimeoutDelaysMs); } private static class InitializableKvClient extends KvClient { @@ -586,6 +925,64 @@ void initializeTransport() throws PDException { } } + private static class FailingStubCreationKvClient extends KvClient { + + FailingStubCreationKvClient(PDConfig pdConfig) { + super(pdConfig); + } + + @Override + protected AbstractStub createStub() { + throw new IllegalStateException("stub creation failed"); + } + + void initializeTransport() throws PDException { + getStub(); + } + } + + private static class MutableLeaderService extends PDGrpc.PDImplBase { + + private final AtomicReference leaderAddress; + private final AtomicInteger calls = new AtomicInteger(); + + MutableLeaderService(AtomicReference leaderAddress) { + this.leaderAddress = leaderAddress; + } + + @Override + public void getMembers(Pdpb.GetMembersRequest request, + StreamObserver observer) { + this.calls.incrementAndGet(); + Metapb.Member leader = Metapb.Member.newBuilder() + .setGrpcUrl(this.leaderAddress.get()) + .build(); + observer.onNext(Pdpb.GetMembersResponse.newBuilder().setLeader(leader).build()); + observer.onCompleted(); + } + } + + private static class RecordingWatchService extends KvServiceGrpc.KvServiceImplBase { + + private final long clientId; + private final AtomicInteger calls = new AtomicInteger(); + private final AtomicReference> observer = + new AtomicReference<>(); + private final CountDownLatch started = new CountDownLatch(1); + + RecordingWatchService(long clientId) { + this.clientId = clientId; + } + + @Override + public void watch(WatchRequest request, StreamObserver observer) { + this.calls.incrementAndGet(); + this.observer.set(observer); + observer.onNext(startingResponse(this.clientId)); + this.started.countDown(); + } + } + private static class SequentialStreamingKvClient extends KvClient { private final Deque channels = new ArrayDeque<>(); @@ -605,62 +1002,80 @@ protected AbstractStub getStub() { } } - private static class FailingChannel extends Channel { - - private final StatusRuntimeException failure; + private static class FailFirstStreamingKvClient extends KvClient { - FailingChannel(StatusRuntimeException failure) { - this.failure = failure; - } + private final StatusRuntimeException firstFailure; + private final AtomicInteger stubCreations = new AtomicInteger(); - @Override - public String authority() { - return "test"; + FailFirstStreamingKvClient(PDConfig pdConfig, + ScheduledExecutorService reconnectExecutor, + StatusRuntimeException firstFailure) { + super(pdConfig, reconnectExecutor); + this.firstFailure = firstFailure; } @Override - public ClientCall newCall( - MethodDescriptor method, CallOptions callOptions) { - return new FailingClientCall<>(this.failure); + protected AbstractStub createStub() { + if (this.stubCreations.getAndIncrement() == 0) { + return KvServiceGrpc.newStub(new FailingChannel(this.firstFailure)); + } + return super.createStub(); } } - private static class FailingClientCall extends ClientCall { - - private final StatusRuntimeException failure; - - FailingClientCall(StatusRuntimeException failure) { - this.failure = failure; - } - - @Override - public void start(Listener responseListener, Metadata headers) { - throw this.failure; - } + private static class CloseRaceKvClient extends KvClient { - @Override - public void request(int numMessages) { - } + private final ManagedChannel reconnectChannel; + private final AtomicInteger calls = new AtomicInteger(); + private final AtomicReference> observer = + new AtomicReference<>(); + private final CountDownLatch reconnectStarted = new CountDownLatch(1); + private final CountDownLatch allowReconnectStart = new CountDownLatch(1); + private final CountDownLatch closeEntered = new CountDownLatch(1); + private final CountDownLatch closeReturned = new CountDownLatch(1); - @Override - public void cancel(String message, Throwable cause) { + CloseRaceKvClient(PDConfig pdConfig, + ScheduledExecutorService reconnectExecutor, + ManagedChannel reconnectChannel) { + super(pdConfig, reconnectExecutor); + this.reconnectChannel = reconnectChannel; } @Override - public void halfClose() { + protected void streamingCall(MethodDescriptor method, + ReqT request, + StreamObserver responseObserver, + int retry) { + this.observer.set((StreamObserver) responseObserver); + if (this.calls.incrementAndGet() == 1) { + return; + } + this.reconnectStarted.countDown(); + try { + if (!this.allowReconnectStart.await(1L, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to start reconnect"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Reconnect was interrupted", e); + } + this.channel = this.reconnectChannel; } @Override - public void sendMessage(ReqT message) { + public void close() { + this.closeEntered.countDown(); + super.close(); + this.closeReturned.countDown(); } } - private static class SuccessfulChannel extends Channel { + private static class FailingChannel extends Channel { - private final AtomicBoolean called; + private final StatusRuntimeException failure; - SuccessfulChannel(AtomicBoolean called) { - this.called = called; + FailingChannel(StatusRuntimeException failure) { + this.failure = failure; } @Override @@ -671,21 +1086,21 @@ public String authority() { @Override public ClientCall newCall( MethodDescriptor method, CallOptions callOptions) { - return new SuccessfulClientCall<>(this.called); + return new FailingClientCall<>(this.failure); } } - private static class SuccessfulClientCall extends ClientCall { + private static class FailingClientCall extends ClientCall { - private final AtomicBoolean called; + private final StatusRuntimeException failure; - SuccessfulClientCall(AtomicBoolean called) { - this.called = called; + FailingClientCall(StatusRuntimeException failure) { + this.failure = failure; } @Override public void start(Listener responseListener, Metadata headers) { - this.called.set(true); + throw this.failure; } @Override @@ -710,16 +1125,22 @@ private static class WatchTestContext implements AutoCloseable { private final TestKvClient client; private final ScheduledExecutorService reconnectExecutor; private final Deque reconnectTasks; + private final Deque watchTimeoutTasks; private final List reconnectDelaysMs; + private final List watchTimeoutDelaysMs; WatchTestContext(TestKvClient client, ScheduledExecutorService reconnectExecutor, Deque reconnectTasks, - List reconnectDelaysMs) { + Deque watchTimeoutTasks, + List reconnectDelaysMs, + List watchTimeoutDelaysMs) { this.client = client; this.reconnectExecutor = reconnectExecutor; this.reconnectTasks = reconnectTasks; + this.watchTimeoutTasks = watchTimeoutTasks; this.reconnectDelaysMs = reconnectDelaysMs; + this.watchTimeoutDelaysMs = watchTimeoutDelaysMs; } Runnable takeNextReconnect() { @@ -731,6 +1152,11 @@ void runNextReconnect() { takeNextReconnect().run(); } + void runNextWatchTimeout() { + assertThat(watchTimeoutTasks).isNotEmpty(); + watchTimeoutTasks.removeFirst().run(); + } + @Override public void close() { client.close(); @@ -742,6 +1168,9 @@ private static class TestKvClient extends KvClient { private final List calls = new ArrayList<>(); private final List lockCalls = new ArrayList<>(); private int remainingFailures; + private PDException nextFailure; + private RuntimeException nextRuntimeFailure; + private int stubInvalidations; private long lockClientId; TestKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { @@ -752,6 +1181,20 @@ void failNextCalls(int failures) { this.remainingFailures = failures; } + void failNextCallWith(PDException failure) { + this.nextFailure = failure; + } + + void failNextCallWith(RuntimeException failure) { + this.nextRuntimeFailure = failure; + } + + int subscriptionCount() throws Exception { + Field field = KvClient.class.getDeclaredField("subscriptions"); + field.setAccessible(true); + return ((java.util.Set) field.get(this)).size(); + } + WatchCall call(int index) { return this.calls.get(index); } @@ -784,12 +1227,27 @@ protected void streamingCall(MethodDescriptor method, this.calls.add(new WatchCall(method.getFullMethodName(), (WatchRequest) request, (StreamObserver) responseObserver)); + if (this.nextRuntimeFailure != null) { + RuntimeException failure = this.nextRuntimeFailure; + this.nextRuntimeFailure = null; + throw failure; + } + if (this.nextFailure != null) { + PDException failure = this.nextFailure; + this.nextFailure = null; + throw failure; + } if (this.remainingFailures > 0) { this.remainingFailures--; throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, "PD is still unreachable"); } } + + @Override + protected void invalidateAsyncStub() { + this.stubInvalidations++; + } } private static class WatchCall { diff --git a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java index 16273485c7..8326dba611 100644 --- a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java +++ b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java @@ -24,6 +24,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Timer; import java.util.TimerTask; import java.util.concurrent.ConcurrentHashMap; @@ -84,7 +85,12 @@ public class SchemaDriver { private SchemaDriver(PDConfig pdConfig, int cacheSize, long expiration) { - this.client = new KvClient<>(pdConfig); + this(new KvClient<>(pdConfig), cacheSize, expiration); + } + + SchemaDriver(KvClient client, int cacheSize, + long expiration) { + this.client = Objects.requireNonNull(client, "client"); this.caches = new SchemaCaches(cacheSize, expiration); this.listenMetaChanges(); log.info(String.format( @@ -111,11 +117,14 @@ public static void init(PDConfig pdConfig, int cacheSize, long expiration) { } public static void destroy() { - SchemaDriver instance = INSTANCE.get(); + SchemaDriver instance = INSTANCE.getAndSet(null); if (instance != null) { - instance.caches.cancelScheduleCacheClean(); - instance.caches.destroyAll(); - INSTANCE.set(null); + try { + instance.client.close(); + } finally { + instance.caches.cancelScheduleCacheClean(); + instance.caches.destroyAll(); + } } } diff --git a/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java new file mode 100644 index 0000000000..b62f6c075f --- /dev/null +++ b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.hugegraph; + +import java.lang.reflect.Field; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; + +import org.apache.hugegraph.pd.client.KvClient; +import org.apache.hugegraph.pd.client.PDConfig; +import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.grpc.kv.WatchResponse; +import org.junit.Assert; +import org.junit.Test; + +public class SchemaDriverTest { + + @Test + public void testDestroyClosesOwnedKvClient() throws Exception { + TrackingKvClient client = new TrackingKvClient(); + SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); + instanceReference().set(driver); + + try { + SchemaDriver.destroy(); + + Assert.assertTrue(client.closed); + Assert.assertNull(SchemaDriver.getInstance()); + } finally { + instanceReference().set(null); + client.close(); + } + } + + @SuppressWarnings("unchecked") + private static AtomicReference instanceReference() throws Exception { + Field field = SchemaDriver.class.getDeclaredField("INSTANCE"); + field.setAccessible(true); + return (AtomicReference) field.get(null); + } + + private static class TrackingKvClient extends KvClient { + + private boolean closed; + + TrackingKvClient() { + super(PDConfig.of("127.0.0.1:8686")); + } + + @Override + public void listen(String key, Consumer consumer) throws PDException { + // Avoid opening a real PD stream while constructing SchemaDriver. + } + + @Override + public void close() { + this.closed = true; + super.close(); + } + } +} From 0a77c55ac6fa0baaf0e97d98eb0d15a1c28a1531 Mon Sep 17 00:00:00 2001 From: contrueCT Date: Mon, 31 Aug 2026 18:57:24 +0800 Subject: [PATCH 5/6] fix(pd): address watch reconnect review gaps --- .../apache/hugegraph/pd/client/KvClient.java | 29 +++- .../hugegraph/pd/client/KvClientTest.java | 147 +++++++++++++++++- .../org/apache/hugegraph/SchemaDriver.java | 34 ++-- .../apache/hugegraph/SchemaDriverTest.java | 75 ++++++++- 4 files changed, 266 insertions(+), 19 deletions(-) diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index 023d3684e3..f058770e58 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -54,6 +54,8 @@ import org.apache.hugegraph.pd.grpc.kv.WatchType; import io.grpc.Status; +import io.grpc.StatusException; +import io.grpc.StatusRuntimeException; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.AbstractStub; import io.grpc.stub.StreamObserver; @@ -349,15 +351,23 @@ private void onStartTimeout(WatchSubscription subscription, } log.warn("Watch for key {} did not receive its first response in {} ms", subscription.key, WATCH_START_TIMEOUT_MS); - requestReconnect(subscription, sourceObserver, true); + requestReconnect(subscription, sourceObserver, true, true); } private void requestReconnect(WatchSubscription subscription, StreamObserver sourceObserver, boolean rotateTransport) { + requestReconnect(subscription, sourceObserver, rotateTransport, false); + } + + private void requestReconnect(WatchSubscription subscription, + StreamObserver sourceObserver, + boolean rotateTransport, + boolean requireMissingFirstFrame) { long attemptGeneration; synchronized (subscription) { - if (closed.get() || subscription.observer.get() != sourceObserver) { + if (closed.get() || subscription.observer.get() != sourceObserver || + (requireMissingFirstFrame && subscription.firstFrameReceived)) { return; } subscription.observer.set(null); @@ -392,7 +402,20 @@ private static boolean isRetryableWatchError(Throwable throwable) { } private static boolean shouldRotateWatchTransport(Throwable throwable) { - return Status.fromThrowable(throwable).getCode() == Status.Code.UNAVAILABLE; + Throwable cause = throwable; + while (cause != null) { + Status.Code code = null; + if (cause instanceof StatusException) { + code = ((StatusException) cause).getStatus().getCode(); + } else if (cause instanceof StatusRuntimeException) { + code = ((StatusRuntimeException) cause).getStatus().getCode(); + } + if (code != null) { + return code == Status.Code.UNAVAILABLE || code == Status.Code.UNKNOWN; + } + cause = cause.getCause(); + } + return false; } private void stopWatch(WatchSubscription subscription, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index da0cd24604..6fdcee1340 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -27,6 +27,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.io.Serializable; import java.lang.reflect.Field; import java.util.ArrayDeque; import java.util.ArrayList; @@ -57,6 +58,15 @@ import org.apache.hugegraph.pd.grpc.kv.WatchResponse; import org.apache.hugegraph.pd.grpc.kv.WatchState; import org.apache.hugegraph.pd.grpc.kv.WatchType; +import org.apache.logging.log4j.Level; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.core.Filter; +import org.apache.logging.log4j.core.Layout; +import org.apache.logging.log4j.core.LogEvent; +import org.apache.logging.log4j.core.LoggerContext; +import org.apache.logging.log4j.core.appender.AbstractAppender; +import org.apache.logging.log4j.core.config.LoggerConfig; +import org.apache.logging.log4j.core.config.Property; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -373,11 +383,7 @@ public void testRetryableErrorSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); - if (status.getCode() == Status.Code.UNAVAILABLE) { - assertThat(context.client.stubInvalidations).isEqualTo(1); - } else { - assertThat(context.client.stubInvalidations).isZero(); - } + assertThat(context.client.stubInvalidations).isEqualTo(1); } } } @@ -456,6 +462,55 @@ public void testLeaderChangedReconnectsToNewPeer() throws Exception { } } + @Test + public void testFollowerErrorReconnectsToNewPeer() throws Exception { + AtomicReference firstAddress = new AtomicReference<>(); + AtomicReference secondAddress = new AtomicReference<>(); + RejectingWatchService firstWatchService = new RejectingWatchService(); + RecordingWatchService secondWatchService = new RecordingWatchService(12L); + Server firstServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(firstAddress)) + .addService(firstWatchService) + .build() + .start(); + Server secondServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(secondAddress)) + .addService(secondWatchService) + .build() + .start(); + firstAddress.set("127.0.0.1:" + firstServer.getPort()); + secondAddress.set("127.0.0.1:" + secondServer.getPort()); + ScheduledExecutorService executor = mock(ScheduledExecutorService.class); + Deque reconnectTasks = new ArrayDeque<>(); + CountDownLatch reconnectScheduled = new CountDownLatch(1); + doAnswer(invocation -> { + long delay = invocation.getArgument(1); + if (invocation.getArgument(2).toMillis(delay) == 1000L) { + reconnectTasks.addLast(invocation.getArgument(0)); + reconnectScheduled.countDown(); + } + return mock(ScheduledFuture.class); + }).when(executor).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + PDConfig config = PDConfig.of(firstAddress.get() + "," + secondAddress.get()) + .setAuthority(user, pwd); + + try (KvClient testClient = new KvClient<>(config, executor)) { + testClient.listen("key", response -> { }); + Assert.assertTrue(firstWatchService.failed.await(2L, TimeUnit.SECONDS)); + Assert.assertTrue(reconnectScheduled.await(2L, TimeUnit.SECONDS)); + assertThat(reconnectTasks).hasSize(1); + + reconnectTasks.removeFirst().run(); + + Assert.assertTrue(secondWatchService.started.await(2L, TimeUnit.SECONDS)); + assertThat(firstWatchService.calls).hasValue(1); + assertThat(secondWatchService.calls).hasValue(1); + } finally { + firstServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + secondServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + @Test public void testMissingFirstFrameSchedulesReconnect() throws Exception { try (WatchTestContext context = newWatchTestContext()) { @@ -486,6 +541,50 @@ public void testFirstFramePreventsWatchTimeoutReconnect() throws Exception { } } + @Test + public void testFirstFrameDuringTimeoutTransitionPreventsReconnect() throws Exception { + String loggerName = KvClient.class.getName(); + LoggerContext loggerContext = (LoggerContext) LogManager.getContext(false); + org.apache.logging.log4j.core.config.Configuration loggerConfiguration = + loggerContext.getConfiguration(); + LoggerConfig existingConfig = loggerConfiguration.getLoggerConfig(loggerName); + LoggerConfig originalConfig = + loggerName.equals(existingConfig.getName()) ? existingConfig : null; + AtomicReference timeoutAction = new AtomicReference<>(); + TimeoutRaceAppender appender = new TimeoutRaceAppender(timeoutAction); + appender.start(); + if (originalConfig != null) { + loggerConfiguration.removeLogger(loggerName); + } + LoggerConfig testConfig = new LoggerConfig(loggerName, Level.WARN, false); + testConfig.addAppender(appender, Level.WARN, null); + loggerConfiguration.addLogger(loggerName, testConfig); + loggerContext.updateLoggers(); + + try (WatchTestContext context = newWatchTestContext()) { + Consumer consumer = mock(Consumer.class); + context.client.listen("key", consumer); + StreamObserver observer = context.client.call(0).observer; + timeoutAction.set(() -> observer.onNext(startingResponse(9L))); + + context.runNextWatchTimeout(); + + assertThat(appender.triggered).isTrue(); + assertThat(context.reconnectTasks).isEmpty(); + assertThat(context.client.stubInvalidations).isZero(); + WatchResponse event = eventResponse(9L); + observer.onNext(event); + verify(consumer).accept(event); + } finally { + loggerConfiguration.removeLogger(loggerName); + if (originalConfig != null) { + loggerConfiguration.addLogger(loggerName, originalConfig); + } + loggerContext.updateLoggers(); + appender.stop(); + } + } + @Test public void testStaleWatchTimeoutDoesNotReplaceCurrentObserver() throws Exception { try (WatchTestContext context = newWatchTestContext()) { @@ -577,7 +676,7 @@ public void testUncheckedReconnectFailureIsRetried() throws Exception { context.runNextReconnect(); assertThat(context.reconnectTasks).hasSize(1); - assertThat(context.client.stubInvalidations).isZero(); + assertThat(context.client.stubInvalidations).isEqualTo(1); WatchResponse activeEvent = eventResponse(21L); context.client.call(1).observer.onNext(activeEvent); verify(activeConsumer).accept(activeEvent); @@ -983,6 +1082,42 @@ public void watch(WatchRequest request, StreamObserver observer) } } + private static class RejectingWatchService extends KvServiceGrpc.KvServiceImplBase { + + private static final String NOT_LEADER = + "node is not leader,it is necessary to redirect to the leader on the client"; + + private final AtomicInteger calls = new AtomicInteger(); + private final CountDownLatch failed = new CountDownLatch(1); + + @Override + public void watch(WatchRequest request, StreamObserver observer) { + this.calls.incrementAndGet(); + observer.onError(new PDException(-1, NOT_LEADER)); + this.failed.countDown(); + } + } + + private static class TimeoutRaceAppender extends AbstractAppender { + + private final AtomicReference action; + private final AtomicBoolean triggered = new AtomicBoolean(false); + + TimeoutRaceAppender(AtomicReference action) { + super("KvClientTimeoutRaceAppender", (Filter) null, + (Layout) null, false, Property.EMPTY_ARRAY); + this.action = action; + } + + @Override + public void append(LogEvent event) { + Runnable callback = this.action.get(); + if (callback != null && this.triggered.compareAndSet(false, true)) { + callback.run(); + } + } + } + private static class SequentialStreamingKvClient extends KvClient { private final Deque channels = new ArrayDeque<>(); diff --git a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java index 8326dba611..5f8f96c274 100644 --- a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java +++ b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java @@ -92,7 +92,16 @@ private SchemaDriver(PDConfig pdConfig, int cacheSize, long expiration) { this.client = Objects.requireNonNull(client, "client"); this.caches = new SchemaCaches(cacheSize, expiration); - this.listenMetaChanges(); + try { + this.listenMetaChanges(); + } catch (RuntimeException e) { + try { + this.closeResources(); + } catch (RuntimeException closeException) { + e.addSuppressed(closeException); + } + throw e; + } log.info(String.format( "The SchemaDriver initialized successfully, cacheSize = %s," + " expiration = %s s", cacheSize, expiration / 1000)); @@ -103,7 +112,7 @@ public static void init(PDConfig pdConfig) { init(pdConfig, 300, 300 * 1000); } - public static void init(PDConfig pdConfig, int cacheSize, long expiration) { + public static synchronized void init(PDConfig pdConfig, int cacheSize, long expiration) { SchemaDriver instance = INSTANCE.get(); if (instance != null) { throw new NotAllowException( @@ -112,22 +121,29 @@ public static void init(PDConfig pdConfig, int cacheSize, long expiration) { "allowed to be initialized again", instance.caches.limit(), instance.caches.expiration(), instance.client); } - INSTANCE.compareAndSet(null, new SchemaDriver(pdConfig, cacheSize, - expiration)); + INSTANCE.set(new SchemaDriver(pdConfig, cacheSize, expiration)); } - public static void destroy() { - SchemaDriver instance = INSTANCE.getAndSet(null); + public static synchronized void destroy() { + SchemaDriver instance = INSTANCE.get(); if (instance != null) { try { - instance.client.close(); + instance.closeResources(); } finally { - instance.caches.cancelScheduleCacheClean(); - instance.caches.destroyAll(); + INSTANCE.compareAndSet(instance, null); } } } + private void closeResources() { + try { + this.client.close(); + } finally { + this.caches.cancelScheduleCacheClean(); + this.caches.destroyAll(); + } + } + public SchemaCaches schemaCaches() { return this.caches; } diff --git a/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java index b62f6c075f..6edd7900f9 100644 --- a/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java +++ b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java @@ -18,9 +18,12 @@ package org.apache.hugegraph; import java.lang.reflect.Field; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; +import org.apache.hugegraph.exception.HugeException; import org.apache.hugegraph.pd.client.KvClient; import org.apache.hugegraph.pd.client.PDConfig; import org.apache.hugegraph.pd.common.PDException; @@ -47,6 +50,44 @@ public void testDestroyClosesOwnedKvClient() throws Exception { } } + @Test(timeout = 3000L) + public void testDestroyKeepsInstanceUntilResourcesAreClosed() throws Exception { + BlockingCloseKvClient client = new BlockingCloseKvClient(); + SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); + AtomicReference instance = instanceReference(); + SchemaDriver previous = instance.getAndSet(driver); + Thread destroyThread = new Thread(SchemaDriver::destroy, "schema-driver-destroy"); + try { + destroyThread.start(); + Assert.assertTrue(client.closeStarted.await(1L, TimeUnit.SECONDS)); + + Assert.assertSame(driver, SchemaDriver.getInstance()); + } finally { + client.allowClose.countDown(); + destroyThread.join(1000L); + instance.set(previous); + client.close(); + } + + Assert.assertFalse(destroyThread.isAlive()); + Assert.assertTrue(client.closed); + } + + @Test + public void testConstructorFailureClosesOwnedKvClient() { + FailingListenKvClient client = new FailingListenKvClient(); + + try { + new SchemaDriver(client, 10, 60_000L); + Assert.fail("SchemaDriver construction should fail"); + } catch (HugeException expected) { + Assert.assertEquals(2, client.listenCalls); + Assert.assertTrue(client.closed); + } finally { + client.close(); + } + } + @SuppressWarnings("unchecked") private static AtomicReference instanceReference() throws Exception { Field field = SchemaDriver.class.getDeclaredField("INSTANCE"); @@ -56,7 +97,7 @@ private static AtomicReference instanceReference() throws Exceptio private static class TrackingKvClient extends KvClient { - private boolean closed; + protected volatile boolean closed; TrackingKvClient() { super(PDConfig.of("127.0.0.1:8686")); @@ -73,4 +114,36 @@ public void close() { super.close(); } } + + private static class BlockingCloseKvClient extends TrackingKvClient { + + private final CountDownLatch closeStarted = new CountDownLatch(1); + private final CountDownLatch allowClose = new CountDownLatch(1); + + @Override + public void close() { + this.closeStarted.countDown(); + try { + if (!this.allowClose.await(1L, TimeUnit.SECONDS)) { + throw new IllegalStateException("Timed out waiting to close client"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Client close was interrupted", e); + } + super.close(); + } + } + + private static class FailingListenKvClient extends TrackingKvClient { + + private int listenCalls; + + @Override + public void listen(String key, Consumer consumer) throws PDException { + if (++this.listenCalls == 2) { + throw new PDException(-1, "listener startup failed"); + } + } + } } From f99b6bdddf7554760f2920f719238bb1811bcdd7 Mon Sep 17 00:00:00 2001 From: contrueCT Date: Tue, 1 Sep 2026 01:34:32 +0800 Subject: [PATCH 6/6] fix(pd): isolate watch reconnect lifecycle --- .../hugegraph/pd/client/AbstractClient.java | 162 +++++++-- .../apache/hugegraph/pd/client/KvClient.java | 182 +++++++--- .../hugegraph/pd/client/KvClientTest.java | 335 +++++++++++++++++- .../org/apache/hugegraph/SchemaDriver.java | 73 +++- .../apache/hugegraph/SchemaDriverTest.java | 108 +++++- 5 files changed, 744 insertions(+), 116 deletions(-) diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java index 24d26ac3df..3943b8cb6e 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/AbstractClient.java @@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.function.Predicate; import java.util.stream.Stream; @@ -57,8 +58,9 @@ public abstract class AbstractClient implements Closeable { protected final Pdpb.RequestHeader header; protected final AbstractClientStubProxy proxy; protected final PDConfig config; - protected ManagedChannel channel = null; + protected volatile ManagedChannel channel = null; protected ConcurrentMap stubs = null; + private final ThreadLocal> streamingAttemptConsumer = new ThreadLocal<>(); protected AbstractClient(PDConfig config) { String[] hosts = config.getServerHost().split(","); @@ -84,7 +86,12 @@ protected static void handleErrors(Pdpb.ResponseHeader header) throws PDExceptio } public static T setBlockingParams(T stub, PDConfig config) { - stub = (T) stub.withDeadlineAfter(config.getGrpcTimeOut(), TimeUnit.MILLISECONDS) + return setBlockingParams(stub, config, config.getGrpcTimeOut()); + } + + private static T setBlockingParams(T stub, PDConfig config, + long timeoutMillis) { + stub = (T) stub.withDeadlineAfter(timeoutMillis, TimeUnit.MILLISECONDS) .withMaxInboundMessageSize(PDConfig.getInboundMessageSize()); return (T) stub.withInterceptors( new Authentication(config.getUserName(), config.getAuthority())); @@ -97,31 +104,23 @@ public static T setAsyncParams(T stub, PDConfig config) new Authentication(config.getUserName(), config.getAuthority())); } - protected AbstractBlockingStub getBlockingStub() throws PDException { + protected synchronized AbstractBlockingStub getBlockingStub() throws PDException { if (proxy.getBlockingStub() == null) { - synchronized (this) { - if (proxy.getBlockingStub() == null) { - String host = resetStub(); - if (host.isEmpty()) { - throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, - "PD unreachable, pd.peers=" + config.getServerHost()); - } - } + String host = resetStub(); + if (host.isEmpty()) { + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "PD unreachable, pd.peers=" + config.getServerHost()); } } return setBlockingParams(proxy.getBlockingStub(), config); } - protected AbstractStub getStub() throws PDException { + protected synchronized AbstractStub getStub() throws PDException { if (proxy.getStub() == null) { - synchronized (this) { - if (proxy.getStub() == null) { - String host = resetStub(); - if (host.isEmpty()) { - throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, - "PD unreachable, pd.peers=" + config.getServerHost()); - } - } + String host = resetStub(); + if (host.isEmpty()) { + throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, + "PD unreachable, pd.peers=" + config.getServerHost()); } } return setAsyncParams(proxy.getStub(), config); @@ -131,31 +130,95 @@ protected synchronized void invalidateAsyncStub() { proxy.setStub(null); } + protected synchronized boolean invalidateAsyncStub(Channel expectedChannel) { + AbstractStub stub = proxy.getStub(); + if (stub == null || stub.getChannel() != expectedChannel) { + return false; + } + proxy.setStub(null); + return true; + } + + protected long stubResetTimeoutMillis() { + return (long) config.getGrpcTimeOut() * Math.max(1, proxy.getHostCount()); + } + + protected boolean isShutdown() { + return false; + } + protected abstract AbstractStub createStub(); protected abstract AbstractBlockingStub createBlockingStub(); private String resetStub() { Exception ex = null; + long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(stubResetTimeoutMillis()); + long deadlineNanos = System.nanoTime() + timeoutNanos; for (int i = 0; i < proxy.getHostCount(); i++) { + if (isShutdown() || remainingMillis(deadlineNanos) <= 0L) { + break; + } String host = proxy.nextHost(); closeConnections(); - channel = ManagedChannelBuilder.forTarget(host).usePlaintext().build(); + if (isShutdown()) { + break; + } + ManagedChannel candidate = + ManagedChannelBuilder.forTarget(host).usePlaintext().build(); + if (isShutdown()) { + closeChannel(candidate); + break; + } + channel = candidate; + if (isShutdown()) { + closeChannel(candidate); + break; + } + long remaining = remainingMillis(deadlineNanos); + if (remaining <= 0L) { + break; + } + int remainingHosts = proxy.getHostCount() - i; + long peerTimeout = Math.max(1L, Math.min(config.getGrpcTimeOut(), + remaining / remainingHosts)); PDBlockingStub blockingStub = - setBlockingParams(PDGrpc.newBlockingStub(channel), config); + setBlockingParams(PDGrpc.newBlockingStub(channel), config, + peerTimeout); try { GetMembersRequest request = Pdpb.GetMembersRequest.newBuilder() .setHeader(header).build(); GetMembersResponse members = blockingStub.getMembers(request); + if (isShutdown()) { + break; + } Metapb.Member leader = members.getLeader(); String leaderHost = leader.getGrpcUrl(); if (!host.equals(leaderHost)) { closeConnections(); - channel = ManagedChannelBuilder.forTarget(leaderHost).usePlaintext().build(); + if (isShutdown()) { + break; + } + candidate = ManagedChannelBuilder.forTarget(leaderHost).usePlaintext().build(); + if (isShutdown()) { + closeChannel(candidate); + break; + } + channel = candidate; + if (isShutdown()) { + closeChannel(candidate); + break; + } + } + AbstractBlockingStub newBlockingStub = + setBlockingParams(createBlockingStub(), config); + AbstractStub newStub = setAsyncParams(createStub(), config); + if (isShutdown()) { + break; } - proxy.setBlockingStub(setBlockingParams(createBlockingStub(), config)); - proxy.setStub(setAsyncParams(createStub(), config)); + proxy.setBlockingStub(newBlockingStub); + proxy.setStub(newStub); log.info("AbstractClient connect to host = {} success", leaderHost); return leaderHost; } catch (StatusRuntimeException se) { @@ -170,6 +233,8 @@ private String resetStub() { proxy.setBlockingStub(null); proxy.setStub(null); } + proxy.setBlockingStub(null); + proxy.setStub(null); closeConnections(); if (ex != null) { log.error(String.format("connect to %s with error: ", config.getServerHost()), ex); @@ -177,6 +242,14 @@ private String resetStub() { return ""; } + private static long remainingMillis(long deadlineNanos) { + long remainingNanos = deadlineNanos - System.nanoTime(); + if (remainingNanos <= 0L) { + return 0L; + } + return Math.max(1L, TimeUnit.NANOSECONDS.toMillis(remainingNanos)); + } + protected RespT blockingUnaryCall( MethodDescriptor method, ReqT req) throws PDException { return blockingUnaryCall(method, req, 0); @@ -253,7 +326,17 @@ protected KVPair concurrentBlockingUnaryCall( protected void streamingCall(MethodDescriptor method, ReqT request, StreamObserver responseObserver, int retry) throws PDException { - AbstractStub stub = getStub(); + AbstractStub stub; + Channel attemptChannel; + synchronized (this) { + stub = getStub(); + AbstractStub currentStub = proxy.getStub(); + attemptChannel = currentStub == null ? stub.getChannel() : currentStub.getChannel(); + Consumer attemptConsumer = this.streamingAttemptConsumer.get(); + if (attemptConsumer != null) { + attemptConsumer.accept(attemptChannel); + } + } try { ClientCall call = stub.getChannel().newCall(method, stub.getCallOptions()); ClientCalls.asyncServerStreamingCall(call, request, responseObserver); @@ -261,9 +344,7 @@ protected void streamingCall(MethodDescriptor method, log.error("rpc call with exception :", e); if (e instanceof StatusRuntimeException) { if (retry < proxy.getHostCount()) { - synchronized (this) { - invalidateAsyncStub(); - } + invalidateAsyncStub(attemptChannel); streamingCall(method, request, responseObserver, ++retry); return; } @@ -273,6 +354,24 @@ protected void streamingCall(MethodDescriptor method, } } + protected void streamingCall(MethodDescriptor method, ReqT request, + StreamObserver responseObserver, + int retry, + Consumer attemptConsumer) + throws PDException { + Consumer previous = this.streamingAttemptConsumer.get(); + this.streamingAttemptConsumer.set(attemptConsumer); + try { + streamingCall(method, request, responseObserver, retry); + } finally { + if (previous == null) { + this.streamingAttemptConsumer.remove(); + } else { + this.streamingAttemptConsumer.set(previous); + } + } + } + @Override public void close() { closeConnections(); @@ -289,9 +388,8 @@ private void closeConnections() { private void closeChannel(ManagedChannel channel) { try { - while (channel != null && - !channel.shutdownNow().awaitTermination(100, TimeUnit.MILLISECONDS)) { - continue; + if (channel != null) { + channel.shutdownNow().awaitTermination(100, TimeUnit.MILLISECONDS); } } catch (Exception e) { log.info("Close channel with error :.", e); diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java index f058770e58..bd95e7370f 100644 --- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java +++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java @@ -25,6 +25,8 @@ import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -53,6 +55,7 @@ import org.apache.hugegraph.pd.grpc.kv.WatchResponse; import org.apache.hugegraph.pd.grpc.kv.WatchType; +import io.grpc.Channel; import io.grpc.Status; import io.grpc.StatusException; import io.grpc.StatusRuntimeException; @@ -66,6 +69,7 @@ public class KvClient extends AbstractClient implements private static final long RECONNECT_DELAY_MS = 1000L; private static final long WATCH_START_TIMEOUT_MS = 5000L; + private static final int RECONNECT_WORKER_THREADS = 4; private static final Set NON_RETRYABLE_WATCH_ERRORS = Set.of(Status.Code.INVALID_ARGUMENT, Status.Code.NOT_FOUND, @@ -81,19 +85,39 @@ public class KvClient extends AbstractClient implements private final Semaphore lockSemaphore = new Semaphore(1); private final AtomicBoolean closed = new AtomicBoolean(false); private final Set subscriptions = ConcurrentHashMap.newKeySet(); - private final ScheduledExecutorService reconnectExecutor; - private long transportGeneration; + private final ScheduledExecutorService reconnectScheduler; + private final Executor reconnectExecutor; public KvClient(PDConfig pdConfig) { - this(pdConfig, Executors.newSingleThreadScheduledExecutor(runnable -> { + this(pdConfig, newReconnectScheduler(), newReconnectExecutor()); + } + + private static ScheduledExecutorService newReconnectScheduler() { + return Executors.newSingleThreadScheduledExecutor(runnable -> { Thread thread = new Thread(runnable, "pd-kv-watch-reconnect"); thread.setDaemon(true); return thread; - })); + }); + } + + private static ExecutorService newReconnectExecutor() { + AtomicLong threadId = new AtomicLong(); + return Executors.newFixedThreadPool(RECONNECT_WORKER_THREADS, runnable -> { + Thread thread = new Thread(runnable, + "pd-kv-watch-worker-" + threadId.incrementAndGet()); + thread.setDaemon(true); + return thread; + }); + } + + KvClient(PDConfig pdConfig, ScheduledExecutorService reconnectScheduler) { + this(pdConfig, reconnectScheduler, Runnable::run); } - KvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { + KvClient(PDConfig pdConfig, ScheduledExecutorService reconnectScheduler, + Executor reconnectExecutor) { super(pdConfig); + this.reconnectScheduler = reconnectScheduler; this.reconnectExecutor = reconnectExecutor; } @@ -107,6 +131,16 @@ protected AbstractBlockingStub createBlockingStub() { return KvServiceGrpc.newBlockingStub(channel); } + @Override + protected long stubResetTimeoutMillis() { + return Math.max(1L, Math.min(config.getGrpcTimeOut(), WATCH_START_TIMEOUT_MS)); + } + + @Override + protected boolean isShutdown() { + return closed.get(); + } + public KvResponse put(String key, String value) throws PDException { Kv kv = Kv.newBuilder().setKey(key).setValue(value).build(); KvResponse response = blockingUnaryCall(KvServiceGrpc.getPutMethod(), kv); @@ -276,27 +310,37 @@ private boolean startWatch(WatchSubscription subscription) throws PDException { .build(); StreamObserver observer = getObserver(subscription); try { - synchronized (this) { - synchronized (subscription) { - if (closed.get() || !subscriptions.contains(subscription)) { - return false; - } - if (subscription.observer.get() != null) { - return true; - } - subscription.firstFrameReceived = false; - subscription.attemptGeneration = this.transportGeneration; - subscription.observer.set(observer); + synchronized (subscription) { + if (closed.get() || !subscriptions.contains(subscription)) { + return false; } - if (subscription.prefix) { - streamingCall(KvServiceGrpc.getWatchPrefixMethod(), request, observer, 1); - } else { - streamingCall(KvServiceGrpc.getWatchMethod(), request, observer, 1); + if (subscription.observer.get() != null) { + return true; } + subscription.firstFrameReceived = false; + subscription.attemptChannel = null; + subscription.observer.set(observer); + } + if (subscription.prefix) { + streamingCall(KvServiceGrpc.getWatchPrefixMethod(), request, observer, 1, + channel -> recordAttemptChannel(subscription, observer, channel)); + } else { + streamingCall(KvServiceGrpc.getWatchMethod(), request, observer, 1, + channel -> recordAttemptChannel(subscription, observer, channel)); + } + if (closed.get()) { + cleanupFailedStart(subscription); + super.close(); + return false; } scheduleStartTimeout(subscription, observer); return true; } catch (Exception e) { + if (closed.get()) { + cleanupFailedStart(subscription); + super.close(); + return false; + } if (e instanceof PDException) { throw (PDException) e; } @@ -304,6 +348,16 @@ private boolean startWatch(WatchSubscription subscription) throws PDException { } } + private void recordAttemptChannel(WatchSubscription subscription, + StreamObserver sourceObserver, + Channel channel) { + synchronized (subscription) { + if (subscription.observer.get() == sourceObserver) { + subscription.attemptChannel = channel; + } + } + } + private boolean acceptFirstFrame(WatchSubscription subscription, StreamObserver sourceObserver) { synchronized (subscription) { @@ -321,7 +375,7 @@ private void scheduleStartTimeout(WatchSubscription subscription, throws PDException { ScheduledFuture timeout; try { - timeout = reconnectExecutor.schedule( + timeout = reconnectScheduler.schedule( () -> onStartTimeout(subscription, sourceObserver), WATCH_START_TIMEOUT_MS, TimeUnit.MILLISECONDS); } catch (RuntimeException e) { @@ -364,7 +418,6 @@ private void requestReconnect(WatchSubscription subscription, StreamObserver sourceObserver, boolean rotateTransport, boolean requireMissingFirstFrame) { - long attemptGeneration; synchronized (subscription) { if (closed.get() || subscription.observer.get() != sourceObserver || (requireMissingFirstFrame && subscription.firstFrameReceived)) { @@ -373,20 +426,18 @@ private void requestReconnect(WatchSubscription subscription, subscription.observer.set(null); cancelStartTimeout(subscription); subscription.clientId.set(0L); - attemptGeneration = subscription.attemptGeneration; - } - if (rotateTransport) { - invalidateAttemptStub(attemptGeneration); + if (rotateTransport) { + subscription.rotateTransport = true; + subscription.reconnectChannel = subscription.attemptChannel; + } } scheduleReconnect(subscription); } - private synchronized void invalidateAttemptStub(long attemptGeneration) { - if (attemptGeneration != this.transportGeneration) { - return; + private void invalidateAttemptStub(Channel attemptChannel) { + if (attemptChannel != null) { + invalidateAsyncStub(attemptChannel); } - this.transportGeneration++; - invalidateAsyncStub(); } private void cancelStartTimeout(WatchSubscription subscription) { @@ -459,8 +510,8 @@ private void scheduleReconnect(WatchSubscription subscription) { return; } try { - reconnectExecutor.schedule(() -> reconnect(subscription), RECONNECT_DELAY_MS, - TimeUnit.MILLISECONDS); + reconnectScheduler.schedule(() -> submitReconnect(subscription), + RECONNECT_DELAY_MS, TimeUnit.MILLISECONDS); } catch (RuntimeException e) { subscription.reconnectScheduled.set(false); if (!closed.get()) { @@ -471,12 +522,40 @@ private void scheduleReconnect(WatchSubscription subscription) { } } + private void submitReconnect(WatchSubscription subscription) { + if (closed.get() || !subscriptions.contains(subscription)) { + subscription.reconnectScheduled.set(false); + return; + } + try { + reconnectExecutor.execute(() -> reconnect(subscription)); + } catch (RuntimeException e) { + subscription.reconnectScheduled.set(false); + if (!closed.get()) { + log.warn("Failed to execute watch reconnect for key {}", subscription.key, e); + subscriptions.remove(subscription); + notifyWatchStopped(subscription, e); + } + } + } + private void reconnect(WatchSubscription subscription) { subscription.reconnectScheduled.set(false); if (closed.get() || !subscriptions.contains(subscription)) { return; } try { + Channel reconnectChannel = null; + synchronized (subscription) { + if (subscription.rotateTransport) { + reconnectChannel = subscription.reconnectChannel; + subscription.rotateTransport = false; + subscription.reconnectChannel = null; + } + } + if (reconnectChannel != null) { + invalidateAttemptStub(reconnectChannel); + } if (!startWatch(subscription)) { scheduleReconnect(subscription); } @@ -624,26 +703,27 @@ public void close() { if (!closed.compareAndSet(false, true)) { return; } - reconnectExecutor.shutdownNow(); + reconnectScheduler.shutdownNow(); + if (reconnectExecutor instanceof ExecutorService) { + ((ExecutorService) reconnectExecutor).shutdownNow(); + } release(lockSemaphore); - synchronized (this) { - for (WatchSubscription subscription : subscriptions) { - try { - StreamObserver observer; - synchronized (subscription) { - observer = subscription.observer.getAndSet(null); - cancelStartTimeout(subscription); - } - if (observer != null) { - observer.onCompleted(); - } - } catch (Exception e) { - log.warn("Failed to close watch for key {}", subscription.key, e); + for (WatchSubscription subscription : subscriptions) { + try { + StreamObserver observer; + synchronized (subscription) { + observer = subscription.observer.getAndSet(null); + cancelStartTimeout(subscription); } + if (observer != null) { + observer.onCompleted(); + } + } catch (Exception e) { + log.warn("Failed to close watch for key {}", subscription.key, e); } - subscriptions.clear(); - super.close(); } + subscriptions.clear(); + super.close(); } private final class WatchSubscription { @@ -658,7 +738,9 @@ private final class WatchSubscription { private final AtomicLong clientId; private final AtomicReference> startTimeout; private boolean firstFrameReceived; - private long attemptGeneration; + private Channel attemptChannel; + private boolean rotateTransport; + private Channel reconnectChannel; private WatchSubscription(String key, Consumer consumer, Consumer errorConsumer, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java index 6fdcee1340..9186185ec3 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientTest.java @@ -34,6 +34,7 @@ import java.util.Deque; import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; @@ -383,6 +384,7 @@ public void testRetryableErrorSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); + context.runNextReconnect(); assertThat(context.client.stubInvalidations).isEqualTo(1); } } @@ -403,8 +405,8 @@ public void testLeaderChangedSchedulesReconnect() throws Exception { assertThat(context.reconnectTasks).hasSize(1); assertThat(context.reconnectDelaysMs).containsExactly(1000L); - assertThat(context.client.stubInvalidations).isEqualTo(1); context.runNextReconnect(); + assertThat(context.client.stubInvalidations).isEqualTo(1); assertThat(context.client.calls).hasSize(2); assertThat(context.client.call(1).request.getClientId()).isZero(); } @@ -521,8 +523,8 @@ public void testMissingFirstFrameSchedulesReconnect() throws Exception { context.runNextWatchTimeout(); assertThat(context.reconnectTasks).hasSize(1); - assertThat(context.client.stubInvalidations).isEqualTo(1); context.runNextReconnect(); + assertThat(context.client.stubInvalidations).isEqualTo(1); assertThat(context.client.calls).hasSize(2); assertThat(context.client.call(1).request.getClientId()).isZero(); } @@ -755,6 +757,152 @@ public void testSubscriptionsReconnectIndependently() throws Exception { } } + @Test + public void testBlockedReconnectDoesNotDelayAnotherSubscription() throws Exception { + BlockingReconnectKvClient testClient = new BlockingReconnectKvClient(getPdConfig()); + try { + testClient.listen("first", response -> { }); + testClient.listen("second", response -> { }); + + testClient.observer(0).onError(Status.UNAVAILABLE.asRuntimeException()); + testClient.observer(1).onError(Status.UNAVAILABLE.asRuntimeException()); + + Assert.assertTrue(testClient.blockedReconnectStarted.await(5L, TimeUnit.SECONDS)); + Assert.assertTrue(testClient.otherReconnectStarted.await(5L, TimeUnit.SECONDS)); + } finally { + testClient.allowBlockedReconnect.countDown(); + testClient.close(); + } + } + + @Test + public void testUnreachablePeerDoesNotBlockOtherSubscriptionReconnect() throws Exception { + AtomicReference liveAddress = new AtomicReference<>(); + MultiWatchService watchService = new MultiWatchService(); + HangingMembersService hangingService = new HangingMembersService(); + Server liveServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(liveAddress)) + .addService(watchService) + .build() + .start(); + Server unreachableServer = ServerBuilder.forPort(0) + .addService(hangingService) + .build() + .start(); + liveAddress.set("127.0.0.1:" + liveServer.getPort()); + String unreachableAddress = "127.0.0.1:" + unreachableServer.getPort(); + PDConfig config = PDConfig.of(liveAddress.get() + "," + unreachableAddress) + .setAuthority(user, pwd); + config.setGrpcTimeOut(60_000L); + try (BoundedResetKvClient testClient = new BoundedResetKvClient(config)) { + testClient.listen("first", response -> { }); + testClient.listen("second", response -> { }); + Assert.assertTrue(watchService.initialWatchesStarted.await(5L, TimeUnit.SECONDS)); + + watchService.observer(0).onError(Status.UNAVAILABLE.asRuntimeException()); + watchService.observer(1).onError(Status.UNAVAILABLE.asRuntimeException()); + + Assert.assertTrue(hangingService.called.await(5L, TimeUnit.SECONDS)); + Assert.assertTrue(watchService.bothWatchesReconnected.await(5L, TimeUnit.SECONDS)); + assertThat(watchService.calls).hasValue(4); + assertThat(watchService.firstCalls).hasValue(2); + assertThat(watchService.secondCalls).hasValue(2); + } finally { + liveServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + unreachableServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + + @Test + public void testInitialConnectTriesHealthyPeerAfterUnresponsivePeer() throws Exception { + AtomicReference liveAddress = new AtomicReference<>(); + RecordingWatchService watchService = new RecordingWatchService(11L); + HangingMembersService hangingService = new HangingMembersService(); + Server liveServer = ServerBuilder.forPort(0) + .addService(new MutableLeaderService(liveAddress)) + .addService(watchService) + .build() + .start(); + Server unresponsiveServer = ServerBuilder.forPort(0) + .addService(hangingService) + .build() + .start(); + liveAddress.set("127.0.0.1:" + liveServer.getPort()); + String unresponsiveAddress = "127.0.0.1:" + unresponsiveServer.getPort(); + PDConfig config = PDConfig.of(unresponsiveAddress + "," + liveAddress.get()) + .setAuthority(user, pwd); + config.setGrpcTimeOut(60_000L); + try (BoundedResetKvClient testClient = new BoundedResetKvClient(config)) { + testClient.listen("key", response -> { }); + + Assert.assertTrue(hangingService.called.await(5L, TimeUnit.SECONDS)); + Assert.assertTrue(watchService.started.await(5L, TimeUnit.SECONDS)); + assertThat(watchService.calls).hasValue(1); + } finally { + liveServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + unresponsiveServer.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS); + } + } + + @Test + public void testAttemptTracksTransportUsedAfterConcurrentRotation() throws Exception { + ScheduledExecutorService reconnectScheduler = mock(ScheduledExecutorService.class); + Deque reconnectTasks = new ArrayDeque<>(); + doAnswer(invocation -> { + long delay = invocation.getArgument(1); + if (invocation.getArgument(2).toMillis(delay) == 1000L) { + reconnectTasks.addLast(invocation.getArgument(0)); + } + return mock(ScheduledFuture.class); + }).when(reconnectScheduler).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + AttemptRaceKvClient testClient = + new AttemptRaceKvClient(getPdConfig(), reconnectScheduler); + AtomicReference listenFailure = new AtomicReference<>(); + Thread listenThread = new Thread(() -> { + try { + testClient.listen("second", response -> { }); + } catch (Throwable throwable) { + listenFailure.set(throwable); + } + }, "test-watch-listen"); + try { + testClient.listen("first", response -> { }); + listenThread.start(); + Assert.assertTrue(testClient.secondAttemptStarted.await(5L, TimeUnit.SECONDS)); + + testClient.observer(0).onError(Status.UNAVAILABLE.asRuntimeException()); + reconnectTasks.removeFirst().run(); + assertThat(testClient.stubInvalidations).isEqualTo(1); + + testClient.allowSecondAttempt.countDown(); + listenThread.join(5000L); + Assert.assertFalse(listenThread.isAlive()); + assertThat(listenFailure.get()).isNull(); + + testClient.observer(1).onError(Status.UNAVAILABLE.asRuntimeException()); + reconnectTasks.removeFirst().run(); + assertThat(testClient.stubInvalidations).isEqualTo(2); + } finally { + testClient.allowSecondAttempt.countDown(); + listenThread.join(5000L); + testClient.close(); + } + } + + @Test + public void testLegacyStreamingCallOverrideRemainsDispatchTarget() throws Exception { + ScheduledExecutorService reconnectScheduler = mock(ScheduledExecutorService.class); + LegacyStreamingOverrideKvClient testClient = + new LegacyStreamingOverrideKvClient(getPdConfig(), reconnectScheduler); + try { + testClient.listen("key", response -> { }); + + assertThat(testClient.calls).hasValue(1); + } finally { + testClient.close(); + } + } + @Test public void testReconnectsDoNotSharePermit() throws Exception { try (WatchTestContext context = newWatchTestContext()) { @@ -872,7 +1020,7 @@ public void testCloseStopsScheduledReconnect() throws Exception { } @Test(timeout = 3000L) - public void testCloseWaitsForRunningReconnectAndClosesItsChannel() throws Exception { + public void testCloseDoesNotWaitForRunningReconnectAndClosesItsChannel() throws Exception { ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class); Deque reconnectTasks = new ArrayDeque<>(); doAnswer(invocation -> { @@ -901,7 +1049,7 @@ public void testCloseWaitsForRunningReconnectAndClosesItsChannel() throws Except closeThread = new Thread(testClient::close, "test-kv-client-close"); closeThread.start(); Assert.assertTrue(testClient.closeEntered.await(1L, TimeUnit.SECONDS)); - assertThat(testClient.closeReturned.await(200L, TimeUnit.MILLISECONDS)).isFalse(); + assertThat(testClient.closeReturned.await(1L, TimeUnit.SECONDS)).isTrue(); } finally { testClient.allowReconnectStart.countDown(); if (reconnectThread != null) { @@ -1061,6 +1209,60 @@ public void getMembers(Pdpb.GetMembersRequest request, } } + private static class HangingMembersService extends PDGrpc.PDImplBase { + + private final CountDownLatch called = new CountDownLatch(1); + + @Override + public void getMembers(Pdpb.GetMembersRequest request, + StreamObserver observer) { + this.called.countDown(); + // Keep the call open until the client-side reconnect deadline expires. + } + } + + private static class MultiWatchService extends KvServiceGrpc.KvServiceImplBase { + + private final AtomicInteger calls = new AtomicInteger(); + private final List> observers = + new CopyOnWriteArrayList<>(); + private final AtomicInteger firstCalls = new AtomicInteger(); + private final AtomicInteger secondCalls = new AtomicInteger(); + private final CountDownLatch initialWatchesStarted = new CountDownLatch(2); + private final CountDownLatch bothWatchesReconnected = new CountDownLatch(2); + + StreamObserver observer(int index) { + return this.observers.get(index); + } + + @Override + public void watch(WatchRequest request, StreamObserver observer) { + int call = this.calls.incrementAndGet(); + this.observers.add(observer); + observer.onNext(startingResponse(call)); + AtomicInteger keyCalls = request.getKey().equals("first") ? + this.firstCalls : this.secondCalls; + int keyCall = keyCalls.incrementAndGet(); + if (keyCall == 1) { + this.initialWatchesStarted.countDown(); + } else if (keyCall == 2) { + this.bothWatchesReconnected.countDown(); + } + } + } + + private static class BoundedResetKvClient extends KvClient { + + BoundedResetKvClient(PDConfig pdConfig) { + super(pdConfig); + } + + @Override + protected long stubResetTimeoutMillis() { + return 300L; + } + } + private static class RecordingWatchService extends KvServiceGrpc.KvServiceImplBase { private final long clientId; @@ -1168,6 +1370,7 @@ private static class CloseRaceKvClient extends KvClient { private final CountDownLatch allowReconnectStart = new CountDownLatch(1); private final CountDownLatch closeEntered = new CountDownLatch(1); private final CountDownLatch closeReturned = new CountDownLatch(1); + private final Channel initialChannel = mock(Channel.class); CloseRaceKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor, @@ -1180,9 +1383,11 @@ private static class CloseRaceKvClient extends KvClient { protected void streamingCall(MethodDescriptor method, ReqT request, StreamObserver responseObserver, - int retry) { + int retry, + Consumer attemptConsumer) { this.observer.set((StreamObserver) responseObserver); if (this.calls.incrementAndGet() == 1) { + attemptConsumer.accept(this.initialChannel); return; } this.reconnectStarted.countDown(); @@ -1195,6 +1400,7 @@ protected void streamingCall(MethodDescriptor method, throw new IllegalStateException("Reconnect was interrupted", e); } this.channel = this.reconnectChannel; + attemptConsumer.accept(this.reconnectChannel); } @Override @@ -1205,6 +1411,112 @@ public void close() { } } + private static class BlockingReconnectKvClient extends KvClient { + + private final List> observers = new ArrayList<>(); + private final AtomicInteger calls = new AtomicInteger(); + private final CountDownLatch blockedReconnectStarted = new CountDownLatch(1); + private final CountDownLatch otherReconnectStarted = new CountDownLatch(1); + private final CountDownLatch allowBlockedReconnect = new CountDownLatch(1); + private final Channel attemptChannel = mock(Channel.class); + + BlockingReconnectKvClient(PDConfig pdConfig) { + super(pdConfig); + } + + StreamObserver observer(int index) { + return this.observers.get(index); + } + + @Override + protected void streamingCall(MethodDescriptor method, + ReqT request, + StreamObserver responseObserver, + int retry, + Consumer attemptConsumer) { + synchronized (this.observers) { + this.observers.add((StreamObserver) responseObserver); + } + attemptConsumer.accept(this.attemptChannel); + int call = this.calls.incrementAndGet(); + if (call == 3) { + this.blockedReconnectStarted.countDown(); + try { + this.allowBlockedReconnect.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } else if (call == 4) { + this.otherReconnectStarted.countDown(); + } + } + } + + private static class AttemptRaceKvClient extends KvClient { + + private final List> observers = new ArrayList<>(); + private final AtomicInteger calls = new AtomicInteger(); + private final CountDownLatch secondAttemptStarted = new CountDownLatch(1); + private final CountDownLatch allowSecondAttempt = new CountDownLatch(1); + private Channel attemptChannel = mock(Channel.class); + private int stubInvalidations; + + AttemptRaceKvClient(PDConfig pdConfig, + ScheduledExecutorService reconnectScheduler) { + super(pdConfig, reconnectScheduler); + } + + StreamObserver observer(int index) { + return this.observers.get(index); + } + + @Override + protected void streamingCall(MethodDescriptor method, + ReqT request, + StreamObserver responseObserver, + int retry, + Consumer attemptConsumer) { + this.observers.add((StreamObserver) responseObserver); + if (this.calls.incrementAndGet() == 2) { + this.secondAttemptStarted.countDown(); + try { + this.allowSecondAttempt.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + attemptConsumer.accept(this.attemptChannel); + } + + @Override + protected boolean invalidateAsyncStub(Channel expectedChannel) { + if (expectedChannel != this.attemptChannel) { + return false; + } + this.stubInvalidations++; + this.attemptChannel = mock(Channel.class); + return true; + } + } + + private static class LegacyStreamingOverrideKvClient extends KvClient { + + private final AtomicInteger calls = new AtomicInteger(); + + LegacyStreamingOverrideKvClient(PDConfig pdConfig, + ScheduledExecutorService reconnectScheduler) { + super(pdConfig, reconnectScheduler); + } + + @Override + protected void streamingCall(MethodDescriptor method, + ReqT request, + StreamObserver responseObserver, + int retry) { + this.calls.incrementAndGet(); + } + } + private static class FailingChannel extends Channel { private final StatusRuntimeException failure; @@ -1307,6 +1619,7 @@ private static class TestKvClient extends KvClient { private RuntimeException nextRuntimeFailure; private int stubInvalidations; private long lockClientId; + private Channel attemptChannel = mock(Channel.class); TestKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) { super(pdConfig, reconnectExecutor); @@ -1358,7 +1671,10 @@ protected RespT blockingUnaryCall( protected void streamingCall(MethodDescriptor method, ReqT request, StreamObserver responseObserver, - int retry) throws PDException { + int retry, + Consumer attemptConsumer) + throws PDException { + attemptConsumer.accept(this.attemptChannel); this.calls.add(new WatchCall(method.getFullMethodName(), (WatchRequest) request, (StreamObserver) responseObserver)); @@ -1380,8 +1696,13 @@ protected void streamingCall(MethodDescriptor method, } @Override - protected void invalidateAsyncStub() { + protected boolean invalidateAsyncStub(Channel expectedChannel) { + if (expectedChannel != this.attemptChannel) { + return false; + } this.stubInvalidations++; + this.attemptChannel = mock(Channel.class); + return true; } } diff --git a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java index 5f8f96c274..f4ad3790d2 100644 --- a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java +++ b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java @@ -28,6 +28,7 @@ import java.util.Timer; import java.util.TimerTask; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -78,6 +79,9 @@ public class SchemaDriver { private static final AtomicReference INSTANCE = new AtomicReference<>(); + private static final Object LIFECYCLE_LOCK = new Object(); + private static boolean destroying; + private static CountDownLatch destroyCompletion = new CountDownLatch(0); // Client for accessing PD private final KvClient client; @@ -112,26 +116,65 @@ public static void init(PDConfig pdConfig) { init(pdConfig, 300, 300 * 1000); } - public static synchronized void init(PDConfig pdConfig, int cacheSize, long expiration) { - SchemaDriver instance = INSTANCE.get(); - if (instance != null) { - throw new NotAllowException( - "The SchemaDriver [cacheSize=%s, expiration=%s, " + - "client=%s] has already been initialized and is not " + - "allowed to be initialized again", instance.caches.limit(), - instance.caches.expiration(), instance.client); + public static void init(PDConfig pdConfig, int cacheSize, long expiration) { + synchronized (LIFECYCLE_LOCK) { + SchemaDriver instance = INSTANCE.get(); + if (instance != null) { + throw new NotAllowException( + "The SchemaDriver [cacheSize=%s, expiration=%s, " + + "client=%s] has already been initialized and is not " + + "allowed to be initialized again", instance.caches.limit(), + instance.caches.expiration(), instance.client); + } + INSTANCE.set(new SchemaDriver(pdConfig, cacheSize, expiration)); } - INSTANCE.set(new SchemaDriver(pdConfig, cacheSize, expiration)); } - public static synchronized void destroy() { - SchemaDriver instance = INSTANCE.get(); - if (instance != null) { - try { - instance.closeResources(); - } finally { + public static void destroy() { + SchemaDriver instance = null; + CountDownLatch completion; + boolean closeResources = false; + synchronized (LIFECYCLE_LOCK) { + if (destroying) { + completion = destroyCompletion; + } else { + instance = INSTANCE.get(); + if (instance == null) { + return; + } + destroying = true; + completion = new CountDownLatch(1); + destroyCompletion = completion; + closeResources = true; + } + } + if (!closeResources) { + awaitDestroy(completion); + return; + } + try { + instance.closeResources(); + } finally { + synchronized (LIFECYCLE_LOCK) { INSTANCE.compareAndSet(instance, null); + destroying = false; } + completion.countDown(); + } + } + + private static void awaitDestroy(CountDownLatch completion) { + boolean interrupted = false; + while (true) { + try { + completion.await(); + break; + } catch (InterruptedException ignored) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); } } diff --git a/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java index 6edd7900f9..2c8bd5bc75 100644 --- a/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java +++ b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java @@ -24,20 +24,37 @@ import java.util.function.Consumer; import org.apache.hugegraph.exception.HugeException; +import org.apache.hugegraph.exception.NotAllowException; import org.apache.hugegraph.pd.client.KvClient; import org.apache.hugegraph.pd.client.PDConfig; import org.apache.hugegraph.pd.common.PDException; import org.apache.hugegraph.pd.grpc.kv.WatchResponse; import org.junit.Assert; +import org.junit.After; +import org.junit.Before; import org.junit.Test; public class SchemaDriverTest { + private AtomicReference instance; + private SchemaDriver previousInstance; + + @Before + public void setUp() throws Exception { + this.instance = instanceReference(); + this.previousInstance = this.instance.getAndSet(null); + } + + @After + public void tearDown() { + this.instance.set(this.previousInstance); + } + @Test - public void testDestroyClosesOwnedKvClient() throws Exception { + public void testDestroyClosesOwnedKvClient() { TrackingKvClient client = new TrackingKvClient(); SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); - instanceReference().set(driver); + this.instance.set(driver); try { SchemaDriver.destroy(); @@ -45,27 +62,24 @@ public void testDestroyClosesOwnedKvClient() throws Exception { Assert.assertTrue(client.closed); Assert.assertNull(SchemaDriver.getInstance()); } finally { - instanceReference().set(null); client.close(); } } - @Test(timeout = 3000L) + @Test public void testDestroyKeepsInstanceUntilResourcesAreClosed() throws Exception { BlockingCloseKvClient client = new BlockingCloseKvClient(); SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); - AtomicReference instance = instanceReference(); - SchemaDriver previous = instance.getAndSet(driver); + this.instance.set(driver); Thread destroyThread = new Thread(SchemaDriver::destroy, "schema-driver-destroy"); try { destroyThread.start(); - Assert.assertTrue(client.closeStarted.await(1L, TimeUnit.SECONDS)); + client.awaitCloseStarted(); Assert.assertSame(driver, SchemaDriver.getInstance()); } finally { client.allowClose.countDown(); - destroyThread.join(1000L); - instance.set(previous); + destroyThread.join(5000L); client.close(); } @@ -73,6 +87,74 @@ public void testDestroyKeepsInstanceUntilResourcesAreClosed() throws Exception { Assert.assertTrue(client.closed); } + @Test + public void testInitDoesNotWaitForDestroyCleanup() throws Exception { + BlockingCloseKvClient client = new BlockingCloseKvClient(); + SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); + this.instance.set(driver); + AtomicReference initFailure = new AtomicReference<>(); + CountDownLatch initFinished = new CountDownLatch(1); + Thread destroyThread = new Thread(SchemaDriver::destroy, "schema-driver-destroy"); + Thread initThread = new Thread(() -> { + try { + SchemaDriver.init(PDConfig.of("127.0.0.1:8686"), 10, 60_000L); + } catch (Throwable throwable) { + initFailure.set(throwable); + } finally { + initFinished.countDown(); + } + }, "schema-driver-init"); + boolean initCompletedDuringCleanup; + try { + destroyThread.start(); + client.awaitCloseStarted(); + initThread.start(); + initCompletedDuringCleanup = initFinished.await(1L, TimeUnit.SECONDS); + } finally { + client.allowClose.countDown(); + destroyThread.join(5000L); + initThread.join(5000L); + client.close(); + } + + Assert.assertTrue(initCompletedDuringCleanup); + Assert.assertTrue(initFailure.get() instanceof NotAllowException); + Assert.assertFalse(destroyThread.isAlive()); + Assert.assertFalse(initThread.isAlive()); + } + + @Test + public void testConcurrentDestroyWaitsForCleanup() throws Exception { + BlockingCloseKvClient client = new BlockingCloseKvClient(); + SchemaDriver driver = new SchemaDriver(client, 10, 60_000L); + this.instance.set(driver); + CountDownLatch secondDestroyStarted = new CountDownLatch(1); + CountDownLatch secondDestroyReturned = new CountDownLatch(1); + Thread firstDestroy = new Thread(SchemaDriver::destroy, "schema-driver-destroy-first"); + Thread secondDestroy = new Thread(() -> { + secondDestroyStarted.countDown(); + SchemaDriver.destroy(); + secondDestroyReturned.countDown(); + }, "schema-driver-destroy-second"); + boolean returnedBeforeCleanup; + try { + firstDestroy.start(); + client.awaitCloseStarted(); + secondDestroy.start(); + Assert.assertTrue(secondDestroyStarted.await(5L, TimeUnit.SECONDS)); + returnedBeforeCleanup = secondDestroyReturned.await(500L, TimeUnit.MILLISECONDS); + } finally { + client.allowClose.countDown(); + firstDestroy.join(5000L); + secondDestroy.join(5000L); + client.close(); + } + + Assert.assertFalse(returnedBeforeCleanup); + Assert.assertFalse(firstDestroy.isAlive()); + Assert.assertFalse(secondDestroy.isAlive()); + } + @Test public void testConstructorFailureClosesOwnedKvClient() { FailingListenKvClient client = new FailingListenKvClient(); @@ -124,15 +206,17 @@ private static class BlockingCloseKvClient extends TrackingKvClient { public void close() { this.closeStarted.countDown(); try { - if (!this.allowClose.await(1L, TimeUnit.SECONDS)) { - throw new IllegalStateException("Timed out waiting to close client"); - } + this.allowClose.await(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IllegalStateException("Client close was interrupted", e); } super.close(); } + + private void awaitCloseStarted() throws InterruptedException { + Assert.assertTrue(this.closeStarted.await(5L, TimeUnit.SECONDS)); + } } private static class FailingListenKvClient extends TrackingKvClient {