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 431f6e6b4 refactor(store): recover retries after store replacement 
(#3130)
431f6e6b4 is described below

commit 431f6e6b431014feb5b50d27ea67180b4610f10b
Author: KAI <[email protected]>
AuthorDate: Sun Aug 16 00:27:26 2026 +0530

    refactor(store): recover retries after store replacement (#3130)
    
    - Remove the custom IP fingerprint, five-second DNS polling, pool 
replacement, and maintenance executors. Kubernetes owns Pod IPs, while gRPC 
owns DNS resolution and transport reconnect.
    - When gRPC reports UNAVAILABLE, discard only the exact failed Store node, 
close that target, and replace its cached session during the same transaction 
retry. CANCELLED and unrelated failures do not evict the node.
    - Rebuild blocking and async stubs only when their channel generation 
changes, with focused tests for stale sessions, Store replacement, and channel 
rebinding.
    
    ---------
    
    Co-authored-by: imbajin <[email protected]>
---
 .../hugegraph/store/client/HgStoreNodeManager.java |  30 ++-
 .../hugegraph/store/client/NodeTxExecutor.java     |   8 +-
 .../store/client/grpc/AbstractGrpcClient.java      | 147 ++++++++-----
 .../store/client/grpc/NotifyingExecutor.java       |  13 +-
 .../hugegraph/store/client/ClientSuiteTest.java    |   3 +-
 .../hugegraph/store/client/NodeTxExecutorTest.java | 118 ++++++++++
 .../store/client/grpc/AbstractGrpcClientTest.java  | 237 ++++++++++++++++++++-
 7 files changed, 486 insertions(+), 70 deletions(-)

diff --git 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/HgStoreNodeManager.java
 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/HgStoreNodeManager.java
index 956bd2b67..269405bcf 100644
--- 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/HgStoreNodeManager.java
+++ 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/HgStoreNodeManager.java
@@ -26,6 +26,7 @@ import java.util.concurrent.ConcurrentHashMap;
 
 import javax.annotation.concurrent.ThreadSafe;
 
+import org.apache.hugegraph.store.client.grpc.AbstractGrpcClient;
 import org.apache.hugegraph.store.client.grpc.GrpcStoreNodeBuilder;
 import org.apache.hugegraph.store.client.type.HgNodeStatus;
 import org.apache.hugegraph.store.client.type.HgStoreClientException;
@@ -150,7 +151,6 @@ public final class HgStoreNodeManager {
      * @throws HgStoreClientException
      */
     public Integer notifying(String graphName, HgStoreNotice notice) {
-
         if (this.nodeNotifier != null) {
 
             synchronized (Thread.currentThread()) {
@@ -170,6 +170,34 @@ public final class HgStoreNodeManager {
         return null;
     }
 
+    public Integer notifying(String graphName, HgStoreNotice notice,
+                             HgStoreNode expectedNode) {
+        if (notice.getNodeStatus() == HgNodeStatus.NOT_ONLINE ||
+            notice.getNodeStatus() == HgNodeStatus.NOT_WORK) {
+            this.evictNode(expectedNode);
+        }
+        return this.notifying(graphName, notice);
+    }
+
+    private void evictNode(HgStoreNode expectedNode) {
+        synchronized (this.nodeIdMap) {
+            HgStoreNode currentNode = 
this.nodeIdMap.get(expectedNode.getNodeId());
+            if (currentNode == expectedNode) {
+                this.nodeIdMap.remove(expectedNode.getNodeId(), expectedNode);
+                this.addressMap.remove(expectedNode.getAddress(), 
expectedNode);
+                AbstractGrpcClient.closeChannel(expectedNode.getAddress());
+            } else if (currentNode == null ||
+                       
!currentNode.getAddress().equals(expectedNode.getAddress())) {
+                AbstractGrpcClient.closeChannel(expectedNode.getAddress());
+            }
+        }
+        synchronized (this.graphNodesMap) {
+            for (List<HgStoreNode> nodes : this.graphNodesMap.values()) {
+                nodes.removeIf(node -> node == expectedNode);
+            }
+        }
+    }
+
     /**
      * Return a collection of HgStoreNode who is in charge of the graph passed 
in the argument.
      *
diff --git 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutor.java
 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutor.java
index 5dae57c4e..939342182 100644
--- 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutor.java
+++ 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutor.java
@@ -263,8 +263,14 @@ final class NodeTxExecutor {
         }
     }
 
-    public HgStoreSession openNodeSession(HgStoreNode node) {
+    public synchronized HgStoreSession openNodeSession(HgStoreNode node) {
         HgStoreSession res = this.sessions.get(node.getNodeId());
+        if (res instanceof HgStoreNodeSession &&
+            ((HgStoreNodeSession) res).getStoreNode() != node) {
+            // A retry can receive the same node ID backed by a new Store 
process.
+            this.sessions.remove(node.getNodeId());
+            res = null;
+        }
         if (res == null) {
             this.sessions.put(node.getNodeId(), (res = 
node.openSession(this.graphName)));
         }
diff --git 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClient.java
 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClient.java
index 693781d19..12eb9c6a7 100644
--- 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClient.java
+++ 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClient.java
@@ -91,39 +91,43 @@ public abstract class AbstractGrpcClient {
     public abstract AbstractBlockingStub getBlockingStub(ManagedChannel 
channel);
 
     public AbstractBlockingStub getBlockingStub(String target) {
-        ManagedChannel[] channels = getChannels(target);
-        HgPair<ManagedChannel, AbstractBlockingStub>[] pairs = 
blockingStubs.get(target);
-        long l = counter.getAndIncrement();
-        if (l >= limit) {
-            counter.set(0);
-        }
-        int index = (int) (l & (concurrency - 1));
-        if (pairs == null) {
+        while (true) {
+            ManagedChannel[] channels = getChannels(target);
+            HgPair<ManagedChannel, AbstractBlockingStub>[] pairs =
+                    blockingStubs.get(target);
+            int index = nextStubIndex();
+            if (usesChannels(pairs, channels)) {
+                return (AbstractBlockingStub) setBlockingStubOption(
+                        pairs[index].getValue());
+            }
             synchronized (blockingStubs) {
+                channels = getChannels(target);
                 pairs = blockingStubs.get(target);
-                if (pairs == null) {
-                    HgPair<ManagedChannel, AbstractBlockingStub>[] value = new 
HgPair[concurrency];
-                    IntStream.range(0, concurrency).forEach(i -> {
-                        ManagedChannel channel = channels[i];
-                        AbstractBlockingStub stub = getBlockingStub(channel);
-                        value[i] = new HgPair<>(channel, stub);
-                        // log.info("create channel for {}",target);
-                    });
-                    blockingStubs.put(target, value);
-                    AbstractBlockingStub stub = value[index].getValue();
-                    return (AbstractBlockingStub) setBlockingStubOption(stub);
+                if (usesChannels(pairs, channels)) {
+                    return (AbstractBlockingStub) setBlockingStubOption(
+                            pairs[index].getValue());
                 }
+                ManagedChannel[] currentChannels = channels;
+                HgPair<ManagedChannel, AbstractBlockingStub>[] value =
+                        new HgPair[concurrency];
+                IntStream.range(0, concurrency).forEach(i -> {
+                    ManagedChannel channel = currentChannels[i];
+                    AbstractBlockingStub stub = getBlockingStub(channel);
+                    value[i] = new HgPair<>(channel, stub);
+                });
+                if (!publishStubs(target, currentChannels, blockingStubs, 
value)) {
+                    continue;
+                }
+                return (AbstractBlockingStub) setBlockingStubOption(
+                        value[index].getValue());
             }
         }
-        return (AbstractBlockingStub) 
setBlockingStubOption(pairs[index].getValue());
     }
 
     private AbstractStub setBlockingStubOption(AbstractBlockingStub stub) {
         return stub.withDeadlineAfter(config.getGrpcTimeoutSeconds(), 
TimeUnit.SECONDS)
-                   .withMaxInboundMessageSize(
-                           config.getGrpcMaxInboundMessageSize())
-                   .withMaxOutboundMessageSize(
-                           config.getGrpcMaxOutboundMessageSize());
+                   
.withMaxInboundMessageSize(config.getGrpcMaxInboundMessageSize())
+                   
.withMaxOutboundMessageSize(config.getGrpcMaxOutboundMessageSize());
     }
 
     public AbstractAsyncStub getAsyncStub(ManagedChannel channel) {
@@ -131,46 +135,85 @@ public abstract class AbstractGrpcClient {
     }
 
     public AbstractAsyncStub getAsyncStub(String target) {
-        ManagedChannel[] channels = getChannels(target);
-        HgPair<ManagedChannel, AbstractAsyncStub>[] pairs = 
asyncStubs.get(target);
-        long l = counter.getAndIncrement();
-        if (l >= limit) {
-            counter.set(0);
-        }
-        int index = (int) (l & (concurrency - 1));
-        if (pairs == null) {
+        while (true) {
+            ManagedChannel[] channels = getChannels(target);
+            HgPair<ManagedChannel, AbstractAsyncStub>[] pairs = 
asyncStubs.get(target);
+            int index = nextStubIndex();
+            if (usesChannels(pairs, channels)) {
+                return (AbstractAsyncStub) 
setStubOption(pairs[index].getValue());
+            }
             synchronized (asyncStubs) {
+                channels = getChannels(target);
                 pairs = asyncStubs.get(target);
-                if (pairs == null) {
-                    HgPair<ManagedChannel, AbstractAsyncStub>[] value = new 
HgPair[concurrency];
-                    IntStream.range(0, concurrency).parallel().forEach(i -> {
-                        ManagedChannel channel = channels[i];
-                        AbstractAsyncStub stub = getAsyncStub(channel);
-                        // 
stub.withMaxInboundMessageSize(config.getGrpcMaxInboundMessageSize())
-                        //    
.withMaxOutboundMessageSize(config.getGrpcMaxOutboundMessageSize());
-                        value[i] = new HgPair<>(channel, stub);
-                        // log.info("create channel for {}",target);
-                    });
-                    asyncStubs.put(target, value);
-                    AbstractAsyncStub stub =
-                            (AbstractAsyncStub) 
setStubOption(value[index].getValue());
-                    return stub;
+                if (usesChannels(pairs, channels)) {
+                    return (AbstractAsyncStub) setStubOption(
+                            pairs[index].getValue());
+                }
+                ManagedChannel[] currentChannels = channels;
+                HgPair<ManagedChannel, AbstractAsyncStub>[] value =
+                        new HgPair[concurrency];
+                IntStream.range(0, concurrency).parallel().forEach(i -> {
+                    ManagedChannel channel = currentChannels[i];
+                    AbstractAsyncStub stub = getAsyncStub(channel);
+                    value[i] = new HgPair<>(channel, stub);
+                });
+                if (!publishStubs(target, currentChannels, asyncStubs, value)) 
{
+                    continue;
                 }
+                return (AbstractAsyncStub) 
setStubOption(value[index].getValue());
             }
         }
-        return (AbstractAsyncStub) setStubOption(pairs[index].getValue());
+    }
 
+    private static boolean usesChannels(HgPair<ManagedChannel, ?>[] pairs,
+                                        ManagedChannel[] channels) {
+        if (pairs == null || pairs.length != channels.length) {
+            return false;
+        }
+        return IntStream.range(0, channels.length)
+                        .allMatch(i -> pairs[i].getKey() == channels[i]);
+    }
+
+    private static int nextStubIndex() {
+        long value = counter.getAndIncrement();
+        if (value >= limit) {
+            counter.set(0);
+        }
+        return (int) (value & (concurrency - 1));
+    }
+
+    private static <S> boolean publishStubs(
+            String target, ManagedChannel[] currentChannels,
+            Map<String, HgPair<ManagedChannel, S>[]> stubs,
+            HgPair<ManagedChannel, S>[] value) {
+        synchronized (channels) {
+            if (channels.get(target) != currentChannels) {
+                return false;
+            }
+            stubs.put(target, value);
+            return true;
+        }
     }
 
     protected AbstractStub setStubOption(AbstractStub value) {
-        return value.withMaxInboundMessageSize(
-                            config.getGrpcMaxInboundMessageSize())
-                    .withMaxOutboundMessageSize(
-                            config.getGrpcMaxOutboundMessageSize());
+        return 
value.withMaxInboundMessageSize(config.getGrpcMaxInboundMessageSize())
+                    
.withMaxOutboundMessageSize(config.getGrpcMaxOutboundMessageSize());
+    }
+
+    public static void closeChannel(String target) {
+        ManagedChannel[] targetChannels;
+        synchronized (channels) {
+            targetChannels = channels.remove(target);
+        }
+        if (targetChannels != null) {
+            IntStream.range(0, targetChannels.length)
+                     .mapToObj(i -> targetChannels[i])
+                     .filter(channel -> channel != null)
+                     .forEach(ManagedChannel::shutdown);
+        }
     }
 
     protected ManagedChannel createChannel(String target) {
         return ManagedChannelBuilder.forTarget(target).usePlaintext().build();
     }
-
 }
diff --git 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/NotifyingExecutor.java
 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/NotifyingExecutor.java
index 491ad94b3..75cbd88c3 100644
--- 
a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/NotifyingExecutor.java
+++ 
b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/NotifyingExecutor.java
@@ -38,6 +38,7 @@ import 
org.apache.hugegraph.store.grpc.session.PartitionLeader;
 
 import com.google.protobuf.util.JsonFormat;
 
+import io.grpc.Status;
 import lombok.extern.slf4j.Slf4j;
 
 /**
@@ -242,11 +243,15 @@ final class NotifyingExecutor {
 
     private Consumer<Throwable> notifyErrConsumer(HgNodeStatus status) {
         return t -> {
-            nodeManager.notifying(
-                    this.graphName,
+            HgStoreNotice notice =
                     
HgStoreNotice.of(this.nodeSession.getStoreNode().getNodeId(), status,
-                                     t.getMessage())
-            );
+                                     t.getMessage());
+            if (Status.fromThrowable(t).getCode() == Status.Code.UNAVAILABLE) {
+                nodeManager.notifying(this.graphName, notice,
+                                      this.nodeSession.getStoreNode());
+            } else {
+                nodeManager.notifying(this.graphName, notice);
+            }
         };
     }
 
diff --git 
a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java
 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java
index 4217a4c1d..c6e785f22 100644
--- 
a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java
+++ 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java
@@ -27,7 +27,8 @@ import org.junit.runners.Suite;
  */
 @RunWith(Suite.class)
 @Suite.SuiteClasses({
-        AbstractGrpcClientTest.class
+        AbstractGrpcClientTest.class,
+        NodeTxExecutorTest.class
 })
 public class ClientSuiteTest {
 }
diff --git 
a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutorTest.java
 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutorTest.java
new file mode 100644
index 000000000..00a05c65c
--- /dev/null
+++ 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxExecutorTest.java
@@ -0,0 +1,118 @@
+/*
+ * 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.store.client;
+
+import static org.junit.Assert.assertSame;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.hugegraph.store.HgStoreSession;
+import org.junit.Test;
+
+public class NodeTxExecutorTest {
+
+    @Test
+    public void testRetryReplacesSessionFromEvictedNode() {
+        long nodeId = 1L;
+        HgStoreNode oldNode = mock(HgStoreNode.class);
+        HgStoreNode currentNode = mock(HgStoreNode.class);
+        HgStoreNodeSession oldSession = mock(HgStoreNodeSession.class);
+        HgStoreNodeSession currentSession = mock(HgStoreNodeSession.class);
+
+        when(oldNode.getNodeId()).thenReturn(nodeId);
+        when(currentNode.getNodeId()).thenReturn(nodeId);
+        when(oldNode.openSession("graph")).thenReturn(oldSession);
+        when(currentNode.openSession("graph")).thenReturn(currentSession);
+        when(oldSession.getStoreNode()).thenReturn(oldNode);
+        when(currentSession.getStoreNode()).thenReturn(currentNode);
+
+        NodeTxExecutor executor = NodeTxExecutor.graphOf("graph", null);
+        executor.setTx(true);
+        AtomicInteger attempts = new AtomicInteger();
+        Optional<HgStoreSession> result = executor.retryingInvoke(() -> {
+            HgStoreNode node = attempts.getAndIncrement() == 0 ? oldNode : 
currentNode;
+            HgStoreSession session = executor.openNodeSession(node);
+            if (node == oldNode) {
+                throw new RuntimeException("simulated transport failure");
+            }
+            return session;
+        });
+
+        assertSame(currentSession, result.get());
+        verify(oldSession).beginTx();
+        verify(currentSession).beginTx();
+    }
+
+    @Test
+    public void testParallelReplacementUsesOneCurrentSession() throws 
Exception {
+        long nodeId = 2L;
+        HgStoreNode oldNode = mock(HgStoreNode.class);
+        HgStoreNode currentNode = mock(HgStoreNode.class);
+        HgStoreNodeSession oldSession = mock(HgStoreNodeSession.class);
+        HgStoreNodeSession firstCurrentSession = 
mock(HgStoreNodeSession.class);
+        HgStoreNodeSession secondCurrentSession = 
mock(HgStoreNodeSession.class);
+
+        when(oldNode.getNodeId()).thenReturn(nodeId);
+        when(currentNode.getNodeId()).thenReturn(nodeId);
+        when(oldNode.openSession("graph")).thenReturn(oldSession);
+        when(oldSession.getStoreNode()).thenReturn(oldNode);
+        when(firstCurrentSession.getStoreNode()).thenReturn(currentNode);
+        when(secondCurrentSession.getStoreNode()).thenReturn(currentNode);
+
+        CountDownLatch concurrentCreations = new CountDownLatch(2);
+        AtomicInteger creations = new AtomicInteger();
+        when(currentNode.openSession("graph")).thenAnswer(invocation -> {
+            int creation = creations.getAndIncrement();
+            concurrentCreations.countDown();
+            concurrentCreations.await(1, TimeUnit.SECONDS);
+            return creation == 0 ? firstCurrentSession : secondCurrentSession;
+        });
+
+        NodeTxExecutor executor = NodeTxExecutor.graphOf("graph", null);
+        executor.openNodeSession(oldNode);
+
+        ExecutorService workers = Executors.newFixedThreadPool(2);
+        CountDownLatch start = new CountDownLatch(1);
+        Future<HgStoreSession> first = workers.submit(() -> {
+            start.await();
+            return executor.openNodeSession(currentNode);
+        });
+        Future<HgStoreSession> second = workers.submit(() -> {
+            start.await();
+            return executor.openNodeSession(currentNode);
+        });
+        try {
+            start.countDown();
+            assertSame(first.get(3, TimeUnit.SECONDS),
+                       second.get(3, TimeUnit.SECONDS));
+            verify(currentNode, times(1)).openSession("graph");
+        } finally {
+            workers.shutdownNow();
+        }
+    }
+}
diff --git 
a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClientTest.java
 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClientTest.java
index 59f1c86ab..dd2f21218 100644
--- 
a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClientTest.java
+++ 
b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/AbstractGrpcClientTest.java
@@ -18,8 +18,15 @@
 package org.apache.hugegraph.store.client.grpc;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -28,9 +35,21 @@ import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.apache.hugegraph.store.HgStoreSession;
+import org.apache.hugegraph.store.client.HgStoreNode;
+import org.apache.hugegraph.store.client.HgStoreNodeManager;
+import org.apache.hugegraph.store.client.HgStoreNodeSession;
+import org.apache.hugegraph.store.client.HgStoreNotice;
+import org.apache.hugegraph.store.client.type.HgNodeStatus;
+import org.apache.hugegraph.store.client.type.HgStoreClientException;
 import org.junit.Test;
 
 import io.grpc.CallOptions;
@@ -38,13 +57,10 @@ import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.ManagedChannel;
 import io.grpc.MethodDescriptor;
+import io.grpc.Status;
 import io.grpc.stub.AbstractAsyncStub;
 import io.grpc.stub.AbstractBlockingStub;
 
-/**
- * Verifies that the stub pools of {@link AbstractGrpcClient} spread their 
entries over every
- * channel created for a target, instead of binding all of them to a single 
channel.
- */
 public class AbstractGrpcClientTest {
 
     private static final AtomicInteger TARGET_SEQ = new AtomicInteger();
@@ -66,13 +82,14 @@ public class AbstractGrpcClientTest {
         ManagedChannel[] channels = client.getChannels(target);
         assertTrue("pool must hold more than one channel", channels.length > 
1);
 
-        // Pool initialisation: one stub per channel, each bound to a 
different channel.
         assertNotNull(client.getBlockingStub(target));
-        assertEquals("one stub per channel", channels.length, 
client.blockingStubChannels.size());
+        assertEquals("one stub per channel", channels.length,
+                     client.blockingStubChannels.size());
         Set<ManagedChannel> bound = identitySet(client.blockingStubChannels);
         assertEquals("stubs must not share a channel", channels.length, 
bound.size());
         assertTrue("stubs must cover the channels of the target",
                    bound.containsAll(Arrays.asList(channels)));
+        AbstractGrpcClient.closeChannel(target);
     }
 
     @Test
@@ -88,11 +105,172 @@ public class AbstractGrpcClientTest {
         assertEquals("stubs must not share a channel", channels.length, 
bound.size());
         assertTrue("stubs must cover the channels of the target",
                    bound.containsAll(Arrays.asList(channels)));
+        AbstractGrpcClient.closeChannel(target);
+    }
+
+    @Test
+    public void testClosedTargetRebindsStubPools() {
+        String target = uniqueTarget("rebind");
+        RecordingGrpcClient client = new RecordingGrpcClient();
+        ManagedChannel[] oldChannels = client.getChannels(target);
+        client.getBlockingStub(target);
+        client.getAsyncStub(target);
+        int oldBlockingCount = client.blockingStubChannels.size();
+        int oldAsyncCount = client.asyncStubChannels.size();
+
+        AbstractGrpcClient.closeChannel(target);
+        
assertTrue(Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown));
+
+        ManagedChannel[] currentChannels = client.getChannels(target);
+        assertNotSame(oldChannels, currentChannels);
+        client.getBlockingStub(target);
+        client.getAsyncStub(target);
+
+        List<ManagedChannel> currentBlocking =
+                client.blockingStubChannels.subList(oldBlockingCount,
+                                                    
client.blockingStubChannels.size());
+        List<ManagedChannel> currentAsync =
+                client.asyncStubChannels.subList(oldAsyncCount,
+                                                 
client.asyncStubChannels.size());
+        assertEquals(identitySet(Arrays.asList(currentChannels)),
+                     identitySet(currentBlocking));
+        assertEquals(identitySet(Arrays.asList(currentChannels)), 
identitySet(currentAsync));
+        Set<ManagedChannel> retired = identitySet(Arrays.asList(oldChannels));
+        assertFalse(currentBlocking.stream().anyMatch(retired::contains));
+        assertFalse(currentAsync.stream().anyMatch(retired::contains));
+        AbstractGrpcClient.closeChannel(target);
+    }
+
+    @Test
+    public void testCloseDuringStubBuildDoesNotPublishRetiredChannel() throws 
Exception {
+        String target = uniqueTarget("concurrent-rebind");
+        PausingGrpcClient client = new PausingGrpcClient();
+        client.getBlockingStub(target);
+        AbstractGrpcClient.closeChannel(target);
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        client.pauseStubCreation = true;
+        Future<AbstractBlockingStub> future = executor.submit(
+                () -> client.getBlockingStub(target));
+        try {
+            assertTrue(client.stubCreationStarted.await(5, TimeUnit.SECONDS));
+            ManagedChannel[] retired = AbstractGrpcClient.channels.get(target);
+            assertNotNull(retired);
+
+            AbstractGrpcClient.closeChannel(target);
+            ManagedChannel[] current = client.getChannels(target);
+            client.releaseStubCreation.countDown();
+
+            AbstractBlockingStub stub = future.get(5, TimeUnit.SECONDS);
+            
assertTrue(identitySet(Arrays.asList(current)).contains(stub.getChannel()));
+            assertFalse(((ManagedChannel) stub.getChannel()).isShutdown());
+        } finally {
+            client.releaseStubCreation.countDown();
+            executor.shutdownNow();
+            AbstractGrpcClient.closeChannel(target);
+        }
+    }
+
+    @Test
+    public void testUnavailableRpcEvictsExpectedNodeAndTarget() {
+        String graph = uniqueTarget("terminal-graph");
+        String target = uniqueTarget("terminal-node");
+        long nodeId = TARGET_SEQ.incrementAndGet();
+        RecordingGrpcClient client = new RecordingGrpcClient();
+        HgStoreNodeManager manager = HgStoreNodeManager.getInstance();
+        HgStoreNode node = testNode(nodeId, target);
+        HgStoreNodeSession session = mock(HgStoreNodeSession.class);
+        when(session.getStoreNode()).thenReturn(node);
+
+        manager.addNode(graph, node);
+        ManagedChannel[] published = client.getChannels(target);
+        NotifyingExecutor notifier = new NotifyingExecutor(graph, manager, 
session);
+        try {
+            notifier.invoke(() -> {
+                throw Status.UNAVAILABLE.asRuntimeException();
+            }, response -> true);
+            fail("the unavailable RPC must still be reported to the caller");
+        } catch (HgStoreClientException ignored) {
+            // Expected.
+        }
+
+        assertNull(manager.getStoreNode(nodeId));
+        assertFalse(AbstractGrpcClient.channels.containsKey(target));
+        
assertTrue(Arrays.stream(published).allMatch(ManagedChannel::isShutdown));
+    }
+
+    @Test
+    public void testStaleNoticePreservesSameAddressReplacement() {
+        String graph = uniqueTarget("replacement-graph");
+        String target = uniqueTarget("replacement-node");
+        long nodeId = TARGET_SEQ.incrementAndGet();
+        RecordingGrpcClient client = new RecordingGrpcClient();
+        HgStoreNodeManager manager = HgStoreNodeManager.getInstance();
+        HgStoreNode staleNode = testNode(nodeId, target);
+        HgStoreNode replacementNode = testNode(nodeId, target);
+
+        manager.addNode(graph, staleNode);
+        manager.addNode(graph, replacementNode);
+        ManagedChannel[] published = client.getChannels(target);
+        manager.notifying(graph, HgStoreNotice.of(nodeId, 
HgNodeStatus.NOT_WORK), staleNode);
+
+        assertSame(replacementNode, manager.getStoreNode(nodeId));
+        assertSame(published, AbstractGrpcClient.channels.get(target));
+        
assertTrue(Arrays.stream(published).noneMatch(ManagedChannel::isShutdown));
+
+        manager.notifying(graph, HgStoreNotice.of(nodeId, 
HgNodeStatus.NOT_WORK),
+                          replacementNode);
+    }
+
+    @Test
+    public void testCancelledRpcDoesNotEvictCurrentNode() {
+        String graph = uniqueTarget("cancelled-graph");
+        String target = uniqueTarget("cancelled-node");
+        long nodeId = TARGET_SEQ.incrementAndGet();
+        RecordingGrpcClient client = new RecordingGrpcClient();
+        HgStoreNodeManager manager = HgStoreNodeManager.getInstance();
+        HgStoreNode node = testNode(nodeId, target);
+        HgStoreNodeSession session = mock(HgStoreNodeSession.class);
+        when(session.getStoreNode()).thenReturn(node);
+
+        manager.addNode(graph, node);
+        ManagedChannel[] published = client.getChannels(target);
+        NotifyingExecutor notifier = new NotifyingExecutor(graph, manager, 
session);
+        try {
+            notifier.invoke(() -> {
+                throw Status.CANCELLED.asRuntimeException();
+            }, response -> true);
+            fail("the cancelled RPC must still be reported to the caller");
+        } catch (HgStoreClientException ignored) {
+            // Expected.
+        }
+
+        assertSame(node, manager.getStoreNode(nodeId));
+        assertSame(published, AbstractGrpcClient.channels.get(target));
+        
assertTrue(Arrays.stream(published).noneMatch(ManagedChannel::isShutdown));
+
+        manager.notifying(graph, HgStoreNotice.of(nodeId, 
HgNodeStatus.NOT_WORK), node);
+    }
+
+    private static HgStoreNode testNode(long nodeId, String address) {
+        return new HgStoreNode() {
+            @Override
+            public Long getNodeId() {
+                return nodeId;
+            }
+
+            @Override
+            public String getAddress() {
+                return address;
+            }
+
+            @Override
+            public HgStoreSession openSession(String graphName) {
+                return null;
+            }
+        };
     }
 
-    /**
-     * A client whose channels and stubs are local fakes, so the test needs no 
PD or store node.
-     */
     private static class RecordingGrpcClient extends AbstractGrpcClient {
 
         private final AtomicInteger channelSeq = new AtomicInteger();
@@ -101,9 +279,24 @@ public class AbstractGrpcClientTest {
         private final List<ManagedChannel> asyncStubChannels =
                 Collections.synchronizedList(new ArrayList<>());
 
+        @Override
+        public ManagedChannel[] getChannels(String target) {
+            synchronized (channels) {
+                ManagedChannel[] current = channels.get(target);
+                if (current == null) {
+                    current = new ManagedChannel[concurrency];
+                    for (int i = 0; i < current.length; i++) {
+                        current[i] = this.createChannel(target);
+                    }
+                    channels.put(target, current);
+                }
+                return current;
+            }
+        }
+
         @Override
         protected ManagedChannel createChannel(String target) {
-            return new FakeManagedChannel(target + "#" + 
channelSeq.getAndIncrement());
+            return new FakeManagedChannel(target + "#" + 
this.channelSeq.getAndIncrement());
         }
 
         @Override
@@ -119,6 +312,28 @@ public class AbstractGrpcClientTest {
         }
     }
 
+    private static class PausingGrpcClient extends RecordingGrpcClient {
+
+        private final AtomicBoolean paused = new AtomicBoolean();
+        private final CountDownLatch stubCreationStarted = new 
CountDownLatch(1);
+        private final CountDownLatch releaseStubCreation = new 
CountDownLatch(1);
+        private volatile boolean pauseStubCreation;
+
+        @Override
+        public AbstractBlockingStub getBlockingStub(ManagedChannel channel) {
+            if (this.pauseStubCreation && this.paused.compareAndSet(false, 
true)) {
+                this.stubCreationStarted.countDown();
+                try {
+                    this.releaseStubCreation.await();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new RuntimeException(e);
+                }
+            }
+            return super.getBlockingStub(channel);
+        }
+    }
+
     private static class FakeBlockingStub extends 
AbstractBlockingStub<FakeBlockingStub> {
 
         FakeBlockingStub(Channel channel, CallOptions callOptions) {
@@ -171,7 +386,7 @@ public class AbstractGrpcClientTest {
 
         @Override
         public ManagedChannel shutdownNow() {
-            return shutdown();
+            return this.shutdown();
         }
 
         @Override

Reply via email to