imbajin commented on code in PR #3157:
URL: https://github.com/apache/hugegraph/pull/3157#discussion_r3892814413


##########
hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java:
##########
@@ -84,7 +85,12 @@ public class SchemaDriver {
 
     private SchemaDriver(PDConfig pdConfig, int cacheSize,
                          long expiration) {
-        this.client = new KvClient<>(pdConfig);
+        this(new KvClient<>(pdConfig), cacheSize, expiration);

Review Comment:
   ⚠️ Important. This constructor now owns a `KvClient` with reconnect 
resources and starts four watches before construction completes. If a later 
`listenMetaChanges()` call fails after an earlier watch has started, the 
constructor exits without closing that client or its executor because 
`INSTANCE` is assigned only after construction succeeds. Add failure cleanup 
around listener initialization, or make client ownership explicit and close it 
on partial startup.



##########
hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java:
##########
@@ -180,50 +199,286 @@ public void onNext(WatchResponse value) {
 
             @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 (this) {
+                synchronized (subscription) {
+                    if (closed.get() || !subscriptions.contains(subscription)) 
{
+                        return false;
+                    }
+                    if (subscription.observer.get() != null) {
+                        return true;
+                    }
+                    subscription.firstFrameReceived = false;
+                    subscription.attemptGeneration = this.transportGeneration;
+                    subscription.observer.set(observer);
+                }
+                if (subscription.prefix) {
+                    streamingCall(KvServiceGrpc.getWatchPrefixMethod(), 
request, observer, 1);
+                } else {
+                    streamingCall(KvServiceGrpc.getWatchMethod(), request, 
observer, 1);
+                }
+            }
+            scheduleStartTimeout(subscription, observer);
+            return true;
         } catch (Exception e) {
-            release();
+            if (e instanceof PDException) {
+                throw (PDException) e;
+            }
             throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE, e);
         }
     }
 
-    private void acquire() {
+    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 = reconnectExecutor.schedule(
+                    () -> onStartTimeout(subscription, sourceObserver),
+                    WATCH_START_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        } catch (RuntimeException e) {
+            throw new PDException(Pdpb.ErrorType.PD_UNREACHABLE_VALUE,
+                                  "Failed to schedule watch start timeout", e);
+        }
+        synchronized (subscription) {
+            if (closed.get() || subscription.observer.get() != sourceObserver 
||
+                subscription.firstFrameReceived) {
+                timeout.cancel(false);
+                return;
+            }
+            ScheduledFuture<?> previous = 
subscription.startTimeout.getAndSet(timeout);
+            if (previous != null) {
+                previous.cancel(false);
+            }
+        }
+    }
+
+    private void onStartTimeout(WatchSubscription subscription,

Review Comment:
   ⚠️ Important. `onStartTimeout()` checks `firstFrameReceived` under the 
subscription lock, releases it, and then calls `requestReconnect()`, which does 
not recheck that flag. If the `Starting` frame arrives in this window, the 
timeout still clears the now-healthy observer and schedules a duplicate watch. 
Make the timeout check and observer invalidation one atomic transition, or 
recheck `firstFrameReceived` inside `requestReconnect()`.



##########
hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java:
##########
@@ -56,13 +62,37 @@
 @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 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 reconnectExecutor;
+    private long transportGeneration;

Review Comment:
   🧹 Design suggestion. The reconnect behavior is now spread across 
per-subscription atomics, nested locks, a shared `transportGeneration`, a 
scheduler, and a first-frame watchdog. Before adding another guard for each 
race, consider one explicit per-watch state/attempt object with an atomic 
`ACTIVE -> RECONNECTING -> STOPPED` transition, the observer/client ID, and the 
cancellable timeout; keep transport refresh behind one small helper and make 
the server's follower response an explicit redirectable status. Then timeout, 
error, leader change, and close can share the same transition instead of 
accumulating independent flags and patches.



##########
hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java:
##########
@@ -180,50 +199,286 @@ public void onNext(WatchResponse value) {
 
             @Override
             public void onError(Throwable t) {
-                release();
-                if (!closed.get()) {
-                    clientId.set(0);
-                    listenWrapper.accept(key, consumer);
+                if (isRetryableWatchError(t)) {

Review Comment:
   ‼️ Critical. `shouldRotateWatchTransport()` only handles `UNAVAILABLE`, but 
the server reports a non-leader watch through `responseObserver.onError(new 
PDException(-1, msg))` without a redirect status. This reaches the retry path 
as an `UNKNOWN`-style error, so the cached stub is retained and a watch that 
lands on a follower can reconnect to that same follower forever. Return an 
explicit redirectable status from the server or invalidate the transport for 
this response so leader discovery runs again, and add a two-peer regression 
test for this path.



##########
hugegraph-struct/src/main/java/org/apache/hugegraph/SchemaDriver.java:
##########
@@ -111,11 +117,14 @@ public static void init(PDConfig pdConfig, int cacheSize, 
long expiration) {
     }
 
     public static void destroy() {
-        SchemaDriver instance = INSTANCE.get();
+        SchemaDriver instance = INSTANCE.getAndSet(null);

Review Comment:
   ⚠️ Important. `INSTANCE` is cleared before the old client and caches finish 
closing. A concurrent `SchemaDriver.init()` can observe `null` and publish a 
second driver while the first driver's watches, reconnect executor, and cache 
cleanup are still active, producing overlapping metadata listeners. Serialize 
init/destroy with one lifecycle lock, or clear `INSTANCE` only after cleanup 
completes.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to