You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hugegraph.apache.org by ji...@apache.org on 2026/10/23 05:28:14 UTC
(hugegraph) branch master updated: fix: preserve traversal and PD client contracts (#3395)
This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hugegraph.git
The following commit(s) were added to refs/heads/master by this push:
new 833840282 fix: preserve traversal and PD client contracts (#3395)
833840282 is described below
commit 8338402826e5f9d05ed6b87784fe334659e381c8
Author: imbajin <ji...@apache.org>
AuthorDate: Fri Oct 23 13:28:09 2026 +0800
fix: preserve traversal and PD client contracts (#3395)
Algorithm results and filtered queries can lose data or leave scans open. PD lock initialization and leader discovery can use the wrong failure or routing path, and one failing watch listener can stop notifications for every listener. This cleanup preserves API signatures, configuration, retry budgets and stored formats.
---
docs/storage-lifecycle.md | 26 ++
.../org/apache/hugegraph/pd/client/KvClient.java | 10 +-
.../org/apache/hugegraph/pd/client/PDClient.java | 27 +-
.../hugegraph/pd/client/KvClientLockTest.java | 191 ++++++++++++
.../hugegraph/pd/client/PDClientLeaderTest.java | 224 ++++++++++++++
.../hugegraph/pd/client/PDClientSuiteTest.java | 3 +
.../pd/client/PDClientWatchDispatchTest.java | 329 +++++++++++++++++++++
.../hugegraph/job/algorithm/AbstractAlgorithm.java | 50 +++-
.../job/algorithm/CountEdgeAlgorithm.java | 26 +-
.../job/algorithm/CountVertexAlgorithm.java | 26 +-
.../algorithm/records/record/RecordFactory.java | 20 +-
.../job/algorithm/AlgorithmExecutionTest.java | 129 ++++++++
.../job/algorithm/AlgorithmJsonMapTest.java | 114 +++++++
.../job/algorithm/CountAlgorithmTest.java | 196 ++++++++++++
.../org/apache/hugegraph/unit/UnitTestSuite.java | 6 +
.../business/itrv2/IntersectionFilterIterator.java | 78 +++--
.../business/itrv2/MapValueFilterIterator.java | 13 +-
.../store/business/itrv2/UnionFilterIterator.java | 19 +-
.../grpc/query/FilterIteratorContractTest.java | 307 +++++++++++++++++++
.../grpc/query/IntersectionSpillIteratorTest.java | 251 ++++++++++++++++
.../hugegraph/store/service/ServerSuiteTest.java | 4 +
21 files changed, 1938 insertions(+), 111 deletions(-)
diff --git a/docs/storage-lifecycle.md b/docs/storage-lifecycle.md
index 94a1987e3..ddf7731a0 100644
--- a/docs/storage-lifecycle.md
+++ b/docs/storage-lifecycle.md
@@ -99,6 +99,14 @@ the interrupt flag; input that never entered the queue cannot be reported as acc
Ephemeral jobs retain their negative task IDs and existing wraparound rule. Concurrent
schedulers allocate these IDs atomically, avoiding duplicate IDs caused by lost updates.
+Algorithm traversal helpers restore the caller's query capacity and close the traversal
+on success or failure. If calculation and close both fail, the calculation error stays
+primary and the close error is suppressed. String IDs and values in algorithm JSON
+results preserve quotes, backslashes and control characters through JSON encoding;
+numeric result values remain JSON numbers.
+Vertex and edge counting also close their scan iterator when reading or reporting
+progress fails, preserving the calculation error if closing fails as well.
+
### Store sort and shuffle iterators
Sort/shuffle applies the requested comparator both in memory and after spilling to
@@ -116,6 +124,14 @@ them is terminal and discards pending values and completion markers. The source
closed once; a failed close remains visible on later close calls. ORDER BY propagates
sort failures to the query's existing error and cleanup path with the original cause.
+Union and intersection filters preserve a pending record across
+repeated probes, including the extra probe at RPC batch boundaries. Reading with
+`next()` consumes that record exactly once. Union results also support reading without
+a preceding probe. Closing either filter discards pending results and stops further
+probes of its source.
+For a spilled intersection, the requested group size also applies to the final group;
+the default mode continues to select groups containing more than one record.
+
### Callback executor configuration
The gRPC application callback queue is always unbounded (`Integer.MAX_VALUE`), matching
@@ -138,6 +154,16 @@ the RPC method and original cause. Existing retry budgets and response-header er
codes are preserved, so registration and metadata callers receive a PD error rather
than a missing response.
+An interrupted wait for the first lock client's initialization stops before issuing a
+lock RPC, preserves the interrupt flag and leaves the initialization permit with its
+owner. Successful leader discovery binds subsequent PD calls to the discovered leader,
+including when the configured peer used for discovery is a follower.
+
+PD watch notifications isolate runtime failures from individual business listeners.
+The failure is logged and the remaining listeners still receive the event on the same
+callback thread. Cache updates retain their existing order and errors still propagate;
+failed listener work is not retried automatically.
+
PD snapshot compression finishes the ZIP central directory before flushing and syncing
the archive. An unreadable source directory fails with an I/O error. Archive entries
and the checksum format remain compatible with existing snapshots.
diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java
index e513fa62e..fd34f491b 100644
--- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java
+++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/KvClient.java
@@ -580,7 +580,7 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
}
}
- private void acquire(AtomicLong clientId, Semaphore semaphore) {
+ private void acquire(AtomicLong clientId, Semaphore semaphore) throws PDException {
if (clientId.get() == 0L) {
try {
semaphore.acquire();
@@ -589,7 +589,9 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
}
log.info("wait for client starting....");
} catch (InterruptedException e) {
- log.error("get semaphore with error:", e);
+ Thread.currentThread().interrupt();
+ throw new PDException(Pdpb.ErrorType.UNKNOWN_VALUE,
+ "Interrupted while waiting for lock client initialization", e);
}
}
}
@@ -704,8 +706,8 @@ public class KvClient<T extends WatchResponse> extends AbstractClient implements
return;
}
reconnectScheduler.shutdownNow();
- if (reconnectExecutor instanceof ExecutorService) {
- ((ExecutorService) reconnectExecutor).shutdownNow();
+ if (reconnectExecutor instanceof ExecutorService executor) {
+ executor.shutdownNow();
}
release(lockSemaphore);
for (WatchSubscription subscription : subscriptions) {
diff --git a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java
index 99b5f3772..5d0963b6d 100644
--- a/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java
+++ b/hugegraph-pd/hg-pd-client/src/main/java/org/apache/hugegraph/pd/client/PDClient.java
@@ -28,6 +28,7 @@ import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
import org.apache.hugegraph.pd.client.impl.PDPulseImpl2;
import org.apache.hugegraph.pd.client.interceptor.Authentication;
@@ -52,6 +53,7 @@ import org.apache.hugegraph.pd.grpc.Pdpb.GetPartitionRequest;
import org.apache.hugegraph.pd.grpc.Pdpb.GetPartitionResponse;
import org.apache.hugegraph.pd.grpc.Pdpb.GraphStatsResponse;
import org.apache.hugegraph.pd.grpc.watch.WatchResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchType;
import org.apache.hugegraph.pd.watch.NodeEvent;
import org.apache.hugegraph.pd.watch.PartitionEvent;
import org.apache.hugegraph.pd.watch.PartitionEvent.ChangeType;
@@ -165,7 +167,8 @@ public class PDClient {
if (response.getChangeType() == ChangeType.DEL) {
cache.removeAll(response.getGraph());
}
- listeners.forEach(listener -> listener.onPartitionChanged(response));
+ notifyListeners(WatchType.WATCH_TYPE_PARTITION_CHANGE,
+ listener -> listener.onPartitionChanged(response));
}
@Override
@@ -200,7 +203,8 @@ public class PDClient {
}
}
- listeners.forEach(listener -> listener.onStoreChanged(response));
+ notifyListeners(WatchType.WATCH_TYPE_STORE_NODE_CHANGE,
+ listener -> listener.onStoreChanged(response));
}
@Override
@@ -214,7 +218,8 @@ public class PDClient {
graphWatcher = pdWatch.watchGraph(new PDWatch.Listener<>() {
@Override
public void onNext(WatchResponse response) {
- listeners.forEach(listener -> listener.onGraphChanged(response));
+ notifyListeners(WatchType.WATCH_TYPE_GRAPH_CHANGE,
+ listener -> listener.onGraphChanged(response));
}
@Override
@@ -245,7 +250,8 @@ public class PDClient {
break;
}
}
- listeners.forEach(listener -> listener.onShardGroupChanged(response));
+ notifyListeners(WatchType.WATCH_TYPE_SHARD_GROUP_CHANGE,
+ listener -> listener.onShardGroupChanged(response));
}
@Override
@@ -256,6 +262,17 @@ public class PDClient {
}
+ private void notifyListeners(WatchType eventType, Consumer<PDEventListener> action) {
+ for (PDEventListener listener : this.listeners) {
+ try {
+ action.accept(listener);
+ } catch (RuntimeException e) {
+ String listenerClass = listener == null ? "null" : listener.getClass().getName();
+ log.error("PD watch listener {} failed to handle {} event", listenerClass, eventType, e);
+ }
+ }
+ }
+
private synchronized void closeStub(boolean closeWatcher) {
// TODO ManagedChannel Did not close properly
proxy.set(null);
@@ -314,7 +331,7 @@ public class PDClient {
var leaderIp = getLeaderIp(stub);
if (!leaderIp.equalsIgnoreCase(host)) {
leaderHost = leaderIp;
- proxy.set(getStub(PDGrpc.newBlockingStub(channel)));
+ proxy.set(getStub(PDGrpc.newBlockingStub(getChannel(leaderIp))));
} else {
proxy.set(stub);
leaderHost = host;
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientLockTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientLockTest.java
new file mode 100644
index 000000000..f07d63a0b
--- /dev/null
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/KvClientLockTest.java
@@ -0,0 +1,191 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.client;
+
+import java.lang.reflect.Field;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.hugegraph.pd.common.PDException;
+import org.apache.hugegraph.pd.grpc.Pdpb;
+import org.apache.hugegraph.pd.grpc.kv.LockRequest;
+import org.apache.hugegraph.pd.grpc.kv.LockResponse;
+import org.apache.hugegraph.pd.grpc.kv.WatchResponse;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.mockito.Mockito;
+
+import io.grpc.MethodDescriptor;
+
+@RunWith(Parameterized.class)
+public class KvClientLockTest {
+
+ private final boolean reentrant;
+
+ public KvClientLockTest(boolean reentrant) {
+ this.reentrant = reentrant;
+ }
+
+ @Parameterized.Parameters(name = "reentrant={0}")
+ public static Object[][] lockMethods() {
+ return new Object[][]{{true}, {false}};
+ }
+
+ @Test(timeout = 10000)
+ public void testPreInterruptedLockDoesNotSendRequest() throws Exception {
+ try (ControlledKvClient client = new ControlledKvClient(false)) {
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ AtomicBoolean interrupted = new AtomicBoolean();
+ Thread caller = new Thread(() -> {
+ Thread.currentThread().interrupt();
+ callAndCapture(client, failure, interrupted);
+ });
+ try {
+ caller.start();
+ caller.join(5000);
+ Assert.assertFalse(caller.isAlive());
+ assertInterrupted(failure.get(), interrupted.get());
+ Assert.assertTrue(client.requests.isEmpty());
+
+ Assert.assertTrue(callLock(client).getSucceed());
+ Assert.assertEquals(0L, client.requests.get(0).getClientId());
+ Assert.assertTrue(callLock(client).getSucceed());
+ Assert.assertEquals(42L, client.requests.get(1).getClientId());
+ } finally {
+ caller.interrupt();
+ caller.join(5000);
+ }
+ }
+ }
+
+ @Test(timeout = 15000)
+ public void testInterruptedWaitDoesNotReleaseAnotherCallersPermit() throws Exception {
+ try (ControlledKvClient client = new ControlledKvClient(true)) {
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ AtomicBoolean interrupted = new AtomicBoolean();
+ Thread waiter = new Thread(() -> callAndCapture(client, failure, interrupted));
+ try {
+ Future<LockResponse> first = executor.submit(() -> callLock(client));
+ Assert.assertTrue(client.firstRequest.await(5, TimeUnit.SECONDS));
+ Semaphore semaphore = lockSemaphore(client);
+ waiter.start();
+ awaitQueuedWaiter(semaphore);
+
+ waiter.interrupt();
+ waiter.join(5000);
+ Assert.assertFalse(waiter.isAlive());
+ assertInterrupted(failure.get(), interrupted.get());
+ Assert.assertEquals(1, client.requests.size());
+ Assert.assertEquals("The first caller must still own the permit", 0,
+ semaphore.availablePermits());
+
+ client.releaseFirst.countDown();
+ Assert.assertEquals(42L, first.get(5, TimeUnit.SECONDS).getClientId());
+ Assert.assertEquals(1, semaphore.availablePermits());
+ Assert.assertTrue(callLock(client).getSucceed());
+ Assert.assertEquals(42L, client.requests.get(1).getClientId());
+ } finally {
+ client.releaseFirst.countDown();
+ waiter.interrupt();
+ waiter.join(5000);
+ executor.shutdownNow();
+ Assert.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+ }
+
+ private LockResponse callLock(KvClient<?> client) throws PDException {
+ return this.reentrant ? client.lock("key", 1000L) :
+ client.lockWithoutReentrant("key", 1000L);
+ }
+
+ private void callAndCapture(KvClient<?> client, AtomicReference<Throwable> failure,
+ AtomicBoolean interrupted) {
+ try {
+ callLock(client);
+ } catch (Throwable e) {
+ failure.set(e);
+ interrupted.set(Thread.currentThread().isInterrupted());
+ }
+ }
+
+ private static void assertInterrupted(Throwable failure, boolean interrupted) {
+ Assert.assertTrue(failure instanceof PDException);
+ Assert.assertEquals(Pdpb.ErrorType.UNKNOWN_VALUE, ((PDException) failure).getErrorCode());
+ Assert.assertTrue(failure.getCause() instanceof InterruptedException);
+ Assert.assertTrue(interrupted);
+ }
+
+ private static Semaphore lockSemaphore(KvClient<?> client) throws Exception {
+ Field field = KvClient.class.getDeclaredField("lockSemaphore");
+ field.setAccessible(true);
+ return (Semaphore) field.get(client);
+ }
+
+ private static void awaitQueuedWaiter(Semaphore semaphore) throws InterruptedException {
+ long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (!semaphore.hasQueuedThreads() && System.nanoTime() < deadline) {
+ Thread.sleep(1);
+ }
+ Assert.assertTrue("The second caller must be waiting for initialization",
+ semaphore.hasQueuedThreads());
+ }
+
+ private static class ControlledKvClient extends KvClient<WatchResponse> {
+
+ private final boolean holdFirst;
+ private final CountDownLatch firstRequest = new CountDownLatch(1);
+ private final CountDownLatch releaseFirst = new CountDownLatch(1);
+ private final List<LockRequest> requests = new CopyOnWriteArrayList<>();
+
+ ControlledKvClient(boolean holdFirst) {
+ super(PDConfig.of("unused:8686"), Mockito.mock(ScheduledExecutorService.class));
+ this.holdFirst = holdFirst;
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ protected <ReqT, RespT> RespT blockingUnaryCall(
+ MethodDescriptor<ReqT, RespT> method, ReqT request) throws PDException {
+ this.requests.add((LockRequest) request);
+ if (this.holdFirst && this.requests.size() == 1) {
+ this.firstRequest.countDown();
+ try {
+ Assert.assertTrue(this.releaseFirst.await(10, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new PDException(Pdpb.ErrorType.UNKNOWN_VALUE, e);
+ }
+ }
+ return (RespT) LockResponse.newBuilder().setHeader(AbstractClient.okHeader)
+ .setClientId(42L).setSucceed(true).build();
+ }
+ }
+}
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientLeaderTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientLeaderTest.java
new file mode 100644
index 000000000..468ff759f
--- /dev/null
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientLeaderTest.java
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.client;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.hugegraph.pd.grpc.Metapb;
+import org.apache.hugegraph.pd.grpc.PDGrpc;
+import org.apache.hugegraph.pd.grpc.Pdpb;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import io.grpc.CallOptions;
+import io.grpc.ClientCall;
+import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
+
+public class PDClientLeaderTest {
+
+ private final String prefix = "pd-leader-test-" + UUID.randomUUID();
+ private final List<TestChannel> peers = new ArrayList<>();
+ private ConcurrentMap<String, ManagedChannel> channels;
+ private ManagedChannel previousChannel;
+
+ @SuppressWarnings("unchecked")
+ @Before
+ public void setUp() throws Exception {
+ this.channels = (ConcurrentMap<String, ManagedChannel>) field("channels").get(null);
+ this.previousChannel = (ManagedChannel) field("channel").get(null);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ for (TestChannel peer : this.peers) {
+ this.channels.remove(peer.address, peer);
+ peer.shutdownNow();
+ }
+ field("channel").set(null, this.previousChannel);
+ }
+
+ @Test
+ public void testRegistrationUsesDiscoveredLeader() throws Exception {
+ TestChannel leader = addPeer("leader");
+ TestChannel follower = addPeer("follower");
+ follower.leaderAddress = leader.address;
+ PDClient client = newClient(follower.address, leader.address);
+
+ Assert.assertEquals(17L, client.registerStore(Metapb.Store.getDefaultInstance()));
+
+ Assert.assertEquals(1, follower.discoveryCalls);
+ Assert.assertEquals(0, follower.registrationCalls);
+ Assert.assertEquals(1, leader.registrationCalls);
+ }
+
+ @Test
+ public void testRegistrationReusesPeerWhenAlreadyLeader() throws Exception {
+ TestChannel leader = addPeer("leader");
+ PDClient client = newClient(leader.address, leader.address);
+
+ Assert.assertEquals(17L, client.registerStore(Metapb.Store.getDefaultInstance()));
+
+ Assert.assertEquals(1, leader.discoveryCalls);
+ Assert.assertEquals(1, leader.registrationCalls);
+ }
+
+ @Test
+ public void testDiscoveryFailureTriesNextPeer() throws Exception {
+ TestChannel unavailable = addPeer("unavailable");
+ unavailable.failDiscovery = true;
+ TestChannel leader = addPeer("leader");
+ PDClient client = newClient(unavailable.address + "," + leader.address, leader.address);
+
+ Assert.assertEquals(17L, client.registerStore(Metapb.Store.getDefaultInstance()));
+
+ Assert.assertEquals(1, unavailable.discoveryCalls);
+ Assert.assertEquals(0, unavailable.registrationCalls);
+ Assert.assertEquals(1, leader.discoveryCalls);
+ Assert.assertEquals(1, leader.registrationCalls);
+ }
+
+ private TestChannel addPeer(String name) {
+ TestChannel channel = new TestChannel(this.prefix + "-" + name + ":8686");
+ Assert.assertNull(this.channels.putIfAbsent(channel.address, channel));
+ this.peers.add(channel);
+ return channel;
+ }
+
+ private PDClient newClient(String hosts, String leader) throws Exception {
+ PDClient client = PDClient.create(PDConfig.of(hosts).setAuthority("store", "secret"));
+ // Isolate unrelated watch streams; discovery and registration still use the real client path.
+ PDWatch watch = Mockito.mock(PDWatch.class);
+ Mockito.when(watch.getCurrentHost()).thenReturn(leader);
+ Mockito.when(watch.checkChannel()).thenReturn(true);
+ field("pdWatch").set(client, watch);
+ return client;
+ }
+
+ private static Field field(String name) throws NoSuchFieldException {
+ Field field = PDClient.class.getDeclaredField(name);
+ field.setAccessible(true);
+ return field;
+ }
+
+ private static class TestChannel extends ManagedChannel {
+
+ private final String address;
+ private String leaderAddress;
+ private boolean failDiscovery;
+ private boolean shutdown;
+ private int discoveryCalls;
+ private int registrationCalls;
+
+ TestChannel(String address) {
+ this.address = address;
+ this.leaderAddress = address;
+ }
+
+ @Override
+ public ManagedChannel shutdown() {
+ this.shutdown = true;
+ return this;
+ }
+
+ @Override
+ public ManagedChannel shutdownNow() {
+ return shutdown();
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return this.shutdown;
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return this.shutdown;
+ }
+
+ @Override
+ public boolean awaitTermination(long timeout, TimeUnit unit) {
+ return this.shutdown;
+ }
+
+ @Override
+ public String authority() {
+ return this.address;
+ }
+
+ @Override
+ public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
+ MethodDescriptor<ReqT, RespT> method, CallOptions callOptions) {
+ return new ClientCall<>() {
+
+ private Listener<RespT> listener;
+
+ @Override
+ public void start(Listener<RespT> listener, Metadata headers) {
+ this.listener = listener;
+ }
+
+ @Override
+ public void request(int numMessages) {
+ }
+
+ @Override
+ public void cancel(String message, Throwable cause) {
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public void halfClose() {
+ Object response;
+ if (method.getFullMethodName().equals(PDGrpc.getGetMembersMethod().getFullMethodName())) {
+ discoveryCalls++;
+ if (failDiscovery) {
+ this.listener.onClose(Status.UNAVAILABLE, new Metadata());
+ return;
+ }
+ response = Pdpb.GetMembersResponse.newBuilder()
+ .setLeader(Metapb.Member.newBuilder()
+ .setGrpcUrl(leaderAddress))
+ .build();
+ } else {
+ Assert.assertEquals(PDGrpc.getRegisterStoreMethod().getFullMethodName(),
+ method.getFullMethodName());
+ registrationCalls++;
+ response = Pdpb.RegisterStoreResponse.newBuilder().setStoreId(17L).build();
+ }
+ this.listener.onMessage((RespT) response);
+ this.listener.onClose(Status.OK, new Metadata());
+ }
+
+ @Override
+ public void sendMessage(ReqT message) {
+ }
+ };
+ }
+ }
+}
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientSuiteTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientSuiteTest.java
index fd4bac8a0..8948fef16 100644
--- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientSuiteTest.java
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientSuiteTest.java
@@ -25,6 +25,9 @@ import lombok.extern.slf4j.Slf4j;
@RunWith(Suite.class)
@Suite.SuiteClasses({
BlockingUnaryCallTest.class,
+ KvClientLockTest.class,
+ PDClientLeaderTest.class,
+ PDClientWatchDispatchTest.class,
PDConnectionManagerTest.class,
PDClientTest.class,
KvClientTest.class,
diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientWatchDispatchTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientWatchDispatchTest.java
new file mode 100644
index 000000000..609f9f6b5
--- /dev/null
+++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/client/PDClientWatchDispatchTest.java
@@ -0,0 +1,329 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.pd.client;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.EnumMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.hugegraph.pd.client.listener.PDEventListener;
+import org.apache.hugegraph.pd.common.KVPair;
+import org.apache.hugegraph.pd.grpc.Metapb;
+import org.apache.hugegraph.pd.grpc.watch.HgPdWatchGrpc;
+import org.apache.hugegraph.pd.grpc.watch.NodeEventType;
+import org.apache.hugegraph.pd.grpc.watch.WatchChangeType;
+import org.apache.hugegraph.pd.grpc.watch.WatchGraphResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchNodeResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchPartitionResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchRequest;
+import org.apache.hugegraph.pd.grpc.watch.WatchResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchShardGroupResponse;
+import org.apache.hugegraph.pd.grpc.watch.WatchType;
+import org.apache.hugegraph.pd.watch.NodeEvent;
+import org.apache.hugegraph.pd.watch.PartitionEvent;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.mockito.Mockito;
+
+import io.grpc.CallOptions;
+import io.grpc.ClientCall;
+import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+
+@RunWith(Parameterized.class)
+public class PDClientWatchDispatchTest {
+
+ private static final String GRAPH = "test-graph";
+ private static final int PARTITION_ID = 3;
+ private static final long STORE_ID = 7L;
+ private static final Metapb.ShardGroup SHARD_GROUP =
+ Metapb.ShardGroup.newBuilder().setId(PARTITION_ID).build();
+
+ private final WatchType watchType;
+ private final String address = "pd-watch-test-" + UUID.randomUUID() + ":8686";
+ private final List<String> events = new ArrayList<>();
+ private final WatchChannel channel = new WatchChannel();
+ private ConcurrentMap<String, ManagedChannel> channels;
+ private ClientCache cache;
+ private PDClient client;
+ private Thread dispatchThread;
+
+ public PDClientWatchDispatchTest(WatchType watchType) {
+ this.watchType = watchType;
+ }
+
+ @Parameterized.Parameters(name = "{0}")
+ public static Object[][] watchTypes() {
+ return new Object[][]{
+ {WatchType.WATCH_TYPE_PARTITION_CHANGE},
+ {WatchType.WATCH_TYPE_STORE_NODE_CHANGE},
+ {WatchType.WATCH_TYPE_GRAPH_CHANGE},
+ {WatchType.WATCH_TYPE_SHARD_GROUP_CHANGE}
+ };
+ }
+
+ @SuppressWarnings("unchecked")
+ @Before
+ public void setUp() throws Exception {
+ Field channelsField = Channels.class.getDeclaredField("chs");
+ channelsField.setAccessible(true);
+ this.channels = (ConcurrentMap<String, ManagedChannel>) channelsField.get(null);
+ Assert.assertNull(this.channels.putIfAbsent(this.address, this.channel));
+ this.client = PDClient.create(PDConfig.of(this.address).setEnableCache(true));
+ this.cache = Mockito.mock(ClientCache.class);
+ Field cacheField = PDClient.class.getDeclaredField("cache");
+ cacheField.setAccessible(true);
+ cacheField.set(this.client, this.cache);
+ prepareCache();
+ this.dispatchThread = Thread.currentThread();
+
+ // Exercise the public setup and the real PDWatchImpl response adapters.
+ this.client.startWatch(this.address);
+ Assert.assertEquals(4, this.channel.listeners.size());
+ }
+
+ @After
+ public void tearDown() {
+ this.channels.remove(this.address, this.channel);
+ this.channel.shutdownNow();
+ }
+
+ @Test
+ public void testRuntimeFailureDoesNotStopOtherListenersOrLaterFrames() {
+ RuntimeException failure = new IllegalStateException("listener failed");
+ this.client.addEventListener(listener("first", () -> { throw failure; }));
+ this.client.addEventListener(listener("second", () -> { }));
+
+ emit();
+ emit();
+
+ List<String> expected = expectedEvents("first", "second");
+ expected.addAll(expectedEvents("first", "second"));
+ Assert.assertEquals(expected, this.events);
+ Assert.assertEquals(4, this.channel.listeners.size());
+ }
+
+ @Test
+ public void testErrorPropagatesWithoutCallingLaterListeners() {
+ AssertionError failure = new AssertionError("fatal listener failure");
+ this.client.addEventListener(listener("first", () -> { throw failure; }));
+ this.client.addEventListener(listener("second", () -> { }));
+
+ Assert.assertSame(failure, Assert.assertThrows(AssertionError.class, this::emit));
+
+ Assert.assertEquals(expectedEvents("first"), this.events);
+ }
+
+ @Test
+ public void testInternalCacheBoundaryIsPreserved() {
+ RuntimeException failure = new IllegalStateException("cache update failed");
+ this.client.addEventListener(listener("listener", () -> { }));
+ switch (this.watchType) {
+ case WATCH_TYPE_PARTITION_CHANGE ->
+ Mockito.doThrow(failure).when(this.cache).getPartitionById(GRAPH, PARTITION_ID);
+ case WATCH_TYPE_STORE_NODE_CHANGE ->
+ Mockito.doThrow(failure).when(this.cache).removeStore(STORE_ID);
+ case WATCH_TYPE_SHARD_GROUP_CHANGE ->
+ Mockito.doThrow(failure).when(this.cache).updateShardGroup(SHARD_GROUP);
+ case WATCH_TYPE_GRAPH_CHANGE -> {
+ emit();
+ Assert.assertEquals(List.of("listener"), this.events);
+ Mockito.verifyNoInteractions(this.cache);
+ return;
+ }
+ default -> throw new AssertionError(this.watchType);
+ }
+ Assert.assertSame(failure, Assert.assertThrows(IllegalStateException.class, this::emit));
+
+ Assert.assertTrue(this.events.isEmpty());
+ }
+
+ private void prepareCache() {
+ Mockito.when(this.cache.getPartitionById(GRAPH, PARTITION_ID)).thenAnswer(invocation -> {
+ this.events.add("lookup-partition");
+ return new KVPair<>(Metapb.Partition.newBuilder().setId(PARTITION_ID).build(),
+ Metapb.Shard.getDefaultInstance());
+ });
+ Mockito.doAnswer(invocation -> {
+ this.events.add("remove-partition");
+ return null;
+ }).when(this.cache).removePartition(GRAPH, PARTITION_ID);
+ Mockito.doAnswer(invocation -> {
+ this.events.add("remove-graph");
+ return null;
+ }).when(this.cache).removeAll(GRAPH);
+ Mockito.doAnswer(invocation -> {
+ this.events.add("remove-store");
+ return null;
+ }).when(this.cache).removeStore(STORE_ID);
+ Mockito.when(this.cache.updateShardGroup(SHARD_GROUP)).thenAnswer(invocation -> {
+ this.events.add("update-shard-group");
+ return true;
+ });
+ }
+
+ private PDEventListener listener(String name, Runnable callback) {
+ return new PDEventListener() {
+ private void accept(WatchType type) {
+ Assert.assertSame(watchType, type);
+ Assert.assertSame(dispatchThread, Thread.currentThread());
+ events.add(name);
+ callback.run();
+ }
+
+ @Override
+ public void onStoreChanged(NodeEvent event) {
+ accept(WatchType.WATCH_TYPE_STORE_NODE_CHANGE);
+ }
+
+ @Override
+ public void onPartitionChanged(PartitionEvent event) {
+ accept(WatchType.WATCH_TYPE_PARTITION_CHANGE);
+ }
+
+ @Override
+ public void onGraphChanged(WatchResponse event) {
+ accept(WatchType.WATCH_TYPE_GRAPH_CHANGE);
+ }
+
+ @Override
+ public void onShardGroupChanged(WatchResponse event) {
+ accept(WatchType.WATCH_TYPE_SHARD_GROUP_CHANGE);
+ }
+ };
+ }
+
+ private List<String> expectedEvents(String... listeners) {
+ List<String> cacheEvents = switch (this.watchType) {
+ case WATCH_TYPE_PARTITION_CHANGE ->
+ List.of("lookup-partition", "remove-partition", "remove-graph");
+ case WATCH_TYPE_STORE_NODE_CHANGE -> List.of("remove-store");
+ case WATCH_TYPE_SHARD_GROUP_CHANGE -> List.of("update-shard-group");
+ case WATCH_TYPE_GRAPH_CHANGE -> List.of();
+ default -> throw new AssertionError(this.watchType);
+ };
+ List<String> expected = new ArrayList<>(cacheEvents);
+ expected.addAll(Arrays.asList(listeners));
+ return expected;
+ }
+
+ private void emit() {
+ WatchResponse.Builder response = WatchResponse.newBuilder().setWatchType(this.watchType);
+ switch (this.watchType) {
+ case WATCH_TYPE_PARTITION_CHANGE -> response.setPartitionResponse(
+ WatchPartitionResponse.newBuilder().setGraph(GRAPH).setPartitionId(PARTITION_ID)
+ .setChangeType(WatchChangeType.WATCH_CHANGE_TYPE_DEL));
+ case WATCH_TYPE_STORE_NODE_CHANGE -> response.setNodeResponse(
+ WatchNodeResponse.newBuilder().setGraph(GRAPH).setNodeId(STORE_ID)
+ .setNodeEventType(NodeEventType.NODE_EVENT_TYPE_NODE_OFFLINE));
+ case WATCH_TYPE_GRAPH_CHANGE -> response.setGraphResponse(
+ WatchGraphResponse.newBuilder().setGraph(Metapb.Graph.newBuilder().setGraphName(GRAPH)));
+ case WATCH_TYPE_SHARD_GROUP_CHANGE -> response.setShardGroupResponse(
+ WatchShardGroupResponse.newBuilder().setShardGroupId(PARTITION_ID)
+ .setShardGroup(SHARD_GROUP)
+ .setType(WatchChangeType.WATCH_CHANGE_TYPE_ALTER));
+ default -> throw new AssertionError(this.watchType);
+ }
+ this.channel.listeners.get(this.watchType).onMessage(response.build());
+ }
+
+ private static class WatchChannel extends ManagedChannel {
+
+ private final Map<WatchType, ClientCall.Listener<WatchResponse>> listeners =
+ new EnumMap<>(WatchType.class);
+ private boolean shutdown;
+
+ @Override
+ public ManagedChannel shutdown() {
+ this.shutdown = true;
+ this.listeners.clear();
+ return this;
+ }
+
+ @Override
+ public ManagedChannel shutdownNow() {
+ return shutdown();
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return this.shutdown;
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return this.shutdown;
+ }
+
+ @Override
+ public boolean awaitTermination(long timeout, TimeUnit unit) {
+ return this.shutdown;
+ }
+
+ @Override
+ public String authority() {
+ return "test";
+ }
+
+ @Override
+ public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
+ MethodDescriptor<ReqT, RespT> method, CallOptions callOptions) {
+ Assert.assertEquals(HgPdWatchGrpc.getWatchMethod().getFullMethodName(), method.getFullMethodName());
+ return new ClientCall<>() {
+
+ private Listener<RespT> listener;
+
+ @Override
+ public void start(Listener<RespT> listener, Metadata headers) {
+ this.listener = listener;
+ }
+
+ @Override
+ public void request(int numMessages) {
+ }
+
+ @Override
+ public void cancel(String message, Throwable cause) {
+ }
+
+ @Override
+ public void halfClose() {
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public void sendMessage(ReqT message) {
+ WatchRequest request = (WatchRequest) message;
+ listeners.put(request.getCreateRequest().getWatchType(),
+ (Listener<WatchResponse>) this.listener);
+ }
+ };
+ }
+ }
+}
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/AbstractAlgorithm.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/AbstractAlgorithm.java
index e757bdb81..a2630cf97 100644
--- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/AbstractAlgorithm.java
+++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/AbstractAlgorithm.java
@@ -27,7 +27,6 @@ import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.function.Consumer;
-import org.apache.commons.lang3.StringEscapeUtils;
import org.apache.commons.lang3.mutable.MutableLong;
import org.apache.hugegraph.exception.HugeException;
import org.apache.hugegraph.id.Id;
@@ -57,7 +56,6 @@ import org.apache.tinkerpop.gremlin.structure.util.CloseableIterator;
import com.google.common.base.Objects;
-@SuppressWarnings("deprecation") // StringEscapeUtils
public abstract class AbstractAlgorithm implements Algorithm {
public static final long MAX_RESULT_SIZE = 100L * Bytes.MB;
@@ -307,6 +305,26 @@ public abstract class AbstractAlgorithm implements Algorithm {
}
}
+ protected Object countElements(Iterator<? extends Element> elements) {
+ Throwable failure = null;
+ try {
+ Map<String, MutableLong> counts = new HashMap<>();
+ long total = 0L;
+ while (elements.hasNext()) {
+ String label = elements.next().label();
+ counts.computeIfAbsent(label, key -> new MutableLong()).increment();
+ this.updateProgress(++total);
+ }
+ counts.put("*", new MutableLong(total));
+ return JsonUtil.asJson(counts);
+ } catch (RuntimeException | Error e) {
+ failure = e;
+ throw e;
+ } finally {
+ closeIterator(elements, failure);
+ }
+ }
+
protected long traverse(String sourceLabel, String sourceCLabel,
Consumer<Vertex> consumer) {
return this.traverse(sourceLabel, sourceCLabel, consumer, null);
@@ -472,17 +490,31 @@ public abstract class AbstractAlgorithm implements Algorithm {
protected <V> V execute(GraphTraversal<?, ?> traversal,
Callable<V> callback) {
long capacity = Query.defaultCapacity(MAX_QUERY_LIMIT);
+ Throwable failure = null;
try {
return callback.call();
} catch (Exception e) {
- throw new HugeException("Failed to execute algorithm: %s",
- e, e.getMessage());
+ HugeException error = new HugeException("Failed to execute algorithm: %s",
+ e, e.getMessage());
+ failure = error;
+ throw error;
+ } catch (Error e) {
+ failure = e;
+ throw e;
} finally {
Query.defaultCapacity(capacity);
try {
traversal.close();
- } catch (Exception e) {
- throw new HugeException("Can't close traversal", e);
+ } catch (Exception | Error e) {
+ if (failure == null) {
+ if (e instanceof Error error) {
+ throw error;
+ }
+ throw new HugeException("Can't close traversal", e);
+ }
+ if (failure != e) {
+ failure.addSuppressed(e);
+ }
}
}
}
@@ -655,10 +687,8 @@ public abstract class AbstractAlgorithm implements Algorithm {
}
private StringBuilder appendString(String str) {
- if (str.indexOf('"') >= 0) {
- str = StringEscapeUtils.escapeJson(str);
- }
- return this.json.append('"').append(str).append('"');
+ E.checkNotNull(str, "string");
+ return this.json.append(JsonUtil.toJson(str));
}
public void checkSizeLimit() {
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountEdgeAlgorithm.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountEdgeAlgorithm.java
index fd20f8a82..7f6732ab3 100644
--- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountEdgeAlgorithm.java
+++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountEdgeAlgorithm.java
@@ -17,14 +17,9 @@
package org.apache.hugegraph.job.algorithm;
-import java.util.HashMap;
-import java.util.Iterator;
import java.util.Map;
-import org.apache.commons.lang3.mutable.MutableLong;
import org.apache.hugegraph.job.UserJob;
-import org.apache.hugegraph.util.JsonUtil;
-import org.apache.tinkerpop.gremlin.structure.Edge;
public class CountEdgeAlgorithm extends AbstractAlgorithm {
@@ -52,26 +47,7 @@ public class CountEdgeAlgorithm extends AbstractAlgorithm {
}
public Object count() {
- Iterator<Edge> edges = this.edges(null);
-
- Map<String, MutableLong> counts = new HashMap<>();
- long total = 0L;
-
- while (edges.hasNext()) {
- Edge edge = edges.next();
- String label = edge.label();
- MutableLong count = counts.get(label);
- if (count != null) {
- count.increment();
- } else {
- counts.put(label, new MutableLong(1L));
- }
- total++;
- this.updateProgress(total);
- }
- counts.put("*", new MutableLong(total));
-
- return JsonUtil.asJson(counts);
+ return this.countElements(this.edges(null));
}
}
}
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountVertexAlgorithm.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountVertexAlgorithm.java
index c325370b3..9a33b45be 100644
--- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountVertexAlgorithm.java
+++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/job/algorithm/CountVertexAlgorithm.java
@@ -17,14 +17,9 @@
package org.apache.hugegraph.job.algorithm;
-import java.util.HashMap;
-import java.util.Iterator;
import java.util.Map;
-import org.apache.commons.lang3.mutable.MutableLong;
import org.apache.hugegraph.job.UserJob;
-import org.apache.hugegraph.util.JsonUtil;
-import org.apache.tinkerpop.gremlin.structure.Vertex;
public class CountVertexAlgorithm extends AbstractAlgorithm {
@@ -52,26 +47,7 @@ public class CountVertexAlgorithm extends AbstractAlgorithm {
}
public Object count() {
- Iterator<Vertex> vertices = this.vertices();
-
- Map<String, MutableLong> counts = new HashMap<>();
- long total = 0L;
-
- while (vertices.hasNext()) {
- Vertex vertex = vertices.next();
- String label = vertex.label();
- MutableLong count = counts.get(label);
- if (count != null) {
- count.increment();
- } else {
- counts.put(label, new MutableLong(1L));
- }
- total++;
- this.updateProgress(total);
- }
- counts.put("*", new MutableLong(total));
-
- return JsonUtil.asJson(counts);
+ return this.countElements(this.vertices());
}
}
}
diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/records/record/RecordFactory.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/records/record/RecordFactory.java
index 9c41a93bf..6b510cc0f 100644
--- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/records/record/RecordFactory.java
+++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/records/record/RecordFactory.java
@@ -24,20 +24,12 @@ public class RecordFactory {
}
public static Record newRecord(RecordType type, boolean concurrent) {
- Record record;
- switch (type) {
- case INT:
- record = new Int2IntRecord();
- break;
- case SET:
- record = new Int2SetRecord();
- break;
- case ARRAY:
- record = new Int2ArrayRecord();
- break;
- default:
- throw new AssertionError("Unsupported record type: " + type);
- }
+ Record record = switch (type) {
+ case INT -> new Int2IntRecord();
+ case SET -> new Int2SetRecord();
+ case ARRAY -> new Int2ArrayRecord();
+ default -> throw new AssertionError("Unsupported record type: " + type);
+ };
if (concurrent && !record.concurrent()) {
record = new SyncRecord(record);
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmExecutionTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmExecutionTest.java
new file mode 100644
index 000000000..61d5301d1
--- /dev/null
+++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmExecutionTest.java
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.job.algorithm;
+
+import org.apache.hugegraph.HugeGraph;
+import org.apache.hugegraph.config.CoreOptions;
+import org.apache.hugegraph.exception.HugeException;
+import org.apache.hugegraph.job.UserJob;
+import org.apache.hugegraph.job.algorithm.AbstractAlgorithm.AlgoTraverser;
+import org.apache.hugegraph.query.Query;
+import org.apache.hugegraph.testutil.Assert;
+import org.apache.hugegraph.type.define.CollectionType;
+import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class AlgorithmExecutionTest {
+
+ private static final long CAPACITY = 1234L;
+
+ private long previousCapacity;
+ private AlgoTraverser traverser;
+ private GraphTraversal<?, ?> traversal;
+
+ @Before
+ @SuppressWarnings("unchecked")
+ public void setUp() throws Exception {
+ this.previousCapacity = Query.defaultCapacity(CAPACITY);
+ HugeGraph graph = Mockito.mock(HugeGraph.class);
+ Mockito.when(graph.option(CoreOptions.OLTP_COLLECTION_TYPE)).thenReturn(CollectionType.JCF);
+ UserJob<Object> job = Mockito.mock(UserJob.class);
+ Mockito.when(job.graph()).thenReturn(graph);
+ this.traverser = new AlgoTraverser(job);
+ this.traversal = Mockito.mock(GraphTraversal.class);
+ Mockito.doAnswer(invocation -> {
+ Assert.assertEquals(CAPACITY, Query.defaultCapacity());
+ return null;
+ }).when(this.traversal).close();
+ }
+
+ @After
+ public void tearDown() {
+ try {
+ Assert.assertEquals(CAPACITY, Query.defaultCapacity());
+ } finally {
+ Query.defaultCapacity(this.previousCapacity);
+ }
+ }
+
+ @Test
+ public void testSuccessfulExecutionReturnsResultAndClosesTraversal() throws Exception {
+ Object expected = new Object();
+ Assert.assertSame(expected, this.traverser.execute(this.traversal, () -> {
+ Assert.assertEquals(AbstractAlgorithm.MAX_QUERY_LIMIT, Query.defaultCapacity());
+ return expected;
+ }));
+ Mockito.verify(this.traversal).close();
+ }
+
+ @Test
+ public void testCheckedComputationFailureRetainsCloseFailure() throws Exception {
+ Exception computation = new Exception("query failed");
+ Exception close = new Exception("close failed");
+ Mockito.doThrow(close).when(this.traversal).close();
+ HugeException thrown = Assert.assertThrows(HugeException.class,
+ () -> this.traverser.execute(this.traversal, () -> { throw computation; }));
+ Assert.assertSame(computation, thrown.getCause());
+ Assert.assertEquals("Failed to execute algorithm: query failed", thrown.getMessage());
+ Assert.assertArrayEquals(new Throwable[]{close}, thrown.getSuppressed());
+ Mockito.verify(this.traversal).close();
+ }
+
+ @Test
+ public void testRuntimeComputationFailureWithoutMessageRetainsCloseError() throws Exception {
+ RuntimeException computation = new IllegalStateException();
+ Error close = new AssertionError("close failed");
+ Mockito.doThrow(close).when(this.traversal).close();
+ HugeException thrown = Assert.assertThrows(HugeException.class,
+ () -> this.traverser.execute(this.traversal, () -> { throw computation; }));
+ Assert.assertSame(computation, thrown.getCause());
+ Assert.assertArrayEquals(new Throwable[]{close}, thrown.getSuppressed());
+ }
+
+ @Test
+ public void testFatalComputationErrorRetainsCloseFailure() throws Exception {
+ Error computation = new AssertionError("computation failed");
+ Exception close = new Exception("close failed");
+ Mockito.doThrow(close).when(this.traversal).close();
+ Error thrown = Assert.assertThrows(AssertionError.class,
+ () -> this.traverser.execute(this.traversal, () -> { throw computation; }));
+ Assert.assertSame(computation, thrown);
+ Assert.assertArrayEquals(new Throwable[]{close}, thrown.getSuppressed());
+ }
+
+ @Test
+ public void testCloseFailureAfterSuccessKeepsExistingWrapper() throws Exception {
+ Exception close = new Exception("close failed");
+ Mockito.doThrow(close).when(this.traversal).close();
+ HugeException thrown = Assert.assertThrows(HugeException.class,
+ () -> this.traverser.execute(this.traversal, () -> 1));
+ Assert.assertSame(close, thrown.getCause());
+ Assert.assertEquals("Can't close traversal", thrown.getMessage());
+ }
+
+ @Test
+ public void testCloseErrorAfterSuccessIsPropagatedUnchanged() throws Exception {
+ Error close = new AssertionError("close failed");
+ Mockito.doThrow(close).when(this.traversal).close();
+ Assert.assertSame(close, Assert.assertThrows(AssertionError.class,
+ () -> this.traverser.execute(this.traversal, () -> 1)));
+ }
+}
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmJsonMapTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmJsonMapTest.java
new file mode 100644
index 000000000..be0dcb387
--- /dev/null
+++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/AlgorithmJsonMapTest.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.job.algorithm;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.hugegraph.id.IdGenerator;
+import org.apache.hugegraph.job.algorithm.AbstractAlgorithm.JsonMap;
+import org.apache.hugegraph.testutil.Assert;
+import org.apache.hugegraph.util.JsonUtil;
+import org.junit.Test;
+
+public class AlgorithmJsonMapTest {
+
+ @Test
+ public void testStringVertexIdsRoundTripAsValues() {
+ List<String> ids = stringIds();
+ JsonMap json = new JsonMap();
+ json.startList();
+ for (String id : ids) {
+ json.append(IdGenerator.of(id).toString());
+ }
+ json.endList();
+ Assert.assertEquals(ids, JsonUtil.fromJson(JsonUtil.toJson(json.asJson()), List.class));
+ }
+
+ @Test
+ public void testStringVertexIdsRoundTripAsNumericResultKeys() {
+ List<String> ids = stringIds();
+ JsonMap json = new JsonMap();
+ json.startObject();
+ for (int i = 0; i < ids.size(); i++) {
+ json.append(IdGenerator.of(ids.get(i)), (long) i);
+ }
+ json.endObject();
+ Map<?, ?> results = JsonUtil.fromJson(JsonUtil.toJson(json.asJson()), Map.class);
+ Assert.assertEquals(ids.size(), results.size());
+ for (int i = 0; i < ids.size(); i++) {
+ Assert.assertEquals(i, ((Number) results.get(ids.get(i))).intValue());
+ }
+ }
+
+ @Test
+ public void testStringKeysAndValuesRoundTripTogether() {
+ JsonMap json = new JsonMap();
+ json.startObject();
+ for (String id : stringIds()) {
+ json.append(id, id);
+ }
+ json.endObject();
+ Map<?, ?> results = JsonUtil.fromJson(JsonUtil.toJson(json.asJson()), Map.class);
+ for (String id : stringIds()) {
+ Assert.assertEquals(id, results.get(id));
+ }
+ }
+
+ @Test
+ public void testAppendKeyEscapesNestedResultName() {
+ String key = "label\\name\n";
+ JsonMap json = new JsonMap();
+ json.startObject();
+ json.appendKey(key);
+ json.startList();
+ json.append("vertex\\name\t");
+ json.endList();
+ json.endObject();
+ Map<?, ?> results = JsonUtil.fromJson(JsonUtil.toJson(json.asJson()), Map.class);
+ Assert.assertEquals(List.of("vertex\\name\t"), results.get(key));
+ }
+
+ @Test
+ public void testPlainNumericAndRawResultsKeepTheirShape() {
+ JsonMap json = new JsonMap();
+ json.startObject();
+ json.append("degree", 3L);
+ json.append("rank", 0.5D);
+ json.appendRaw("stats", "{\"vertices\":2}");
+ json.endObject();
+ Assert.assertEquals("{\"degree\":3,\"rank\":0.5,\"stats\":{\"vertices\":2}}",
+ JsonUtil.toJson(json.asJson()));
+ }
+
+ @Test
+ public void testNullStringInputsRemainInvalid() {
+ JsonMap json = new JsonMap();
+ Assert.assertThrows(NullPointerException.class, () -> json.append((String) null));
+ Assert.assertThrows(NullPointerException.class, () -> json.appendKey(null));
+ }
+
+ private static List<String> stringIds() {
+ List<String> ids = new ArrayList<>(List.of("plain", "path\\name", "quoted\"id", "中文顶点😀"));
+ for (char control = 0; control < 32; control++) {
+ ids.add("vertex" + control + "id");
+ }
+ return ids;
+ }
+}
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/CountAlgorithmTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/CountAlgorithmTest.java
new file mode 100644
index 000000000..53cde6870
--- /dev/null
+++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/job/algorithm/CountAlgorithmTest.java
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.job.algorithm;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+
+import org.apache.hugegraph.HugeGraph;
+import org.apache.hugegraph.config.CoreOptions;
+import org.apache.hugegraph.exception.HugeException;
+import org.apache.hugegraph.job.UserJob;
+import org.apache.hugegraph.query.Query;
+import org.apache.hugegraph.testutil.Assert;
+import org.apache.hugegraph.type.define.CollectionType;
+import org.apache.hugegraph.util.JsonUtil;
+import org.apache.tinkerpop.gremlin.structure.Edge;
+import org.apache.tinkerpop.gremlin.structure.Element;
+import org.apache.tinkerpop.gremlin.structure.Vertex;
+import org.apache.tinkerpop.gremlin.structure.util.CloseableIterator;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.mockito.Mockito;
+
+@RunWith(Parameterized.class)
+public class CountAlgorithmTest {
+
+ private final boolean vertices;
+ private final Algorithm algorithm;
+ private UserJob<Object> job;
+ private TestCursor cursor;
+
+ public CountAlgorithmTest(String elements) {
+ this.vertices = "vertices".equals(elements);
+ this.algorithm = this.vertices ? new CountVertexAlgorithm() : new CountEdgeAlgorithm();
+ }
+
+ @Parameterized.Parameters(name = "{0}")
+ public static Collection<Object[]> elements() {
+ return List.of(new Object[]{"vertices"}, new Object[]{"edges"});
+ }
+
+ @Before
+ @SuppressWarnings("unchecked")
+ public void setUp() {
+ HugeGraph graph = Mockito.mock(HugeGraph.class);
+ Mockito.when(graph.option(CoreOptions.OLTP_COLLECTION_TYPE)).thenReturn(CollectionType.JCF);
+ this.job = Mockito.mock(UserJob.class);
+ Mockito.when(this.job.graph()).thenReturn(graph);
+ this.cursor = new TestCursor();
+ if (this.vertices) {
+ Mockito.doReturn(this.cursor).when(graph).vertices(Mockito.any(Query.class));
+ } else {
+ Mockito.doReturn(this.cursor).when(graph).edges(Mockito.any(Query.class));
+ }
+ }
+
+ @Test
+ public void testCountsLabelsAndTotalAndClosesScan() {
+ this.addElement("person");
+ this.addElement("software");
+ this.addElement("person");
+ Object result = this.algorithm.call(this.job, Map.of());
+ Assert.assertEquals(Map.of("person", 2, "software", 1, "*", 3),
+ JsonUtil.fromJson(JsonUtil.toJson(result), Map.class));
+ Assert.assertEquals(1, this.cursor.closes);
+ Mockito.verify(this.job).updateProgress(1);
+ Mockito.verify(this.job).updateProgress(2);
+ Mockito.verify(this.job).updateProgress(3);
+ }
+
+ @Test
+ public void testEmptyScanReturnsZeroAndCloses() {
+ Object result = this.algorithm.call(this.job, Map.of());
+ Assert.assertEquals(Map.of("*", 0), JsonUtil.fromJson(JsonUtil.toJson(result), Map.class));
+ Assert.assertEquals(1, this.cursor.closes);
+ Mockito.verify(this.job, Mockito.never()).updateProgress(Mockito.anyInt());
+ }
+
+ @Test
+ public void testReadFailureRemainsPrimaryWhenClosingFails() {
+ this.addElement("person");
+ RuntimeException read = new IllegalStateException("scan failed");
+ RuntimeException close = new IllegalStateException("close failed");
+ this.cursor.readFailure = read;
+ this.cursor.failAfter = 1;
+ this.cursor.closeFailure = close;
+ Throwable thrown = Assert.assertThrows(IllegalStateException.class,
+ () -> this.algorithm.call(this.job, Map.of()));
+ this.assertFailureAndClosed(read, close, thrown);
+ Mockito.verify(this.job).updateProgress(1);
+ }
+
+ @Test
+ public void testProgressFailureRemainsPrimaryWhenClosingFails() {
+ this.addElement("person");
+ RuntimeException progress = new IllegalStateException("progress save failed");
+ RuntimeException close = new IllegalStateException("close failed");
+ Mockito.doThrow(progress).when(this.job).updateProgress(1);
+ this.cursor.closeFailure = close;
+ Throwable thrown = Assert.assertThrows(IllegalStateException.class,
+ () -> this.algorithm.call(this.job, Map.of()));
+ this.assertFailureAndClosed(progress, close, thrown);
+ }
+
+ @Test
+ public void testCloseFailureAfterCountingIsReported() {
+ this.addElement("person");
+ RuntimeException close = new IllegalStateException("close failed");
+ this.cursor.closeFailure = close;
+ Throwable thrown = Assert.assertThrows(RuntimeException.class,
+ () -> this.algorithm.call(this.job, Map.of()));
+ Assert.assertSame(close, HugeException.rootCause(thrown));
+ Assert.assertEquals(1, this.cursor.closes);
+ }
+
+ @Test
+ public void testFatalReadErrorRemainsPrimaryWhenClosingFails() {
+ Error read = new AssertionError("fatal scan failure");
+ RuntimeException close = new IllegalStateException("close failed");
+ this.cursor.readFailure = read;
+ this.cursor.closeFailure = close;
+ Throwable thrown = Assert.assertThrows(AssertionError.class,
+ () -> this.algorithm.call(this.job, Map.of()));
+ this.assertFailureAndClosed(read, close, thrown);
+ }
+
+ private void addElement(String label) {
+ Element element = this.vertices ? Mockito.mock(Vertex.class) : Mockito.mock(Edge.class);
+ Mockito.when(element.label()).thenReturn(label);
+ this.cursor.elements.add(element);
+ }
+
+ private void assertFailureAndClosed(Throwable primary, Throwable close, Throwable thrown) {
+ Assert.assertSame(primary, thrown);
+ Assert.assertEquals(1, thrown.getSuppressed().length);
+ Assert.assertSame(close, HugeException.rootCause(thrown.getSuppressed()[0]));
+ Assert.assertEquals(1, this.cursor.closes);
+ }
+
+ private static final class TestCursor implements CloseableIterator<Element> {
+
+ private final List<Element> elements = new ArrayList<>();
+ private Throwable readFailure;
+ private RuntimeException closeFailure;
+ private int position;
+ private int failAfter;
+ private int closes;
+
+ @Override
+ public boolean hasNext() {
+ if (this.position == this.failAfter && this.readFailure != null) {
+ if (this.readFailure instanceof Error error) {
+ throw error;
+ }
+ throw (RuntimeException) this.readFailure;
+ }
+ return this.position < this.elements.size();
+ }
+
+ @Override
+ public Element next() {
+ if (!this.hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return this.elements.get(this.position++);
+ }
+
+ @Override
+ public void close() {
+ this.closes++;
+ if (this.closeFailure != null) {
+ throw this.closeFailure;
+ }
+ }
+ }
+}
diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java
index e9609752d..1a699defa 100644
--- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java
+++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java
@@ -39,6 +39,9 @@ import org.apache.hugegraph.backend.tx.GraphTransactionTest;
import org.apache.hugegraph.core.memory.MemoryManageTest;
import org.apache.hugegraph.job.EphemeralJobBuilderTest;
import org.apache.hugegraph.job.algorithm.AlgorithmConsumerLifecycleTest;
+import org.apache.hugegraph.job.algorithm.AlgorithmExecutionTest;
+import org.apache.hugegraph.job.algorithm.AlgorithmJsonMapTest;
+import org.apache.hugegraph.job.algorithm.CountAlgorithmTest;
import org.apache.hugegraph.job.system.DeleteExpiredJobTest;
import org.apache.hugegraph.meta.EtcdMetaDriverTest;
import org.apache.hugegraph.meta.MetaManagerSchemaCacheClearEventTest;
@@ -299,6 +302,9 @@ import org.junit.runners.Suite;
ConsumersTest.class,
AlgorithmConsumersTest.class,
AlgorithmConsumerLifecycleTest.class,
+ AlgorithmExecutionTest.class,
+ AlgorithmJsonMapTest.class,
+ CountAlgorithmTest.class,
VersionTest.class,
JsonUtilTest.class,
StringEncodingTest.class,
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/IntersectionFilterIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/IntersectionFilterIterator.java
index 2487ce810..bdcae813e 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/IntersectionFilterIterator.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/IntersectionFilterIterator.java
@@ -23,6 +23,7 @@ import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
+import java.util.NoSuchElementException;
import org.apache.hugegraph.rocksdb.access.RocksDBSession;
import org.apache.hugegraph.rocksdb.access.ScanIterator;
@@ -40,6 +41,7 @@ public class IntersectionFilterIterator implements ScanIterator {
private final ScanIterator iterator;
private final IntersectionWrapper wrapper;
private boolean processed = false;
+ private boolean closed;
private Iterator innerIterator;
private SortShuffle<RocksDBSession.BackendColumn> sortShuffle;
@@ -69,6 +71,9 @@ public class IntersectionFilterIterator implements ScanIterator {
@Override
public boolean hasNext() {
+ if (this.closed) {
+ return false;
+ }
if (!processed) {
try {
dedup();
@@ -137,7 +142,7 @@ public class IntersectionFilterIterator implements ScanIterator {
@Override
public boolean isValid() {
- if (this.processed) {
+ if (this.closed || this.processed) {
return false;
}
return iterator.isValid();
@@ -145,15 +150,20 @@ public class IntersectionFilterIterator implements ScanIterator {
@Override
public <T> T next() {
+ if (this.closed) {
+ throw new NoSuchElementException();
+ }
return (T) this.innerIterator.next();
}
@Override
public void close() {
+ this.closed = true;
try (SortShuffle<?> shuffle = this.sortShuffle) {
this.iterator.close();
} finally {
this.sortShuffle = null;
+ this.innerIterator = null;
this.map.clear();
}
}
@@ -188,6 +198,10 @@ public class IntersectionFilterIterator implements ScanIterator {
private int count = 0;
+ private boolean ready;
+
+ private boolean exhausted;
+
private final Iterator<E> iterator;
private final Comparator<E> comparator;
@@ -202,37 +216,48 @@ public class IntersectionFilterIterator implements ScanIterator {
}
/**
- * Consecutive duplicate elimination. When prev == current, record data. When not equal,
- * return previous data. Note: Final result may contain duplicates.
+ * Return the last element of each matching group. A group must contain exactly
+ * {@code adjacent} elements, or at least two when {@code adjacent == -1}.
*/
@Override
public boolean hasNext() {
+ if (this.ready) {
+ return true;
+ }
+ if (this.exhausted) {
+ return false;
+ }
while (iterator.hasNext()) {
+ E next = iterator.next();
if (prev == null) {
- prev = iterator.next();
+ prev = next;
+ current = next;
continue;
}
- current = iterator.next();
- if (comparator.compare(prev, current) == 0) {
- data = current;
+ if (comparator.compare(prev, next) == 0) {
+ current = next;
count += 1;
- } else {
- // count starts from 0, so the size is count + 1
- if (count > 0 && this.adjacent == -1 || count + 1 == this.adjacent) {
- count = 0;
- prev = current;
- return true;
- } else {
- count = 0;
- prev = current;
- }
+ continue;
}
- }
- // last result
- if (count > 0) {
+ boolean matches = matchesGroup();
+ if (matches) {
+ data = current;
+ }
count = 0;
+ prev = next;
+ current = next;
+ if (matches) {
+ this.ready = true;
+ return true;
+ }
+ }
+
+ this.exhausted = true;
+ if (prev != null && matchesGroup()) {
+ data = current;
+ this.ready = true;
return true;
}
@@ -241,7 +266,18 @@ public class IntersectionFilterIterator implements ScanIterator {
@Override
public E next() {
- return data;
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ E result = this.data;
+ this.data = null;
+ this.ready = false;
+ return result;
+ }
+
+ private boolean matchesGroup() {
+ // count excludes the first element of the group.
+ return this.adjacent == -1 ? this.count > 0 : this.count + 1 == this.adjacent;
}
}
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/MapValueFilterIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/MapValueFilterIterator.java
index bccd126f5..e91511ce7 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/MapValueFilterIterator.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/MapValueFilterIterator.java
@@ -27,6 +27,7 @@ public class MapValueFilterIterator<K> implements Iterator<K> {
Iterator<Map.Entry<K, Integer>> mapIterator;
private final IntPredicate filter;
private K value;
+ private boolean ready;
public MapValueFilterIterator(Map<K, Integer> map, IntPredicate filter) {
this.mapIterator = map.entrySet().iterator();
@@ -35,10 +36,14 @@ public class MapValueFilterIterator<K> implements Iterator<K> {
@Override
public boolean hasNext() {
+ if (this.ready) {
+ return true;
+ }
while (mapIterator.hasNext()) {
Map.Entry<K, Integer> entry = mapIterator.next();
if (filter.test(entry.getValue())) {
value = entry.getKey();
+ this.ready = true;
return true;
}
}
@@ -48,10 +53,12 @@ public class MapValueFilterIterator<K> implements Iterator<K> {
@Override
public K next() {
- if (value == null) {
+ if (!hasNext()) {
throw new NoSuchElementException();
}
-
- return value;
+ K result = this.value;
+ this.value = null;
+ this.ready = false;
+ return result;
}
}
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/UnionFilterIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/UnionFilterIterator.java
index 8395fddd9..5e471fd45 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/UnionFilterIterator.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/itrv2/UnionFilterIterator.java
@@ -44,6 +44,7 @@ public class UnionFilterIterator<T extends Serializable> implements ScanIterator
private final SortShuffleSerializer<T> serializer;
private Object current;
private boolean isProcessed = false;
+ private boolean closed;
public UnionFilterIterator(ScanIterator iterator, IntersectionWrapper<T> wrapper,
Comparator<T> comparator, SortShuffleSerializer<T> serializer) {
@@ -74,6 +75,12 @@ public class UnionFilterIterator<T extends Serializable> implements ScanIterator
@Override
public boolean hasNext() {
+ if (this.closed) {
+ return false;
+ }
+ if (this.current != null) {
+ return true;
+ }
while (this.iterator.hasNext()) {
var obj = (T) this.iterator.next();
// batch get or index lookup may generate null
@@ -131,24 +138,28 @@ public class UnionFilterIterator<T extends Serializable> implements ScanIterator
@Override
public boolean isValid() {
// todo: check logic
- return this.iterator.isValid() || hasNext();
+ return !this.closed && (this.iterator.isValid() || hasNext());
}
@Override
public <X> X next() {
- if (current == null) {
+ if (!hasNext() || this.current == null) {
throw new NoSuchElementException();
}
-
- return (X) current;
+ X result = (X) this.current;
+ this.current = null;
+ return result;
}
@Override
public void close() {
+ this.closed = true;
+ this.current = null;
try (SortShuffle<T> shuffle = this.sortShuffle) {
this.iterator.close();
} finally {
this.sortShuffle = null;
+ this.innerIterator = null;
this.map.clear();
}
}
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/FilterIteratorContractTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/FilterIteratorContractTest.java
new file mode 100644
index 000000000..6a956a77f
--- /dev/null
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/FilterIteratorContractTest.java
@@ -0,0 +1,307 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.store.node.grpc.query;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn;
+import org.apache.hugegraph.rocksdb.access.ScanIterator;
+import org.apache.hugegraph.store.business.itrv2.IntersectionFilterIterator;
+import org.apache.hugegraph.store.business.itrv2.IntersectionWrapper;
+import org.apache.hugegraph.store.business.itrv2.MapValueFilterIterator;
+import org.apache.hugegraph.store.business.itrv2.UnionFilterIterator;
+import org.apache.hugegraph.store.business.itrv2.io.SortShuffleSerializer;
+import org.apache.hugegraph.store.grpc.query.QueryRequest;
+import org.apache.hugegraph.store.grpc.query.QueryResponse;
+import org.apache.hugegraph.store.node.grpc.query.model.QueryPlan;
+import org.junit.Test;
+
+import io.grpc.stub.StreamObserver;
+
+public class FilterIteratorContractTest extends AggregativeQueryTestSupport {
+
+ @Test
+ public void testMapFilterSupportsDirectNextAndRepeatedProbes() {
+ Map<String, Integer> counts = new LinkedHashMap<>();
+ counts.put("skip-first", 1);
+ counts.put("first", 2);
+ counts.put("skip-middle", 3);
+ counts.put("second", 2);
+ counts.put("third", 2);
+ counts.put("skip-last", 0);
+ Iterator<String> iterator = new MapValueFilterIterator<>(counts, count -> count == 2);
+
+ assertEquals("first", iterator.next());
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.hasNext());
+ assertEquals("second", iterator.next());
+ assertEquals("third", iterator.next());
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+
+ @Test
+ public void testMapFilterDoesNotRepeatConsumedValue() {
+ Iterator<String> iterator = new MapValueFilterIterator<>(Map.of("only", 1), count -> count > 0);
+
+ assertTrue(iterator.hasNext());
+ assertEquals("only", iterator.next());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ assertFalse(iterator.hasNext());
+ }
+
+ @Test
+ public void testEmptyAndRejectedMapsRemainExhausted() {
+ for (Map<String, Integer> counts : List.of(Map.<String, Integer>of(), Map.of("rejected", 1))) {
+ Iterator<String> iterator = new MapValueFilterIterator<>(counts, count -> count == 2);
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+ }
+
+ @Test
+ public void testUnionUniqueStreamPreservesPendingResults() {
+ assertFilterContract(FilterKind.UNION_UNIQUE);
+ }
+
+ @Test
+ public void testUnionBufferedResultsPreservePendingResults() {
+ assertFilterContract(FilterKind.UNION_BUFFERED);
+ }
+
+ @Test
+ public void testIntersectionInMemoryPreservesPendingResults() {
+ assertFilterContract(FilterKind.INTERSECTION);
+ }
+
+ @Test
+ public void testEmptyUnionRemainsExhausted() {
+ IntersectionWrapper<BackendColumn> wrapper = new IntersectionWrapper<>(
+ new RowIterator(List.of()), column -> column.name[0]);
+ wrapper.proc();
+ try (ScanIterator iterator = new UnionFilterIterator<>(new RowIterator(List.of()), wrapper,
+ BackendColumn::compareTo,
+ SortShuffleSerializer.ofBackendColumnSerializer())) {
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+ }
+
+ @Test
+ public void testCloseDiscardsPendingResultsWithoutProbingClosedSource() {
+ for (FilterKind kind : FilterKind.values()) {
+ FilterFixture fixture = filter(kind);
+ assertTrue(fixture.iterator.hasNext());
+
+ fixture.iterator.close();
+
+ assertClosedWithoutProbingSource(fixture);
+ }
+ }
+
+ @Test
+ public void testCloseFailureStillDiscardsPendingResultsAndStopsProbing() {
+ for (FilterKind kind : FilterKind.values()) {
+ FilterFixture fixture = filter(kind);
+ IllegalStateException failure = new IllegalStateException("source close failed");
+ fixture.source.closeFailure = failure;
+ assertTrue(fixture.iterator.hasNext());
+
+ assertSame(failure, assertThrows(IllegalStateException.class, fixture.iterator::close));
+
+ assertClosedWithoutProbingSource(fixture);
+ }
+ }
+
+ @Test(timeout = 5000)
+ public void testUnionUniqueStreamDoesNotLoseRecordsBetweenSingleRecordBatches() throws Exception {
+ assertSingleRecordBatches(FilterKind.UNION_UNIQUE);
+ }
+
+ @Test(timeout = 5000)
+ public void testUnionBufferedResultsDoNotLoseRecordsBetweenSingleRecordBatches() throws Exception {
+ assertSingleRecordBatches(FilterKind.UNION_BUFFERED);
+ }
+
+ @Test(timeout = 5000)
+ public void testIntersectionDoesNotLoseRecordsBetweenSingleRecordBatches() throws Exception {
+ assertSingleRecordBatches(FilterKind.INTERSECTION);
+ }
+
+ private static void assertFilterContract(FilterKind kind) {
+ FilterFixture fixture = filter(kind);
+ List<Integer> keys = new ArrayList<>();
+ try (ScanIterator iterator = fixture.iterator) {
+ // Intersection initializes its reduction on the first probe; keep that API unchanged.
+ if (kind == FilterKind.INTERSECTION) {
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.hasNext());
+ }
+ BackendColumn first = iterator.next();
+ keys.add((int) first.name[0]);
+ while (iterator.hasNext()) {
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.hasNext());
+ BackendColumn column = iterator.next();
+ keys.add((int) column.name[0]);
+ }
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+ keys.sort(Integer::compareTo);
+ assertEquals(List.of(1, 2, 3), keys);
+ assertTrue(fixture.source.closed);
+ }
+
+ private static void assertClosedWithoutProbingSource(FilterFixture fixture) {
+ int probes = fixture.source.probes;
+ assertFalse(fixture.iterator.hasNext());
+ assertFalse(fixture.iterator.hasNext());
+ assertFalse(fixture.iterator.isValid());
+ assertThrows(NoSuchElementException.class, fixture.iterator::next);
+ assertEquals(probes, fixture.source.probes);
+ assertTrue(fixture.source.closed);
+ }
+
+ private static void assertSingleRecordBatches(FilterKind kind) throws Exception {
+ FilterFixture fixture = filter(kind);
+ ThreadPoolExecutor pool = pool(1);
+ QueryPlan plan = new QueryPlan();
+ AggregativeQueryService service = new AggregativeQueryService(pool, 500, 1) {
+ @Override
+ AggregativeQueryObserver newObserver(StreamObserver<QueryResponse> sender) {
+ return fixture(super.newObserver(sender), fixture.iterator, plan, null);
+ }
+ };
+ ResponseRecorder response = new ResponseRecorder();
+ try {
+ service.query(response).onNext(QueryRequest.newBuilder().setQueryId(kind.name()).build());
+ assertTrue(response.finished.await(2, TimeUnit.SECONDS));
+
+ List<Integer> keys = new ArrayList<>();
+ assertEquals(3, response.responses.size());
+ for (int i = 0; i < response.responses.size(); i++) {
+ QueryResponse batch = response.responses.get(i);
+ assertTrue(batch.getIsOk());
+ assertEquals(kind.name(), batch.getQueryId());
+ assertEquals(1, batch.getDataCount());
+ assertEquals(i == response.responses.size() - 1, batch.getIsFinished());
+ keys.add((int) batch.getData(0).getKey().byteAt(0));
+ }
+ keys.sort(Integer::compareTo);
+ assertEquals(List.of(1, 2, 3), keys);
+ assertEquals(1, response.completed.get());
+ assertEquals(0, response.errors.get());
+ assertTrue(fixture.source.closed);
+ service.shutdownQueries();
+ } finally {
+ pool.shutdownNow();
+ pool.awaitTermination(1, TimeUnit.SECONDS);
+ fixture.iterator.close();
+ }
+ }
+
+ private static FilterFixture filter(FilterKind kind) {
+ BackendColumn first = BackendColumn.of(new byte[]{1}, new byte[]{11});
+ BackendColumn second = BackendColumn.of(new byte[]{2}, new byte[]{12});
+ BackendColumn third = BackendColumn.of(new byte[]{3}, new byte[]{13});
+ List<BackendColumn> rows = kind == FilterKind.UNION_UNIQUE ? List.of(first, second, third) :
+ Arrays.asList(first, second, third, first, second, third);
+ IntersectionWrapper<BackendColumn> wrapper = new IntersectionWrapper<>(
+ new RowIterator(rows), column -> column.name[0]);
+ wrapper.proc();
+ RowIterator source = new RowIterator(rows);
+ ScanIterator iterator = kind == FilterKind.INTERSECTION ?
+ new IntersectionFilterIterator(source, wrapper, 2) :
+ new UnionFilterIterator<>(source, wrapper, BackendColumn::compareTo,
+ SortShuffleSerializer.ofBackendColumnSerializer());
+ return new FilterFixture(iterator, source);
+ }
+
+ private enum FilterKind {
+ UNION_UNIQUE, UNION_BUFFERED, INTERSECTION
+ }
+
+ private static class FilterFixture {
+
+ private final ScanIterator iterator;
+ private final RowIterator source;
+
+ private FilterFixture(ScanIterator iterator, RowIterator source) {
+ this.iterator = iterator;
+ this.source = source;
+ }
+ }
+
+ private static class RowIterator implements ScanIterator {
+
+ private final Iterator<BackendColumn> rows;
+ private boolean closed;
+ private int probes;
+ private RuntimeException closeFailure;
+
+ private RowIterator(List<BackendColumn> rows) {
+ this.rows = rows.iterator();
+ }
+
+ @Override
+ public boolean hasNext() {
+ this.probes++;
+ return !this.closed && this.rows.hasNext();
+ }
+
+ @Override
+ public boolean isValid() {
+ return hasNext();
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public <T> T next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return (T) this.rows.next();
+ }
+
+ @Override
+ public void close() {
+ this.closed = true;
+ if (this.closeFailure != null) {
+ throw this.closeFailure;
+ }
+ }
+ }
+}
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/IntersectionSpillIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/IntersectionSpillIteratorTest.java
new file mode 100644
index 000000000..c2a90d363
--- /dev/null
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/IntersectionSpillIteratorTest.java
@@ -0,0 +1,251 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hugegraph.store.node.grpc.query;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import java.io.File;
+import java.io.OutputStream;
+import java.lang.reflect.Field;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Deque;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn;
+import org.apache.hugegraph.rocksdb.access.ScanIterator;
+import org.apache.hugegraph.store.business.itrv2.IntersectionFilterIterator;
+import org.apache.hugegraph.store.business.itrv2.IntersectionFilterIterator.ReduceIterator;
+import org.apache.hugegraph.store.business.itrv2.IntersectionWrapper;
+import org.apache.hugegraph.store.business.itrv2.io.SortShuffleSerializer;
+import org.apache.hugegraph.store.grpc.query.QueryRequest;
+import org.apache.hugegraph.store.grpc.query.QueryResponse;
+import org.apache.hugegraph.store.node.grpc.query.model.QueryPlan;
+import org.apache.hugegraph.store.util.SortShuffle;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import io.grpc.stub.StreamObserver;
+
+public class IntersectionSpillIteratorTest extends AggregativeQueryTestSupport {
+
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+
+ private String previousBasePath;
+
+ @Before
+ public void setUp() {
+ this.previousBasePath = SortShuffle.getBasePath();
+ SortShuffle.setBasePath(this.folder.getRoot().getAbsolutePath() + File.separator);
+ }
+
+ @After
+ public void tearDown() {
+ SortShuffle.setBasePath(this.previousBasePath);
+ }
+
+ @Test
+ public void testReduceIteratorSupportsDirectNextAndRepeatedProbes() {
+ Iterator<Integer> iterator = new ReduceIterator<>(List.of(1, 1, 2, 2, 3, 3).iterator(),
+ Integer::compareTo, 2);
+
+ assertEquals(Integer.valueOf(1), iterator.next());
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.hasNext());
+ assertEquals(Integer.valueOf(2), iterator.next());
+ assertEquals(Integer.valueOf(3), iterator.next());
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+
+ @Test
+ public void testEmptyReduceIteratorRemainsExhausted() {
+ for (int size : new int[]{-1, 1, 2, 3}) {
+ Iterator<Integer> iterator = new ReduceIterator<>(List.<Integer>of().iterator(),
+ Integer::compareTo, size);
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+ }
+
+ @Test
+ public void testSpillKeepsLastRepresentativeAtGroupBoundaryAndEof() throws Exception {
+ assertEquals(List.of(List.of(1, 12), List.of(2, 22)),
+ readSpill(2, column(1, 11), column(1, 12), column(2, 21), column(2, 22)));
+ }
+
+ @Test
+ public void testExactSizeRejectsShortFinalGroup() throws Exception {
+ assertEquals(List.of(List.of(1, 13)),
+ readSpill(3, column(1, 11), column(1, 12), column(1, 13),
+ column(2, 21), column(2, 22)));
+ }
+
+ @Test
+ public void testExactSizeRejectsOversizedFinalGroup() throws Exception {
+ assertEquals(List.of(List.of(1, 12)),
+ readSpill(2, column(1, 11), column(1, 12), column(2, 21),
+ column(2, 22), column(2, 23)));
+ }
+
+ @Test
+ public void testExactSizeUsesSamePredicateForEveryGroup() throws Exception {
+ assertEquals(List.of(List.of(2, 23)),
+ readSpill(3, column(1, 11), column(1, 12),
+ column(2, 21), column(2, 22), column(2, 23),
+ column(3, 31), column(3, 32), column(3, 33), column(3, 34)));
+ }
+
+ @Test
+ public void testDefaultDuplicatePredicateSkipsSingletonsAtBothEnds() throws Exception {
+ assertEquals(List.of(List.of(2, 22), List.of(4, 43)),
+ readSpill(-1, column(1, 11), column(2, 21), column(2, 22), column(3, 31),
+ column(4, 41), column(4, 42), column(4, 43), column(5, 51)));
+ assertEquals(List.of(List.of(2, 22)),
+ readSpill(-1, column(1, 11), column(2, 21), column(2, 22)));
+ }
+
+ @Test
+ public void testExactOneReturnsSingletonWithoutNullOrStaleRepresentative() throws Exception {
+ assertEquals(List.of(List.of(1, 11), List.of(3, 31)),
+ readSpill(1, column(1, 11), column(2, 21), column(2, 22), column(3, 31)));
+ }
+
+ @Test(timeout = 5000)
+ public void testSingleRecordRpcBatchesPreserveEverySpilledGroup() throws Exception {
+ ScanIterator iterator = spill(2, column(1, 11), column(1, 12), column(2, 21), column(2, 22),
+ column(3, 31), column(3, 32));
+ ThreadPoolExecutor pool = pool(1);
+ QueryPlan plan = new QueryPlan();
+ AggregativeQueryService service = new AggregativeQueryService(pool, 500, 1) {
+ @Override
+ AggregativeQueryObserver newObserver(StreamObserver<QueryResponse> sender) {
+ return fixture(super.newObserver(sender), iterator, plan, null);
+ }
+ };
+ ResponseRecorder response = new ResponseRecorder();
+ try {
+ service.query(response).onNext(QueryRequest.newBuilder().setQueryId("intersection-spill").build());
+ assertTrue(response.finished.await(2, TimeUnit.SECONDS));
+
+ assertEquals(3, response.responses.size());
+ for (int i = 0; i < response.responses.size(); i++) {
+ QueryResponse batch = response.responses.get(i);
+ assertTrue(batch.getIsOk());
+ assertEquals("intersection-spill", batch.getQueryId());
+ assertEquals(1, batch.getDataCount());
+ assertEquals(i + 1, batch.getData(0).getKey().byteAt(0));
+ assertEquals((i + 1) * 10 + 2, batch.getData(0).getValue().byteAt(0));
+ assertEquals(i == 2, batch.getIsFinished());
+ }
+ assertEquals(1, response.completed.get());
+ assertEquals(0, response.errors.get());
+ assertEquals(0, this.folder.getRoot().list().length);
+ service.shutdownQueries();
+ } finally {
+ pool.shutdownNow();
+ pool.awaitTermination(1, TimeUnit.SECONDS);
+ iterator.close();
+ }
+ }
+
+ private List<List<Integer>> readSpill(int size, BackendColumn... columns) throws Exception {
+ List<List<Integer>> result = new ArrayList<>();
+ try (ScanIterator iterator = spill(size, columns)) {
+ while (iterator.hasNext()) {
+ assertTrue(iterator.hasNext());
+ assertTrue(iterator.hasNext());
+ BackendColumn column = iterator.next();
+ result.add(List.of((int) column.name[0], (int) column.value[0]));
+ }
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.hasNext());
+ assertThrows(NoSuchElementException.class, iterator::next);
+ }
+ assertEquals(0, this.folder.getRoot().list().length);
+ return result;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static ScanIterator spill(int size, BackendColumn... columns) throws Exception {
+ SortShuffle<BackendColumn> shuffle = new SortShuffle<>(
+ (left, right) -> Arrays.compare(left.name, right.name),
+ SortShuffleSerializer.ofBackendColumnSerializer());
+ try {
+ // Seed a real sorted spill without changing the production spill thresholds.
+ File spill = new File((String) field(SortShuffle.class, "path").get(shuffle), "seed");
+ try (OutputStream output = Files.newOutputStream(spill.toPath())) {
+ for (BackendColumn column : columns) {
+ SortShuffleSerializer.ofBackendColumnSerializer().write(output, column);
+ }
+ }
+ ((Deque<String>) field(SortShuffle.class, "files").get(shuffle)).add(spill.getAbsolutePath());
+ ScanIterator source = new ScanIterator() {
+ @Override
+ public boolean hasNext() {
+ return false;
+ }
+
+ @Override
+ public boolean isValid() {
+ return false;
+ }
+
+ @Override
+ public <T> T next() {
+ throw new NoSuchElementException();
+ }
+
+ @Override
+ public void close() {
+ }
+ };
+ IntersectionWrapper<BackendColumn> wrapper = new IntersectionWrapper<>(source, column -> column.name[0]);
+ IntersectionFilterIterator iterator = new IntersectionFilterIterator(source, wrapper, size);
+ field(IntersectionFilterIterator.class, "sortShuffle").set(iterator, shuffle);
+ return iterator;
+ } catch (Exception | Error failure) {
+ shuffle.close();
+ throw failure;
+ }
+ }
+
+ private static BackendColumn column(int key, int value) {
+ return BackendColumn.of(new byte[]{(byte) key}, new byte[]{(byte) value});
+ }
+
+ private static Field field(Class<?> owner, String name) throws Exception {
+ Field field = owner.getDeclaredField(name);
+ field.setAccessible(true);
+ return field;
+ }
+}
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java
index aaf3679a7..a40320919 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java
@@ -21,6 +21,8 @@ import org.apache.hugegraph.rocksdb.access.RocksDBScanIteratorCloseTest;
import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest;
import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryLifecycleTest;
import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryShutdownTest;
+import org.apache.hugegraph.store.node.grpc.query.FilterIteratorContractTest;
+import org.apache.hugegraph.store.node.grpc.query.IntersectionSpillIteratorTest;
import org.apache.hugegraph.store.node.grpc.query.OrderByStageTest;
import org.apache.hugegraph.store.node.grpc.query.UnaryQueryLifecycleTest;
import org.apache.hugegraph.store.node.grpc.BatchScanHalfCloseTest;
@@ -40,6 +42,8 @@ import lombok.extern.slf4j.Slf4j;
PartitionEngineShutdownTest.class,
AggregativeQueryLifecycleTest.class,
AggregativeQueryShutdownTest.class,
+ FilterIteratorContractTest.class,
+ IntersectionSpillIteratorTest.class,
OrderByStageTest.class,
UnaryQueryLifecycleTest.class,
BatchScanHalfCloseTest.class,