You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hugegraph.apache.org by ji...@apache.org on 2026/09/17 20:38:38 UTC
(hugegraph) branch master updated: fix(pd): keep KvClient watches alive after reconnect failures (#3157)
This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hugegraph.git
The following commit(s) were added to refs/heads/master by this push:
new 4f5202b81 fix(pd): keep KvClient watches alive after reconnect failures (#3157)
4f5202b81 is described below
commit 4f5202b81bbe5d2704cb20b425f7ce36b5c1adc6
Author: contrueCT <co...@outlook.com>
AuthorDate: Fri Sep 18 04:38:33 2026 +0800
fix(pd): keep KvClient watches alive after reconnect failures (#3157)
KvClient previously retried a failed watch only once. If that reconnect also failed, the
watch stopped permanently, and unexpected stream completion did not trigger recovery.
The reconnect path also exposed a lifecycle issue: AbstractClient.resetStub() invoked the
virtual close() method, so transport initialization on a KvClient dispatched to
KvClient.close() and marked the whole client closed.
---------
Co-authored-by: imbajin <ji...@apache.org>
---
hugegraph-pd/docs/api-reference.md | 19 +
.../apache/hugegraph/pd/client/AbstractClient.java | 194 ++-
.../org/apache/hugegraph/pd/client/KvClient.java | 565 +++++--
.../apache/hugegraph/pd/client/KvClientTest.java | 1650 ++++++++++++++++++++
.../java/org/apache/hugegraph/SchemaDriver.java | 102 +-
.../org/apache/hugegraph/SchemaDriverTest.java | 233 +++
6 files changed, 2611 insertions(+), 152 deletions(-)
diff --git a/hugegraph-pd/docs/api-reference.md b/hugegraph-pd/docs/api-reference.md
index f7c37be9d..8fbd28a65 100644
--- a/hugegraph-pd/docs/api-reference.md
+++ b/hugegraph-pd/docs/api-reference.md
@@ -756,6 +756,25 @@ for (Map.Entry<String, byte[]> entry : results.entrySet()) {
}
```
+### KV watch recovery
+
+`KvClient.listen()` and `listenPrefix()` restore subscriptions after retryable
+stream failures, including repeated reconnect failures and stream completion.
+Recovery resumes future notifications; events emitted while disconnected are
+not replayed. Consumers that require convergence must reconcile against durable
+PD state (see [graph metadata reconciliation #3151](https://github.com/apache/hugegraph/issues/3151)
+and [schema cache discussion #3205](https://github.com/apache/hugegraph/discussions/3205)).
+
+The overloads accepting an `errorConsumer` report permanent subscription failure,
+such as authentication or permission errors. The two-argument overloads log these
+failures. An initial synchronous registration failure still throws `PDException`.
+Returning from `listen()` does not mean the server has acknowledged registration.
+Close the owning `KvClient` to stop its watches and reconnect workers.
+
+Watch discovery uses at most a five-second budget, capped by `grpcTimeOut`; a
+separate five-second timer retries streams that receive no first response. These
+watch limits do not replace the configured timeout for blocking KV or lock calls.
+
## REST API
PD exposes a REST API for management and monitoring (default port: 8620).
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 570bacc6b..44ca5dcce 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,10 +23,10 @@ import java.util.concurrent.ConcurrentHashMap;
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;
-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;
@@ -58,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<String, AbstractBlockingStub> stubs = null;
+ private final ThreadLocal<Consumer<Channel>> streamingAttemptConsumer = new ThreadLocal<>();
protected AbstractClient(PDConfig config) {
String[] hosts = config.getServerHost().split(",");
@@ -85,7 +86,12 @@ public abstract class AbstractClient implements Closeable {
}
public static <T extends AbstractStub> T setBlockingParams(T stub, PDConfig config) {
- stub = (T) stub.withDeadlineAfter(config.getGrpcTimeOut(), TimeUnit.MILLISECONDS)
+ return setBlockingParams(stub, config, config.getGrpcTimeOut());
+ }
+
+ private static <T extends AbstractStub> 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()));
@@ -98,67 +104,125 @@ public abstract class AbstractClient implements Closeable {
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(stubResetTimeoutMillis());
+ 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(asyncStubResetTimeoutMillis());
+ if (host.isEmpty()) {
+ throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE,
+ "PD unreachable, pd.peers=" + config.getServerHost());
}
}
return setAsyncParams(proxy.getStub(), config);
}
+ 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 long asyncStubResetTimeoutMillis() {
+ return stubResetTimeoutMillis();
+ }
+
+ protected boolean isShutdown() {
+ return false;
+ }
+
protected abstract AbstractStub createStub();
protected abstract AbstractBlockingStub createBlockingStub();
- private String resetStub() {
- String leaderHost = "";
+ private String resetStub(long timeoutMillis) {
Exception ex = null;
+ long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(timeoutMillis);
+ long deadlineNanos = System.nanoTime() + timeoutNanos;
for (int i = 0; i < proxy.getHostCount(); i++) {
+ if (isShutdown() || remainingMillis(deadlineNanos) <= 0L) {
+ break;
+ }
String host = proxy.nextHost();
- close();
+ 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();
- leaderHost = leader.getGrpcUrl();
+ String leaderHost = leader.getGrpcUrl();
if (!host.equals(leaderHost)) {
- close();
- channel = ManagedChannelBuilder.forTarget(leaderHost).usePlaintext().build();
+ closeConnections();
+ 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);
- break;
+ return leaderHost;
} catch (StatusRuntimeException se) {
ex = se;
- continue;
} catch (Exception e) {
ex = e;
String msg =
@@ -166,11 +230,24 @@ public abstract class AbstractClient implements Closeable {
e.getMessage());
log.error(msg, e);
}
+ proxy.setBlockingStub(null);
+ proxy.setStub(null);
}
- if (StringUtils.isEmpty(leaderHost) && ex != null) {
+ proxy.setBlockingStub(null);
+ proxy.setStub(null);
+ closeConnections();
+ if (ex != null) {
log.error(String.format("connect to %s with error: ", config.getServerHost()), ex);
}
- return leaderHost;
+ 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 <ReqT, RespT> RespT blockingUnaryCall(
@@ -249,7 +326,17 @@ public abstract class AbstractClient implements Closeable {
protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method, ReqT request,
StreamObserver<RespT> 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<Channel> attemptConsumer = this.streamingAttemptConsumer.get();
+ if (attemptConsumer != null) {
+ attemptConsumer.accept(attemptChannel);
+ }
+ }
try {
ClientCall<ReqT, RespT> call = stub.getChannel().newCall(method, stub.getCallOptions());
ClientCalls.asyncServerStreamingCall(call, request, responseObserver);
@@ -257,31 +344,52 @@ public abstract class AbstractClient implements Closeable {
log.error("rpc call with exception :", e);
if (e instanceof StatusRuntimeException) {
if (retry < proxy.getHostCount()) {
- synchronized (this) {
- proxy.setStub(null);
- }
+ invalidateAsyncStub(attemptChannel);
streamingCall(method, request, responseObserver, ++retry);
+ return;
}
}
+ throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE,
+ "RPC streaming call failed", e);
+ }
+ }
+
+ protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method, ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry,
+ Consumer<Channel> attemptConsumer)
+ throws PDException {
+ Consumer<Channel> 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();
+ }
+
+ private void closeConnections() {
closeChannel(channel);
if (stubs != null) {
for (AbstractBlockingStub stub : stubs.values()) {
closeChannel((ManagedChannel) stub.getChannel());
}
}
-
}
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 cbc15bfe3..e513fa62e 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,12 +22,19 @@ import java.util.HashMap;
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.Executor;
+import java.util.concurrent.ExecutorService;
+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;
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;
@@ -48,6 +55,10 @@ import org.apache.hugegraph.pd.grpc.kv.WatchRequest;
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;
import io.grpc.stub.AbstractBlockingStub;
import io.grpc.stub.AbstractStub;
import io.grpc.stub.StreamObserver;
@@ -56,13 +67,58 @@ import lombok.extern.slf4j.Slf4j;
@Slf4j
public class KvClient<T extends WatchResponse> extends AbstractClient implements Closeable {
- private AtomicLong clientId = new AtomicLong(0);
- private Semaphore semaphore = new Semaphore(1);
- private AtomicBoolean closed = new AtomicBoolean(false);
- private Set<StreamObserver> observers = ConcurrentHashMap.newKeySet();
+ 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<Status.Code> NON_RETRYABLE_WATCH_ERRORS =
+ Set.of(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 Semaphore lockSemaphore = new Semaphore(1);
+ private final AtomicBoolean closed = new AtomicBoolean(false);
+ private final Set<WatchSubscription> subscriptions = ConcurrentHashMap.newKeySet();
+ private final ScheduledExecutorService reconnectScheduler;
+ private final Executor reconnectExecutor;
public KvClient(PDConfig pdConfig) {
+ 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 reconnectScheduler,
+ Executor reconnectExecutor) {
super(pdConfig);
+ this.reconnectScheduler = reconnectScheduler;
+ this.reconnectExecutor = reconnectExecutor;
}
@Override
@@ -75,6 +131,16 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
return KvServiceGrpc.newBlockingStub(channel);
}
+ @Override
+ protected long asyncStubResetTimeoutMillis() {
+ 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);
@@ -127,7 +193,6 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
private void onEvent(WatchResponse value, Consumer<T> consumer) {
log.debug("receive message for {},event Count:{}", value, value.getEventsCount());
- clientId.compareAndSet(0L, value.getClientId());
if (value.getEventsCount() != 0) {
try {
consumer.accept((T) value);
@@ -140,35 +205,25 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
}
}
- private StreamObserver<WatchResponse> getObserver(String key, Consumer<T> consumer,
- BiConsumer<String, Consumer> listenWrapper,
- long client) {
- StreamObserver<WatchResponse> observer = getObserver(key, consumer, listenWrapper);
- observers.add(observer);
- return observer;
- }
-
- private StreamObserver<WatchResponse> getObserver(String key, Consumer<T> consumer,
- BiConsumer<String, Consumer> listenWrapper) {
+ private StreamObserver<WatchResponse> getObserver(WatchSubscription subscription) {
return new StreamObserver<WatchResponse>() {
@Override
public void onNext(WatchResponse value) {
+ if (!acceptFirstFrame(subscription, this)) {
+ return;
+ }
switch (value.getState()) {
case Starting:
- boolean b = clientId.compareAndSet(0, value.getClientId());
+ boolean b = subscription.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, true);
break;
case Alive:
// only for check client is alive, do nothing
@@ -180,50 +235,352 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
@Override
public void onError(Throwable t) {
- release();
- if (!closed.get()) {
- clientId.set(0);
- listenWrapper.accept(key, consumer);
+ if (isRetryableWatchError(t)) {
+ requestReconnect(subscription, this, shouldRotateWatchTransport(t));
+ } else {
+ stopWatch(subscription, this, t);
}
}
@Override
public void onCompleted() {
-
+ requestReconnect(subscription, this, false);
}
};
}
public void listen(String key, Consumer<T> consumer) throws PDException {
- long value = clientId.get();
- StreamObserver<WatchResponse> observer = getObserver(key, consumer, listenWrapper, value);
- acquire();
+ listen(key, consumer, throwable -> { }, false);
+ }
+
+ public void listen(String key, Consumer<T> consumer,
+ Consumer<Throwable> errorConsumer) throws PDException {
+ listen(key, consumer, errorConsumer, false);
+ }
+
+ public void listenPrefix(String prefix, Consumer<T> consumer) throws PDException {
+ listen(prefix, consumer, throwable -> { }, true);
+ }
+
+ public void listenPrefix(String prefix, Consumer<T> consumer,
+ Consumer<Throwable> errorConsumer) throws PDException {
+ listen(prefix, consumer, errorConsumer, true);
+ }
+
+ private void listen(String key, Consumer<T> consumer,
+ Consumer<Throwable> 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 {
- 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) {
+ cleanupFailedStart(subscription);
+ subscriptions.remove(subscription);
+ throw e;
+ } catch (RuntimeException e) {
+ cleanupFailedStart(subscription);
+ subscriptions.remove(subscription);
+ throw e;
}
}
- public void listenPrefix(String prefix, Consumer<T> consumer) throws PDException {
- long value = clientId.get();
- StreamObserver<WatchResponse> observer =
- getObserver(prefix, consumer, prefixListenWrapper, value);
- acquire();
+ private void cleanupFailedStart(WatchSubscription subscription) {
+ synchronized (subscription) {
+ subscription.observer.set(null);
+ cancelStartTimeout(subscription);
+ }
+ }
+
+ private boolean startWatch(WatchSubscription subscription) throws PDException {
+ if (closed.get() || !subscriptions.contains(subscription)) {
+ return false;
+ }
+
+ WatchRequest request = WatchRequest.newBuilder()
+ .setClientId(subscription.clientId.get())
+ .setKey(subscription.key)
+ .build();
+ StreamObserver<WatchResponse> observer = getObserver(subscription);
try {
- WatchRequest k =
- WatchRequest.newBuilder().setClientId(clientId.get()).setKey(prefix).build();
- streamingCall(KvServiceGrpc.getWatchPrefixMethod(), k, observer, 1);
+ synchronized (subscription) {
+ if (closed.get() || !subscriptions.contains(subscription)) {
+ return false;
+ }
+ 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) {
- release();
+ if (closed.get()) {
+ cleanupFailedStart(subscription);
+ super.close();
+ return false;
+ }
+ if (e instanceof PDException) {
+ throw (PDException) e;
+ }
throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, e);
}
}
- private void acquire() {
+ private void recordAttemptChannel(WatchSubscription subscription,
+ StreamObserver<WatchResponse> sourceObserver,
+ Channel channel) {
+ synchronized (subscription) {
+ if (subscription.observer.get() == sourceObserver) {
+ subscription.attemptChannel = channel;
+ }
+ }
+ }
+
+ private boolean acceptFirstFrame(WatchSubscription subscription,
+ StreamObserver<WatchResponse> sourceObserver) {
+ synchronized (subscription) {
+ if (subscription.observer.get() != sourceObserver) {
+ return false;
+ }
+ subscription.firstFrameReceived = true;
+ cancelStartTimeout(subscription);
+ return true;
+ }
+ }
+
+ private void scheduleStartTimeout(WatchSubscription subscription,
+ StreamObserver<WatchResponse> sourceObserver)
+ throws PDException {
+ ScheduledFuture<?> timeout;
+ try {
+ timeout = reconnectScheduler.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<WatchResponse> 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, true);
+ }
+
+ private void requestReconnect(WatchSubscription subscription,
+ StreamObserver<WatchResponse> sourceObserver,
+ boolean rotateTransport) {
+ requestReconnect(subscription, sourceObserver, rotateTransport, false);
+ }
+
+ private void requestReconnect(WatchSubscription subscription,
+ StreamObserver<WatchResponse> sourceObserver,
+ boolean rotateTransport,
+ boolean requireMissingFirstFrame) {
+ synchronized (subscription) {
+ if (closed.get() || subscription.observer.get() != sourceObserver ||
+ (requireMissingFirstFrame && subscription.firstFrameReceived)) {
+ return;
+ }
+ subscription.observer.set(null);
+ cancelStartTimeout(subscription);
+ subscription.clientId.set(0L);
+ if (rotateTransport) {
+ subscription.rotateTransport = true;
+ subscription.reconnectChannel = subscription.attemptChannel;
+ }
+ }
+ scheduleReconnect(subscription);
+ }
+
+ private void invalidateAttemptStub(Channel attemptChannel) {
+ if (attemptChannel != null) {
+ invalidateAsyncStub(attemptChannel);
+ }
+ }
+
+ 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) {
+ 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,
+ StreamObserver<WatchResponse> sourceObserver,
+ Throwable throwable) {
+ synchronized (subscription) {
+ if (subscription.observer.get() != sourceObserver) {
+ return;
+ }
+ subscription.observer.set(null);
+ cancelStartTimeout(subscription);
+ }
+ 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() || !subscriptions.contains(subscription) ||
+ !subscription.reconnectScheduled.compareAndSet(false, true)) {
+ return;
+ }
+ try {
+ reconnectScheduler.schedule(() -> submitReconnect(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);
+ subscriptions.remove(subscription);
+ notifyWatchStopped(subscription, e);
+ }
+ }
+ }
+
+ 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);
+ }
+ } catch (RuntimeException | PDException e) {
+ StreamObserver<WatchResponse> observer = subscription.observer.get();
+ 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 void acquire(AtomicLong clientId, Semaphore semaphore) {
if (clientId.get() == 0L) {
try {
semaphore.acquire();
@@ -237,7 +594,7 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
}
}
- private void release() {
+ private void release(Semaphore semaphore) {
try {
if (semaphore.availablePermits() == 0) {
semaphore.release();
@@ -276,105 +633,127 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
}
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;
}
@Override
public void close() {
- for (StreamObserver o : observers) {
+ if (!closed.compareAndSet(false, true)) {
+ return;
+ }
+ reconnectScheduler.shutdownNow();
+ if (reconnectExecutor instanceof ExecutorService) {
+ ((ExecutorService) reconnectExecutor).shutdownNow();
+ }
+ release(lockSemaphore);
+ for (WatchSubscription subscription : subscriptions) {
try {
- if (o != null) {
- o.onCompleted();
+ StreamObserver<WatchResponse> 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);
}
}
- observers.clear();
- closed.set(true);
+ subscriptions.clear();
super.close();
}
- BiConsumer<String, Consumer> listenWrapper = (key, consumer) -> {
- try {
- listen(key, consumer);
- } catch (PDException e) {
- try {
- log.warn("start listen with warning:", e);
- Thread.sleep(1000);
- } catch (InterruptedException ex) {
- }
- }
- };
-
- BiConsumer<String, Consumer> 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 class WatchSubscription {
+
+ private final String key;
+ private final Consumer<T> consumer;
+ private final Consumer<Throwable> errorConsumer;
+ private final boolean prefix;
+ private final AtomicReference<StreamObserver<WatchResponse>> observer;
+ private final AtomicBoolean reconnectScheduled;
+ private final AtomicBoolean terminated;
+ private final AtomicLong clientId;
+ private final AtomicReference<ScheduledFuture<?>> startTimeout;
+ private boolean firstFrameReceived;
+ private Channel attemptChannel;
+ private boolean rotateTransport;
+ private Channel reconnectChannel;
+
+ private WatchSubscription(String key, Consumer<T> consumer,
+ Consumer<Throwable> 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 ae44fb234..8c28f854e 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,74 @@
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 static org.mockito.Mockito.when;
+import java.io.Serializable;
+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.CountDownLatch;
+import java.util.concurrent.CopyOnWriteArrayList;
+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;
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.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;
+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.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;
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;
+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;
public class KvClientTest extends BaseClientTest {
@@ -45,6 +96,11 @@ public class KvClientTest extends BaseClientTest {
client = new KvClient<>(getPdConfig());
}
+ @After
+ public void tearDown() {
+ client.close();
+ }
+
@Test
public void testCreateStub() {
// Setup
@@ -67,6 +123,984 @@ public class KvClientTest extends BaseClientTest {
}
}
+ @Test
+ public void testTransportInitializationDoesNotCloseClient() throws Exception {
+ AtomicReference<String> grpcAddress = new AtomicReference<>();
+ Server server = ServerBuilder.forPort(0)
+ .addService(new PDGrpc.PDImplBase() {
+ @Override
+ public void getMembers(
+ Pdpb.GetMembersRequest request,
+ StreamObserver<Pdpb.GetMembersResponse> 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 testStubCreationFailureIsReportedAsPdException() throws Exception {
+ AtomicReference<String> grpcAddress = new AtomicReference<>();
+ Server server = ServerBuilder.forPort(0)
+ .addService(new PDGrpc.PDImplBase() {
+ @Override
+ public void getMembers(
+ Pdpb.GetMembersRequest request,
+ StreamObserver<Pdpb.GetMembersResponse> 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);
+ 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 {
+ AtomicReference<String> firstAddress = new AtomicReference<>();
+ AtomicReference<String> 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();
+ PDConfig config = PDConfig.of(firstAddress.get() + "," + secondAddress.get())
+ .setAuthority(user, pwd);
+ try (FailFirstStreamingKvClient testClient =
+ new FailFirstStreamingKvClient(config, reconnectExecutor, failure)) {
+ testClient.listen("key", response -> { });
+
+ 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);
+ }
+ }
+
+ @Test
+ public void testReconnectRetriesAfterFirstFailure() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ Consumer<WatchResponse> 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.reconnectDelaysMs).containsExactly(1000L, 1000L);
+ 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();
+ 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<WatchResponse> 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.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<WatchResponse> 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 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)) {
+ 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);
+ context.runNextReconnect();
+ assertThat(context.client.stubInvalidations).isEqualTo(1);
+ }
+ }
+ }
+
+ @Test
+ public void testLeaderChangedSchedulesReconnect() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ context.client.listen("key", response -> { });
+ StreamObserver<WatchResponse> 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).containsExactly(1000L);
+ context.runNextReconnect();
+ assertThat(context.client.stubInvalidations).isEqualTo(1);
+ assertThat(context.client.calls).hasSize(2);
+ assertThat(context.client.call(1).request.getClientId()).isZero();
+ }
+ }
+
+ @Test
+ public void testLeaderChangedReconnectsToNewPeer() throws Exception {
+ AtomicReference<String> 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<Runnable> reconnectTasks = new ArrayDeque<>();
+ CountDownLatch reconnectScheduled = new CountDownLatch(1);
+ doAnswer(invocation -> {
+ long delay = invocation.getArgument(1);
+ if (invocation.<TimeUnit>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<WatchResponse> 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 testFollowerErrorReconnectsToNewPeer() throws Exception {
+ AtomicReference<String> firstAddress = new AtomicReference<>();
+ AtomicReference<String> 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<Runnable> reconnectTasks = new ArrayDeque<>();
+ CountDownLatch reconnectScheduled = new CountDownLatch(1);
+ doAnswer(invocation -> {
+ long delay = invocation.getArgument(1);
+ if (invocation.<TimeUnit>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<WatchResponse> 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()) {
+ context.client.listen("key", response -> { });
+
+ assertThat(context.watchTimeoutTasks).hasSize(1);
+ assertThat(context.watchTimeoutDelaysMs).containsExactly(5000L);
+ context.runNextWatchTimeout();
+
+ assertThat(context.reconnectTasks).hasSize(1);
+ context.runNextReconnect();
+ assertThat(context.client.stubInvalidations).isEqualTo(1);
+ 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 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<Runnable> 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<WatchResponse> consumer = mock(Consumer.class);
+ context.client.listen("key", consumer);
+ StreamObserver<WatchResponse> 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()) {
+ 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<Throwable> 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<Throwable> terminalErrors = new ArrayList<>();
+ context.client.listenPrefix("prefix", response -> { }, terminalErrors::add);
+ StatusRuntimeException failure = Status.PERMISSION_DENIED.asRuntimeException();
+ StreamObserver<WatchResponse> 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<WatchResponse> 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).isEqualTo(1);
+ 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()) {
+ context.client.listen("key", response -> { });
+
+ context.client.call(0).observer.onCompleted();
+
+ assertThat(context.reconnectTasks).hasSize(1);
+ assertThat(context.reconnectDelaysMs).containsExactly(1000L);
+ assertThat(context.client.stubInvalidations).isZero();
+ context.runNextReconnect();
+ assertThat(context.client.calls).hasSize(2);
+ }
+ }
+
+ @Test
+ public void testObserverSchedulesOnlyOneReconnect() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ context.client.listen("key", response -> { });
+ StreamObserver<WatchResponse> observer = context.client.call(0).observer;
+
+ observer.onError(new RuntimeException("disconnected"));
+ observer.onCompleted();
+
+ assertThat(context.reconnectTasks).hasSize(1);
+ assertThat(context.reconnectDelaysMs).containsExactly(1000L);
+ }
+ }
+
+ @Test
+ public void testSubscriptionsReconnectIndependently() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ Consumer<WatchResponse> firstConsumer = mock(Consumer.class);
+ Consumer<WatchResponse> secondConsumer = mock(Consumer.class);
+ 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());
+
+ 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");
+ 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);
+
+ verify(firstConsumer).accept(firstEvent);
+ verify(secondConsumer).accept(secondEvent);
+ verify(firstConsumer, never()).accept(secondEvent);
+ verify(secondConsumer, never()).accept(firstEvent);
+ }
+ }
+
+ @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<String> 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<String> 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<Runnable> reconnectTasks = new ArrayDeque<>();
+ doAnswer(invocation -> {
+ long delay = invocation.getArgument(1);
+ if (invocation.<TimeUnit>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<Throwable> 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 testBlockingDiscoveryKeepsConfiguredTimeout() throws Exception {
+ AtomicReference<String> address = new AtomicReference<>();
+ AtomicReference<Long> remainingMillis = new AtomicReference<>();
+ Server server = ServerBuilder.forPort(0)
+ .addService(new PDGrpc.PDImplBase() {
+ @Override
+ public void getMembers(Pdpb.GetMembersRequest request,
+ StreamObserver<Pdpb.GetMembersResponse> observer) {
+ remainingMillis.set(io.grpc.Context.current()
+ .getDeadline()
+ .timeRemaining(TimeUnit.MILLISECONDS));
+ observer.onNext(Pdpb.GetMembersResponse.newBuilder()
+ .setLeader(Metapb.Member.newBuilder()
+ .setGrpcUrl(address.get()))
+ .build());
+ observer.onCompleted();
+ }
+ }).build().start();
+ address.set("127.0.0.1:" + server.getPort());
+ PDConfig config = PDConfig.of(address.get() + "," + address.get());
+ config.setGrpcTimeOut(60_000L);
+ try (KvClient<WatchResponse> blockingClient = new KvClient<>(config);
+ KvClient<WatchResponse> watchClient = new KvClient<>(config)) {
+ blockingClient.getBlockingStub();
+ assertThat(remainingMillis.get()).isGreaterThan(30_000L);
+
+ watchClient.getStub();
+ assertThat(remainingMillis.get()).isBetween(1L, 2500L);
+ } finally {
+ server.shutdownNow().awaitTermination(5L, TimeUnit.SECONDS);
+ }
+ }
+
+ @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()) {
+ 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();
+ context.runNextReconnect();
+
+ 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.client.call(2).request.getClientId()).isZero();
+ }
+ }
+
+ @Test
+ public void testStaleChannelErrorDoesNotInvalidateFreshStub() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ context.client.listen("first", response -> { });
+ context.client.listen("second", response -> { });
+ StreamObserver<WatchResponse> 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);
+ }
+ }
+
+ @Test
+ public void testStaleObserverDoesNotScheduleReconnect() throws Exception {
+ try (WatchTestContext context = newWatchTestContext()) {
+ Consumer<WatchResponse> consumer = mock(Consumer.class);
+ context.client.listen("key", consumer);
+ StreamObserver<WatchResponse> 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<WatchResponse> currentObserver = context.client.call(1).observer;
+ currentObserver.onNext(WatchResponse.newBuilder()
+ .setState(WatchState.Starting)
+ .setClientId(8L)
+ .build());
+ currentObserver.onNext(response);
+ verify(consumer).accept(response);
+ }
+ }
+
+ @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();
+ }
+ }
+
+ @Test(timeout = 3000L)
+ public void testCloseDoesNotWaitForRunningReconnectAndClosesItsChannel() throws Exception {
+ ScheduledExecutorService reconnectExecutor = mock(ScheduledExecutorService.class);
+ Deque<Runnable> 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(1L, TimeUnit.SECONDS)).isTrue();
+ } 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";
@@ -114,4 +1148,620 @@ public class KvClientTest extends BaseClientTest {
}
}
+
+ private static boolean isClosed(KvClient<?> client) throws Exception {
+ Field field = KvClient.class.getDeclaredField("closed");
+ field.setAccessible(true);
+ 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<Runnable> reconnectTasks = new ArrayDeque<>();
+ Deque<Runnable> watchTimeoutTasks = new ArrayDeque<>();
+ List<Long> reconnectDelaysMs = new ArrayList<>();
+ List<Long> watchTimeoutDelaysMs = new ArrayList<>();
+ doAnswer(invocation -> {
+ long delay = invocation.getArgument(1);
+ TimeUnit unit = invocation.getArgument(2);
+ 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, watchTimeoutTasks,
+ reconnectDelaysMs, watchTimeoutDelaysMs);
+ }
+
+ private static class InitializableKvClient extends KvClient<WatchResponse> {
+
+ InitializableKvClient(PDConfig pdConfig) {
+ super(pdConfig);
+ }
+
+ void initializeTransport() throws PDException {
+ getStub();
+ }
+ }
+
+ private static class FailingStubCreationKvClient extends KvClient<WatchResponse> {
+
+ 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<String> leaderAddress;
+ private final AtomicInteger calls = new AtomicInteger();
+
+ MutableLeaderService(AtomicReference<String> leaderAddress) {
+ this.leaderAddress = leaderAddress;
+ }
+
+ @Override
+ public void getMembers(Pdpb.GetMembersRequest request,
+ StreamObserver<Pdpb.GetMembersResponse> 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 HangingMembersService extends PDGrpc.PDImplBase {
+
+ private final CountDownLatch called = new CountDownLatch(1);
+
+ @Override
+ public void getMembers(Pdpb.GetMembersRequest request,
+ StreamObserver<Pdpb.GetMembersResponse> 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<StreamObserver<WatchResponse>> 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<WatchResponse> observer(int index) {
+ return this.observers.get(index);
+ }
+
+ @Override
+ public void watch(WatchRequest request, StreamObserver<WatchResponse> 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<WatchResponse> {
+
+ BoundedResetKvClient(PDConfig pdConfig) {
+ super(pdConfig);
+ }
+
+ @Override
+ protected long asyncStubResetTimeoutMillis() {
+ return 300L;
+ }
+ }
+
+ private static class RecordingWatchService extends KvServiceGrpc.KvServiceImplBase {
+
+ private final long clientId;
+ private final AtomicInteger calls = new AtomicInteger();
+ private final AtomicReference<StreamObserver<WatchResponse>> observer =
+ new AtomicReference<>();
+ private final CountDownLatch started = new CountDownLatch(1);
+
+ RecordingWatchService(long clientId) {
+ this.clientId = clientId;
+ }
+
+ @Override
+ public void watch(WatchRequest request, StreamObserver<WatchResponse> observer) {
+ this.calls.incrementAndGet();
+ this.observer.set(observer);
+ observer.onNext(startingResponse(this.clientId));
+ this.started.countDown();
+ }
+ }
+
+ 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<WatchResponse> observer) {
+ this.calls.incrementAndGet();
+ observer.onError(new PDException(-1, NOT_LEADER));
+ this.failed.countDown();
+ }
+ }
+
+ private static class TimeoutRaceAppender extends AbstractAppender {
+
+ private final AtomicReference<Runnable> action;
+ private final AtomicBoolean triggered = new AtomicBoolean(false);
+
+ TimeoutRaceAppender(AtomicReference<Runnable> action) {
+ super("KvClientTimeoutRaceAppender", (Filter) null,
+ (Layout<? extends Serializable>) 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<WatchResponse> {
+
+ private final Deque<Channel> 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 FailFirstStreamingKvClient extends KvClient<WatchResponse> {
+
+ private final StatusRuntimeException firstFailure;
+ private final AtomicInteger stubCreations = new AtomicInteger();
+
+ FailFirstStreamingKvClient(PDConfig pdConfig,
+ ScheduledExecutorService reconnectExecutor,
+ StatusRuntimeException firstFailure) {
+ super(pdConfig, reconnectExecutor);
+ this.firstFailure = firstFailure;
+ }
+
+ @Override
+ protected AbstractStub createStub() {
+ if (this.stubCreations.getAndIncrement() == 0) {
+ return KvServiceGrpc.newStub(new FailingChannel(this.firstFailure));
+ }
+ return super.createStub();
+ }
+ }
+
+ private static class CloseRaceKvClient extends KvClient<WatchResponse> {
+
+ private final ManagedChannel reconnectChannel;
+ private final AtomicInteger calls = new AtomicInteger();
+ private final AtomicReference<StreamObserver<WatchResponse>> 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);
+ private final Channel initialChannel = mock(Channel.class);
+
+ CloseRaceKvClient(PDConfig pdConfig,
+ ScheduledExecutorService reconnectExecutor,
+ ManagedChannel reconnectChannel) {
+ super(pdConfig, reconnectExecutor);
+ this.reconnectChannel = reconnectChannel;
+ }
+
+ @Override
+ protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method,
+ ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry,
+ Consumer<Channel> attemptConsumer) {
+ this.observer.set((StreamObserver<WatchResponse>) responseObserver);
+ if (this.calls.incrementAndGet() == 1) {
+ attemptConsumer.accept(this.initialChannel);
+ 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;
+ attemptConsumer.accept(this.reconnectChannel);
+ }
+
+ @Override
+ public void close() {
+ this.closeEntered.countDown();
+ super.close();
+ this.closeReturned.countDown();
+ }
+ }
+
+ private static class BlockingReconnectKvClient extends KvClient<WatchResponse> {
+
+ private final List<StreamObserver<WatchResponse>> 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<WatchResponse> observer(int index) {
+ return this.observers.get(index);
+ }
+
+ @Override
+ protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method,
+ ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry,
+ Consumer<Channel> attemptConsumer) {
+ synchronized (this.observers) {
+ this.observers.add((StreamObserver<WatchResponse>) 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<WatchResponse> {
+
+ private final List<StreamObserver<WatchResponse>> 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<WatchResponse> observer(int index) {
+ return this.observers.get(index);
+ }
+
+ @Override
+ protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method,
+ ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry,
+ Consumer<Channel> attemptConsumer) {
+ this.observers.add((StreamObserver<WatchResponse>) 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<WatchResponse> {
+
+ private final AtomicInteger calls = new AtomicInteger();
+
+ LegacyStreamingOverrideKvClient(PDConfig pdConfig,
+ ScheduledExecutorService reconnectScheduler) {
+ super(pdConfig, reconnectScheduler);
+ }
+
+ @Override
+ protected <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method,
+ ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry) {
+ this.calls.incrementAndGet();
+ }
+ }
+
+ private static class FailingChannel extends Channel {
+
+ private final StatusRuntimeException failure;
+
+ FailingChannel(StatusRuntimeException failure) {
+ this.failure = failure;
+ }
+
+ @Override
+ public String authority() {
+ return "test";
+ }
+
+ @Override
+ public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
+ MethodDescriptor<ReqT, RespT> method, CallOptions callOptions) {
+ return new FailingClientCall<>(this.failure);
+ }
+ }
+
+ private static class FailingClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
+
+ private final StatusRuntimeException failure;
+
+ FailingClientCall(StatusRuntimeException failure) {
+ this.failure = failure;
+ }
+
+ @Override
+ public void start(Listener<RespT> 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 WatchTestContext implements AutoCloseable {
+
+ private final TestKvClient client;
+ private final ScheduledExecutorService reconnectExecutor;
+ private final Deque<Runnable> reconnectTasks;
+ private final Deque<Runnable> watchTimeoutTasks;
+ private final List<Long> reconnectDelaysMs;
+ private final List<Long> watchTimeoutDelaysMs;
+
+ WatchTestContext(TestKvClient client,
+ ScheduledExecutorService reconnectExecutor,
+ Deque<Runnable> reconnectTasks,
+ Deque<Runnable> watchTimeoutTasks,
+ List<Long> reconnectDelaysMs,
+ List<Long> watchTimeoutDelaysMs) {
+ this.client = client;
+ this.reconnectExecutor = reconnectExecutor;
+ this.reconnectTasks = reconnectTasks;
+ this.watchTimeoutTasks = watchTimeoutTasks;
+ this.reconnectDelaysMs = reconnectDelaysMs;
+ this.watchTimeoutDelaysMs = watchTimeoutDelaysMs;
+ }
+
+ Runnable takeNextReconnect() {
+ assertThat(reconnectTasks).isNotEmpty();
+ return reconnectTasks.removeFirst();
+ }
+
+ void runNextReconnect() {
+ takeNextReconnect().run();
+ }
+
+ void runNextWatchTimeout() {
+ assertThat(watchTimeoutTasks).isNotEmpty();
+ watchTimeoutTasks.removeFirst().run();
+ }
+
+ @Override
+ public void close() {
+ client.close();
+ }
+ }
+
+ private static class TestKvClient extends KvClient<WatchResponse> {
+
+ private final List<WatchCall> calls = new ArrayList<>();
+ private final List<LockCall> lockCalls = new ArrayList<>();
+ private int remainingFailures;
+ private PDException nextFailure;
+ private RuntimeException nextRuntimeFailure;
+ private int stubInvalidations;
+ private long lockClientId;
+ private Channel attemptChannel = mock(Channel.class);
+
+ TestKvClient(PDConfig pdConfig, ScheduledExecutorService reconnectExecutor) {
+ super(pdConfig, reconnectExecutor);
+ }
+
+ 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);
+ }
+
+ void respondWithLockClientId(long clientId) {
+ this.lockClientId = clientId;
+ }
+
+ LockCall lockCall(int index) {
+ return this.lockCalls.get(index);
+ }
+
+ @Override
+ protected <ReqT, RespT> RespT blockingUnaryCall(
+ MethodDescriptor<ReqT, RespT> 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 <ReqT, RespT> void streamingCall(MethodDescriptor<ReqT, RespT> method,
+ ReqT request,
+ StreamObserver<RespT> responseObserver,
+ int retry,
+ Consumer<Channel> attemptConsumer)
+ throws PDException {
+ attemptConsumer.accept(this.attemptChannel);
+ this.calls.add(new WatchCall(method.getFullMethodName(),
+ (WatchRequest) request,
+ (StreamObserver<WatchResponse>) 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 boolean invalidateAsyncStub(Channel expectedChannel) {
+ if (expectedChannel != this.attemptChannel) {
+ return false;
+ }
+ this.stubInvalidations++;
+ this.attemptChannel = mock(Channel.class);
+ return true;
+ }
+ }
+
+ private static class WatchCall {
+
+ private final String methodName;
+ private final WatchRequest request;
+ private final StreamObserver<WatchResponse> observer;
+
+ WatchCall(String methodName, WatchRequest request,
+ StreamObserver<WatchResponse> observer) {
+ this.methodName = methodName;
+ this.request = request;
+ 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;
+ }
+ }
}
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 16273485c..11849b80a 100644
--- a/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java
+++ b/hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java
@@ -24,9 +24,11 @@ import java.util.ArrayList;
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;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
@@ -77,6 +79,9 @@ public class SchemaDriver {
private static final AtomicReference<SchemaDriver> 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<WatchResponse> client;
@@ -84,9 +89,23 @@ public class SchemaDriver {
private SchemaDriver(PDConfig pdConfig, int cacheSize,
long expiration) {
- this.client = new KvClient<>(pdConfig);
+ this(new KvClient<>(pdConfig), cacheSize, expiration);
+ }
+
+ SchemaDriver(KvClient<WatchResponse> client, 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));
@@ -98,24 +117,75 @@ public class SchemaDriver {
}
public static 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);
+ synchronized (LIFECYCLE_LOCK) {
+ if (destroying) {
+ throw new NotAllowException("The SchemaDriver is being destroyed");
+ }
+ 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.compareAndSet(null, new SchemaDriver(pdConfig, cacheSize,
- expiration));
}
public static void destroy() {
- SchemaDriver instance = INSTANCE.get();
- if (instance != null) {
- instance.caches.cancelScheduleCacheClean();
- instance.caches.destroyAll();
- INSTANCE.set(null);
+ SchemaDriver instance = null;
+ CountDownLatch completion;
+ boolean closeResources = false;
+ synchronized (LIFECYCLE_LOCK) {
+ if (destroying) {
+ completion = destroyCompletion;
+ } else {
+ instance = INSTANCE.getAndSet(null);
+ 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) {
+ 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();
+ }
+ }
+
+ private void closeResources() {
+ try {
+ this.client.close();
+ } finally {
+ this.caches.cancelScheduleCacheClean();
+ this.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 000000000..e5398c2fb
--- /dev/null
+++ b/hugegraph-struct/src/test/java/org/apache/hugegraph/SchemaDriverTest.java
@@ -0,0 +1,233 @@
+/*
+ * 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.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.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<SchemaDriver> 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() {
+ TrackingKvClient client = new TrackingKvClient();
+ SchemaDriver driver = new SchemaDriver(client, 10, 60_000L);
+ this.instance.set(driver);
+
+ try {
+ SchemaDriver.destroy();
+
+ Assert.assertTrue(client.closed);
+ Assert.assertNull(SchemaDriver.getInstance());
+ } finally {
+ client.close();
+ }
+ }
+
+ @Test
+ public void testDestroyUnpublishesInstanceBeforeClosingResources() throws Exception {
+ BlockingCloseKvClient client = new BlockingCloseKvClient();
+ SchemaDriver driver = new SchemaDriver(client, 10, 60_000L);
+ this.instance.set(driver);
+ Thread destroyThread = new Thread(SchemaDriver::destroy, "schema-driver-destroy");
+ try {
+ destroyThread.start();
+ client.awaitCloseStarted();
+
+ Assert.assertNull(SchemaDriver.getInstance());
+ } finally {
+ client.allowClose.countDown();
+ destroyThread.join(5000L);
+ client.close();
+ }
+
+ Assert.assertFalse(destroyThread.isAlive());
+ 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<Throwable> 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();
+
+ 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<SchemaDriver> instanceReference() throws Exception {
+ Field field = SchemaDriver.class.getDeclaredField("INSTANCE");
+ field.setAccessible(true);
+ return (AtomicReference<SchemaDriver>) field.get(null);
+ }
+
+ private static class TrackingKvClient extends KvClient<WatchResponse> {
+
+ protected volatile boolean closed;
+
+ TrackingKvClient() {
+ super(PDConfig.of("127.0.0.1:8686"));
+ }
+
+ @Override
+ public void listen(String key, Consumer<WatchResponse> consumer) throws PDException {
+ // Avoid opening a real PD stream while constructing SchemaDriver.
+ }
+
+ @Override
+ public void close() {
+ this.closed = true;
+ 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 {
+ 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 {
+
+ private int listenCalls;
+
+ @Override
+ public void listen(String key, Consumer<WatchResponse> consumer) throws PDException {
+ if (++this.listenCalls == 2) {
+ throw new PDException(-1, "listener startup failed");
+ }
+ }
+ }
+}