From 4665c078fc268b457089ff878d0cdf8838eb1992 Mon Sep 17 00:00:00 2001 From: Himanshu Verma Date: Thu, 30 Jul 2026 18:46:33 +0530 Subject: [PATCH 1/3] fix(store): refresh gRPC channels after address changes Refresh cached channel pools when a Store target resolves to a new address. Rebuild stale blocking and async stub pools, and guard concurrent resolution and publication races. Fixes #3124 --- .../store/client/grpc/AbstractGrpcClient.java | 196 ++++++++++++---- .../store/client/ClientSuiteTest.java | 4 - .../client/grpc/AbstractGrpcClientTest.java | 210 +++++++++++++++--- 3 files changed, 324 insertions(+), 86 deletions(-) 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 693781d19d..1f7c87b51c 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 @@ -17,12 +17,17 @@ package org.apache.hugegraph.store.client.grpc; +import java.net.InetAddress; +import java.net.URI; +import java.net.UnknownHostException; +import java.util.Arrays; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.hugegraph.store.client.util.ExecutorPool; @@ -38,6 +43,9 @@ public abstract class AbstractGrpcClient { protected static Map channels = new ConcurrentHashMap<>(); + private static final Map resolvedTargets = new ConcurrentHashMap<>(); + private static final Map resolutionRequests = new ConcurrentHashMap<>(); + private static final Map appliedResolutions = new ConcurrentHashMap<>(); private static final int n = 5; protected static int concurrency = 1 << n; private static final AtomicLong counter = new AtomicLong(0); @@ -58,6 +66,7 @@ public AbstractGrpcClient() { } public ManagedChannel[] getChannels(String target) { + this.refreshChannelsIfAddressChanged(target); ManagedChannel[] tc; if ((tc = channels.get(target)) == null) { synchronized (channels) { @@ -91,31 +100,44 @@ public ManagedChannel[] getChannels(String target) { public abstract AbstractBlockingStub getBlockingStub(ManagedChannel channel); public AbstractBlockingStub getBlockingStub(String target) { - ManagedChannel[] channels = getChannels(target); - HgPair[] pairs = blockingStubs.get(target); - long l = counter.getAndIncrement(); - if (l >= limit) { - counter.set(0); - } - int index = (int) (l & (concurrency - 1)); - if (pairs == null) { - synchronized (blockingStubs) { - pairs = blockingStubs.get(target); - if (pairs == null) { - HgPair[] 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); + while (true) { + ManagedChannel[] targetChannels = getChannels(target); + HgPair[] pairs = blockingStubs.get(target); + long l = counter.getAndIncrement(); + if (l >= limit) { + counter.set(0); + } + int index = (int) (l & (concurrency - 1)); + if (!usesChannels(pairs, targetChannels)) { + synchronized (blockingStubs) { + pairs = blockingStubs.get(target); + if (!usesChannels(pairs, targetChannels)) { + HgPair[] value = + new HgPair[concurrency]; + IntStream.range(0, concurrency).forEach(i -> { + ManagedChannel channel = targetChannels[index]; + AbstractBlockingStub stub = getBlockingStub(channel); + value[i] = new HgPair<>(channel, stub); + // log.info("create channel for {}",target); + }); + synchronized (channels) { + if (channels.get(target) != targetChannels) { + continue; + } + blockingStubs.put(target, value); + AbstractBlockingStub stub = value[index].getValue(); + return (AbstractBlockingStub) setBlockingStubOption(stub); + } + } } } + synchronized (channels) { + if (channels.get(target) != targetChannels) { + continue; + } + return (AbstractBlockingStub) setBlockingStubOption(pairs[index].getValue()); + } } - return (AbstractBlockingStub) setBlockingStubOption(pairs[index].getValue()); } private AbstractStub setBlockingStubOption(AbstractBlockingStub stub) { @@ -131,35 +153,49 @@ public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { } public AbstractAsyncStub getAsyncStub(String target) { - ManagedChannel[] channels = getChannels(target); - HgPair[] pairs = asyncStubs.get(target); - long l = counter.getAndIncrement(); - if (l >= limit) { - counter.set(0); - } - int index = (int) (l & (concurrency - 1)); - if (pairs == null) { - synchronized (asyncStubs) { - pairs = asyncStubs.get(target); - if (pairs == null) { - HgPair[] 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; + while (true) { + ManagedChannel[] targetChannels = getChannels(target); + HgPair[] pairs = asyncStubs.get(target); + long l = counter.getAndIncrement(); + if (l >= limit) { + counter.set(0); + } + int index = (int) (l & (concurrency - 1)); + if (!usesChannels(pairs, targetChannels)) { + synchronized (asyncStubs) { + pairs = asyncStubs.get(target); + if (!usesChannels(pairs, targetChannels)) { + HgPair[] value = + new HgPair[concurrency]; + IntStream.range(0, concurrency).parallel().forEach(i -> { + ManagedChannel channel = targetChannels[index]; + AbstractAsyncStub stub = getAsyncStub(channel); + // stub.withMaxInboundMessageSize( + // config.getGrpcMaxInboundMessageSize()) + // .withMaxOutboundMessageSize( + // config.getGrpcMaxOutboundMessageSize()); + value[i] = new HgPair<>(channel, stub); + // log.info("create channel for {}",target); + }); + synchronized (channels) { + if (channels.get(target) != targetChannels) { + continue; + } + asyncStubs.put(target, value); + AbstractAsyncStub stub = + (AbstractAsyncStub) setStubOption(value[index].getValue()); + return stub; + } + } + } + } + synchronized (channels) { + if (channels.get(target) != targetChannels) { + continue; } + return (AbstractAsyncStub) setStubOption(pairs[index].getValue()); } } - return (AbstractAsyncStub) setStubOption(pairs[index].getValue()); - } protected AbstractStub setStubOption(AbstractStub value) { @@ -169,6 +205,70 @@ protected AbstractStub setStubOption(AbstractStub value) { config.getGrpcMaxOutboundMessageSize()); } + private static boolean usesChannels(HgPair[] pairs, + ManagedChannel[] channels) { + if (pairs == null || pairs.length != channels.length) { + return false; + } + for (HgPair pair : pairs) { + if (pair == null || pair.getKey() == null || + !containsChannel(channels, pair.getKey())) { + return false; + } + } + return true; + } + + private static boolean containsChannel(ManagedChannel[] channels, + ManagedChannel expected) { + return Arrays.stream(channels).anyMatch(channel -> channel == expected); + } + + private void refreshChannelsIfAddressChanged(String target) { + long resolutionRequest = resolutionRequests.computeIfAbsent(target, + key -> new AtomicLong()) + .incrementAndGet(); + String resolvedTarget = this.resolveTarget(target); + if (resolvedTarget.isEmpty()) { + return; + } + synchronized (channels) { + Long appliedResolution = appliedResolutions.get(target); + if (appliedResolution != null && appliedResolution >= resolutionRequest) { + return; + } + appliedResolutions.put(target, resolutionRequest); + String previousTarget = resolvedTargets.put(target, resolvedTarget); + if (previousTarget == null && !channels.containsKey(target)) { + return; + } + if (resolvedTarget.equals(previousTarget)) { + return; + } + ManagedChannel[] staleChannels = channels.remove(target); + if (staleChannels != null) { + Arrays.stream(staleChannels) + .filter(channel -> channel != null && !channel.isShutdown()) + .forEach(ManagedChannel::shutdownNow); + } + } + } + + protected String resolveTarget(String target) { + try { + String host = URI.create("dns://" + target).getHost(); + if (host == null) { + return ""; + } + return Arrays.stream(InetAddress.getAllByName(host)) + .map(InetAddress::getHostAddress) + .sorted() + .collect(Collectors.joining(",")); + } catch (IllegalArgumentException | UnknownHostException ignored) { + return ""; + } + } + protected ManagedChannel createChannel(String target) { return ManagedChannelBuilder.forTarget(target).usePlaintext().build(); } 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 4217a4c1de..885d5a46ad 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 @@ -21,10 +21,6 @@ import org.junit.runner.RunWith; import org.junit.runners.Suite; -/** - * Entry point of the {@code store-client-test} profile. Cluster-dependent tests are deliberately - * excluded. - */ @RunWith(Suite.class) @Suite.SuiteClasses({ AbstractGrpcClientTest.class 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 59f1c86ab1..6b717929ff 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 @@ -19,15 +19,18 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collection; 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.AtomicInteger; @@ -42,8 +45,7 @@ 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. + * Verifies that Store address changes replace channels and their cached stubs safely. */ public class AbstractGrpcClientTest { @@ -53,49 +55,135 @@ private static String uniqueTarget(String prefix) { return prefix + "-" + TARGET_SEQ.incrementAndGet() + ":8500"; } - private static Set identitySet(Collection channels) { - Set set = Collections.newSetFromMap(new IdentityHashMap<>()); - set.addAll(channels); - return set; + private static boolean belongsToPool(Channel channel, + ManagedChannel[] channels) { + return Arrays.stream(channels).anyMatch(current -> current == channel); } @Test - public void testBlockingStubPoolCoversEveryChannel() { - String target = uniqueTarget("blocking"); + public void testAddressChangeReplacesChannelAndStubPools() { + String target = uniqueTarget("address-change"); RecordingGrpcClient client = new RecordingGrpcClient(); - ManagedChannel[] channels = client.getChannels(target); - assertTrue("pool must hold more than one channel", channels.length > 1); + ManagedChannel[] oldChannels = client.getChannels(target); + assertNotNull(client.getBlockingStub(target)); + assertNotNull(client.getAsyncStub(target)); - // Pool initialisation: one stub per channel, each bound to a different channel. + client.resolvedTarget = "10.0.0.2"; + ManagedChannel[] newChannels = client.getChannels(target); + assertNotSame("an address change must replace the channel pool", + oldChannels, newChannels); + assertTrue("every stale channel must be shut down", + Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); + + client.blockingStubChannels.clear(); + client.asyncStubChannels.clear(); assertNotNull(client.getBlockingStub(target)); - assertEquals("one stub per channel", channels.length, client.blockingStubChannels.size()); - Set 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))); + assertNotNull(client.getAsyncStub(target)); + assertEquals("the blocking stub pool must be rebuilt", + newChannels.length, client.blockingStubChannels.size()); + assertTrue("replacement blocking stubs must use the new channel pool", + client.blockingStubChannels.stream() + .allMatch(channel -> + belongsToPool(channel, newChannels))); + assertEquals("the async stub pool must be rebuilt", + newChannels.length, client.asyncStubChannels.size()); + assertTrue("replacement async stubs must use the new channel pool", + client.asyncStubChannels.stream() + .allMatch(channel -> + belongsToPool(channel, newChannels))); } @Test - public void testAsyncStubPoolCoversEveryChannel() { - String target = uniqueTarget("async"); + public void testFirstSuccessfulResolutionReplacesUnknownChannels() { + String target = uniqueTarget("first-successful-resolution"); RecordingGrpcClient client = new RecordingGrpcClient(); - ManagedChannel[] channels = client.getChannels(target); - assertTrue("pool must hold more than one channel", channels.length > 1); + client.resolvedTarget = ""; + ManagedChannel[] unknownChannels = client.getChannels(target); + + client.resolvedTarget = "10.0.0.1"; + ManagedChannel[] resolvedChannels = client.getChannels(target); + assertNotSame("a pool with unknown addresses must be replaced", + unknownChannels, resolvedChannels); + assertTrue("every channel from the unknown pool must be shut down", + Arrays.stream(unknownChannels).allMatch(ManagedChannel::isShutdown)); + assertTrue("the resolved channel pool must remain live", + Arrays.stream(resolvedChannels).noneMatch(ManagedChannel::isShutdown)); + } - assertNotNull(client.getAsyncStub(target)); - assertEquals("one stub per channel", channels.length, client.asyncStubChannels.size()); - Set bound = identitySet(client.asyncStubChannels); - 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))); + @Test + public void testOlderResolutionCannotReplaceNewerChannels() throws Exception { + String target = uniqueTarget("concurrent-address-change"); + OutOfOrderResolverGrpcClient client = new OutOfOrderResolverGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future staleResolution = + executor.submit(() -> client.getChannels(target)); + assertTrue("the stale resolution must be in flight", + client.staleResolutionStarted.await(5, TimeUnit.SECONDS)); + Future freshResolution = + executor.submit(() -> client.getChannels(target)); + ManagedChannel[] freshChannels = freshResolution.get(5, TimeUnit.SECONDS); + + assertNotSame("the newer address must replace the old channel pool", + oldChannels, freshChannels); + client.releaseStaleResolution.countDown(); + assertSame("the late stale result must retain the newer channel pool", + freshChannels, staleResolution.get(5, TimeUnit.SECONDS)); + assertTrue("the replaced channel pool must be shut down", + Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); + assertTrue("the newer channel pool must remain live", + Arrays.stream(freshChannels).noneMatch(ManagedChannel::isShutdown)); + } finally { + client.releaseStaleResolution.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void testStubBuildRetriesAfterChannelRefresh() throws Exception { + String target = uniqueTarget("concurrent-stub-refresh"); + StubInterleavingGrpcClient client = new StubInterleavingGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future staleStub = + executor.submit(() -> client.getBlockingStub(target)); + assertTrue("the old stub pool build must be in flight", + client.staleStubBuildStarted.await(5, TimeUnit.SECONDS)); + client.resolvedTarget = "10.0.0.2"; + Future freshStub = + executor.submit(() -> client.getBlockingStub(target)); + for (int i = 0; i < 500 && + Arrays.stream(oldChannels).anyMatch(channel -> !channel.isShutdown()); + i++) { + Thread.sleep(10L); + } + assertTrue("refresh must shut down the old channel pool", + Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); + client.releaseStaleStubBuild.countDown(); + + AbstractBlockingStub staleResult = staleStub.get(5, TimeUnit.SECONDS); + AbstractBlockingStub freshResult = freshStub.get(5, TimeUnit.SECONDS); + ManagedChannel[] currentChannels = client.getChannels(target); + assertTrue("the stale build must retry against the current pool", + belongsToPool(staleResult.getChannel(), currentChannels)); + assertTrue("the concurrent build must use the current pool", + belongsToPool(freshResult.getChannel(), currentChannels)); + assertTrue("the current channel pool must remain live", + Arrays.stream(currentChannels).noneMatch(ManagedChannel::isShutdown)); + } finally { + client.releaseStaleStubBuild.countDown(); + executor.shutdownNow(); + } } - /** - * 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(); + protected volatile String resolvedTarget = "10.0.0.1"; private final List blockingStubChannels = Collections.synchronizedList(new ArrayList<>()); private final List asyncStubChannels = @@ -103,7 +191,12 @@ private static class RecordingGrpcClient extends AbstractGrpcClient { @Override protected ManagedChannel createChannel(String target) { - return new FakeManagedChannel(target + "#" + channelSeq.getAndIncrement()); + return new FakeManagedChannel(target + "#" + this.channelSeq.getAndIncrement()); + } + + @Override + protected String resolveTarget(String target) { + return this.resolvedTarget; } @Override @@ -119,6 +212,55 @@ public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { } } + private static class OutOfOrderResolverGrpcClient extends RecordingGrpcClient { + + private final AtomicInteger resolutionSeq = new AtomicInteger(); + private final CountDownLatch staleResolutionStarted = new CountDownLatch(1); + private final CountDownLatch releaseStaleResolution = new CountDownLatch(1); + + @Override + protected String resolveTarget(String target) { + int sequence = this.resolutionSeq.incrementAndGet(); + if (sequence == 1) { + return "10.0.0.1"; + } + if (sequence == 2) { + this.staleResolutionStarted.countDown(); + try { + assertTrue("the stale resolution must be released", + this.releaseStaleResolution.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + return "10.0.0.1"; + } + return "10.0.0.2"; + } + } + + private static class StubInterleavingGrpcClient extends RecordingGrpcClient { + + private final AtomicInteger blockingStubSeq = new AtomicInteger(); + private final CountDownLatch staleStubBuildStarted = new CountDownLatch(1); + private final CountDownLatch releaseStaleStubBuild = new CountDownLatch(1); + + @Override + public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { + if (this.blockingStubSeq.incrementAndGet() == 1) { + this.staleStubBuildStarted.countDown(); + try { + assertTrue("the stale stub build must be released", + this.releaseStaleStubBuild.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + return super.getBlockingStub(channel); + } + } + private static class FakeBlockingStub extends AbstractBlockingStub { FakeBlockingStub(Channel channel, CallOptions callOptions) { @@ -171,7 +313,7 @@ public ManagedChannel shutdown() { @Override public ManagedChannel shutdownNow() { - return shutdown(); + return this.shutdown(); } @Override From 21e31f1a52eba709b1cf8568a378171d64522cfa Mon Sep 17 00:00:00 2001 From: Himanshu Verma Date: Fri, 31 Jul 2026 14:14:26 +0530 Subject: [PATCH 2/3] fix(store): address channel refresh review --- .../store/client/grpc/AbstractGrpcClient.java | 238 ++++++++-- .../client/grpc/AbstractGrpcClientTest.java | 442 +++++++++++++++--- 2 files changed, 565 insertions(+), 115 deletions(-) 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 1f7c87b51c..9454e32a26 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 @@ -18,7 +18,6 @@ package org.apache.hugegraph.store.client.grpc; import java.net.InetAddress; -import java.net.URI; import java.net.UnknownHostException; import java.util.Arrays; import java.util.Map; @@ -27,6 +26,8 @@ import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -44,8 +45,10 @@ public abstract class AbstractGrpcClient { protected static Map channels = new ConcurrentHashMap<>(); private static final Map resolvedTargets = new ConcurrentHashMap<>(); - private static final Map resolutionRequests = new ConcurrentHashMap<>(); - private static final Map appliedResolutions = new ConcurrentHashMap<>(); + private static final Map nextResolutions = new ConcurrentHashMap<>(); + private static final Map refreshLocks = new ConcurrentHashMap<>(); + private static final long DEFAULT_CHANNEL_REFRESH_INTERVAL_NANOS = + TimeUnit.SECONDS.toNanos(5L); private static final int n = 5; protected static int concurrency = 1 << n; private static final AtomicLong counter = new AtomicLong(0); @@ -71,26 +74,7 @@ public ManagedChannel[] getChannels(String target) { if ((tc = channels.get(target)) == null) { synchronized (channels) { if ((tc = channels.get(target)) == null) { - try { - ManagedChannel[] value = new ManagedChannel[concurrency]; - CountDownLatch latch = new CountDownLatch(concurrency); - for (int i = 0; i < concurrency; i++) { - int fi = i; - executor.execute(() -> { - try { - value[fi] = createChannel(target); - } catch (Exception e) { - throw new RuntimeException(e); - } finally { - latch.countDown(); - } - }); - } - latch.await(); - channels.put(target, tc = value); - } catch (Exception e) { - throw new RuntimeException(e); - } + channels.put(target, tc = this.createChannels(target)); } } } @@ -115,7 +99,7 @@ public AbstractBlockingStub getBlockingStub(String target) { HgPair[] value = new HgPair[concurrency]; IntStream.range(0, concurrency).forEach(i -> { - ManagedChannel channel = targetChannels[index]; + ManagedChannel channel = targetChannels[i]; AbstractBlockingStub stub = getBlockingStub(channel); value[i] = new HgPair<>(channel, stub); // log.info("create channel for {}",target); @@ -168,7 +152,7 @@ public AbstractAsyncStub getAsyncStub(String target) { HgPair[] value = new HgPair[concurrency]; IntStream.range(0, concurrency).parallel().forEach(i -> { - ManagedChannel channel = targetChannels[index]; + ManagedChannel channel = targetChannels[i]; AbstractAsyncStub stub = getAsyncStub(channel); // stub.withMaxInboundMessageSize( // config.getGrpcMaxInboundMessageSize()) @@ -225,46 +209,208 @@ private static boolean containsChannel(ManagedChannel[] channels, } private void refreshChannelsIfAddressChanged(String target) { - long resolutionRequest = resolutionRequests.computeIfAbsent(target, - key -> new AtomicLong()) - .incrementAndGet(); - String resolvedTarget = this.resolveTarget(target); - if (resolvedTarget.isEmpty()) { + if (!this.shouldRefreshChannels(target)) { + return; + } + + ReentrantLock refreshLock = refreshLocks.computeIfAbsent(target, + key -> new ReentrantLock()); + if (!refreshLock.tryLock()) { return; } - synchronized (channels) { - Long appliedResolution = appliedResolutions.get(target); - if (appliedResolution != null && appliedResolution >= resolutionRequest) { + + try { + if (!this.shouldRefreshChannels(target)) { + return; + } + + String resolvedTarget = this.resolveTarget(target); + this.postponeNextRefresh(target); + if (resolvedTarget.isEmpty()) { return; } - appliedResolutions.put(target, resolutionRequest); - String previousTarget = resolvedTargets.put(target, resolvedTarget); - if (previousTarget == null && !channels.containsKey(target)) { + + ManagedChannel[] staleChannels = channels.get(target); + String previousTarget = resolvedTargets.get(target); + if (previousTarget == null && staleChannels == null) { + resolvedTargets.put(target, resolvedTarget); return; } if (resolvedTarget.equals(previousTarget)) { return; } - ManagedChannel[] staleChannels = channels.remove(target); - if (staleChannels != null) { - Arrays.stream(staleChannels) - .filter(channel -> channel != null && !channel.isShutdown()) - .forEach(ManagedChannel::shutdownNow); + if (staleChannels == null) { + resolvedTargets.put(target, resolvedTarget); + return; + } + + ManagedChannel[] replacementChannels; + try { + replacementChannels = this.createChannels(target); + } catch (RuntimeException ignored) { + return; } + + boolean replaced = false; + synchronized (channels) { + if (channels.get(target) == staleChannels) { + channels.put(target, replacementChannels); + resolvedTargets.put(target, resolvedTarget); + replaced = true; + } + } + + if (replaced) { + this.retireChannels(staleChannels); + } else { + this.retireChannels(replacementChannels); + } + } finally { + refreshLock.unlock(); } } - protected String resolveTarget(String target) { + private boolean shouldRefreshChannels(String target) { + AtomicLong nextResolution = nextResolutions.computeIfAbsent(target, + key -> new AtomicLong()); + return System.nanoTime() - nextResolution.get() >= 0L; + } + + private void postponeNextRefresh(String target) { + long interval = Math.max(0L, this.channelRefreshIntervalNanos()); + nextResolutions.computeIfAbsent(target, key -> new AtomicLong()) + .set(System.nanoTime() + interval); + } + + protected long channelRefreshIntervalNanos() { + return DEFAULT_CHANNEL_REFRESH_INTERVAL_NANOS; + } + + protected long channelDrainTimeoutNanos() { + return TimeUnit.SECONDS.toNanos(config.getGrpcTimeoutSeconds()); + } + + private ManagedChannel[] createChannels(String target) { try { - String host = URI.create("dns://" + target).getHost(); - if (host == null) { + ManagedChannel[] value = new ManagedChannel[concurrency]; + CountDownLatch latch = new CountDownLatch(concurrency); + AtomicReference failure = new AtomicReference<>(); + for (int i = 0; i < concurrency; i++) { + int fi = i; + executor.execute(() -> { + try { + value[fi] = createChannel(target); + } catch (Exception e) { + failure.compareAndSet(null, new RuntimeException(e)); + } finally { + latch.countDown(); + } + }); + } + latch.await(); + if (failure.get() != null) { + throw failure.get(); + } + return value; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private void retireChannels(ManagedChannel[] retiredChannels) { + Arrays.stream(retiredChannels) + .filter(channel -> channel != null && !channel.isShutdown()) + .forEach(ManagedChannel::shutdown); + + executor.execute(() -> forceTerminateChannels(retiredChannels)); + } + + private void forceTerminateChannels(ManagedChannel[] retiredChannels) { + long deadline = System.nanoTime() + Math.max(0L, this.channelDrainTimeoutNanos()); + boolean interrupted = false; + + for (ManagedChannel channel : retiredChannels) { + if (channel == null || channel.isTerminated()) { + continue; + } + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0L || + !channel.awaitTermination(remaining, TimeUnit.NANOSECONDS)) { + channel.shutdownNow(); + } + } catch (InterruptedException e) { + interrupted = true; + channel.shutdownNow(); + } + } + + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private static String targetHost(String target) { + if (target == null || target.isEmpty()) { + return ""; + } + + String endpoint = target; + if (target.startsWith("dns://")) { + endpoint = target.substring("dns://".length()); + while (endpoint.startsWith("/")) { + endpoint = endpoint.substring(1); + } + int pathStart = endpoint.indexOf('/'); + if (pathStart >= 0) { + endpoint = endpoint.substring(pathStart + 1); + } + } else if (target.contains("://")) { + return ""; + } + + return endpointHost(endpoint); + } + + private static String endpointHost(String endpoint) { + if (endpoint == null || endpoint.isEmpty()) { + return ""; + } + + if (endpoint.charAt(0) == '[') { + int hostEnd = endpoint.indexOf(']'); + if (hostEnd <= 1) { return ""; } - return Arrays.stream(InetAddress.getAllByName(host)) + return endpoint.substring(1, hostEnd); + } + + int lastColon = endpoint.lastIndexOf(':'); + if (lastColon < 0) { + return endpoint; + } + if (endpoint.indexOf(':') != lastColon) { + return endpoint; + } + return endpoint.substring(0, lastColon); + } + + protected InetAddress[] resolveHost(String host) throws UnknownHostException { + return InetAddress.getAllByName(host); + } + + protected String resolveTarget(String target) { + String host = targetHost(target); + if (host.isEmpty()) { + return ""; + } + try { + return Arrays.stream(this.resolveHost(host)) .map(InetAddress::getHostAddress) .sorted() .collect(Collectors.joining(",")); - } catch (IllegalArgumentException | UnknownHostException ignored) { + } catch (UnknownHostException ignored) { return ""; } } 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 6b717929ff..5875ca6db0 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,22 +18,29 @@ 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.assertSame; import static org.junit.Assert.assertTrue; +import java.lang.reflect.Field; +import java.net.InetAddress; +import java.net.UnknownHostException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Map; 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 java.util.stream.Collectors; +import org.apache.hugegraph.store.term.HgPair; import org.junit.Test; import io.grpc.CallOptions; @@ -45,7 +52,7 @@ import io.grpc.stub.AbstractBlockingStub; /** - * Verifies that Store address changes replace channels and their cached stubs safely. + * Verifies that Store address changes replace channels and cached stubs safely. */ public class AbstractGrpcClientTest { @@ -60,6 +67,57 @@ private static boolean belongsToPool(Channel channel, return Arrays.stream(channels).anyMatch(current -> current == channel); } + private static boolean allChannelsAreShutdown(ManagedChannel[] channels) { + return Arrays.stream(channels).allMatch(ManagedChannel::isShutdown); + } + + private static boolean allChannelsAreLive(ManagedChannel[] channels) { + return Arrays.stream(channels).noneMatch(ManagedChannel::isShutdown); + } + + private static List fakeChannels(ManagedChannel[] channels) { + return Arrays.stream(channels) + .map(channel -> (FakeManagedChannel) channel) + .collect(Collectors.toList()); + } + + private static void assertUsesEveryChannel(String message, + List stubChannels, + ManagedChannel[] channels) { + assertEquals(message, channels.length, new HashSet<>(stubChannels).size()); + } + + private static void assertCachedChannelsCurrentAndLive(String message, + List cached, + ManagedChannel[] current) { + assertEquals(message, current.length, cached.size()); + assertTrue(message, cached.stream().allMatch(channel -> + belongsToPool(channel, current) && !channel.isShutdown())); + } + + private static void awaitCondition(String message, Condition condition) throws Exception { + for (int i = 0; i < 500; i++) { + if (condition.isTrue()) { + return; + } + Thread.sleep(10L); + } + assertTrue(message, condition.isTrue()); + } + + @SuppressWarnings("unchecked") + private static List cachedAsyncStubChannels(AbstractGrpcClient client, + String target) + throws Exception { + Field field = AbstractGrpcClient.class.getDeclaredField("asyncStubs"); + field.setAccessible(true); + Map[]> stubs = + (Map[]>) field.get(client); + HgPair[] pairs = stubs.get(target); + assertNotNull("the async stub cache must exist", pairs); + return Arrays.stream(pairs).map(HgPair::getKey).collect(Collectors.toList()); + } + @Test public void testAddressChangeReplacesChannelAndStubPools() { String target = uniqueTarget("address-change"); @@ -72,8 +130,11 @@ public void testAddressChangeReplacesChannelAndStubPools() { ManagedChannel[] newChannels = client.getChannels(target); assertNotSame("an address change must replace the channel pool", oldChannels, newChannels); - assertTrue("every stale channel must be shut down", - Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); + assertTrue("every stale channel must be gracefully shut down", + allChannelsAreShutdown(oldChannels)); + assertFalse("refresh must not force close stale channels immediately", + fakeChannels(oldChannels).stream() + .anyMatch(FakeManagedChannel::isForceShutdown)); client.blockingStubChannels.clear(); client.asyncStubChannels.clear(); @@ -85,12 +146,16 @@ public void testAddressChangeReplacesChannelAndStubPools() { client.blockingStubChannels.stream() .allMatch(channel -> belongsToPool(channel, newChannels))); + assertUsesEveryChannel("blocking stubs must be spread across the pool", + client.blockingStubChannels, newChannels); assertEquals("the async stub pool must be rebuilt", newChannels.length, client.asyncStubChannels.size()); assertTrue("replacement async stubs must use the new channel pool", client.asyncStubChannels.stream() .allMatch(channel -> belongsToPool(channel, newChannels))); + assertUsesEveryChannel("async stubs must be spread across the pool", + client.asyncStubChannels, newChannels); } @Test @@ -104,46 +169,112 @@ public void testFirstSuccessfulResolutionReplacesUnknownChannels() { ManagedChannel[] resolvedChannels = client.getChannels(target); assertNotSame("a pool with unknown addresses must be replaced", unknownChannels, resolvedChannels); - assertTrue("every channel from the unknown pool must be shut down", - Arrays.stream(unknownChannels).allMatch(ManagedChannel::isShutdown)); + assertTrue("every channel from the unknown pool must be gracefully shut down", + allChannelsAreShutdown(unknownChannels)); + assertFalse("unknown channels must not be force closed immediately", + fakeChannels(unknownChannels).stream() + .anyMatch(FakeManagedChannel::isForceShutdown)); assertTrue("the resolved channel pool must remain live", - Arrays.stream(resolvedChannels).noneMatch(ManagedChannel::isShutdown)); + allChannelsAreLive(resolvedChannels)); + } + + @Test + public void testStubAcquisitionReusesResolutionWithinRefreshInterval() { + String target = uniqueTarget("throttled-refresh"); + CountingResolverGrpcClient client = new CountingResolverGrpcClient(); + client.refreshIntervalNanos = TimeUnit.HOURS.toNanos(1L); + + assertNotNull(client.getBlockingStub(target)); + assertNotNull(client.getAsyncStub(target)); + for (int i = 0; i < 10; i++) { + assertNotNull(client.getBlockingStub(target)); + assertNotNull(client.getAsyncStub(target)); + } + + assertEquals("stub acquisition must not resolve again inside the refresh interval", + 1, client.resolutionCount.get()); } @Test - public void testOlderResolutionCannotReplaceNewerChannels() throws Exception { - String target = uniqueTarget("concurrent-address-change"); - OutOfOrderResolverGrpcClient client = new OutOfOrderResolverGrpcClient(); + public void testConcurrentStubAcquisitionRetainsHealthyPoolDuringDelayedRefresh() + throws Exception { + String target = uniqueTarget("delayed-refresh"); + DelayedResolverGrpcClient client = new DelayedResolverGrpcClient(); + client.refreshIntervalNanos = 0L; ManagedChannel[] oldChannels = client.getChannels(target); - ExecutorService executor = Executors.newFixedThreadPool(2); + assertNotNull(client.getBlockingStub(target)); + int resolutionsBeforeConcurrentCalls = client.resolutionCount.get(); + client.resolvedTarget = "10.0.0.2"; + client.delayChangedResolution = true; + client.refreshIntervalNanos = TimeUnit.HOURS.toNanos(1L); + + ExecutorService executor = Executors.newFixedThreadPool(6); + List> futures = new ArrayList<>(); try { - Future staleResolution = - executor.submit(() -> client.getChannels(target)); - assertTrue("the stale resolution must be in flight", - client.staleResolutionStarted.await(5, TimeUnit.SECONDS)); - Future freshResolution = - executor.submit(() -> client.getChannels(target)); - ManagedChannel[] freshChannels = freshResolution.get(5, TimeUnit.SECONDS); - - assertNotSame("the newer address must replace the old channel pool", - oldChannels, freshChannels); - client.releaseStaleResolution.countDown(); - assertSame("the late stale result must retain the newer channel pool", - freshChannels, staleResolution.get(5, TimeUnit.SECONDS)); - assertTrue("the replaced channel pool must be shut down", - Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); - assertTrue("the newer channel pool must remain live", - Arrays.stream(freshChannels).noneMatch(ManagedChannel::isShutdown)); + for (int i = 0; i < 6; i++) { + futures.add(executor.submit(() -> client.getBlockingStub(target))); + } + assertTrue("one refresh should be waiting in the delayed resolver", + client.delayedResolutionStarted.await(5, TimeUnit.SECONDS)); + awaitCondition("callers that miss the refresh lock must keep using the cache", + () -> futures.stream().anyMatch(Future::isDone)); + assertTrue("the existing healthy pool must stay live during refresh", + allChannelsAreLive(oldChannels)); + + client.releaseDelayedResolution.countDown(); + for (Future future : futures) { + assertNotNull(future.get(5, TimeUnit.SECONDS)); + } + + ManagedChannel[] currentChannels = client.getChannels(target); + assertNotSame("the completed refresh must publish a new channel pool", + oldChannels, currentChannels); + assertTrue("the previous pool must be retired after replacement is published", + allChannelsAreShutdown(oldChannels)); + assertEquals("concurrent callers must share a single refresh resolution", + resolutionsBeforeConcurrentCalls + 1, + client.resolutionCount.get()); } finally { - client.releaseStaleResolution.countDown(); + client.releaseDelayedResolution.countDown(); executor.shutdownNow(); } } @Test - public void testStubBuildRetriesAfterChannelRefresh() throws Exception { - String target = uniqueTarget("concurrent-stub-refresh"); + public void testRefreshGracefullyRetiresActiveStreamChannels() throws Exception { + String target = uniqueTarget("active-stream-refresh"); + ActiveRetirementGrpcClient client = new ActiveRetirementGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + AbstractAsyncStub activeStreamStub = client.getAsyncStub(target); + assertTrue("the simulated active stream must be on the old pool", + belongsToPool(activeStreamStub.getChannel(), oldChannels)); + + client.resolvedTarget = "10.0.0.2"; + ManagedChannel[] newChannels = client.getChannels(target); + assertNotSame("an address change must publish a replacement pool first", + oldChannels, newChannels); + List retiredChannels = fakeChannels(oldChannels); + assertTrue("the retired pool must receive graceful shutdown", + retiredChannels.stream().allMatch(FakeManagedChannel::isShutdown)); + assertFalse("active streams must not be force closed immediately", + retiredChannels.stream().anyMatch(FakeManagedChannel::isForceShutdown)); + assertTrue("retirement should wait for in-flight calls to drain", + retiredChannels.get(0) + .awaitTerminationStarted(5, TimeUnit.SECONDS)); + + client.finishActiveCalls(); + awaitCondition("retired channels should terminate after active calls finish", + () -> retiredChannels.stream().allMatch(FakeManagedChannel::isTerminated)); + assertFalse("drained channels must not need forced shutdown", + retiredChannels.stream().anyMatch(FakeManagedChannel::isForceShutdown)); + assertTrue("the replacement pool must remain live", + allChannelsAreLive(newChannels)); + } + + @Test + public void testBlockingStubBuildRetriesAfterChannelRefresh() throws Exception { + String target = uniqueTarget("concurrent-blocking-stub-refresh"); StubInterleavingGrpcClient client = new StubInterleavingGrpcClient(); ManagedChannel[] oldChannels = client.getChannels(target); ExecutorService executor = Executors.newFixedThreadPool(2); @@ -151,19 +282,14 @@ public void testStubBuildRetriesAfterChannelRefresh() throws Exception { try { Future staleStub = executor.submit(() -> client.getBlockingStub(target)); - assertTrue("the old stub pool build must be in flight", - client.staleStubBuildStarted.await(5, TimeUnit.SECONDS)); + assertTrue("the old blocking stub pool build must be in flight", + client.staleBlockingStubBuildStarted.await(5, TimeUnit.SECONDS)); client.resolvedTarget = "10.0.0.2"; Future freshStub = executor.submit(() -> client.getBlockingStub(target)); - for (int i = 0; i < 500 && - Arrays.stream(oldChannels).anyMatch(channel -> !channel.isShutdown()); - i++) { - Thread.sleep(10L); - } - assertTrue("refresh must shut down the old channel pool", - Arrays.stream(oldChannels).allMatch(ManagedChannel::isShutdown)); - client.releaseStaleStubBuild.countDown(); + awaitCondition("refresh must retire the old channel pool", + () -> allChannelsAreShutdown(oldChannels)); + client.releaseStaleBlockingStubBuild.countDown(); AbstractBlockingStub staleResult = staleStub.get(5, TimeUnit.SECONDS); AbstractBlockingStub freshResult = freshStub.get(5, TimeUnit.SECONDS); @@ -173,25 +299,94 @@ public void testStubBuildRetriesAfterChannelRefresh() throws Exception { assertTrue("the concurrent build must use the current pool", belongsToPool(freshResult.getChannel(), currentChannels)); assertTrue("the current channel pool must remain live", - Arrays.stream(currentChannels).noneMatch(ManagedChannel::isShutdown)); + allChannelsAreLive(currentChannels)); } finally { - client.releaseStaleStubBuild.countDown(); + client.releaseStaleBlockingStubBuild.countDown(); executor.shutdownNow(); } } + @Test + public void testAsyncStubBuildRetriesAfterChannelRefresh() throws Exception { + String target = uniqueTarget("concurrent-async-stub-refresh"); + StubInterleavingGrpcClient client = new StubInterleavingGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future staleStub = + executor.submit(() -> client.getAsyncStub(target)); + assertTrue("the old async stub pool build must be in flight", + client.staleAsyncStubBuildStarted.await(5, TimeUnit.SECONDS)); + client.resolvedTarget = "10.0.0.2"; + Future freshStub = + executor.submit(() -> client.getAsyncStub(target)); + awaitCondition("refresh must retire the old channel pool", + () -> allChannelsAreShutdown(oldChannels)); + client.releaseStaleAsyncStubBuild.countDown(); + + AbstractAsyncStub staleResult = staleStub.get(5, TimeUnit.SECONDS); + AbstractAsyncStub freshResult = freshStub.get(5, TimeUnit.SECONDS); + ManagedChannel[] currentChannels = client.getChannels(target); + assertTrue("the stale async build must retry against the current pool", + belongsToPool(staleResult.getChannel(), currentChannels)); + assertTrue("the concurrent async build must use the current pool", + belongsToPool(freshResult.getChannel(), currentChannels)); + assertCachedChannelsCurrentAndLive( + "the final async cache must only reference current live channels", + cachedAsyncStubChannels(client, target), currentChannels); + } finally { + client.releaseStaleAsyncStubBuild.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void testResolveTargetSupportsDnsUriAndBracketedIpv6Targets() { + HostCapturingGrpcClient dnsClient = new HostCapturingGrpcClient(); + assertEquals("10.0.0.1", dnsClient.resolveTarget("dns:///store.example.com:8500")); + assertEquals("store.example.com", dnsClient.capturedHost); + + HostCapturingGrpcClient ipv6Client = new HostCapturingGrpcClient(); + assertEquals("10.0.0.1", ipv6Client.resolveTarget("[2001:db8::1]:8500")); + assertEquals("2001:db8::1", ipv6Client.capturedHost); + } + + @Test + public void testResolveTargetSkipsUnsupportedGrpcSchemes() { + HostCapturingGrpcClient client = new HostCapturingGrpcClient(); + assertEquals("", client.resolveTarget("unix:///var/run/store.sock")); + assertEquals("unsupported schemes must not invoke DNS resolution", + 0, client.resolutionCount.get()); + } + + private interface Condition { + + boolean isTrue(); + } + private static class RecordingGrpcClient extends AbstractGrpcClient { private final AtomicInteger channelSeq = new AtomicInteger(); protected volatile String resolvedTarget = "10.0.0.1"; - private final List blockingStubChannels = + protected volatile long refreshIntervalNanos = 0L; + protected final List blockingStubChannels = Collections.synchronizedList(new ArrayList<>()); - private final List asyncStubChannels = + protected final List asyncStubChannels = Collections.synchronizedList(new ArrayList<>()); + @Override + protected long channelRefreshIntervalNanos() { + return this.refreshIntervalNanos; + } + @Override protected ManagedChannel createChannel(String target) { - return new FakeManagedChannel(target + "#" + this.channelSeq.getAndIncrement()); + return this.newFakeChannel(target + "#" + this.channelSeq.getAndIncrement()); + } + + protected FakeManagedChannel newFakeChannel(String authority) { + return new FakeManagedChannel(authority); } @Override @@ -212,46 +407,75 @@ public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { } } - private static class OutOfOrderResolverGrpcClient extends RecordingGrpcClient { + private static class CountingResolverGrpcClient extends RecordingGrpcClient { - private final AtomicInteger resolutionSeq = new AtomicInteger(); - private final CountDownLatch staleResolutionStarted = new CountDownLatch(1); - private final CountDownLatch releaseStaleResolution = new CountDownLatch(1); + protected final AtomicInteger resolutionCount = new AtomicInteger(); @Override protected String resolveTarget(String target) { - int sequence = this.resolutionSeq.incrementAndGet(); - if (sequence == 1) { - return "10.0.0.1"; - } - if (sequence == 2) { - this.staleResolutionStarted.countDown(); + this.resolutionCount.incrementAndGet(); + return super.resolveTarget(target); + } + } + + private static class DelayedResolverGrpcClient extends CountingResolverGrpcClient { + + private final CountDownLatch delayedResolutionStarted = new CountDownLatch(1); + private final CountDownLatch releaseDelayedResolution = new CountDownLatch(1); + private volatile boolean delayChangedResolution; + + @Override + protected String resolveTarget(String target) { + this.resolutionCount.incrementAndGet(); + if (this.delayChangedResolution) { + this.delayedResolutionStarted.countDown(); try { - assertTrue("the stale resolution must be released", - this.releaseStaleResolution.await(5, TimeUnit.SECONDS)); + assertTrue("the delayed resolution must be released", + this.releaseDelayedResolution.await(5, TimeUnit.SECONDS)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); } - return "10.0.0.1"; } - return "10.0.0.2"; + return this.resolvedTarget; + } + } + + private static class ActiveRetirementGrpcClient extends RecordingGrpcClient { + + private final CountDownLatch activeCallsFinished = new CountDownLatch(1); + + @Override + protected long channelDrainTimeoutNanos() { + return TimeUnit.SECONDS.toNanos(5L); + } + + @Override + protected FakeManagedChannel newFakeChannel(String authority) { + return new FakeManagedChannel(authority, this.activeCallsFinished); + } + + private void finishActiveCalls() { + this.activeCallsFinished.countDown(); } } private static class StubInterleavingGrpcClient extends RecordingGrpcClient { private final AtomicInteger blockingStubSeq = new AtomicInteger(); - private final CountDownLatch staleStubBuildStarted = new CountDownLatch(1); - private final CountDownLatch releaseStaleStubBuild = new CountDownLatch(1); + private final AtomicInteger asyncStubSeq = new AtomicInteger(); + private final CountDownLatch staleBlockingStubBuildStarted = new CountDownLatch(1); + private final CountDownLatch releaseStaleBlockingStubBuild = new CountDownLatch(1); + private final CountDownLatch staleAsyncStubBuildStarted = new CountDownLatch(1); + private final CountDownLatch releaseStaleAsyncStubBuild = new CountDownLatch(1); @Override public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { if (this.blockingStubSeq.incrementAndGet() == 1) { - this.staleStubBuildStarted.countDown(); + this.staleBlockingStubBuildStarted.countDown(); try { - assertTrue("the stale stub build must be released", - this.releaseStaleStubBuild.await(5, TimeUnit.SECONDS)); + assertTrue("the stale blocking stub build must be released", + this.releaseStaleBlockingStubBuild.await(5, TimeUnit.SECONDS)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); @@ -259,6 +483,49 @@ public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { } return super.getBlockingStub(channel); } + + @Override + public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { + if (this.asyncStubSeq.incrementAndGet() == 1) { + this.staleAsyncStubBuildStarted.countDown(); + try { + assertTrue("the stale async stub build must be released", + this.releaseStaleAsyncStubBuild.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + return super.getAsyncStub(channel); + } + } + + private static class HostCapturingGrpcClient extends AbstractGrpcClient { + + private final AtomicInteger resolutionCount = new AtomicInteger(); + private volatile String capturedHost; + + @Override + protected ManagedChannel createChannel(String target) { + return new FakeManagedChannel(target); + } + + @Override + public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { + return new FakeBlockingStub(channel, CallOptions.DEFAULT); + } + + @Override + public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { + return new FakeAsyncStub(channel, CallOptions.DEFAULT); + } + + @Override + protected InetAddress[] resolveHost(String host) throws UnknownHostException { + this.resolutionCount.incrementAndGet(); + this.capturedHost = host; + return new InetAddress[]{InetAddress.getByName("10.0.0.1")}; + } } private static class FakeBlockingStub extends AbstractBlockingStub { @@ -288,10 +555,19 @@ protected FakeAsyncStub build(Channel channel, CallOptions callOptions) { private static class FakeManagedChannel extends ManagedChannel { private final String authority; + private final CountDownLatch activeCallsFinished; + private final CountDownLatch awaitTerminationStarted = new CountDownLatch(1); private volatile boolean shutdown; + private volatile boolean forceShutdown; + private volatile boolean terminated; FakeManagedChannel(String authority) { + this(authority, null); + } + + FakeManagedChannel(String authority, CountDownLatch activeCallsFinished) { this.authority = authority; + this.activeCallsFinished = activeCallsFinished; } @Override @@ -308,12 +584,18 @@ public ClientCall newCall(MethodDescriptor method, @Override public ManagedChannel shutdown() { this.shutdown = true; + if (this.activeCallsFinished == null) { + this.terminated = true; + } return this; } @Override public ManagedChannel shutdownNow() { - return this.shutdown(); + this.shutdown = true; + this.forceShutdown = true; + this.terminated = true; + return this; } @Override @@ -321,14 +603,36 @@ public boolean isShutdown() { return this.shutdown; } + boolean isForceShutdown() { + return this.forceShutdown; + } + @Override public boolean isTerminated() { - return this.shutdown; + return this.terminated; + } + + boolean awaitTerminationStarted(long timeout, TimeUnit unit) + throws InterruptedException { + return this.awaitTerminationStarted.await(timeout, unit); } @Override - public boolean awaitTermination(long timeout, TimeUnit unit) { - return this.shutdown; + public boolean awaitTermination(long timeout, TimeUnit unit) + throws InterruptedException { + this.awaitTerminationStarted.countDown(); + if (this.terminated) { + return true; + } + if (this.activeCallsFinished == null) { + this.terminated = this.shutdown; + return this.terminated; + } + if (this.activeCallsFinished.await(timeout, unit)) { + this.terminated = true; + return true; + } + return false; } } } From ddeef7a9adbbb19b28f45c48229a399ffa920008 Mon Sep 17 00:00:00 2001 From: imbajin Date: Fri, 31 Jul 2026 21:13:21 +0800 Subject: [PATCH 3/3] fix(store): harden channel refresh cleanup - move graceful retirement to a cleanup scheduler - force-close partial pools after creation failures - validate cached stubs against channels by index - cover saturation, interruption, and drain deadlines --- .../store/client/grpc/AbstractGrpcClient.java | 94 ++++--- .../client/grpc/AbstractGrpcClientTest.java | 238 +++++++++++++++++- 2 files changed, 279 insertions(+), 53 deletions(-) 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 9454e32a26..8b069d5a80 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 @@ -23,6 +23,7 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; @@ -47,6 +48,9 @@ public abstract class AbstractGrpcClient { private static final Map resolvedTargets = new ConcurrentHashMap<>(); private static final Map nextResolutions = new ConcurrentHashMap<>(); private static final Map refreshLocks = new ConcurrentHashMap<>(); + private static final ScheduledThreadPoolExecutor CHANNEL_CLEANUP_EXECUTOR = + new ScheduledThreadPoolExecutor( + 1, ExecutorPool.newThreadFactory("channel-cleanup")); private static final long DEFAULT_CHANNEL_REFRESH_INTERVAL_NANOS = TimeUnit.SECONDS.toNanos(5L); private static final int n = 5; @@ -194,20 +198,15 @@ private static boolean usesChannels(HgPair[] pairs, if (pairs == null || pairs.length != channels.length) { return false; } - for (HgPair pair : pairs) { - if (pair == null || pair.getKey() == null || - !containsChannel(channels, pair.getKey())) { + for (int i = 0; i < pairs.length; i++) { + HgPair pair = pairs[i]; + if (pair == null || pair.getKey() != channels[i]) { return false; } } return true; } - private static boolean containsChannel(ManagedChannel[] channels, - ManagedChannel expected) { - return Arrays.stream(channels).anyMatch(channel -> channel == expected); - } - private void refreshChannelsIfAddressChanged(String target) { if (!this.shouldRefreshChannels(target)) { return; @@ -291,31 +290,42 @@ protected long channelDrainTimeoutNanos() { } private ManagedChannel[] createChannels(String target) { - try { - ManagedChannel[] value = new ManagedChannel[concurrency]; - CountDownLatch latch = new CountDownLatch(concurrency); - AtomicReference failure = new AtomicReference<>(); - for (int i = 0; i < concurrency; i++) { - int fi = i; - executor.execute(() -> { - try { - value[fi] = createChannel(target); - } catch (Exception e) { - failure.compareAndSet(null, new RuntimeException(e)); - } finally { - latch.countDown(); - } - }); - } - latch.await(); - if (failure.get() != null) { - throw failure.get(); + ManagedChannel[] value = new ManagedChannel[concurrency]; + CountDownLatch latch = new CountDownLatch(concurrency); + AtomicReference failure = new AtomicReference<>(); + for (int i = 0; i < concurrency; i++) { + int fi = i; + executor.execute(() -> { + try { + value[fi] = createChannel(target); + } catch (Exception e) { + failure.compareAndSet(null, new RuntimeException(e)); + } finally { + latch.countDown(); + } + }); + } + + InterruptedException interruption = null; + while (latch.getCount() > 0L) { + try { + latch.await(); + } catch (InterruptedException e) { + interruption = e; } - return value; - } catch (InterruptedException e) { + } + + if (failure.get() != null || interruption != null) { + forceTerminateChannels(value); + } + if (interruption != null) { Thread.currentThread().interrupt(); - throw new RuntimeException(e); + throw new RuntimeException(interruption); + } + if (failure.get() != null) { + throw failure.get(); } + return value; } private void retireChannels(ManagedChannel[] retiredChannels) { @@ -323,32 +333,18 @@ private void retireChannels(ManagedChannel[] retiredChannels) { .filter(channel -> channel != null && !channel.isShutdown()) .forEach(ManagedChannel::shutdown); - executor.execute(() -> forceTerminateChannels(retiredChannels)); + long timeout = Math.max(0L, this.channelDrainTimeoutNanos()); + CHANNEL_CLEANUP_EXECUTOR.schedule( + () -> forceTerminateChannels(retiredChannels), timeout, + TimeUnit.NANOSECONDS); } private void forceTerminateChannels(ManagedChannel[] retiredChannels) { - long deadline = System.nanoTime() + Math.max(0L, this.channelDrainTimeoutNanos()); - boolean interrupted = false; - for (ManagedChannel channel : retiredChannels) { - if (channel == null || channel.isTerminated()) { - continue; - } - try { - long remaining = deadline - System.nanoTime(); - if (remaining <= 0L || - !channel.awaitTermination(remaining, TimeUnit.NANOSECONDS)) { - channel.shutdownNow(); - } - } catch (InterruptedException e) { - interrupted = true; + if (channel != null && !channel.isTerminated()) { channel.shutdownNow(); } } - - if (interrupted) { - Thread.currentThread().interrupt(); - } } private static String targetHost(String target) { 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 5875ca6db0..fbcefa5845 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 @@ -36,8 +36,11 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.ThreadPoolExecutor; 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.stream.Collectors; import org.apache.hugegraph.store.term.HgPair; @@ -118,6 +121,25 @@ private static List cachedAsyncStubChannels(AbstractGrpcClient c return Arrays.stream(pairs).map(HgPair::getKey).collect(Collectors.toList()); } + @SuppressWarnings("unchecked") + private static HgPair[] cachedBlockingStubs( + AbstractGrpcClient client, String target) throws Exception { + Field field = AbstractGrpcClient.class.getDeclaredField("blockingStubs"); + field.setAccessible(true); + Map[]> stubs = + (Map[]>) field.get(client); + HgPair[] pairs = stubs.get(target); + assertNotNull("the blocking stub cache must exist", pairs); + return pairs; + } + + private static ThreadPoolExecutor channelCreationExecutor(AbstractGrpcClient client) + throws Exception { + Field field = AbstractGrpcClient.class.getDeclaredField("executor"); + field.setAccessible(true); + return (ThreadPoolExecutor) field.get(client); + } + @Test public void testAddressChangeReplacesChannelAndStubPools() { String target = uniqueTarget("address-change"); @@ -178,6 +200,153 @@ public void testFirstSuccessfulResolutionReplacesUnknownChannels() { allChannelsAreLive(resolvedChannels)); } + @Test + public void testRetirementDoesNotBlockWhenCreationExecutorIsSaturated() + throws Exception { + String target = uniqueTarget("saturated-retirement"); + ActiveRetirementGrpcClient client = new ActiveRetirementGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + ThreadPoolExecutor channelExecutor = channelCreationExecutor(client); + CountDownLatch workersStarted = new CountDownLatch(AbstractGrpcClient.concurrency); + CountDownLatch releaseWorkers = new CountDownLatch(1); + ExecutorService caller = Executors.newSingleThreadExecutor(); + + try { + for (int i = 0; i < AbstractGrpcClient.concurrency; i++) { + channelExecutor.execute(() -> { + workersStarted.countDown(); + try { + releaseWorkers.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + } + assertTrue("the shared creation executor must be fully occupied", + workersStarted.await(5, TimeUnit.SECONDS)); + + client.resolvedTarget = "10.0.0.2"; + Future refreshed = + caller.submit(() -> client.getChannels(target)); + ManagedChannel[] newChannels = refreshed.get(1, TimeUnit.SECONDS); + + assertNotSame("refresh must publish the replacement pool", + oldChannels, newChannels); + assertTrue("the retired pool must be shut down before refresh returns", + allChannelsAreShutdown(oldChannels)); + assertFalse("the scheduled cleanup must preserve the drain window", + fakeChannels(oldChannels).stream() + .anyMatch(FakeManagedChannel::isForceShutdown)); + } finally { + releaseWorkers.countDown(); + client.finishActiveCalls(); + caller.shutdownNow(); + } + } + + @Test + public void testPartialChannelsAreRetiredAfterMixedCreationFailure() + throws Exception { + String target = uniqueTarget("partial-creation-failure"); + FailingCreationGrpcClient client = new FailingCreationGrpcClient(5); + + try { + client.getChannels(target); + assertTrue("channel creation must propagate the injected failure", false); + } catch (RuntimeException ignored) { + // Expected. + } + + assertEquals("all creation tasks must converge before failure is returned", + AbstractGrpcClient.concurrency - 1, client.createdChannels.size()); + assertTrue("every partial channel must be force terminated before failure returns", + client.createdChannels.stream().allMatch(channel -> + channel.isTerminated() && + ((FakeManagedChannel) channel).isForceShutdown())); + } + + @Test + public void testInterruptedCreationWaitsAndRetiresPartialChannels() + throws Exception { + String target = uniqueTarget("interrupted-creation"); + DelayedCreationGrpcClient client = new DelayedCreationGrpcClient(); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean interrupted = new AtomicBoolean(); + Thread caller = new Thread(() -> { + try { + client.getChannels(target); + } catch (Throwable e) { + failure.set(e); + interrupted.set(Thread.currentThread().isInterrupted()); + } + }); + + caller.start(); + assertTrue("all channel creation tasks must start", + client.creationStarted.await(5, TimeUnit.SECONDS)); + caller.interrupt(); + try { + Thread.sleep(100L); + assertTrue("an interrupted caller must wait for creation tasks to converge", + caller.isAlive()); + } finally { + client.releaseCreation.countDown(); + } + caller.join(TimeUnit.SECONDS.toMillis(5L)); + + assertFalse("the interrupted creation call must finish", caller.isAlive()); + assertTrue("interruption must be reported as a runtime failure", + failure.get() instanceof RuntimeException); + assertTrue("the caller interrupt status must be restored", interrupted.get()); + assertEquals("every creation task must finish before interruption is reported", + AbstractGrpcClient.concurrency, client.createdChannels.size()); + assertTrue("all partial channels must be force terminated before interruption returns", + client.createdChannels.stream().allMatch(channel -> + channel.isTerminated() && + ((FakeManagedChannel) channel).isForceShutdown())); + } + + @Test + public void testDrainDeadlineForceTerminatesRetiredChannels() throws Exception { + String target = uniqueTarget("drain-deadline"); + ImmediateRetirementGrpcClient client = new ImmediateRetirementGrpcClient(); + ManagedChannel[] oldChannels = client.getChannels(target); + + client.resolvedTarget = "10.0.0.2"; + ManagedChannel[] newChannels = client.getChannels(target); + + assertNotSame("refresh must publish a replacement pool", oldChannels, newChannels); + awaitCondition("expired drain deadline must force terminate the retired pool", + () -> fakeChannels(oldChannels).stream().allMatch(channel -> + channel.isTerminated() && channel.isForceShutdown())); + assertTrue("the replacement pool must remain live", allChannelsAreLive(newChannels)); + } + + @Test + public void testStubCacheRequiresIndexIdentityMapping() throws Exception { + String target = uniqueTarget("index-mapping"); + RecordingGrpcClient client = new RecordingGrpcClient(); + ManagedChannel[] targetChannels = client.getChannels(target); + assertNotNull(client.getBlockingStub(target)); + HgPair[] pairs = + cachedBlockingStubs(client, target); + HgPair first = pairs[0]; + pairs[0] = pairs[1]; + pairs[1] = first; + client.blockingStubChannels.clear(); + + assertNotNull(client.getBlockingStub(target)); + + assertEquals("a permuted cache must be rebuilt instead of reused", + AbstractGrpcClient.concurrency, client.blockingStubChannels.size()); + HgPair[] rebuilt = + cachedBlockingStubs(client, target); + for (int i = 0; i < targetChannels.length; i++) { + assertTrue("each cached stub must map to the channel at the same index", + rebuilt[i].getKey() == targetChannels[i]); + } + } + @Test public void testStubAcquisitionReusesResolutionWithinRefreshInterval() { String target = uniqueTarget("throttled-refresh"); @@ -259,10 +428,6 @@ public void testRefreshGracefullyRetiresActiveStreamChannels() throws Exception retiredChannels.stream().allMatch(FakeManagedChannel::isShutdown)); assertFalse("active streams must not be force closed immediately", retiredChannels.stream().anyMatch(FakeManagedChannel::isForceShutdown)); - assertTrue("retirement should wait for in-flight calls to drain", - retiredChannels.get(0) - .awaitTerminationStarted(5, TimeUnit.SECONDS)); - client.finishActiveCalls(); awaitCondition("retired channels should terminate after active calls finish", () -> retiredChannels.stream().allMatch(FakeManagedChannel::isTerminated)); @@ -460,6 +625,67 @@ private void finishActiveCalls() { } } + private static class ImmediateRetirementGrpcClient extends RecordingGrpcClient { + + private final CountDownLatch activeCallsFinished = new CountDownLatch(1); + + @Override + protected long channelDrainTimeoutNanos() { + return 0L; + } + + @Override + protected FakeManagedChannel newFakeChannel(String authority) { + return new FakeManagedChannel(authority, this.activeCallsFinished); + } + } + + private static class FailingCreationGrpcClient extends RecordingGrpcClient { + + private final int failedAttempt; + private final AtomicInteger attempt = new AtomicInteger(); + private final List createdChannels = + Collections.synchronizedList(new ArrayList<>()); + + private FailingCreationGrpcClient(int failedAttempt) { + this.failedAttempt = failedAttempt; + } + + @Override + protected ManagedChannel createChannel(String target) { + int current = this.attempt.getAndIncrement(); + if (current == this.failedAttempt) { + throw new IllegalStateException("injected channel creation failure"); + } + ManagedChannel channel = this.newFakeChannel(target + "#" + current); + this.createdChannels.add(channel); + return channel; + } + } + + private static class DelayedCreationGrpcClient extends RecordingGrpcClient { + + private final CountDownLatch creationStarted = + new CountDownLatch(AbstractGrpcClient.concurrency); + private final CountDownLatch releaseCreation = new CountDownLatch(1); + private final List createdChannels = + Collections.synchronizedList(new ArrayList<>()); + + @Override + protected ManagedChannel createChannel(String target) { + this.creationStarted.countDown(); + try { + this.releaseCreation.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + ManagedChannel channel = this.newFakeChannel(target); + this.createdChannels.add(channel); + return channel; + } + } + private static class StubInterleavingGrpcClient extends RecordingGrpcClient { private final AtomicInteger blockingStubSeq = new AtomicInteger(); @@ -609,6 +835,10 @@ boolean isForceShutdown() { @Override public boolean isTerminated() { + if (this.shutdown && this.activeCallsFinished != null && + this.activeCallsFinished.getCount() == 0L) { + this.terminated = true; + } return this.terminated; }