From 4665c078fc268b457089ff878d0cdf8838eb1992 Mon Sep 17 00:00:00 2001 From: Himanshu Verma Date: Thu, 30 Jul 2026 18:46:33 +0530 Subject: [PATCH 1/4] 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/4] 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/4] 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; } From 198de19e9b0fd277752359c5874b03c51917d831 Mon Sep 17 00:00:00 2001 From: Himanshu Verma Date: Sun, 2 Aug 2026 20:58:54 +0530 Subject: [PATCH 4/4] fix(store): resolve store addresses off the request thread Channel refresh resolved DNS on whichever thread asked for a stub. Under the default launcher that thread can be a Gremlin worker, and HugeSecurityManager denies it socket access, so an HStore-backed request could fail with a SecurityException instead of using the healthy pool. - run resolution, replacement creation and retirement on a channel maintenance executor, keeping the last healthy pool when resolution fails or times out - build the first pool for a target once its address is known, so a cold start no longer creates and immediately retires a pool - replace the per-target refresh lock with a single-flight task map that the cold path can also wait on, and throttle from both submission and completion - route QueryV2Client through the guarded async stub path instead of taking a channel straight from the pool, and restrict getChannels to subclasses - drop the channel monitor from stub acquisition: publishing a pool before retiring the previous one already orders the check, and the monitor is static - log refresh failures and pool replacements, which the executor otherwise discards, and never let a denied thread creation wedge refresh for a target - parse targets with URI, rejecting resolver schemes such as unix:/path that were resolved as a host named after the scheme - fold the blocking and async stub acquisition loops into one implementation --- .../store/client/grpc/AbstractGrpcClient.java | 377 ++++++---- .../store/client/query/QueryV2Client.java | 12 +- .../store/client/ClientSuiteTest.java | 4 + .../client/grpc/AbstractGrpcClientTest.java | 664 ++++++++++-------- 4 files changed, 597 insertions(+), 460 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 8b069d5a80..cf277880ff 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,9 +18,12 @@ package org.apache.hugegraph.store.client.grpc; import java.net.InetAddress; +import java.net.URI; +import java.net.URISyntaxException; import java.net.UnknownHostException; import java.util.Arrays; import java.util.Map; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ScheduledThreadPoolExecutor; @@ -28,7 +31,7 @@ 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.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -41,18 +44,42 @@ import io.grpc.stub.AbstractAsyncStub; import io.grpc.stub.AbstractBlockingStub; import io.grpc.stub.AbstractStub; +import lombok.extern.slf4j.Slf4j; +@Slf4j public abstract class AbstractGrpcClient { protected static Map channels = new ConcurrentHashMap<>(); 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 = + private static final Map> refreshTasks = + new ConcurrentHashMap<>(); + /* + * Refresh runs here rather than on a request thread: a caller of getChannels() may hold a + * Gremlin worker stack, which HugeSecurityManager denies socket access to. Creating the very + * first pool for a target is still done by the caller, so that path stays exposed. + */ + private static final ScheduledThreadPoolExecutor CHANNEL_MAINTENANCE_EXECUTOR = new ScheduledThreadPoolExecutor( - 1, ExecutorPool.newThreadFactory("channel-cleanup")); + 2, ExecutorPool.newThreadFactory("channel-maintenance")); private static final long DEFAULT_CHANNEL_REFRESH_INTERVAL_NANOS = TimeUnit.SECONDS.toNanos(5L); + private static final long DEFAULT_INITIAL_RESOLUTION_TIMEOUT_NANOS = + TimeUnit.SECONDS.toNanos(1L); + private static final String DNS_SCHEME = "dns:"; + + static { + try { + // Create the maintenance threads eagerly, so no request thread ever creates one. + CHANNEL_MAINTENANCE_EXECUTOR.prestartAllCoreThreads(); + } catch (Throwable e) { + // A denied prestart must not leave this class permanently uninitializable, but a + // request thread has to create the thread instead, where it may be denied again. + log.warn("Failed to start the channel maintenance threads eagerly, " + + "channel refresh may be delayed until a permitted thread submits one", e); + } + } + private static final int n = 5; protected static int concurrency = 1 << n; private static final AtomicLong counter = new AtomicLong(0); @@ -72,14 +99,22 @@ public AbstractGrpcClient() { } - public ManagedChannel[] getChannels(String target) { - this.refreshChannelsIfAddressChanged(target); - ManagedChannel[] tc; - if ((tc = channels.get(target)) == null) { - synchronized (channels) { - if ((tc = channels.get(target)) == null) { - channels.put(target, tc = this.createChannels(target)); - } + protected ManagedChannel[] getChannels(String target) { + CompletableFuture refresh = this.triggerChannelRefresh(target); + ManagedChannel[] tc = channels.get(target); + if (tc != null) { + return tc; + } + + /* + * Only the very first pool for a target waits, and only for a bounded time: building it + * before its address is known makes the resolution that lands next rebuild it. Waiting + * avoids that in the common case; if the wait expires the rebuild still happens. + */ + this.awaitInitialResolution(refresh); + synchronized (channels) { + if ((tc = channels.get(target)) == null) { + channels.put(target, tc = this.createChannels(target)); } } return tc; @@ -88,44 +123,58 @@ public ManagedChannel[] getChannels(String target) { public abstract AbstractBlockingStub getBlockingStub(ManagedChannel channel); public AbstractBlockingStub getBlockingStub(String target) { + return this.acquireStub(target, this.blockingStubs, this::getBlockingStub, + stub -> (AbstractBlockingStub) this.setBlockingStubOption(stub)); + } + + /** + * Returns a cached stub bound to a channel of the target's current pool, rebuilding the + * cache when the pool has been replaced. The stub comes from the pool that was published at + * the last check; a refresh landing immediately afterwards can still retire that pool, so + * callers are not shielded from an in-flight replacement. + * + *

The pool check needs no lock: the pool is published before the previous one is retired, + * so reading the current pool from the map is enough to know retirement has not started. + */ + @SuppressWarnings("unchecked") + private S acquireStub(String target, + Map[]> stubCache, + Function stubFactory, + Function stubOption) { 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)); + ManagedChannel[] targetChannels = this.getChannels(target); + HgPair[] pairs = stubCache.get(target); + int index = nextStubIndex(); if (!usesChannels(pairs, targetChannels)) { - synchronized (blockingStubs) { - pairs = blockingStubs.get(target); + synchronized (stubCache) { + pairs = stubCache.get(target); if (!usesChannels(pairs, targetChannels)) { - HgPair[] value = - new HgPair[concurrency]; + HgPair[] value = new HgPair[concurrency]; IntStream.range(0, concurrency).forEach(i -> { ManagedChannel channel = targetChannels[i]; - AbstractBlockingStub stub = getBlockingStub(channel); - value[i] = new HgPair<>(channel, stub); - // log.info("create channel for {}",target); + value[i] = new HgPair<>(channel, stubFactory.apply(channel)); }); - synchronized (channels) { - if (channels.get(target) != targetChannels) { - continue; - } - blockingStubs.put(target, value); - AbstractBlockingStub stub = value[index].getValue(); - return (AbstractBlockingStub) setBlockingStubOption(stub); + if (channels.get(target) != targetChannels) { + continue; } + stubCache.put(target, value); + return stubOption.apply(value[index].getValue()); } } } - synchronized (channels) { - if (channels.get(target) != targetChannels) { - continue; - } - return (AbstractBlockingStub) setBlockingStubOption(pairs[index].getValue()); + if (channels.get(target) != targetChannels) { + continue; } + return stubOption.apply(pairs[index].getValue()); + } + } + + private static int nextStubIndex() { + long l = counter.getAndIncrement(); + if (l >= limit) { + counter.set(0); } + return (int) (l & (concurrency - 1)); } private AbstractStub setBlockingStubOption(AbstractBlockingStub stub) { @@ -141,49 +190,8 @@ public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { } public AbstractAsyncStub getAsyncStub(String target) { - 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[i]; - 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 this.acquireStub(target, this.asyncStubs, this::getAsyncStub, + stub -> (AbstractAsyncStub) this.setStubOption(stub)); } protected AbstractStub setStubOption(AbstractStub value) { @@ -207,66 +215,125 @@ private static boolean usesChannels(HgPair[] pairs, return true; } - private void refreshChannelsIfAddressChanged(String target) { + /** + * Submits a refresh for the target unless one is already in flight or the refresh interval + * has not elapsed. Returns the in-flight refresh, or null when none is running. + */ + private CompletableFuture triggerChannelRefresh(String target) { + CompletableFuture inFlight = refreshTasks.get(target); + if (inFlight != null) { + return inFlight; + } if (!this.shouldRefreshChannels(target)) { - return; + return null; } - ReentrantLock refreshLock = refreshLocks.computeIfAbsent(target, - key -> new ReentrantLock()); - if (!refreshLock.tryLock()) { - return; + CompletableFuture refresh = new CompletableFuture<>(); + CompletableFuture running = refreshTasks.putIfAbsent(target, refresh); + if (running != null) { + return running; } + // Throttle before submitting, so that a failing resolver cannot be retried in a loop. + this.postponeNextRefresh(target); try { - if (!this.shouldRefreshChannels(target)) { - return; - } + this.submitChannelRefresh(() -> { + try { + this.refreshChannelsIfAddressChanged(target); + } catch (Throwable e) { + // The executor discards what a task throws, so report it here. + log.warn("Failed to refresh channels of target {}", target, e); + } finally { + this.completeRefresh(target, refresh); + } + }); + } catch (Throwable e) { + // Includes a thread creation denied on this thread; never leave the entry behind. + log.warn("Failed to submit a channel refresh for target {}", target, e); + this.completeRefresh(target, refresh); + } + return refresh; + } - String resolvedTarget = this.resolveTarget(target); - this.postponeNextRefresh(target); - if (resolvedTarget.isEmpty()) { - return; - } + private void completeRefresh(String target, CompletableFuture refresh) { + /* + * Throttle from completion as well as from submission: a resolver that is slow rather + * than failing can outlast its own interval, which would let every later call queue + * another lookup behind it. + */ + this.postponeNextRefresh(target); + refreshTasks.remove(target, refresh); + refresh.complete(null); + } - 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; - } - if (staleChannels == null) { - resolvedTargets.put(target, resolvedTarget); - return; - } + private void submitChannelRefresh(Runnable task) { + CHANNEL_MAINTENANCE_EXECUTOR.execute(task); + } - ManagedChannel[] replacementChannels; - try { - replacementChannels = this.createChannels(target); - } catch (RuntimeException ignored) { - return; - } + private void awaitInitialResolution(CompletableFuture refresh) { + if (refresh == null) { + return; + } + try { + refresh.get(Math.max(0L, this.initialResolutionTimeoutNanos()), + TimeUnit.NANOSECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Exception ignored) { + // A slow or failing resolver must not delay the first pool any further. + } + } - boolean replaced = false; - synchronized (channels) { - if (channels.get(target) == staleChannels) { - channels.put(target, replacementChannels); - resolvedTargets.put(target, resolvedTarget); - replaced = true; - } - } + /** + * Runs on a maintenance thread, never on a request thread. At most one runs per target at a + * time — that comes from the refreshTasks entry, not from the size of the executor. Replaces + * the target's pool when its resolved address set has changed, publishing the replacement + * before retiring the previous pool. + */ + private void refreshChannelsIfAddressChanged(String target) { + String resolvedTarget = this.resolveTarget(target); + if (resolvedTarget.isEmpty()) { + return; + } + + ManagedChannel[] staleChannels = channels.get(target); + String previousTarget = resolvedTargets.get(target); + if (resolvedTarget.equals(previousTarget)) { + return; + } + if (staleChannels == null) { + /* + * Nothing to replace yet. Recording the address here is what lets the common path + * build its first pool already knowing the address, instead of rebuilding it. + */ + resolvedTargets.put(target, resolvedTarget); + return; + } + + ManagedChannel[] replacementChannels; + try { + replacementChannels = this.createChannels(target); + } catch (RuntimeException e) { + // Keep serving from the last healthy pool. + log.warn("Failed to create replacement channels of target {}, " + + "keeping the current pool", target, e); + return; + } - if (replaced) { - this.retireChannels(staleChannels); - } else { - this.retireChannels(replacementChannels); + boolean replaced = false; + synchronized (channels) { + if (channels.get(target) == staleChannels) { + channels.put(target, replacementChannels); + resolvedTargets.put(target, resolvedTarget); + replaced = true; } - } finally { - refreshLock.unlock(); } + if (replaced) { + log.info("Replaced the channel pool of target {}, address changed from {} to {}", + target, previousTarget, resolvedTarget); + } + + this.retireChannels(replaced ? staleChannels : replacementChannels); } private boolean shouldRefreshChannels(String target) { @@ -285,6 +352,10 @@ protected long channelRefreshIntervalNanos() { return DEFAULT_CHANNEL_REFRESH_INTERVAL_NANOS; } + private long initialResolutionTimeoutNanos() { + return DEFAULT_INITIAL_RESOLUTION_TIMEOUT_NANOS; + } + protected long channelDrainTimeoutNanos() { return TimeUnit.SECONDS.toNanos(config.getGrpcTimeoutSeconds()); } @@ -334,7 +405,7 @@ private void retireChannels(ManagedChannel[] retiredChannels) { .forEach(ManagedChannel::shutdown); long timeout = Math.max(0L, this.channelDrainTimeoutNanos()); - CHANNEL_CLEANUP_EXECUTOR.schedule( + CHANNEL_MAINTENANCE_EXECUTOR.schedule( () -> forceTerminateChannels(retiredChannels), timeout, TimeUnit.NANOSECONDS); } @@ -347,14 +418,20 @@ private void forceTerminateChannels(ManagedChannel[] retiredChannels) { } } + /** + * Extracts the host that a gRPC target resolves through, covering the plain {@code host:port} + * form and the {@code dns:} scheme in both its {@code dns:host:port} and + * {@code dns://authority/host:port} spellings. Any other resolver scheme returns an empty + * host, leaving that target to gRPC instead of monitoring the wrong endpoint. + */ private static String targetHost(String target) { if (target == null || target.isEmpty()) { return ""; } String endpoint = target; - if (target.startsWith("dns://")) { - endpoint = target.substring("dns://".length()); + if (target.regionMatches(true, 0, DNS_SCHEME, 0, DNS_SCHEME.length())) { + endpoint = target.substring(DNS_SCHEME.length()); while (endpoint.startsWith("/")) { endpoint = endpoint.substring(1); } @@ -362,34 +439,30 @@ private static String targetHost(String target) { if (pathStart >= 0) { endpoint = endpoint.substring(pathStart + 1); } - } else if (target.contains("://")) { + } else if (hasResolverScheme(target)) { 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) { + try { + // The authority parser handles ports and bracketed IPv6 literals. + String host = new URI("//" + endpoint).getHost(); + if (host == null) { return ""; } - return endpoint.substring(1, hostEnd); + return host.startsWith("[") ? host.substring(1, host.length() - 1) : host; + } catch (URISyntaxException ignored) { + return ""; } + } - int lastColon = endpoint.lastIndexOf(':'); - if (lastColon < 0) { - return endpoint; - } - if (endpoint.indexOf(':') != lastColon) { - return endpoint; - } - return endpoint.substring(0, lastColon); + /** + * Tells a resolver scheme from the port of a plain {@code host:port} target: a scheme is + * followed by a path, so {@code unix:/var/run/store.sock} is a scheme while + * {@code store:8500} is not. + */ + private static boolean hasResolverScheme(String target) { + int scheme = target.indexOf(':'); + return scheme >= 0 && scheme + 1 < target.length() && target.charAt(scheme + 1) == '/'; } protected InetAddress[] resolveHost(String host) throws UnknownHostException { diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryV2Client.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryV2Client.java index 4a35e46f73..d8ce443d39 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryV2Client.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryV2Client.java @@ -17,8 +17,6 @@ package org.apache.hugegraph.store.client.query; -import java.util.concurrent.atomic.AtomicInteger; - import org.apache.hugegraph.store.client.grpc.AbstractGrpcClient; import org.apache.hugegraph.store.grpc.query.QueryServiceGrpc; @@ -31,8 +29,6 @@ public class QueryV2Client extends AbstractGrpcClient { private volatile static ManagedChannel channel = null; - private final AtomicInteger seq = new AtomicInteger(0); - @Override public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { return QueryServiceGrpc.newBlockingStub(channel); @@ -48,13 +44,7 @@ public QueryServiceGrpc.QueryServiceBlockingStub getQueryServiceBlockingStub(Str } public QueryServiceGrpc.QueryServiceStub getQueryServiceStub(String target) { - return (QueryServiceGrpc.QueryServiceStub) setStubOption( - QueryServiceGrpc.newStub(getManagedChannel(target))); - // return (QueryServiceGrpc.QueryServiceStub) getAsyncStub(target); - } - - private ManagedChannel getManagedChannel(String target) { - return getChannels(target)[Math.abs(seq.getAndIncrement() % concurrency)]; + return (QueryServiceGrpc.QueryServiceStub) getAsyncStub(target); } public static void setTestChannel(ManagedChannel directChannel) { 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 885d5a46ad..4217a4c1de 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,6 +21,10 @@ 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 fbcefa5845..4b1f4d2a3c 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 @@ -21,6 +21,7 @@ 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; @@ -43,6 +44,8 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; +import org.apache.hugegraph.store.client.query.QueryV2Client; +import org.apache.hugegraph.store.grpc.query.QueryServiceGrpc; import org.apache.hugegraph.store.term.HgPair; import org.junit.Test; @@ -55,18 +58,19 @@ import io.grpc.stub.AbstractBlockingStub; /** - * Verifies that Store address changes replace channels and cached stubs safely. + * Verifies that Store address changes replace channels and cached stubs safely, and that the + * refresh never resolves or creates channels on the thread that asked for a stub. */ public class AbstractGrpcClientTest { + private static final String MAINTENANCE_THREAD_PREFIX = "channel-maintenance"; private static final AtomicInteger TARGET_SEQ = new AtomicInteger(); private static String uniqueTarget(String prefix) { return prefix + "-" + TARGET_SEQ.incrementAndGet() + ":8500"; } - private static boolean belongsToPool(Channel channel, - ManagedChannel[] channels) { + private static boolean belongsToPool(Channel channel, ManagedChannel[] channels) { return Arrays.stream(channels).anyMatch(current -> current == channel); } @@ -108,6 +112,30 @@ private static void awaitCondition(String message, Condition condition) throws E assertTrue(message, condition.isTrue()); } + /** + * Refresh runs on the maintenance thread, so a replacement pool becomes visible some time + * after the address changes rather than on the call that observes the change. Retirement + * deliberately trails publication, so waiting for both is what marks a refresh complete. + */ + private static ManagedChannel[] awaitPoolReplacement(RecordingGrpcClient client, + String target, + ManagedChannel[] staleChannels) + throws Exception { + awaitCondition("refresh must publish a replacement pool", + () -> client.getChannels(target) != staleChannels); + awaitCondition("refresh must retire the previous pool after publishing", + () -> allChannelsAreShutdown(staleChannels)); + return client.getChannels(target); + } + + @SuppressWarnings("unchecked") + private static boolean refreshIsIdle(AbstractGrpcClient client, String target) + throws Exception { + Field field = AbstractGrpcClient.class.getDeclaredField("refreshTasks"); + field.setAccessible(true); + return !((Map) field.get(client)).containsKey(target); + } + @SuppressWarnings("unchecked") private static List cachedAsyncStubChannels(AbstractGrpcClient client, String target) @@ -121,18 +149,6 @@ 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"); @@ -141,7 +157,7 @@ private static ThreadPoolExecutor channelCreationExecutor(AbstractGrpcClient cli } @Test - public void testAddressChangeReplacesChannelAndStubPools() { + public void testAddressChangeReplacesChannelAndStubPools() throws Exception { String target = uniqueTarget("address-change"); RecordingGrpcClient client = new RecordingGrpcClient(); ManagedChannel[] oldChannels = client.getChannels(target); @@ -149,9 +165,7 @@ public void testAddressChangeReplacesChannelAndStubPools() { assertNotNull(client.getAsyncStub(target)); client.resolvedTarget = "10.0.0.2"; - ManagedChannel[] newChannels = client.getChannels(target); - assertNotSame("an address change must replace the channel pool", - oldChannels, newChannels); + ManagedChannel[] newChannels = awaitPoolReplacement(client, target, oldChannels); assertTrue("every stale channel must be gracefully shut down", allChannelsAreShutdown(oldChannels)); assertFalse("refresh must not force close stale channels immediately", @@ -162,35 +176,57 @@ public void testAddressChangeReplacesChannelAndStubPools() { client.asyncStubChannels.clear(); assertNotNull(client.getBlockingStub(target)); 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))); + assertCachedChannelsCurrentAndLive("the blocking stub pool must be rebuilt on the new pool", + client.blockingStubChannels, 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))); + assertCachedChannelsCurrentAndLive("the async stub pool must be rebuilt on the new pool", + client.asyncStubChannels, newChannels); assertUsesEveryChannel("async stubs must be spread across the pool", client.asyncStubChannels, newChannels); } @Test - public void testFirstSuccessfulResolutionReplacesUnknownChannels() { + public void testStubPoolsCoverEveryChannelOnFirstBuild() { + String target = uniqueTarget("initial-stub-spread"); + RecordingGrpcClient client = new RecordingGrpcClient(); + ManagedChannel[] channels = client.getChannels(target); + + assertNotNull(client.getBlockingStub(target)); + assertNotNull(client.getAsyncStub(target)); + assertUsesEveryChannel("the first blocking stub pool must cover every channel", + client.blockingStubChannels, channels); + assertUsesEveryChannel("the first async stub pool must cover every channel", + client.asyncStubChannels, channels); + } + + @Test + public void testFirstPoolIsBuiltAfterItsAddressIsKnown() throws Exception { + String target = uniqueTarget("initial-resolution"); + RecordingGrpcClient client = new RecordingGrpcClient(); + ManagedChannel[] initialChannels = client.getChannels(target); + + assertEquals("the first pool must be built once its address is known", + 1, client.resolutionCount.get()); + // Let every refresh settle first, otherwise the assertions below race the swap. + awaitCondition("the refresh must settle", () -> refreshIsIdle(client, target)); + assertSame("a pool built with a known address must not be replaced", + initialChannels, client.getChannels(target)); + awaitCondition("the settled refresh must leave no further work", + () -> refreshIsIdle(client, target)); + assertTrue("the first pool must not be retired by its own resolution", + allChannelsAreLive(initialChannels)); + } + + @Test + public void testUnknownAddressPoolIsReplacedOnFirstSuccessfulResolution() throws Exception { String target = uniqueTarget("first-successful-resolution"); RecordingGrpcClient client = new RecordingGrpcClient(); 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); + ManagedChannel[] resolvedChannels = awaitPoolReplacement(client, target, unknownChannels); assertTrue("every channel from the unknown pool must be gracefully shut down", allChannelsAreShutdown(unknownChannels)); assertFalse("unknown channels must not be force closed immediately", @@ -200,16 +236,63 @@ public void testFirstSuccessfulResolutionReplacesUnknownChannels() { allChannelsAreLive(resolvedChannels)); } + /** + * HugeSecurityManager denies socket connection and thread creation on Gremlin worker stacks, + * and InetAddress.getAllByName() performs exactly the checkConnect(host, -1) simulated here. + * A refresh triggered by such a caller must therefore resolve somewhere else entirely. + */ @Test - public void testRetirementDoesNotBlockWhenCreationExecutorIsSaturated() - throws Exception { + public void testRefreshSucceedsWhenTheCallerThreadIsDeniedSocketAccess() throws Exception { + String target = uniqueTarget("denied-caller"); + HostCapturingGrpcClient client = new HostCapturingGrpcClient(); + client.checkSocketPermission = true; + ManagedChannel[] oldChannels = client.getChannels(target); + assertNotNull(client.getBlockingStub(target)); + + SecurityManager previous = System.getSecurityManager(); + System.setSecurityManager(new DenyingWorkerSecurityManager()); + AtomicReference failure = new AtomicReference<>(); + AtomicReference stub = new AtomicReference<>(); + try { + client.resolvedAddress = "10.0.0.2"; + // The name is what HugeSecurityManager keys its Gremlin worker check on. + Thread worker = new Thread(() -> { + try { + stub.set(client.getBlockingStub(target)); + } catch (Throwable e) { + failure.set(e); + } + }, "gremlin-server-exec-1"); + worker.start(); + worker.join(TimeUnit.SECONDS.toMillis(10L)); + + assertFalse("the denied caller must finish", worker.isAlive()); + assertNotNull("a denied caller must still receive a stub", stub.get()); + assertTrue("a denied caller must not observe a security failure: " + failure.get(), + failure.get() == null); + awaitCondition("the refresh must still publish a replacement pool", + () -> client.getChannels(target) != oldChannels); + assertFalse("the assertion below is vacuous unless something resolved", + client.resolutionThreads.isEmpty()); + assertTrue("every resolution must run on the channel maintenance thread", + client.resolutionThreads.stream() + .allMatch(name -> name.startsWith( + MAINTENANCE_THREAD_PREFIX))); + } finally { + System.setSecurityManager(previous); + } + } + + @Test + public void testRetirementDoesNotBlockWhenCreationExecutorIsSaturated() throws Exception { String target = uniqueTarget("saturated-retirement"); - ActiveRetirementGrpcClient client = new ActiveRetirementGrpcClient(); + RecordingGrpcClient client = new RecordingGrpcClient(); + client.activeCallsFinished = new CountDownLatch(1); + client.drainTimeoutNanos = TimeUnit.SECONDS.toNanos(5L); 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++) { @@ -226,29 +309,25 @@ public void testRetirementDoesNotBlockWhenCreationExecutorIsSaturated() 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); + ManagedChannel[] newChannels = awaitPoolReplacement(client, target, oldChannels); - assertNotSame("refresh must publish the replacement pool", - oldChannels, newChannels); - assertTrue("the retired pool must be shut down before refresh returns", + assertNotSame("refresh must publish the replacement pool", oldChannels, newChannels); + assertTrue("the retired pool must be shut down once refresh completes", allChannelsAreShutdown(oldChannels)); assertFalse("the scheduled cleanup must preserve the drain window", fakeChannels(oldChannels).stream() .anyMatch(FakeManagedChannel::isForceShutdown)); } finally { releaseWorkers.countDown(); - client.finishActiveCalls(); - caller.shutdownNow(); + client.activeCallsFinished.countDown(); } } @Test - public void testPartialChannelsAreRetiredAfterMixedCreationFailure() - throws Exception { + public void testPartialChannelsAreRetiredAfterMixedCreationFailure() { String target = uniqueTarget("partial-creation-failure"); - FailingCreationGrpcClient client = new FailingCreationGrpcClient(5); + CreationControlGrpcClient client = new CreationControlGrpcClient(); + client.failedAttempt = 5; try { client.getChannels(target); @@ -266,10 +345,10 @@ public void testPartialChannelsAreRetiredAfterMixedCreationFailure() } @Test - public void testInterruptedCreationWaitsAndRetiresPartialChannels() - throws Exception { + public void testInterruptedCreationWaitsAndRetiresPartialChannels() throws Exception { String target = uniqueTarget("interrupted-creation"); - DelayedCreationGrpcClient client = new DelayedCreationGrpcClient(); + CreationControlGrpcClient client = new CreationControlGrpcClient(); + client.releaseCreation = new CountDownLatch(1); AtomicReference failure = new AtomicReference<>(); AtomicBoolean interrupted = new AtomicBoolean(); Thread caller = new Thread(() -> { @@ -309,13 +388,14 @@ public void testInterruptedCreationWaitsAndRetiresPartialChannels() @Test public void testDrainDeadlineForceTerminatesRetiredChannels() throws Exception { String target = uniqueTarget("drain-deadline"); - ImmediateRetirementGrpcClient client = new ImmediateRetirementGrpcClient(); + RecordingGrpcClient client = new RecordingGrpcClient(); + client.activeCallsFinished = new CountDownLatch(1); + client.drainTimeoutNanos = 0L; ManagedChannel[] oldChannels = client.getChannels(target); client.resolvedTarget = "10.0.0.2"; - ManagedChannel[] newChannels = client.getChannels(target); + ManagedChannel[] newChannels = awaitPoolReplacement(client, target, oldChannels); - 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())); @@ -323,34 +403,9 @@ public void testDrainDeadlineForceTerminatesRetiredChannels() throws Exception { } @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() { + public void testStubAcquisitionReusesResolutionWithinRefreshInterval() throws Exception { String target = uniqueTarget("throttled-refresh"); - CountingResolverGrpcClient client = new CountingResolverGrpcClient(); + RecordingGrpcClient client = new RecordingGrpcClient(); client.refreshIntervalNanos = TimeUnit.HOURS.toNanos(1L); assertNotNull(client.getBlockingStub(target)); @@ -360,6 +415,7 @@ public void testStubAcquisitionReusesResolutionWithinRefreshInterval() { assertNotNull(client.getAsyncStub(target)); } + Thread.sleep(100L); assertEquals("stub acquisition must not resolve again inside the refresh interval", 1, client.resolutionCount.get()); } @@ -368,14 +424,16 @@ public void testStubAcquisitionReusesResolutionWithinRefreshInterval() { public void testConcurrentStubAcquisitionRetainsHealthyPoolDuringDelayedRefresh() throws Exception { String target = uniqueTarget("delayed-refresh"); - DelayedResolverGrpcClient client = new DelayedResolverGrpcClient(); - client.refreshIntervalNanos = 0L; + RecordingGrpcClient client = new RecordingGrpcClient(); ManagedChannel[] oldChannels = client.getChannels(target); assertNotNull(client.getBlockingStub(target)); + // Let any refresh already in flight settle, so the resolution count below is stable. + awaitCondition("the initial refresh must settle before the count is captured", + () -> refreshIsIdle(client, target)); int resolutionsBeforeConcurrentCalls = client.resolutionCount.get(); client.resolvedTarget = "10.0.0.2"; - client.delayChangedResolution = true; + client.delayResolution = true; client.refreshIntervalNanos = TimeUnit.HOURS.toNanos(1L); ExecutorService executor = Executors.newFixedThreadPool(6); @@ -386,24 +444,21 @@ public void testConcurrentStubAcquisitionRetainsHealthyPoolDuringDelayedRefresh( } 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)); + for (Future future : futures) { + assertNotNull("no caller may block on the delayed refresh", + future.get(5, TimeUnit.SECONDS)); + } 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); + ManagedChannel[] currentChannels = + awaitPoolReplacement(client, target, oldChannels); assertTrue("the previous pool must be retired after replacement is published", allChannelsAreShutdown(oldChannels)); + assertTrue("the replacement pool must be live", allChannelsAreLive(currentChannels)); assertEquals("concurrent callers must share a single refresh resolution", - resolutionsBeforeConcurrentCalls + 1, - client.resolutionCount.get()); + resolutionsBeforeConcurrentCalls + 1, client.resolutionCount.get()); } finally { client.releaseDelayedResolution.countDown(); executor.shutdownNow(); @@ -413,128 +468,173 @@ public void testConcurrentStubAcquisitionRetainsHealthyPoolDuringDelayedRefresh( @Test public void testRefreshGracefullyRetiresActiveStreamChannels() throws Exception { String target = uniqueTarget("active-stream-refresh"); - ActiveRetirementGrpcClient client = new ActiveRetirementGrpcClient(); + RecordingGrpcClient client = new RecordingGrpcClient(); + client.activeCallsFinished = new CountDownLatch(1); + client.drainTimeoutNanos = TimeUnit.SECONDS.toNanos(5L); 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); + ManagedChannel[] newChannels = awaitPoolReplacement(client, target, oldChannels); 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)); - client.finishActiveCalls(); + client.activeCallsFinished.countDown(); 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)); + assertTrue("the replacement pool must remain live", allChannelsAreLive(newChannels)); } + /** + * Holds a stub pool build open, refreshes the pool underneath it, and asserts that both the + * interleaved build and a concurrent one return stubs bound to the published pool. Blocking + * and asynchronous acquisition share one implementation, so the asynchronous path stands in + * for both; it is the one that also publishes a stub cache worth asserting on. + */ @Test - public void testBlockingStubBuildRetriesAfterChannelRefresh() throws Exception { - String target = uniqueTarget("concurrent-blocking-stub-refresh"); + 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 blocking stub pool build must be in flight", - client.staleBlockingStubBuildStarted.await(5, TimeUnit.SECONDS)); + Future staleStub = + executor.submit(() -> client.getAsyncStub(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)); + Future freshStub = + executor.submit(() -> client.getAsyncStub(target)); awaitCondition("refresh must retire the old channel pool", () -> allChannelsAreShutdown(oldChannels)); - client.releaseStaleBlockingStubBuild.countDown(); + 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)); + belongsToPool(staleStub.get(5, TimeUnit.SECONDS).getChannel(), + currentChannels)); assertTrue("the concurrent build must use the current pool", - belongsToPool(freshResult.getChannel(), currentChannels)); + belongsToPool(freshStub.get(5, TimeUnit.SECONDS).getChannel(), + currentChannels)); assertTrue("the current channel pool must remain live", allChannelsAreLive(currentChannels)); + assertCachedChannelsCurrentAndLive( + "the final stub cache must only reference current live channels", + cachedAsyncStubChannels(client, target), currentChannels); } finally { - client.releaseStaleBlockingStubBuild.countDown(); + client.releaseStaleStubBuild.countDown(); executor.shutdownNow(); } } + /** + * QueryV2 used to take a channel straight from the pool and build its stub afterwards, which + * let a refresh retire that channel in between. It must now go through the guarded path. + */ @Test - public void testAsyncStubBuildRetriesAfterChannelRefresh() throws Exception { - String target = uniqueTarget("concurrent-async-stub-refresh"); - StubInterleavingGrpcClient client = new StubInterleavingGrpcClient(); + public void testQueryV2StubFollowsPublishedPoolAcrossRefresh() throws Exception { + String target = uniqueTarget("query-v2-refresh"); + QueryV2TestClient client = new QueryV2TestClient(); ManagedChannel[] oldChannels = client.getChannels(target); - ExecutorService executor = Executors.newFixedThreadPool(2); + ExecutorService executor = Executors.newSingleThreadExecutor(); 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)); + /* + * Interleave a refresh with the stub build. Taking a channel from the pool and + * building the stub afterwards would bind it to a channel retired in between. + */ + Future stub = + executor.submit(() -> client.getQueryServiceStub(target)); + assertTrue("the QueryV2 stub build must be in flight", + client.stubBuildStarted.await(5, TimeUnit.SECONDS)); client.resolvedTarget = "10.0.0.2"; - Future freshStub = - executor.submit(() -> client.getAsyncStub(target)); + // getChannels is what triggers a refresh, and the blocked build cannot call it. 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); + () -> client.getChannels(target) != oldChannels && + allChannelsAreShutdown(oldChannels)); + client.releaseStubBuild.countDown(); + + ManagedChannel[] newChannels = client.getChannels(target); + Channel channel = stub.get(5, TimeUnit.SECONDS).getChannel(); + assertTrue("QueryV2 must never return a stub bound to a retired channel", + belongsToPool(channel, newChannels)); + assertFalse("QueryV2 must never return a stub on a shut down channel", + ((ManagedChannel) channel).isShutdown()); + + List stubChannels = new ArrayList<>(); + for (int i = 0; i < AbstractGrpcClient.concurrency; i++) { + stubChannels.add((ManagedChannel) client.getQueryServiceStub(target).getChannel()); + } + assertCachedChannelsCurrentAndLive("QueryV2 stubs must stay on the published pool", + stubChannels, newChannels); + assertUsesEveryChannel("QueryV2 stubs must still spread across the pool", + stubChannels, newChannels); } finally { - client.releaseStaleAsyncStubBuild.countDown(); + client.releaseStubBuild.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 client = new HostCapturingGrpcClient(); + assertEquals("10.0.0.1", client.resolveTarget("store.example.com:8500")); + assertEquals("store.example.com", client.capturedHost); + + assertEquals("10.0.0.1", client.resolveTarget("dns:///store.example.com:8500")); + assertEquals("store.example.com", client.capturedHost); + + // The scheme-only spelling is a legal gRPC dns target too. + assertEquals("10.0.0.1", client.resolveTarget("dns:store.example.com:8500")); + assertEquals("store.example.com", client.capturedHost); - HostCapturingGrpcClient ipv6Client = new HostCapturingGrpcClient(); - assertEquals("10.0.0.1", ipv6Client.resolveTarget("[2001:db8::1]:8500")); - assertEquals("2001:db8::1", ipv6Client.capturedHost); + assertEquals("10.0.0.1", client.resolveTarget("dns://8.8.8.8/store.example.com:8500")); + assertEquals("store.example.com", client.capturedHost); + + assertEquals("10.0.0.1", client.resolveTarget("[2001:db8::1]:8500")); + assertEquals("2001:db8::1", client.capturedHost); } @Test public void testResolveTargetSkipsUnsupportedGrpcSchemes() { HostCapturingGrpcClient client = new HostCapturingGrpcClient(); assertEquals("", client.resolveTarget("unix:///var/run/store.sock")); + // The single-slash spelling is legal too, and must not resolve the literal host "unix". + assertEquals("", client.resolveTarget("unix:/var/run/store.sock")); + assertEquals("", client.resolveTarget("xds:///store.example.com")); assertEquals("unsupported schemes must not invoke DNS resolution", - 0, client.resolutionCount.get()); + 0, client.hostResolutionCount.get()); } private interface Condition { - boolean isTrue(); + boolean isTrue() throws Exception; } private static class RecordingGrpcClient extends AbstractGrpcClient { private final AtomicInteger channelSeq = new AtomicInteger(); + protected final AtomicInteger resolutionCount = new AtomicInteger(); + protected final List resolutionThreads = + Collections.synchronizedList(new ArrayList<>()); + protected final CountDownLatch delayedResolutionStarted = new CountDownLatch(1); + protected final CountDownLatch releaseDelayedResolution = new CountDownLatch(1); protected volatile String resolvedTarget = "10.0.0.1"; protected volatile long refreshIntervalNanos = 0L; + protected volatile boolean delayResolution; + /** Resolves through the real implementation instead of returning resolvedTarget. */ + protected volatile boolean useRealResolution; + /** Null keeps the inherited drain deadline. */ + protected volatile Long drainTimeoutNanos; + /** Null makes channels terminate as soon as they are shut down. */ + protected volatile CountDownLatch activeCallsFinished; protected final List blockingStubChannels = Collections.synchronizedList(new ArrayList<>()); protected final List asyncStubChannels = @@ -546,53 +646,22 @@ protected long channelRefreshIntervalNanos() { } @Override - protected ManagedChannel createChannel(String target) { - return this.newFakeChannel(target + "#" + this.channelSeq.getAndIncrement()); - } - - protected FakeManagedChannel newFakeChannel(String authority) { - return new FakeManagedChannel(authority); - } - - @Override - protected String resolveTarget(String target) { - return this.resolvedTarget; - } - - @Override - public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { - this.blockingStubChannels.add(channel); - return new FakeBlockingStub(channel, CallOptions.DEFAULT); - } - - @Override - public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { - this.asyncStubChannels.add(channel); - return new FakeAsyncStub(channel, CallOptions.DEFAULT); + protected long channelDrainTimeoutNanos() { + Long timeout = this.drainTimeoutNanos; + return timeout == null ? super.channelDrainTimeoutNanos() : timeout; } - } - - private static class CountingResolverGrpcClient extends RecordingGrpcClient { - - protected final AtomicInteger resolutionCount = new AtomicInteger(); @Override - protected String resolveTarget(String target) { - this.resolutionCount.incrementAndGet(); - return super.resolveTarget(target); + protected ManagedChannel createChannel(String target) { + return new FakeManagedChannel(target + "#" + this.channelSeq.getAndIncrement(), + this.activeCallsFinished); } - } - - 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.resolutionThreads.add(Thread.currentThread().getName()); + if (this.delayResolution) { this.delayedResolutionStarted.countDown(); try { assertTrue("the delayed resolution must be released", @@ -602,85 +671,51 @@ protected String resolveTarget(String target) { throw new AssertionError(e); } } - return this.resolvedTarget; + return this.useRealResolution ? super.resolveTarget(target) : 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 ImmediateRetirementGrpcClient extends RecordingGrpcClient { - - private final CountDownLatch activeCallsFinished = new CountDownLatch(1); @Override - protected long channelDrainTimeoutNanos() { - return 0L; + public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { + this.blockingStubChannels.add(channel); + return new FakeBlockingStub(channel, CallOptions.DEFAULT); } @Override - protected FakeManagedChannel newFakeChannel(String authority) { - return new FakeManagedChannel(authority, this.activeCallsFinished); + public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { + this.asyncStubChannels.add(channel); + return new FakeAsyncStub(channel, CallOptions.DEFAULT); } } - private static class FailingCreationGrpcClient extends RecordingGrpcClient { + private static class CreationControlGrpcClient extends RecordingGrpcClient { - private final int failedAttempt; private final AtomicInteger attempt = new AtomicInteger(); + private final CountDownLatch creationStarted = + new CountDownLatch(AbstractGrpcClient.concurrency); private final List createdChannels = Collections.synchronizedList(new ArrayList<>()); - - private FailingCreationGrpcClient(int failedAttempt) { - this.failedAttempt = failedAttempt; - } + /** Negative never fails. */ + private volatile int failedAttempt = -1; + /** Null creates channels without delay. */ + private volatile CountDownLatch releaseCreation; @Override protected ManagedChannel createChannel(String target) { int current = this.attempt.getAndIncrement(); + this.creationStarted.countDown(); 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); + CountDownLatch release = this.releaseCreation; + if (release != null) { + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } } - ManagedChannel channel = this.newFakeChannel(target); + ManagedChannel channel = new FakeManagedChannel(target + "#" + current); this.createdChannels.add(channel); return channel; } @@ -688,35 +723,46 @@ protected ManagedChannel createChannel(String target) { private static class StubInterleavingGrpcClient extends RecordingGrpcClient { - private final AtomicInteger blockingStubSeq = new AtomicInteger(); - 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); + private final AtomicInteger stubSeq = 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.staleBlockingStubBuildStarted.countDown(); + public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { + if (this.stubSeq.incrementAndGet() == 1) { + this.staleStubBuildStarted.countDown(); try { - assertTrue("the stale blocking stub build must be released", - this.releaseStaleBlockingStubBuild.await(5, TimeUnit.SECONDS)); + 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); + return super.getAsyncStub(channel); } + } + + private static class QueryV2TestClient extends QueryV2Client { + + private final AtomicInteger channelSeq = new AtomicInteger(); + private final AtomicInteger stubSeq = new AtomicInteger(); + private final CountDownLatch stubBuildStarted = new CountDownLatch(1); + private final CountDownLatch releaseStubBuild = new CountDownLatch(1); + private volatile String resolvedTarget = "10.0.0.1"; @Override - public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { - if (this.asyncStubSeq.incrementAndGet() == 1) { - this.staleAsyncStubBuildStarted.countDown(); + protected long channelRefreshIntervalNanos() { + return 0L; + } + + @Override + public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { + if (this.stubSeq.incrementAndGet() == 1) { + this.stubBuildStarted.countDown(); try { - assertTrue("the stale async stub build must be released", - this.releaseStaleAsyncStubBuild.await(5, TimeUnit.SECONDS)); + assertTrue("the QueryV2 stub build must be released", + this.releaseStubBuild.await(5, TimeUnit.SECONDS)); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new AssertionError(e); @@ -724,33 +770,76 @@ public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { } return super.getAsyncStub(channel); } + + @Override + protected ManagedChannel createChannel(String target) { + return new FakeManagedChannel(target + "#" + this.channelSeq.getAndIncrement()); + } + + @Override + protected String resolveTarget(String target) { + return this.resolvedTarget; + } } - private static class HostCapturingGrpcClient extends AbstractGrpcClient { + /** + * Resolves through the inherited implementation, optionally reproducing the security check + * that InetAddress.getAllByName() performs on whichever thread resolution runs on. + */ + private static class HostCapturingGrpcClient extends RecordingGrpcClient { - private final AtomicInteger resolutionCount = new AtomicInteger(); + private final AtomicInteger hostResolutionCount = new AtomicInteger(); private volatile String capturedHost; + private volatile String resolvedAddress = "10.0.0.1"; + private volatile boolean checkSocketPermission; + + HostCapturingGrpcClient() { + this.useRealResolution = true; + } @Override - protected ManagedChannel createChannel(String target) { - return new FakeManagedChannel(target); + protected InetAddress[] resolveHost(String host) throws UnknownHostException { + SecurityManager security = System.getSecurityManager(); + if (this.checkSocketPermission && security != null) { + security.checkConnect(host, -1); + } + this.hostResolutionCount.incrementAndGet(); + this.capturedHost = host; + return new InetAddress[]{InetAddress.getByName(this.resolvedAddress)}; + } + } + + /** + * Denies a deliberately narrow subset of what HugeSecurityManager denies on a Gremlin worker + * stack: the two checks this fix is about, opening sockets and creating threads. Everything + * else stays permitted so the test JVM keeps working, including restoring the previous + * manager. Keys on the thread name alone, where HugeSecurityManager also requires a Gremlin + * script engine frame on the stack. Relies on System.setSecurityManager, which is a no-op + * from JDK 18 and removed in JDK 24; this module builds and runs on Java 11. + */ + private static class DenyingWorkerSecurityManager extends SecurityManager { + + private static boolean isDeniedWorker() { + return Thread.currentThread().getName().startsWith("gremlin-server-exec"); } @Override - public AbstractBlockingStub getBlockingStub(ManagedChannel channel) { - return new FakeBlockingStub(channel, CallOptions.DEFAULT); + public void checkConnect(String host, int port) { + if (isDeniedWorker()) { + throw new SecurityException("Not allowed to connect socket via Gremlin"); + } } @Override - public AbstractAsyncStub getAsyncStub(ManagedChannel channel) { - return new FakeAsyncStub(channel, CallOptions.DEFAULT); + public void checkAccess(ThreadGroup threadGroup) { + if (isDeniedWorker()) { + throw new SecurityException("Not allowed to access thread group via Gremlin"); + } } @Override - protected InetAddress[] resolveHost(String host) throws UnknownHostException { - this.resolutionCount.incrementAndGet(); - this.capturedHost = host; - return new InetAddress[]{InetAddress.getByName("10.0.0.1")}; + public void checkPermission(java.security.Permission permission) { + // Everything else stays permitted, including restoring the previous manager. } } @@ -782,7 +871,6 @@ 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; @@ -842,27 +930,9 @@ public boolean isTerminated() { 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) - 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; + public boolean awaitTermination(long timeout, TimeUnit unit) { + return this.isTerminated(); } } }