diff --git a/.github/workflows/pd-store-ci.yml b/.github/workflows/pd-store-ci.yml index 231751e303..89d5f5f4e7 100644 --- a/.github/workflows/pd-store-ci.yml +++ b/.github/workflows/pd-store-ci.yml @@ -8,6 +8,21 @@ permissions: contents: read jobs: + store-stop-shell-test: + permissions: + contents: read + runs-on: ubuntu-24.04 + steps: + - name: Checkout + uses: actions/checkout@v7 + with: + persist-credentials: false + + - name: Verify Store stop failure propagation and PID retention + run: | + docker run --rm -v "$GITHUB_WORKSPACE:/work:ro" -w /work ubuntu:24.04 \ + bash hugegraph-server/hugegraph-dist/src/assembly/travis/test-stop-hugegraph-store.sh + rocksdb-compatibility: permissions: contents: read @@ -431,6 +446,11 @@ jobs: mvn test -pl hugegraph-store/hg-store-test -am \ -P store-core-test -Djacoco.sessionId=store-core-test + - name: Run server shutdown test + run: | + mvn test -pl hugegraph-store/hg-store-test -am \ + -P store-server-test -Djacoco.sessionId=store-server-test + - name: Generate aggregate coverage report run: | mvn verify -pl hugegraph-store/hg-store-test -am -P jacoco \ @@ -450,6 +470,8 @@ jobs: --require-test-report \ "$TEST_REPORT_DIR/TEST-org.apache.hugegraph.store.core.CoreSuiteTest.xml" \ --require-test-report \ + "$TEST_REPORT_DIR/TEST-org.apache.hugegraph.store.service.ServerSuiteTest.xml" \ + --require-test-report \ "hugegraph-store/hg-store-node/target/surefire-reports/TEST-org.apache.hugegraph.store.business.StoredRowIngressTest.xml" \ --require-covered-group hg-store-common \ --require-covered-group hg-store-client \ @@ -460,6 +482,7 @@ jobs: --require-session store-rocksdb-test \ --require-session store-raftcore-test \ --require-session store-core-test \ + --require-session store-server-test \ --require-session store-node-test \ "$REPORT_FILE" \ hg-store-grpc hg-store-common hg-store-client hg-store-rocksdb hg-store-core diff --git a/README.md b/README.md index 00cacdb561..cda80674e5 100644 --- a/README.md +++ b/README.md @@ -316,9 +316,6 @@ gremlin> :> g.V().limit(5) ``` For comprehensive documentation, visit the [HugeGraph Documentation](https://hugegraph.apache.org/docs/). - -See [standalone RocksDB snapshot recovery](docs/rocksdb-recovery.md) before restoring data or mounting store directories. - For an existing deployment, read the [RocksDB upgrade guidance](docs/rocksdb-upgrade.md) before upgrading the storage runtime. diff --git a/docs/transaction-lifecycle.md b/docs/transaction-lifecycle.md new file mode 100644 index 0000000000..0d1fbe6ffe --- /dev/null +++ b/docs/transaction-lifecycle.md @@ -0,0 +1,30 @@ +# Request transactions and Store shutdown + +Server releases the current thread's graph transactions when a REST request +finishes and when an authenticated context task exits. Cleanup visits every +registered graph even when one graph fails to close, and reports failures. +Cleanup explicitly rolls back unfinished writes, including when `onClose(COMMIT)` +was configured, then releases backend transactions and resets thread-local +transaction behavior and listeners. Applications must commit successful writes +before leaving the request/task boundary. Closing an OLTP traverser preserves +the caller's transaction so that the caller can still commit or roll it back. + +The cleanup also applies when schema caches have become cold. Deleting an auth +project or user must remove its associated access or belong edges while +preserving unrelated relationships. The special OLAP vertex path remains +separate; negative schema IDs alone do not identify edges that can be skipped. + +Store shutdown refuses new RPCs, cancels active RPCs and scans, and waits for their callbacks and scan/TTL workers before Spring destroys the Store engine and its databases. A failed aggregate-query response callback is logged without skipping cancellation or cleanup waits for other queries. Cancellation may fail in-flight requests; stop writes and check their outcomes before planned maintenance. This is not a guarantee that every in-flight request drains successfully or that a leader transfers without a failover interval. + +If a callback cannot finish, or an ordinary scan or bidirectional aggregate query fails to release its plan or partition iterators, shutdown remains pending rather than closing its database underneath it. This also includes iterators discarded while advancing past empty partitions, initializing a sequential scan, or counting rows. RocksDB iterators retain the first failure during automatic or explicit close so that later cleanup cannot hide it; concurrent close calls wait for the release attempt to finish. All scan entry points, including one-shot scans, retain failed cleanup independently of RPC termination and executor shutdown. Failed iterator releases are attempted once and remain diagnostic shutdown blockers. Bidirectional aggregate queries report cleanup failures to the client and log that shutdown is blocked; their final success batch is sent only after cleanup succeeds. The distribution stop script waits up to 30 seconds, returns a nonzero status on timeout, and retains the PID file for diagnosis. Inspect logs and thread dumps before retrying. Do not add a concurrent shutdown hook that closes the same databases. + +If shutdown cancellation wins while a one-shot scan releases its iterator, the +response terminates with `CANCELLED`; it does not report successful completion +without its result. Iterator cleanup still finishes before the scan unregisters. + +A normal aggregate-query request half-close ends feedback without cancelling already permitted work. A subsequent transport cancellation or deadline still interrupts the workers and releases their resources. A batch-scan RPC accepts one initial query; repeated query requests are ignored before allocating another iterator, including after its final batch. If the remaining feedback credit cannot finish the query, the server returns an explicit query error. Scan task rejection reports `UNAVAILABLE` during shutdown and `RESOURCE_EXHAUSTED` when the running scan pool is full. + +See the [Store shutdown instructions](../hugegraph-store/README.md#stopping-a-store-node) +and the Server [module test guidance](../hugegraph-server/AGENTS.md#tests). + +Shared schema and element caches retain their invalidation listeners until the graph closes. Request cleanup releases backend leases while preserving those graph caches and their schema identity. diff --git a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeFactoryAuthProxy.java b/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeFactoryAuthProxy.java index 78ef3d72d2..98f9ea8bc4 100644 --- a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeFactoryAuthProxy.java +++ b/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeFactoryAuthProxy.java @@ -185,7 +185,7 @@ private static void registerPrivateActions() { "this$0"); Reflection.registerFieldsToFilter(HugeGraphAuthProxy.Context.class, "ADMIN", "user"); Reflection.registerFieldsToFilter(HugeGraphAuthProxy.ContextTask.class, "runner", - "context"); + "cleanup", "context"); Reflection.registerFieldsToFilter(StandardHugeGraph.class, "LOG", "started", "closed", "mode", "variables", "name", "params", "configuration", "schemaEventHub", "graphEventHub", "indexEventHub", @@ -203,7 +203,8 @@ private static void registerPrivateActions() { "access$14", "access$15", "access$16", "access$17", "access$18", "serializer", "loadSchemaStore", "loadSystemStore", "loadGraphStore", "closeTx", - "analyzer", "serverInfoManager", "reloadRamtable", + "closeCurrentThreadTransaction", "analyzer", + "serverInfoManager", "reloadRamtable", "reloadRamtable", "access$19", "access$20", "access$21"); Reflection.registerFieldsToFilter( loadClass("org.apache.hugegraph.StandardHugeGraph$StandardHugeGraphParams"), @@ -298,8 +299,10 @@ private static void registerPrivateActions() { "autoCommit", "beforeRead", "afterWrite", "afterRead", "commitMutation2Backend", "checkOwnerThread", "doAction", "store", "reset"); - Reflection.registerFieldsToFilter(HugeFactory.class, "LOG", "NAME_REGEX", "graphs"); - Reflection.registerMethodsToFilter(HugeFactory.class, "lambda$0"); + Reflection.registerFieldsToFilter(HugeFactory.class, "LOG", "NAME_REGEX", "graphs", + "GRAPHS"); + Reflection.registerMethodsToFilter(HugeFactory.class, "lambda$0", + "closeCurrentThreadTransactions"); Reflection.registerFieldsToFilter(SchemaElement.class, "graph", "id", "name", "userdata", "status"); Reflection.registerFieldsToFilter(HugeVertex.class, "EMPTY_SET", "id", "label", "edges", diff --git a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeGraphAuthProxy.java b/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeGraphAuthProxy.java index 9d23e82556..2ae3d7312e 100644 --- a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeGraphAuthProxy.java +++ b/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/auth/HugeGraphAuthProxy.java @@ -1291,11 +1291,17 @@ public User user() { static class ContextTask implements Runnable { private final Runnable runner; + private final Runnable cleanup; private final Context context; public ContextTask(Runnable runner) { + this(runner, HugeFactory::closeCurrentThreadTransactions); + } + + ContextTask(Runnable runner, Runnable cleanup) { this.context = getContext(); this.runner = runner; + this.cleanup = cleanup; } @Override @@ -1305,7 +1311,7 @@ public void run() { this.runner.run(); } finally { try { - HugeFactory.closeCurrentThreadTransactions(); + this.cleanup.run(); } catch (Throwable e) { LOG.error("Failed to close Gremlin worker transactions", e); } finally { diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/HugeFactory.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/HugeFactory.java index 451455ba0f..78a6de78fc 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/HugeFactory.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/HugeFactory.java @@ -22,6 +22,7 @@ import java.io.File; import java.net.URL; import java.util.ArrayList; +import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -118,6 +119,11 @@ public static void closeCurrentThreadTransactions() { synchronized (HugeFactory.class) { graphs = new ArrayList<>(GRAPHS.values()); } + closeCurrentThreadTransactions(graphs); + } + + static void closeCurrentThreadTransactions( + Collection graphs) { Throwable failure = null; for (StandardHugeGraph graph : graphs) { try { diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java index d9b838e897..d8127922d5 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/StandardHugeGraph.java @@ -1361,22 +1361,25 @@ public void rollback() { } public void close() { - try { - this.graphTx.close(); - } catch (Exception e) { - LOG.error("Failed to close GraphTransaction", e); + Throwable failure = null; + for (Runnable close : new Runnable[]{this.graphTx::close, + this.systemTx::close, + this.schemaTx::close}) { + try { + close.run(); + } catch (RuntimeException | Error error) { + if (failure == null) { + failure = error; + } else if (failure != error) { + failure.addSuppressed(error); + } + } } - - try { - this.systemTx.close(); - } catch (Exception e) { - LOG.error("Failed to close SystemTransaction", e); + if (failure instanceof Error) { + throw (Error) failure; } - - try { - this.schemaTx.close(); - } catch (Exception e) { - LOG.error("Failed to close SchemaTransaction", e); + if (failure != null) { + throw (RuntimeException) failure; } } diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/store/BackendSessionPool.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/store/BackendSessionPool.java index 521f44556f..44137cd766 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/store/BackendSessionPool.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/store/BackendSessionPool.java @@ -56,14 +56,17 @@ public HugeConfig config() { public final BackendSession getOrNewSession() { BackendSession session = this.threadLocalSession.get(); if (session == null) { - session = this.newSession(); - assert session != null; - this.threadLocalSession.set(session); - assert !this.sessions.containsKey(Thread.currentThread().getId()); - this.sessions.put(Thread.currentThread().getId(), session); - int sessionCount = this.sessionCount.incrementAndGet(); - LOG.debug("Now(after connect({})) session count is: {}", - this, sessionCount); + // Serialize new borrowers with the last-session native close. + synchronized (this) { + session = this.newSession(); + assert session != null; + this.threadLocalSession.set(session); + assert !this.sessions.containsKey(Thread.currentThread().getId()); + this.sessions.put(Thread.currentThread().getId(), session); + int sessionCount = this.sessionCount.incrementAndGet(); + LOG.debug("Now(after connect({})) session count is: {}", + this, sessionCount); + } } else { this.detectSession(session); } @@ -131,7 +134,7 @@ public void forceResetSessions() { } } - public boolean close() { + public synchronized boolean close() { Pair result = Pair.of(-1, -1); try { result = this.closeSession(); diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphTransaction.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphTransaction.java index 70487f159e..e1b5c37d27 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphTransaction.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphTransaction.java @@ -765,7 +765,7 @@ public Iterator queryTaskInfos(Query query) { } public Iterator queryTaskInfos(Object... vertexIds) { - if (this.graph().backendStoreFeatures().supportsTaskAndServerVertex()) { + if (this.storeFeatures().supportsTaskAndServerVertex()) { return this.queryVerticesByIds(vertexIds, false, false, HugeType.TASK); } diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java index 9c0d3e9d18..ed4673feb1 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/task/StandardTaskScheduler.java @@ -596,7 +596,7 @@ private Iterator> queryTask(Map conditions, boolean withResult) { return this.call(() -> { ConditionQuery query; - if (this.graph.backendStoreFeatures().supportsTaskAndServerVertex()) { + if (this.tx().storeFeatures().supportsTaskAndServerVertex()) { query = new ConditionQuery(HugeType.TASK); } else { query = new ConditionQuery(HugeType.VERTEX); diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/OltpTraverser.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/OltpTraverser.java index 6e9bae5397..9d12b0ee9a 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/OltpTraverser.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/traversal/algorithm/OltpTraverser.java @@ -73,7 +73,7 @@ protected OltpTraverser(HugeGraph graph) { @Override public void close() { - // pass + // The graph's thread-local transaction belongs to the caller. } public static void destroy() { diff --git a/hugegraph-server/hugegraph-dist/src/assembly/travis/run-native-runtime-smoke-test.sh b/hugegraph-server/hugegraph-dist/src/assembly/travis/run-native-runtime-smoke-test.sh index e8570eef01..6631c04560 100755 --- a/hugegraph-server/hugegraph-dist/src/assembly/travis/run-native-runtime-smoke-test.sh +++ b/hugegraph-server/hugegraph-dist/src/assembly/travis/run-native-runtime-smoke-test.sh @@ -39,7 +39,7 @@ cleanup() { local status=$? trap - EXIT if [[ "$SERVER_START_ATTEMPTED" == "true" ]]; then - "$SERVER_DIR/bin/stop-hugegraph.sh" -m false >/dev/null 2>&1 || status=1 + "$SERVER_DIR/bin/stop-hugegraph.sh" -m false || status=1 fi rm -rf "$WORK_DIR" exit "$status" diff --git a/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh index 59b4caf32c..bf7fd64530 100755 --- a/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh +++ b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-check-jacoco-report.sh @@ -430,8 +430,8 @@ def validation_command(job): def reports_for_option(job, option): pattern = re.escape(option) + ( - r'\s+\\?\s*"\$TEST_REPORT_DIR/' - r'(TEST-[A-Za-z0-9_.]+SuiteTest[.]xml)"' + r'\s+\\?\s*"[^"]*/' + r'(TEST-[A-Za-z0-9_.]+Test[.]xml)"' ) return set(re.findall(pattern, validation_command(job))) @@ -492,9 +492,12 @@ assert_order(store_job, [ "-P store-client-test -Djacoco.sessionId=store-client-test", "-P store-rocksdb-test -Djacoco.sessionId=store-rocksdb-test", "-P store-raftcore-test -Djacoco.sessionId=store-raftcore-test", + "-P store-core-test -Djacoco.sessionId=store-core-test", + "-P store-server-test -Djacoco.sessionId=store-server-test", "mvn verify", "--require-session store-common-test", "--require-session store-client-test", "--require-session store-rocksdb-test", - "--require-session store-raftcore-test", "codecov/codecov-action", + "--require-session store-raftcore-test", "--require-session store-core-test", + "--require-session store-server-test", "codecov/codecov-action", ]) assert store_job.count("mvn clean") == 1 assert "hugegraph-store/hg-store-test/target/site/jacoco/jacoco.xml" in store_job @@ -504,7 +507,7 @@ assert "mvn verify -pl hugegraph-store/hg-store-test -am -P jacoco \\ " \ "-DskipTests -Deditorconfig.skip=true -ntp" in " ".join(store_job.split()) assert selected_profiles(store_job, "store") == { "store-common-test", "store-client-test", "store-rocksdb-test", - "store-raftcore-test", "store-core-test", + "store-raftcore-test", "store-core-test", "store-server-test", } assert reports_for_option(store_job, "--require-test-report") == { "TEST-org.apache.hugegraph.store.common.CommonSuiteTest.xml", @@ -512,8 +515,11 @@ assert reports_for_option(store_job, "--require-test-report") == { "TEST-org.apache.hugegraph.store.rocksdb.RocksDbSuiteTest.xml", "TEST-org.apache.hugegraph.store.raftcore.RaftSuiteTest.xml", "TEST-org.apache.hugegraph.store.core.CoreSuiteTest.xml", + "TEST-org.apache.hugegraph.store.service.ServerSuiteTest.xml", + "TEST-org.apache.hugegraph.store.business.StoredRowIngressTest.xml", } assert not reports_for_option(store_job, "--require-suite-report") +assert values_for_option(store_job, "--require-session") == (selected_profiles(store_job, "store") | {"store-node-test"}) assert values_for_option(store_job, "--require-covered-group") == { "hg-store-common", "hg-store-client", "hg-store-rocksdb", "hg-store-core", } diff --git a/hugegraph-server/hugegraph-dist/src/assembly/travis/test-stop-hugegraph-store.sh b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-stop-hugegraph-store.sh new file mode 100755 index 0000000000..5d05ab6ef2 --- /dev/null +++ b/hugegraph-server/hugegraph-dist/src/assembly/travis/test-stop-hugegraph-store.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# +# 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. +# +# Verify stop failure propagation and PID retention without signalling real processes. +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../.." && pwd)" +SCRIPT="$ROOT/hugegraph-store/hg-store-dist/src/assembly/static/bin/stop-hugegraph-store.sh" +FIXTURE=$(mktemp -d) +trap 'rm -rf "$FIXTURE"' EXIT +mkdir -p "$FIXTURE/bin" +cp "$SCRIPT" "$FIXTURE/bin/stop-hugegraph-store.sh" +cat > "$FIXTURE/bin/util.sh" <<'UTIL' +kill_process_and_wait() { + [[ "$1" == HugeGraphStoreServer && "$2" == 12345 && "$3" == 30 ]] || return 99 + return "$STOP_RESULT" +} +UTIL +printf '12345\n' > "$FIXTURE/bin/pid" +if STOP_RESULT=1 bash "$FIXTURE/bin/stop-hugegraph-store.sh"; then + echo "Stop timeout incorrectly returned success" >&2 + exit 1 +fi +[[ "$(cat "$FIXTURE/bin/pid")" == 12345 ]] +STOP_RESULT=0 bash "$FIXTURE/bin/stop-hugegraph-store.sh" +[[ ! -e "$FIXTURE/bin/pid" ]] +STOP_RESULT=0 bash "$FIXTURE/bin/stop-hugegraph-store.sh" +echo "store-stop-contract-ok" diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/HugeFactoryTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/HugeFactoryTest.java new file mode 100644 index 0000000000..da4ef196dc --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/HugeFactoryTest.java @@ -0,0 +1,44 @@ +/* + * 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; + +import java.util.Arrays; + +import org.apache.hugegraph.exception.HugeException; +import org.apache.hugegraph.testutil.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +public class HugeFactoryTest { + + @Test + public void testCleanupContinuesAfterGraphFailure() { + StandardHugeGraph failed = Mockito.mock(StandardHugeGraph.class); + StandardHugeGraph cleaned = Mockito.mock(StandardHugeGraph.class); + Mockito.doThrow(new HugeException("test")) + .when(failed).closeCurrentThreadTransaction(); + + Assert.assertThrows(HugeException.class, () -> { + HugeFactory.closeCurrentThreadTransactions( + Arrays.asList(failed, cleaned)); + }); + + Mockito.verify(failed).closeCurrentThreadTransaction(); + Mockito.verify(cleaned).closeCurrentThreadTransaction(); + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/ContextTaskTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/ContextTaskTest.java new file mode 100644 index 0000000000..a5ef4ad8d7 --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/ContextTaskTest.java @@ -0,0 +1,85 @@ +/* + * 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.auth; + +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.auth.HugeGraphAuthProxy.Context; +import org.apache.hugegraph.auth.HugeGraphAuthProxy.ContextTask; +import org.apache.hugegraph.testutil.Assert; +import org.junit.After; +import org.junit.Test; + +public class ContextTaskTest { + + @After + public void teardown() { + HugeGraphAuthProxy.resetContext(); + } + + @Test + public void testCleanupRunsBeforeContextReset() { + Context context = Context.admin(); + AtomicBoolean runnerCalled = new AtomicBoolean(false); + AtomicBoolean cleanupCalled = new AtomicBoolean(false); + HugeGraphAuthProxy.setContext(context); + + ContextTask task = new ContextTask(() -> { + Assert.assertSame(context, HugeGraphAuthProxy.getContext()); + runnerCalled.set(true); + }, () -> { + Assert.assertSame(context, HugeGraphAuthProxy.getContext()); + cleanupCalled.set(true); + }); + HugeGraphAuthProxy.resetContext(); + + task.run(); + + Assert.assertTrue(runnerCalled.get()); + Assert.assertTrue(cleanupCalled.get()); + Assert.assertNull(HugeGraphAuthProxy.getContext()); + } + + @Test + public void testCleanupRunsWhenTaskFails() { + AtomicBoolean cleanupCalled = new AtomicBoolean(false); + ContextTask task = new ContextTask(() -> { + throw new IllegalStateException("test"); + }, () -> cleanupCalled.set(true)); + + Assert.assertThrows(IllegalStateException.class, task::run); + Assert.assertTrue(cleanupCalled.get()); + Assert.assertNull(HugeGraphAuthProxy.getContext()); + } + + @Test + public void testContextResetRunsWhenCleanupFails() { + Context context = Context.admin(); + HugeGraphAuthProxy.setContext(context); + ContextTask task = new ContextTask(() -> { + Assert.assertSame(context, HugeGraphAuthProxy.getContext()); + }, () -> { + throw new IllegalStateException("test"); + }); + HugeGraphAuthProxy.resetContext(); + + task.run(); + + Assert.assertNull(HugeGraphAuthProxy.getContext()); + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/TransactionLifecycleTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/TransactionLifecycleTest.java new file mode 100644 index 0000000000..b940bc7f3e --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/auth/TransactionLifecycleTest.java @@ -0,0 +1,123 @@ +/* + * 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.auth; + +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hugegraph.HugeFactory; +import org.apache.hugegraph.auth.HugeGraphAuthProxy.Context; +import org.apache.hugegraph.auth.HugeGraphAuthProxy.ContextTask; +import org.apache.hugegraph.id.IdGenerator; +import org.apache.hugegraph.core.BaseCoreTest; +import org.apache.hugegraph.testutil.Assert; +import org.apache.hugegraph.traversal.algorithm.KneighborTraverser; +import org.apache.hugegraph.type.define.Directions; +import org.apache.tinkerpop.gremlin.structure.T; +import org.apache.tinkerpop.gremlin.structure.Transaction; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public class TransactionLifecycleTest extends BaseCoreTest { + + @Before + public void initSchema() { + graph().schema().vertexLabel("lifecycle").useCustomizeNumberId().create(); + HugeFactory.closeCurrentThreadTransactions(); + } + + @After + public void cleanup() { + try { + HugeFactory.closeCurrentThreadTransactions(); + } finally { + HugeGraphAuthProxy.resetContext(); + } + } + + @Test + public void testTraverserClosePreservesCallerWrites() { + graph().addVertex(T.label, "lifecycle", T.id, 1L); + graph().tx().commit(); + + try (KneighborTraverser traverser = new KneighborTraverser(graph())) { + Assert.assertTrue(traverser.kneighbor(IdGenerator.of(1L), Directions.BOTH, + null, 1, 100L, 100L).isEmpty()); + graph().addVertex(T.label, "lifecycle", T.id, 2L); + } + + Assert.assertTrue(graph().tx().isOpen()); + graph().tx().commit(); + HugeFactory.closeCurrentThreadTransactions(); + Assert.assertEquals(2L, graph().traversal().V().count().next().longValue()); + } + + @Test + public void testFailedTaskRollsBackDespiteCommitOnClose() { + AtomicInteger rollbackEvents = new AtomicInteger(); + HugeGraphAuthProxy.setContext(Context.admin()); + IllegalStateException failure = new IllegalStateException("task failed after writing"); + ContextTask task = new ContextTask(() -> { + Transaction tx = graph().tx(); + tx.onClose(Transaction.CLOSE_BEHAVIOR.COMMIT); + tx.onReadWrite(Transaction.READ_WRITE_BEHAVIOR.MANUAL); + tx.open(); + tx.addTransactionListener(status -> rollbackEvents.incrementAndGet()); + graph().addVertex(T.label, "lifecycle", T.id, 1L); + throw failure; + }); + HugeGraphAuthProxy.resetContext(); + + Assert.assertSame(failure, Assert.assertThrows(IllegalStateException.class, task::run)); + Assert.assertFalse(graph().tx().isOpen()); + Assert.assertNull(HugeGraphAuthProxy.getContext()); + Assert.assertEquals(1, rollbackEvents.get()); + // Reuse the worker thread: AUTO and default rollback must be restored. + Assert.assertEquals(0L, graph().traversal().V().count().next().longValue()); + graph().addVertex(T.label, "lifecycle", T.id, 2L); + graph().tx().close(); + Assert.assertEquals(0L, graph().traversal().V().count().next().longValue()); + Assert.assertEquals(1, rollbackEvents.get()); + } + + @Test + public void testSuccessfulTaskKeepsExplicitCommitAndResetsThreadState() { + AtomicInteger commitEvents = new AtomicInteger(); + HugeGraphAuthProxy.setContext(Context.admin()); + ContextTask task = new ContextTask(() -> { + Transaction tx = graph().tx(); + tx.onClose(Transaction.CLOSE_BEHAVIOR.COMMIT); + tx.onReadWrite(Transaction.READ_WRITE_BEHAVIOR.MANUAL); + tx.open(); + tx.addTransactionListener(status -> commitEvents.incrementAndGet()); + graph().addVertex(T.label, "lifecycle", T.id, 1L); + tx.commit(); + }); + HugeGraphAuthProxy.resetContext(); + + task.run(); + + Assert.assertFalse(graph().tx().isOpen()); + Assert.assertNull(HugeGraphAuthProxy.getContext()); + Assert.assertEquals(1L, graph().traversal().V().count().next().longValue()); + graph().addVertex(T.label, "lifecycle", T.id, 2L); + graph().tx().close(); + Assert.assertEquals(1L, graph().traversal().V().count().next().longValue()); + Assert.assertEquals(1, commitEvents.get()); + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/CoreTestSuite.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/CoreTestSuite.java index 9b45f40a64..3d6d3149fc 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/CoreTestSuite.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/CoreTestSuite.java @@ -19,11 +19,13 @@ import org.apache.hugegraph.HugeGraph; import org.apache.hugegraph.auth.BackendLeaseCleanupTest; +import org.apache.hugegraph.auth.TransactionLifecycleTest; import org.apache.hugegraph.constant.ServiceConstant; import org.apache.hugegraph.dist.RegisterUtil; import org.apache.hugegraph.masterelection.GlobalMasterInfo; import org.apache.hugegraph.meta.MetaManager; import org.apache.hugegraph.meta.PdMetaDriver; +import org.apache.hugegraph.task.StandardTaskSchedulerTxTest; import org.apache.hugegraph.task.TaskAndResultSchedulerTest; import org.apache.hugegraph.testutil.Utils; import org.apache.hugegraph.util.Log; @@ -52,7 +54,9 @@ RestoreCoreTest.class, TaskCoreTest.class, TaskAndResultSchedulerTest.class, + StandardTaskSchedulerTxTest.class, AuthTest.class, + TransactionLifecycleTest.class, MultiGraphsTest.class, RamTableTest.class }) diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/task/StandardTaskSchedulerTxTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/task/StandardTaskSchedulerTxTest.java new file mode 100644 index 0000000000..6b03f55627 --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/task/StandardTaskSchedulerTxTest.java @@ -0,0 +1,53 @@ +/* + * 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.task; + +import java.util.Collections; + +import org.apache.hugegraph.HugeGraph; +import org.apache.hugegraph.id.IdGenerator; +import org.apache.hugegraph.core.BaseCoreTest; +import org.junit.Assert; +import org.junit.Test; + +public class StandardTaskSchedulerTxTest extends BaseCoreTest { + + @Test + public void testTaskQueryDoesNotOpenUpperGraphTransaction() { + HugeGraph graph = this.graph(); + TaskScheduler scheduler = graph.taskScheduler(); + try { + // The task worker must also be quiescent after graph startup. + Assert.assertFalse(scheduler.call(() -> graph.tx().isOpen())); + + scheduler.tasks(TaskStatus.NEW, 1L, null).hasNext(); + Assert.assertFalse(scheduler.call(() -> graph.tx().isOpen())); + + scheduler.tasks(Collections.singletonList(IdGenerator.of(9999999L))) + .hasNext(); + Assert.assertFalse(scheduler.call(() -> graph.tx().isOpen())); + } finally { + scheduler.call(() -> { + if (graph.tx().isOpen()) { + graph.tx().rollback(); + } + return null; + }); + } + } +} 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 3346f8ad6d..dd456d0e58 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 @@ -17,11 +17,13 @@ package org.apache.hugegraph.unit; +import org.apache.hugegraph.HugeFactoryTest; import org.apache.hugegraph.api.auth.GraphSpaceAuthPayloadTest; import org.apache.hugegraph.api.auth.GraphSpaceGroupAPITest; import org.apache.hugegraph.api.cypher.CypherClientTest; import org.apache.hugegraph.auth.ContextGremlinServerHttpTest; import org.apache.hugegraph.auth.HttpGremlinRequestHandlerTest; +import org.apache.hugegraph.auth.ContextTaskTest; import org.apache.hugegraph.auth.StandardAuthManagerV2Test; import org.apache.hugegraph.auth.WsAndHttpBasicAuthHandlerTest; import org.apache.hugegraph.backend.page.QueryListTest; @@ -54,6 +56,7 @@ import org.apache.hugegraph.unit.config.GraphBinarySchemaCompatibilityTest; import org.apache.hugegraph.unit.config.K8sResourceQuotaYamlTest; import org.apache.hugegraph.unit.core.AnalyzerTest; +import org.apache.hugegraph.unit.core.BackendSessionPoolTest; import org.apache.hugegraph.unit.core.BackendMutationTest; import org.apache.hugegraph.unit.core.BackendProviderFactoryTest; import org.apache.hugegraph.unit.core.BackendStoreInfoTest; @@ -142,6 +145,8 @@ GraphSpaceGroupAPITest.class, GraphSpaceAuthPayloadTest.class, StandardAuthManagerV2Test.class, + ContextTaskTest.class, + HugeFactoryTest.class, AuthMetaManagerTest.class, /* api space */ @@ -179,6 +184,7 @@ RowLockTest.class, AnalyzerTest.class, BackendMutationTest.class, + BackendSessionPoolTest.class, BackendProviderFactoryTest.class, ConditionTest.class, GroovyScriptEngineCompatibilityTest.class, diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/cache/RequestCacheLifetimeTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/cache/RequestCacheLifetimeTest.java index 812343be30..df3547e0e6 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/cache/RequestCacheLifetimeTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/cache/RequestCacheLifetimeTest.java @@ -17,6 +17,8 @@ package org.apache.hugegraph.unit.cache; +import java.lang.reflect.Constructor; +import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Field; import java.lang.reflect.Method; import java.util.Collections; @@ -42,6 +44,9 @@ import org.apache.tinkerpop.gremlin.structure.Edge; import org.apache.tinkerpop.gremlin.structure.T; import org.apache.tinkerpop.gremlin.structure.Vertex; +import org.apache.hugegraph.backend.tx.GraphTransaction; +import org.apache.hugegraph.backend.tx.ISchemaTransaction; +import org.mockito.Mockito; import org.junit.Test; import sun.misc.Unsafe; @@ -184,4 +189,38 @@ private static void updateThroughApi(HugeGraph graph, Class api, String bodyC "DEFAULT", graph.name(), id, "append", body); Assert.assertTrue(response.contains(value)); } + @Test + public void testTransactionClosePreservesFailureAndAttemptsEveryOwner() throws Exception { + Class holderType = Class.forName("org.apache.hugegraph.StandardHugeGraph$Txs"); + Class systemType = Class.forName("org.apache.hugegraph.StandardHugeGraph$SysTransaction"); + GraphTransaction graphTx = + Mockito.mock(GraphTransaction.class); + GraphTransaction systemTx = + (GraphTransaction) Mockito.mock(systemType); + ISchemaTransaction schemaTx = + Mockito.mock(ISchemaTransaction.class); + RuntimeException first = new IllegalStateException("graph close failed"); + RuntimeException second = new IllegalArgumentException("system close failed"); + Mockito.doThrow(first).when(graphTx).close(); + Mockito.doThrow(second).when(systemTx).close(); + Constructor constructor = holderType.getDeclaredConstructor( + ISchemaTransaction.class, systemType, + GraphTransaction.class); + constructor.setAccessible(true); + Object holder = constructor.newInstance(schemaTx, systemTx, graphTx); + Method close = holderType.getDeclaredMethod("close"); + close.setAccessible(true); + try { + close.invoke(holder); + Assert.fail("Failed cleanup must reach the request boundary"); + } catch (InvocationTargetException error) { + Assert.assertSame(first, error.getCause()); + Assert.assertEquals(1, first.getSuppressed().length); + Assert.assertSame(second, first.getSuppressed()[0]); + } + Mockito.verify(graphTx).close(); + Mockito.verify(systemTx).close(); + Mockito.verify(schemaTx).close(); + } + } diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/BackendSessionPoolTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/BackendSessionPoolTest.java new file mode 100644 index 0000000000..d84d782c34 --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/BackendSessionPoolTest.java @@ -0,0 +1,231 @@ +/* + * 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.unit.core; + +import java.lang.management.ManagementFactory; +import java.lang.management.ThreadInfo; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; + +import org.apache.hugegraph.backend.store.BackendSession; +import org.apache.hugegraph.backend.store.BackendSessionPool; +import org.apache.hugegraph.unit.FakeObjects; +import org.junit.Assert; +import org.junit.Test; + +public class BackendSessionPoolTest { + + private static final long TIMEOUT_SECONDS = 10L; + private static final long CLOSE_HOLD_SECONDS = 30L; + + @Test + public void testConcurrentAcquireDoesNotLeaveSessionOnClosedBackend() + throws Exception { + TestSessionPool pool = new TestSessionPool(); + pool.open(); + + AtomicReference closeFailure = new AtomicReference<>(); + CountDownLatch acquireAttempted = new CountDownLatch(1); + FutureTask acquisition = + new FutureTask<>(() -> { + acquireAttempted.countDown(); + return pool.getOrNewSession(); + }); + + Thread closer = new Thread(() -> { + try { + pool.getOrNewSession(); + pool.close(); + } catch (Throwable e) { + closeFailure.set(e); + } + }, "backend-session-closer"); + Thread acquirer = new Thread(acquisition, "backend-session-acquirer"); + + closer.start(); + try { + Assert.assertTrue("last-session close did not enter doClose", + pool.closeEntered.await(TIMEOUT_SECONDS, + TimeUnit.SECONDS)); + + acquirer.start(); + Assert.assertTrue("acquirer did not start", + acquireAttempted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + boolean acquiredWhileClosePaused = + awaitBorrowerInCloseWindow(pool, closer, acquirer); + BackendSession activeSession = acquiredWhileClosePaused ? + acquisition.get(TIMEOUT_SECONDS, + TimeUnit.SECONDS) : null; + pool.allowClose.countDown(); + + closer.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + acquirer.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + Assert.assertFalse("closer did not finish", closer.isAlive()); + Assert.assertFalse("acquirer did not finish", acquirer.isAlive()); + Assert.assertNull(closeFailure.get()); + if (acquiredWhileClosePaused) { + Assert.assertNotNull(activeSession); + Assert.assertTrue("new session should still be open", + activeSession.opened()); + Assert.assertFalse("pool should still count the new active session", + pool.closed()); + + Assert.assertTrue("doClose closed the backend while a newly " + + "acquired session remained active", + pool.opened()); + } else { + try { + acquisition.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + Assert.fail("late acquire should be rejected after backend close"); + } catch (ExecutionException e) { + Assert.assertTrue("late acquire must fail because backend is closed", + e.getCause() instanceof BackendClosedException); + Assert.assertFalse("backend should be closed before rejecting late acquire", + pool.opened()); + Assert.assertTrue("pool should count no active session after rejection", + pool.closed()); + } + } + } finally { + pool.allowClose.countDown(); + closer.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + if (acquirer.getState() != Thread.State.NEW) { + acquirer.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + } + } + } + + private static boolean awaitBorrowerInCloseWindow(TestSessionPool pool, + Thread closer, + Thread acquirer) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (System.nanoTime() < deadline) { + if (pool.acquirerEnteredNewSession.getCount() == 0L) { + return true; + } + ThreadInfo info = ManagementFactory.getThreadMXBean() + .getThreadInfo(acquirer.getId()); + if (info != null && info.getThreadState() == Thread.State.BLOCKED && + info.getLockOwnerId() == closer.getId() && + info.getLockInfo() != null && + info.getLockInfo().getClassName().equals(TestSessionPool.class.getName())) { + return false; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1L)); + } + throw new AssertionError("acquirer neither entered newSession nor blocked on close"); + } + + private static final class TestSessionPool extends BackendSessionPool { + + private final AtomicBoolean opened; + private final AtomicReference session; + private final CountDownLatch closeEntered; + private final CountDownLatch allowClose; + private final CountDownLatch acquirerEnteredNewSession; + + private TestSessionPool() { + super(FakeObjects.newConfig(), "test"); + this.opened = new AtomicBoolean(); + this.session = new AtomicReference<>(); + this.closeEntered = new CountDownLatch(1); + this.allowClose = new CountDownLatch(1); + this.acquirerEnteredNewSession = new CountDownLatch(1); + } + + @Override + public void open() { + this.opened.set(true); + } + + @Override + protected boolean opened() { + return this.opened.get(); + } + + @Override + public BackendSession session() { + return this.session.get(); + } + + @Override + protected BackendSession newSession() { + if (Thread.currentThread().getName().equals("backend-session-acquirer")) { + this.acquirerEnteredNewSession.countDown(); + } + if (!this.opened.get()) { + throw new BackendClosedException(); + } + BackendSession newSession = new TestSession(); + this.session.set(newSession); + return newSession; + } + + @Override + protected void doClose() { + this.closeEntered.countDown(); + try { + if (!this.allowClose.await(CLOSE_HOLD_SECONDS, TimeUnit.SECONDS)) { + throw new AssertionError("timed out waiting to finish doClose"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("interrupted while waiting to finish doClose", e); + } + this.opened.set(false); + } + } + + private static final class TestSession extends BackendSession.AbstractBackendSession { + + @Override + public void open() { + this.opened = true; + } + + @Override + public void close() { + this.opened = false; + } + + @Override + public Object commit() { + return null; + } + + @Override + public void rollback() { + // No transaction state in this test session. + } + + @Override + public boolean hasChanges() { + return false; + } + } + + private static final class BackendClosedException extends RuntimeException { + + private static final long serialVersionUID = 1L; + } +} diff --git a/hugegraph-store/README.md b/hugegraph-store/README.md index 58943bb94b..960162c435 100644 --- a/hugegraph-store/README.md +++ b/hugegraph-store/README.md @@ -114,6 +114,15 @@ The assembled distribution will be available at: hugegraph-store/apache-hugegraph-store-/lib/hg-store-node-.jar ``` +The Node module keeps a regular JAR for development dependencies and attaches its +Spring Boot executable as `hg-store-node--exec.jar`. The distribution +uses the executable JAR with the filename shown above, so the start script is +unchanged. To run the Store server tests through the package phase: + +```bash +mvn clean package -pl hugegraph-store/hg-store-test -am -Pstore-server-test +``` + ### Configuration Extract the distribution package and edit `conf/application.yml`: @@ -388,6 +397,20 @@ For Docker and Kubernetes deployment details, see [Deployment Guide](docs/deploy --- +## Stopping a Store node + +For an unpacked distribution, run `bin/stop-hugegraph-store.sh`. It waits up to +30 seconds for process exit. A timeout returns a nonzero status and retains +`bin/pid` for diagnosis; it does not force-kill the process or remove its data. +Check the Store log and thread dump before taking further action. + +Spring owns shutdown: new RPCs are refused, active RPCs are cancelled, and context +close waits for their callbacks and scan/TTL workers to release resources. The +Store engine then stops and joins partition Raft services before releasing +databases. A stuck callback keeps shutdown pending rather than allowing its +database to close underneath it. Do not add a separate JVM hook that closes +those databases concurrently. + ## Documentation Comprehensive documentation for HugeGraph Store: diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java index 7abcbf88d0..da44468c7d 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserver.java @@ -65,6 +65,8 @@ public class CommonKvStreamObserver implements StreamObserver { * It can be ended by the client to stop receiving redundant data. */ private final AtomicBoolean closed = new AtomicBoolean(false); + // Outgoing RPC calls must not hold the monitor used by incoming callbacks. + private final Object requestLock = new Object(); @Setter private long timeout = 1800 * 1000; @@ -85,19 +87,24 @@ public CommonKvStreamObserver(Function> valueExtractor, * Send requests */ public void sendRequest() { - if (!isServerFinished() && !closed.get()) { + synchronized (this.requestLock) { + synchronized (this) { + if (isServerFinished() || this.closed.get()) { + return; + } + this.watcher.setState(ResultState.WAITING); + } this.requestSender.accept(true); - this.watcher.setState(ResultState.WAITING); } } - public boolean isServerFinished() { + public synchronized boolean isServerFinished() { return this.watcher.getState() == ResultState.FINISHED || this.watcher.getState() == ResultState.ERROR; } @Override - public void onNext(R value) { + public synchronized void onNext(R value) { watcher.setState(ResultState.INNER_BUSY); try { var state = stateWatcher.apply(value); @@ -125,14 +132,21 @@ public void onNext(R value) { public Iterator consume() { try { - while (!Thread.currentThread().isInterrupted() && (!this.queue.isEmpty() || - !isServerFinished())) { + while (!Thread.currentThread().isInterrupted()) { + // Read terminal state before the queue, including a concurrently published final batch. + if (isServerFinished() && this.queue.isEmpty()) { + return null; + } var iterator = this.queue.poll(200, TimeUnit.MILLISECONDS); if (iterator != null) { sendRequest(); return iterator; } + // Read terminal state before the queue: a final batch is published with that state. + if (isServerFinished() && this.queue.isEmpty()) { + return null; + } if ((System.nanoTime() - watcher.current) / 1000_000 > this.timeout) { throw new HgStoreClientException("iterator timeout"); } @@ -146,25 +160,29 @@ public Iterator consume() { } /** - * Send onComplete, stop receiving data + * Stop feedback and invoke the request stream's completion or cancellation callback. */ public void clear() { - if (!this.closed.get()) { - this.closed.set(true); - this.transferComplete.accept(true); + synchronized (this.requestLock) { + try { + if (this.closed.compareAndSet(false, true)) { + this.transferComplete.accept(true); + } + } finally { + this.queue.clear(); + } } - this.queue.clear(); } @Override - public void onError(Throwable t) { + public synchronized void onError(Throwable t) { log.error("StreamObserver got error:", t); this.queue.offer(new ErrorMessageIterator<>(t.getMessage())); this.watcher.setState(ResultState.ERROR); } @Override - public void onCompleted() { + public synchronized void onCompleted() { if (watcher.getState() != ResultState.ERROR) { watcher.setState(ResultState.FINISHED); } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java index 0a4373c7f3..22be95b1ce 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIterator.java @@ -17,7 +17,12 @@ package org.apache.hugegraph.store.client.query; +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.Iterator; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.List; import java.util.NoSuchElementException; @@ -27,40 +32,74 @@ public class MultiStreamIterator implements HgKvIterator { private HgKvIterator currentIterator = null; + private final List> iterators; private final Iterator> listIterator; + private final AtomicBoolean closed = new AtomicBoolean(); public MultiStreamIterator(List> iterators) { - this.listIterator = iterators.iterator(); + this.iterators = new ArrayList<>(iterators); + this.listIterator = this.iterators.iterator(); } @Override public byte[] key() { + checkOpen(); return currentIterator.key(); } @Override public byte[] value() { + checkOpen(); return currentIterator.value(); } @Override public void close() { - //Todo is syntax correct? - if (currentIterator != null && currentIterator.hasNext()) { - currentIterator.close(); + if (!this.closed.compareAndSet(false, true)) { + return; + } + Throwable failure = null; + Set> visited = Collections.newSetFromMap(new IdentityHashMap<>()); + for (HgKvIterator iterator : this.iterators) { + if (!visited.add(iterator)) { + continue; + } + try { + iterator.close(); + } catch (RuntimeException | Error error) { + if (failure == null) { + failure = error; + } else if (failure != error) { + failure.addSuppressed(error); + } + } + } + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw (RuntimeException) failure; } } @Override public byte[] position() { + checkOpen(); return currentIterator.position(); } @Override public void seek(byte[] position) { + checkOpen(); this.currentIterator.seek(position); } + private void checkOpen() { + if (this.closed.get()) { + throw new IllegalStateException("Iterator is closed"); + } + } + private void getNextIterator() { if (currentIterator != null && currentIterator.hasNext()) { return; @@ -76,13 +115,16 @@ private void getNextIterator() { @Override public boolean hasNext() { + if (this.closed.get()) { + return false; + } getNextIterator(); return currentIterator != null && currentIterator.hasNext(); } @Override public E next() { - if (currentIterator == null || !currentIterator.hasNext()) { + if (this.closed.get() || currentIterator == null || !currentIterator.hasNext()) { throw new NoSuchElementException(); } return currentIterator.next(); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java index 9f6fd0dc8b..82fee7f905 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/query/QueryExecutor.java @@ -61,6 +61,7 @@ import com.google.protobuf.ByteString; +import io.grpc.Status; import lombok.extern.slf4j.Slf4j; @Slf4j @@ -240,7 +241,14 @@ public BaseElement next() { var reqStream = stub.query(observer); observer.setWatcherQueryId(request.getQueryId() + '-' + address); observer.setRequestSender(r -> reqStream.onNext(request)); - observer.setTransferComplete(r -> reqStream.onCompleted()); + observer.setTransferComplete(r -> { + if (observer.isServerFinished()) { + reqStream.onCompleted(); + } else { + reqStream.onError(Status.CANCELLED.withDescription("Query iterator closed") + .asRuntimeException()); + } + }); observer.setTimeout(this.timeout); var itr = new StreamKvIterator<>(b -> observer.clear(), observer::consume); diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java index 368032f2ce..75572ec609 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/InnerKeyFilter.java @@ -31,13 +31,13 @@ public class InnerKeyFilter implements ScanIterator { final boolean codeFilter; ScanIterator iterator; T current = null; + private boolean initialized; public InnerKeyFilter(ScanIterator iterator) { this.iterator = iterator; this.codeFrom = Integer.MIN_VALUE; this.codeTo = Integer.MAX_VALUE; this.codeFilter = false; - moveNext(); } public InnerKeyFilter(ScanIterator iterator, boolean codeFilter) { @@ -45,7 +45,6 @@ public InnerKeyFilter(ScanIterator iterator, boolean codeFilter) { this.codeFrom = Integer.MIN_VALUE; this.codeTo = Integer.MAX_VALUE; this.codeFilter = codeFilter; - moveNext(); } public InnerKeyFilter(ScanIterator iterator, int codeFrom, int codeTo) { @@ -53,7 +52,14 @@ public InnerKeyFilter(ScanIterator iterator, int codeFrom, int codeTo) { this.codeFrom = codeFrom; this.codeTo = codeTo; this.codeFilter = true; - moveNext(); + } + + private void init() { + if (!this.initialized) { + // Publish the wrapper to its owner before prefetch can close or fail on the raw iterator. + this.initialized = true; + moveNext(); + } } private void moveNext() { @@ -76,16 +82,19 @@ private void moveNext() { @Override public boolean hasNext() { + init(); return current != null; } @Override public boolean isValid() { + init(); return iterator.isValid(); } @Override public T next() { + init(); T column = current; if (!codeFilter) // Remove the image ID and hash suffix @@ -103,11 +112,13 @@ public T next() { @Override public void close() { + this.initialized = true; iterator.close(); } @Override public long count() { + init(); return iterator.count(); } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java index 72cc472b21..030e724072 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/MultiPartitionIterator.java @@ -18,13 +18,13 @@ package org.apache.hugegraph.store.business; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.Collections; import java.util.LinkedList; import java.util.List; import java.util.NoSuchElementException; import java.util.Queue; import java.util.function.BiFunction; -import java.util.stream.Collectors; import org.apache.hugegraph.rocksdb.access.ScanIterator; @@ -45,6 +45,7 @@ public class MultiPartitionIterator implements ScanIterator { private Integer curPartitionId; private Integer positionPartitionId; private byte[] positionKey; + private RuntimeException cleanupFailure; private MultiPartitionIterator(List partitionIds, BiFunction supplier) { @@ -75,31 +76,27 @@ public static int toInt(byte[] bytes) { } private ScanIterator getIterator() { - if (this.partitions.isEmpty()) { - return null; - } - ScanIterator buf = null; - while (!partitions.isEmpty()) { - this.curPartitionId = partitions.poll(); + while (!this.partitions.isEmpty()) { + this.curPartitionId = this.partitions.poll(); if (!this.inPosition(this.curPartitionId)) { continue; } - buf = supplier.apply(this.curPartitionId, getPositionKey(this.curPartitionId)); - if (buf == null) { + ScanIterator child = this.supplier.apply(this.curPartitionId, + getPositionKey(this.curPartitionId)); + if (child == null) { continue; } - if (buf.hasNext()) { - break; + try { + if (child.hasNext()) { + return child; + } + } catch (RuntimeException | Error failure) { + closeAfterFailure(child, failure); + throw failure; } + closeCreatedIterator(child); } - if (buf == null) { - return null; - } - if (!buf.hasNext()) { - buf.close(); - buf = null; - } - return buf; + return null; } private void init() { @@ -128,8 +125,7 @@ public T next() { } T t = this.iterator.next(); if (!this.iterator.hasNext()) { - this.iterator.close(); - this.iterator = null; + closeCurrentIterator(); } return t; } @@ -137,11 +133,16 @@ public T next() { @Override public long count() { long count = 0; - this.iterator = this.getIterator(); - while (this.iterator != null) { - count += this.iterator.count(); - // this.iterator.close(); - this.iterator = this.getIterator(); + while (this.hasNext()) { + try { + count += this.iterator.count(); + } catch (RuntimeException | Error failure) { + ScanIterator child = this.iterator; + this.iterator = null; + closeAfterFailure(child, failure); + throw failure; + } + closeCurrentIterator(); } return count; } @@ -172,8 +173,32 @@ public void seek(byte[] position) { @Override public void close() { - if (this.iterator != null) { - this.iterator.close(); + try { + closeCurrentIterator(); + } catch (RuntimeException | Error failure) { + // closeCreatedIterator retains the failure even after ownership is released here. + throw this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw this.cleanupFailure; + } + } + + private void closeCurrentIterator() { + ScanIterator child = this.iterator; + this.iterator = null; + if (child != null) { + closeCreatedIterator(child); + } + } + + private void closeAfterFailure(ScanIterator child, Throwable failure) { + try { + closeCreatedIterator(child); + } catch (RuntimeException | Error closeFailure) { + if (closeFailure != failure) { + failure.addSuppressed(closeFailure); + } } } @@ -205,10 +230,47 @@ private byte[] getPositionKey(int partitionId) { * @return iteration list */ public List getIterators() { - return this.partitions.stream() - .map(id -> supplier.apply(id, getPositionKey(id))) - .filter(ScanIterator::hasNext) - .collect(Collectors.toList()); + List opened = new ArrayList<>(); + try { + for (int id : this.partitions) { + ScanIterator child = this.supplier.apply(id, getPositionKey(id)); + if (child == null) { + continue; + } + opened.add(child); + if (!child.hasNext()) { + opened.remove(opened.size() - 1); + closeCreatedIterator(child); + } + } + // Ownership transfers to the caller only when all partitions were opened. + return opened; + } catch (RuntimeException | Error failure) { + for (ScanIterator child : opened) { + try { + closeCreatedIterator(child); + } catch (RuntimeException | Error closeFailure) { + if (closeFailure != failure) { + failure.addSuppressed(closeFailure); + } + } + } + throw failure; + } + } + + private void closeCreatedIterator(ScanIterator child) { + try { + child.close(); + } catch (RuntimeException | Error failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = new IllegalStateException( + "partition iterator cleanup failed", failure); + } else { + this.cleanupFailure.addSuppressed(failure); + } + throw failure; + } } } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java index f7f013b89f..cba162a58b 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java @@ -41,6 +41,7 @@ public final class OrderedMultiPartitionIterator implements ScanIterator { private boolean initialized; private boolean closed; + private Throwable cleanupFailure; private Integer currentPartitionId; private OrderedMultiPartitionIterator(List partitionIds, @@ -142,22 +143,29 @@ public void seek(byte[] position) { @Override public void close() { if (this.closed) { + rethrowCleanupFailure(); return; } this.closed = true; - Throwable failure = null; + Throwable failure = this.cleanupFailure; for (SourceEntry source : this.sources) { try { this.closeSource(source); } catch (RuntimeException | Error e) { if (failure == null) { failure = e; - } else { + } else if (failure != e) { failure.addSuppressed(e); } } } this.queue.clear(); + this.cleanupFailure = failure; + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + Throwable failure = this.cleanupFailure; if (failure instanceof RuntimeException) { throw (RuntimeException) failure; } @@ -197,14 +205,23 @@ private void closeSource(SourceEntry source) { return; } source.closed = true; - source.iterator.close(); + try { + source.iterator.close(); + } catch (RuntimeException | Error failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + throw failure; + } } private void closeAfterFailure(Throwable failure) { try { this.close(); } catch (RuntimeException | Error closeFailure) { - failure.addSuppressed(closeFailure); + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } } } diff --git a/hugegraph-store/hg-store-dist/pom.xml b/hugegraph-store/hg-store-dist/pom.xml index 12acb0ba9e..c75fc2ad6d 100644 --- a/hugegraph-store/hg-store-dist/pom.xml +++ b/hugegraph-store/hg-store-dist/pom.xml @@ -42,6 +42,8 @@ org.apache.hugegraph hg-store-node ${revision} + exec + runtime diff --git a/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml b/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml index 2e761200ef..47170a6c22 100644 --- a/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml +++ b/hugegraph-store/hg-store-dist/src/assembly/descriptor/server-assembly.xml @@ -68,11 +68,12 @@ /lib + ${artifact.artifactId}-${artifact.version}.jar false runtime false - org.apache.hugegraph:${executable.jar.name}:jar:* + org.apache.hugegraph:${executable.jar.name}:jar:exec:* diff --git a/hugegraph-store/hg-store-dist/src/assembly/static/bin/stop-hugegraph-store.sh b/hugegraph-store/hg-store-dist/src/assembly/static/bin/stop-hugegraph-store.sh index 8f898df7da..e9f503b02a 100644 --- a/hugegraph-store/hg-store-dist/src/assembly/static/bin/stop-hugegraph-store.sh +++ b/hugegraph-store/hg-store-dist/src/assembly/static/bin/stop-hugegraph-store.sh @@ -41,8 +41,8 @@ if [ ! -f ${PID_FILE} ]; then fi PID=`cat $PID_FILE` -kill_process_and_wait "HugeGraphStoreServer" "$PID" "$SERVER_SHUTDOWN_TIMEOUT_S" - -if [ $? -eq 0 ]; then +if kill_process_and_wait "HugeGraphStoreServer" "$PID" "$SERVER_SHUTDOWN_TIMEOUT_S"; then rm "$PID_FILE" +else + exit 1 fi diff --git a/hugegraph-store/hg-store-node/pom.xml b/hugegraph-store/hg-store-node/pom.xml index 56cff46dba..9c952507af 100644 --- a/hugegraph-store/hg-store-node/pom.xml +++ b/hugegraph-store/hg-store-node/pom.xml @@ -169,6 +169,7 @@ repackage + exec org.apache.hugegraph.store.node.StoreNodeApplication diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java index f7e02481fe..4d26c7d839 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/BatchScanIterator.java @@ -41,6 +41,7 @@ public final class BatchScanIterator implements ScanIterator { private final AtomicBoolean closed = new AtomicBoolean(); private ScanIterator iterator; private boolean hasNext = false; + private Throwable cleanupFailure; private long curCount; private long curLimit; @@ -65,13 +66,14 @@ private ScanIterator getIterator() { do { buf = this.batchSupplier.get().getValue(); + this.iterator = buf; if (buf == null) { break; } if (!buf.hasNext()) { - buf.close(); + closeIterator(); buf = null; } @@ -99,13 +101,17 @@ private ScanIterator getIterator() { @Override public boolean hasNext() { + rethrowCleanupFailure(); + if (this.closed.get()) { + return false; + } if (this.iterator == null) { this.iterator = this.getIterator(); } else if (!this.iterator.hasNext()) { - this.iterator.close(); + closeIterator(); this.iterator = this.getIterator(); } else if (this.curCount == this.curLimit) { - this.iterator.close(); + closeIterator(); this.iterator = this.getIterator(); } @@ -135,10 +141,31 @@ public T next() { @Override public void close() { if (!this.closed.getAndSet(true)) { - if (this.iterator != null) { - this.iterator.close(); + closeIterator(); + } + rethrowCleanupFailure(); + } + + private void closeIterator() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + this.cleanupFailure = failure; } } + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + if (this.cleanupFailure instanceof RuntimeException) { + throw (RuntimeException) this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw (Error) this.cleanupFailure; + } } @Override diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java index 56a1cc8c58..afdcb2e59c 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/FusingScanIterator.java @@ -33,6 +33,8 @@ final class FusingScanIterator implements ScanIterator { private long accumulator; private Supplier supplier; private ScanIterator iterator; + private Throwable cleanupFailure; + private boolean closed; private byte[] position = EMPTY_BYTES; private FusingScanIterator() { @@ -47,14 +49,13 @@ public static FusingScanIterator maxOf(long maxThreshold, } private ScanIterator getIterator() { - ScanIterator buf = this.supplier.get(); - if (buf == null) { - return null; + while ((this.iterator = this.supplier.get()) != null) { + if (this.iterator.hasNext()) { + return this.iterator; + } + closeCurrent(); } - if (!buf.hasNext()) { - buf = null; - } - return buf; + return null; } private void init() { @@ -65,7 +66,8 @@ private void init() { @Override public boolean hasNext() { - if (this.isThresholdExceeded()) { + rethrowCleanupFailure(); + if (this.closed || this.isThresholdExceeded()) { return false; } if (this.iterator == null) { @@ -104,16 +106,36 @@ public T next() { position = this.iterator.position(); this.accumulator++; if (!this.iterator.hasNext() || this.isThresholdExceeded()) { - this.iterator.close(); - this.iterator = null; + closeCurrent(); } return t; } @Override public void close() { - if (this.iterator != null) { - this.iterator.close(); + this.closed = true; + closeCurrent(); + } + + private void closeCurrent() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + this.cleanupFailure = failure; + } + } + rethrowCleanupFailure(); + } + + private void rethrowCleanupFailure() { + if (this.cleanupFailure instanceof RuntimeException) { + throw (RuntimeException) this.cleanupFailure; + } + if (this.cleanupFailure != null) { + throw (Error) this.cleanupFailure; } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java new file mode 100644 index 0000000000..a9a6a9f95a --- /dev/null +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/GrpcShutdownBarrier.java @@ -0,0 +1,110 @@ +/* + * 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; + +import java.util.HashSet; +import java.util.Set; + +import org.lognet.springboot.grpc.GRpcGlobalInterceptor; +import org.springframework.stereotype.Component; + +import io.grpc.ForwardingServerCallListener; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import lombok.extern.slf4j.Slf4j; + +/** Waits for terminal application callbacks, not just transport termination. */ +@Slf4j +@Component +@GRpcGlobalInterceptor +public class GrpcShutdownBarrier implements ServerInterceptor { + + private final Set calls = new HashSet<>(); + private boolean closing; + + public synchronized void stopAcceptingCalls() { + this.closing = true; + } + + @Override + public ServerCall.Listener interceptCall( + ServerCall call, Metadata headers, + ServerCallHandler next) { + Object token = new Object(); + synchronized (this) { + if (this.closing) { + call.close(Status.UNAVAILABLE.withDescription("Store is stopping"), new Metadata()); + return new ServerCall.Listener() { }; + } + this.calls.add(token); + } + try { + return new ForwardingServerCallListener.SimpleForwardingServerCallListener( + next.startCall(call, headers)) { + @Override + public void onCancel() { + try { + super.onCancel(); + } finally { + finished(token); + } + } + + @Override + public void onComplete() { + try { + super.onComplete(); + } finally { + finished(token); + } + } + }; + } catch (RuntimeException | Error e) { + finished(token); + throw e; + } + } + + private synchronized void finished(Object token) { + this.calls.remove(token); + this.notifyAll(); + } + + public synchronized void awaitCallbacks() { + boolean interrupted = false; + try { + while (!this.calls.isEmpty()) { + try { + this.wait(5000); + if (!this.calls.isEmpty()) { + log.warn("Waiting for {} RPC callbacks before closing databases", this.calls.size()); + } + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } +} diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java index 7d01fa3db4..a1ff494056 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreStreamImpl.java @@ -17,9 +17,18 @@ package org.apache.hugegraph.store.node.grpc; +import java.util.Map; +import java.util.function.Function; +import java.util.function.Consumer; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import org.apache.hugegraph.store.business.BusinessHandler; +import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; +import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; +import org.apache.hugegraph.store.node.grpc.scan.ScanResponseObserver; import org.apache.hugegraph.store.grpc.state.ScanState; import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; import org.apache.hugegraph.store.grpc.stream.KvPageRes; @@ -31,6 +40,8 @@ import org.lognet.springboot.grpc.GRpcService; import org.springframework.beans.factory.annotation.Autowired; +import io.grpc.Context; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -47,6 +58,167 @@ public class HgStoreStreamImpl extends HgStoreStreamGrpc.HgStoreStreamImplBase { private AppConfig appConfig; private HgStoreWrapperEx wrapper; private ThreadPoolExecutor executor; + private boolean closing; + private final Map scans = new ConcurrentHashMap<>(); + + /** Close admission before cancelling the gRPC server. */ + public synchronized void stopAcceptingScans() { + this.closing = true; + } + + /** Cancel streams, then drain queued cleanup tasks. Never discard the queue. */ + public void shutdownScans() { + stopAcceptingScans(); + for (Runnable cancel : this.scans.values().toArray(new Runnable[0])) { + try { + cancel.run(); + } catch (RuntimeException | Error failure) { + log.warn("Failed to cancel scan; continuing other cancellations", failure); + } + } + ThreadPoolExecutor current = getRealExecutor(); + if (current != null) { + current.shutdown(); + } + } + + private synchronized void checkAcceptingScans() { + if (this.closing) { + throw Status.UNAVAILABLE.withDescription("Store scans are stopping") + .asRuntimeException(); + } + } + + /** Wait independently of executor termination: failed native release is sticky. */ + public void awaitScanCleanup() { + boolean interrupted = false; + long nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + synchronized (this.scans) { + while (!this.scans.isEmpty()) { + try { + this.scans.wait(5000L); + } catch (InterruptedException e) { + interrupted = true; + } + if (!this.scans.isEmpty() && System.nanoTime() - nextLog >= 0) { + log.warn("Still waiting for {} scans to clean up before closing databases", + this.scans.size()); + for (ScanLifecycle scan : this.scans.keySet()) { + Throwable failure = scan.cleanupFailure(); + if (failure != null) { + log.warn("Scan cleanup failed; database close stays blocked", failure); + } + } + nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private StreamObserver register( + ScanLifecycle lifecycle, Function> factory) { + Context context = Context.current(); + StreamObserver observer = factory.apply(lifecycle); + Runnable cancel = () -> { + if (lifecycle.tryEnterCancellation()) { + try { + observer.onError(Status.CANCELLED.asRuntimeException()); + } finally { + lifecycle.leave(); + } + } + }; + this.scans.put(lifecycle, cancel); + Context.CancellationListener listener = ignored -> cancel.run(); + lifecycle.onFinished(() -> { + this.scans.remove(lifecycle); + context.removeListener(listener); + synchronized (this.scans) { + this.scans.notifyAll(); + } + }); + context.addListener(listener, Runnable::run); + return new StreamObserver() { + @Override + public void onNext(T value) { + if (lifecycle.tryEnter()) { + try { + observer.onNext(value); + } finally { + lifecycle.leave(); + } + } + } + + @Override + public void onError(Throwable failure) { + if (lifecycle.tryEnterCancellation()) { + try { + observer.onError(failure); + } finally { + lifecycle.leave(); + } + } + } + + @Override + public void onCompleted() { + if (lifecycle.tryEnter()) { + try { + observer.onCompleted(); + } finally { + lifecycle.leave(); + } + } + } + }; + } + + public StreamObserver scanGraphPartition( + StreamObserver response, + BusinessHandler handler) { + synchronized (this) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + new ScanResponseObserver<>( + lifecycle.response(response), handler, + task -> lifecycle.execute(getExecutor(), task), lifecycle::failedCleanup)); + } + } + + private void oneShot(StreamObserver response, Consumer action) { + ScanLifecycle lifecycle = new ScanLifecycle(); + synchronized (this) { + checkAcceptingScans(); + lifecycle.enter(); + this.scans.put(lifecycle, lifecycle::tryCancel); + lifecycle.onFinished(() -> { + this.scans.remove(lifecycle); + synchronized (this.scans) { + this.scans.notifyAll(); + } + }); + } + try { + action.accept(lifecycle); + } finally { + try { + // An iterator may close automatically during hasNext(), after which the + // scan exits without entering its ordinary response-completion path. + if (lifecycle.isCancelled()) { + lifecycle.response(response).onError(Status.CANCELLED + .withDescription("Store scan cancelled before completion") + .asRuntimeException()); + } + } finally { + lifecycle.finishWithoutResponse(); + lifecycle.leave(); + } + } + } private HgStoreWrapperEx getWrapper() { if (this.wrapper == null) { @@ -60,20 +232,16 @@ private HgStoreWrapperEx getWrapper() { return this.wrapper; } - public ThreadPoolExecutor getRealExecutor() { - return executor; + public synchronized ThreadPoolExecutor getRealExecutor() { + return this.executor; } - public ThreadPoolExecutor getExecutor() { + public synchronized ThreadPoolExecutor getExecutor() { + checkAcceptingScans(); if (this.executor == null) { - synchronized (this) { - if (this.executor == null) { - AppConfig.ThreadPoolScan scan = this.appConfig.getThreadPoolScan(); - this.executor = - HgExecutorUtil.createExecutor("hg-scan", scan.getCore(), scan.getMax(), - scan.getQueue()); - } - } + AppConfig.ThreadPoolScan scan = this.appConfig.getThreadPoolScan(); + this.executor = HgExecutorUtil.createExecutor("hg-scan", scan.getCore(), + scan.getMax(), scan.getQueue()); } return this.executor; } @@ -95,27 +263,34 @@ public ScanState getState() { } @Override - public StreamObserver scan(StreamObserver response) { - return ScanStreamResponse.of(response, getWrapper(), getExecutor(), appConfig); + public synchronized StreamObserver scan(StreamObserver response) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + new ScanStreamResponse(response, getWrapper(), getExecutor(), appConfig, lifecycle)); } @Override public void scanOneShot(ScanStreamReq request, StreamObserver response) { - ScanOneShotResponse.scanOneShot(request, response, getWrapper()); + oneShot(response, lifecycle -> ScanOneShotResponse.scanOneShot(request, response, getWrapper(), lifecycle)); } @Override - public StreamObserver scanBatch(StreamObserver response) { - return ScanBatchResponse3.of(response, getWrapper(), getExecutor()); + public synchronized StreamObserver scanBatch(StreamObserver response) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + ScanBatchResponse3.of(response, getWrapper(), getExecutor(), lifecycle)); } @Override - public StreamObserver scanBatch2(StreamObserver response) { - return ScanBatchResponseFactory.of(response, getWrapper(), getExecutor()); + public synchronized StreamObserver scanBatch2(StreamObserver response) { + checkAcceptingScans(); + return register(new ScanLifecycle(), lifecycle -> + ScanBatchResponseFactory.of(response, getWrapper(), getExecutor(), lifecycle)); } @Override public void scanBatchOneShot(ScanStreamBatchReq request, StreamObserver response) { - ScanBatchOneShotResponse.scanOneShot(request, response, getWrapper()); + oneShot(response, lifecycle -> + ScanBatchOneShotResponse.scanOneShot(request, response, getWrapper(), lifecycle)); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java index 56ce1f45f3..02a05292a6 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ParallelScanIterator.java @@ -27,6 +27,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.ReentrantLock; import java.util.function.Supplier; +import java.util.function.Consumer; import org.apache.hugegraph.pd.common.KVPair; import org.apache.hugegraph.rocksdb.access.RocksDBSession; @@ -67,11 +68,33 @@ public class ParallelScanIterator implements ScanIterator { private int maxInQueue = maxWorkThreads * 2; private volatile boolean finished; private List current = null; + private final Consumer cleanupReporter; + private volatile Throwable scanFailure; + private volatile Throwable cleanupFailure; + + private static void rethrow(Throwable failure) { + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure != null) { + throw (Error) failure; + } + } + + private synchronized void failedCleanup(Throwable failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + this.cleanupReporter.accept(failure); + this.scanFailure = this.cleanupFailure; + requestStop(); + } private ParallelScanIterator(Supplier> iteratorSupplier, Supplier limitSupplier, ScanQueryRequest query, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, Consumer cleanupReporter) { + this.cleanupReporter = cleanupReporter; this.executor = executor; this.batchSupplier = iteratorSupplier; this.limitSupplier = limitSupplier; @@ -88,7 +111,12 @@ private ParallelScanIterator(Supplier> iter this.maxInQueue = maxWorkThreads * 2; // Edge sorted requires a larger queue queue = new LinkedBlockingQueue<>(maxInQueue * 2); - createScanner(); + try { + createScanner(); + } catch (RuntimeException e) { + close(); + throw e; + } } public static ParallelScanIterator of( @@ -98,7 +126,13 @@ public static ParallelScanIterator of( ThreadPoolExecutor executor) { HgAssert.isArgumentNotNull(iteratorSupplier, "iteratorSupplier"); HgAssert.isArgumentNotNull(limitSupplier, "limitSupplier"); - return new ParallelScanIterator(iteratorSupplier, limitSupplier, query, executor); + return of(iteratorSupplier, limitSupplier, query, executor, ignored -> { }); + } + + static ParallelScanIterator of(Supplier> iteratorSupplier, + Supplier limitSupplier, ScanQueryRequest query, + ThreadPoolExecutor executor, Consumer cleanupReporter) { + return new ParallelScanIterator(iteratorSupplier, limitSupplier, query, executor, cleanupReporter); } @Override @@ -126,6 +160,7 @@ public boolean hasNext() { log.error("Wait data timeout!!!, scanner is {}/{}", scanners.size(), pauseScanners.size()); } + rethrow(this.scanFailure); return current != null && current != NO_DATA; } @@ -144,20 +179,26 @@ public List next() { return t; } + void requestStop() { + this.finished = true; + } + @Override public void close() { - finished = true; + requestStop(); + List pending; synchronized (scanners) { - scanners.forEach(scanner -> { - scanner.close(); - }); + pending = new ArrayList<>(scanners); + } + // Never hold a registry lock while waiting for a scanner's iterator lock. + for (KVScanner scanner : pending) { + scanner.close(); } synchronized (pauseScanners) { - pauseScanners.forEach(s -> { - s.close(); - }); + pauseScanners.clear(); } queue.clear(); + rethrow(this.cleanupFailure); } /** @@ -179,10 +220,15 @@ private void createScanner() { */ private void wakeUpScanner() { synchronized (pauseScanners) { - if (!pauseScanners.isEmpty()) { + if (!finished && !pauseScanners.isEmpty()) { KVScanner scanner = pauseScanners.poll(); if (scanner != null) { - executor.execute(() -> scanner.scanKV()); + try { + executor.execute(() -> scanner.scanKV()); + } catch (java.util.concurrent.RejectedExecutionException e) { + scanner.close(); + throw e; + } } } } @@ -195,7 +241,11 @@ private void wakeUpScanner() { */ private void suspendScanner(KVScanner scanner) { synchronized (pauseScanners) { - pauseScanners.add(scanner); + if (!finished) { + pauseScanners.add(scanner); + } else { + scanner.close(); + } } } @@ -218,30 +268,43 @@ private void quitScanner(KVScanner scanner) { */ private boolean putData(List data) { try { - this.queue.put(data); + while (!finished && !this.queue.offer(data, 100, TimeUnit.MILLISECONDS)) { + // Cancellation must release a producer whose client stopped consuming. + } } catch (InterruptedException e) { log.error("exception ", e); this.finished = true; return false; } - return this.queue.size() < maxInQueue; + return !finished && this.queue.size() < maxInQueue; } private boolean putData(List data, boolean hasNext) { + boolean locked = false; try { - queueLock.lock(); - this.queue.put(data); + while (!finished && !(locked = queueLock.tryLock(100, TimeUnit.MILLISECONDS))) { + // An ordered producer may be waiting behind a cancelled scanner. + } + if (!locked) { + return false; + } + while (!finished && !this.queue.offer(data, 100, TimeUnit.MILLISECONDS)) { + // Recheck cancellation while the output queue is full. + } } catch (InterruptedException e) { - log.error("exception ", e); + Thread.currentThread().interrupt(); this.finished = true; return false; } finally { - if (!hasNext) { + if (locked && finished) { + while (queueLock.isHeldByCurrentThread()) { + queueLock.unlock(); + } + } else if (locked && !hasNext) { queueLock.unlock(); } } - // Data not ended, thread continues to execute - return hasNext || this.queue.size() < maxInQueue; + return !finished && (hasNext || this.queue.size() < maxInQueue); } private synchronized KVPair getIterator() { @@ -309,9 +372,8 @@ class KVScanner { private ScanIterator getIterator() { // Iterator has no data, or the point has reached the limit, switch to a new iterator. if (iterator == null || !iterator.hasNext() || counter >= limit) { - if (iterator != null) { - iterator.close(); - } + closeIterator(); + rethrow(cleanupFailure); KVPair pair = ParallelScanIterator.this.getIterator(); query = pair.getKey(); iterator = pair.getValue(); @@ -328,12 +390,12 @@ public void scanKV() { iteratorLock.lock(); try { long entriesSize = 0, bodySize = 0; - while (canNext && !closed) { + while (canNext && !closed && !finished) { iterator = this.getIterator(); if (iterator == null) { break; } - while (iterator.hasNext() && entriesSize < batchSize && + while (!closed && !finished && iterator.hasNext() && entriesSize < batchSize && bodySize < maxBodySize && counter < limit && !closed) { KV kv = KV.of(iterator.next()); @@ -363,11 +425,18 @@ public void scanKV() { putData(dataList); } } - } catch (Exception e) { - log.error("exception {}", e); + } catch (RuntimeException | Error e) { + scanFailure = e; + requestStop(); } finally { + // putData(..., true) intentionally keeps this lock across batches to + // serialize ordered output. The scanner must always release it before + // suspending or quitting, including when the next iterator is empty. + while (queueLock.isHeldByCurrentThread()) { + queueLock.unlock(); + } iteratorLock.unlock(); - if (iterator != null && counter < limit && !closed) { + if (iterator != null && counter < limit && !closed && !finished) { suspendScanner(this); } else { quitScanner(this); @@ -379,12 +448,22 @@ public void close() { closed = true; iteratorLock.lock(); try { - if (iterator != null) { - iterator.close(); - } + closeIterator(); } finally { iteratorLock.unlock(); } } + + private void closeIterator() { + ScanIterator current = this.iterator; + this.iterator = null; + if (current != null) { + try { + current.close(); + } catch (RuntimeException | Error failure) { + failedCleanup(failure); + } + } + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java index 1234a5b3de..bedb7eaf86 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchOneShotResponse.java @@ -30,6 +30,7 @@ import com.google.protobuf.ByteString; +import io.grpc.Context; import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -51,10 +52,19 @@ public class ScanBatchOneShotResponse { public static void scanOneShot(ScanStreamBatchReq request, StreamObserver responseObserver, HgStoreWrapperEx wrapper) { + scanOneShot(request, responseObserver, wrapper, new ScanLifecycle()); + } + + static void scanOneShot(ScanStreamBatchReq request, StreamObserver output, + HgStoreWrapperEx wrapper, ScanLifecycle lifecycle) { + StreamObserver responseObserver = lifecycle.response(output); String graph = request.getHeader().getGraph(); ScanQueryRequest queryRequest = request.getQueryRequest(); - ScanIterator iterator = getIterator(graph, queryRequest, wrapper); + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + ScanIterator iterator = getIterator(graph, queryRequest, wrapper, lifecycle::failedCleanup); KvPageRes.Builder resBuilder = KvPageRes.newBuilder(); Kv.Builder kvBuilder = Kv.newBuilder(); @@ -68,9 +78,11 @@ public static void scanOneShot(ScanStreamBatchReq request, } int count = 0; + boolean cleanupAttempted = false; try { - while (iterator.hasNext()) { + while (!(Context.current().isCancelled() || lifecycle.isCancelled()) && + !Thread.currentThread().isInterrupted() && iterator.hasNext()) { if (++count > limit) { break; @@ -86,6 +98,19 @@ public static void scanOneShot(ScanStreamBatchReq request, } + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + if (Thread.currentThread().isInterrupted()) { + responseObserver.onError(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } + cleanupAttempted = true; + if (!lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + return; + } responseObserver.onNext(resBuilder.build()); responseObserver.onCompleted(); @@ -94,7 +119,10 @@ public static void scanOneShot(ScanStreamBatchReq request, responseObserver.onError( HgGrpc.toErr(Status.Code.INTERNAL, msg, t)); } finally { - iterator.close(); + if (!cleanupAttempted && !lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java index f4485e1527..7d014568f2 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse.java @@ -21,6 +21,7 @@ import java.util.List; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; @@ -33,6 +34,7 @@ import org.apache.hugegraph.store.node.util.HgGrpc; import org.apache.hugegraph.store.node.util.PropertyUtil; +import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -59,7 +61,7 @@ public class ScanBatchResponse implements StreamObserver { private final Object stateLock = new Object(); private final Lock iteratorLock = new ReentrantLock(); // Currently traversing iterator - private ScanIterator iterator; + private volatile ScanIterator iterator; // Next send sequence number private volatile int seqNo; // Client consumed sequence number @@ -72,10 +74,19 @@ public class ScanBatchResponse implements StreamObserver { // Last read data time private long activeTime; private volatile State state; + private final AtomicBoolean cancelled = new AtomicBoolean(); + private final ScanLifecycle lifecycle; + private boolean iteratorClosed; public ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - this.sender = response; + this(response, wrapper, executor, new ScanLifecycle()); + } + + ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wrapper, + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; + this.sender = lifecycle.response(response); this.wrapper = wrapper; this.executor = executor; this.iterator = null; @@ -92,40 +103,46 @@ public ScanBatchResponse(StreamObserver response, HgStoreWrapperEx wra */ @Override public void onNext(ScanStreamBatchReq request) { + if (this.cancelled.get()) { + return; + } switch (request.getQueryCase()) { case QUERY_REQUEST: // query conditions - executor.execute(() -> { + submit(() -> { startQuery(request.getHeader().getGraph(), request.getQueryRequest()); }); break; case RECEIPT_REQUEST: // Message asynchronous response this.clientSeqNo = request.getReceiptRequest().getTimes(); if (seqNo - clientSeqNo < maxInFlightCount) { + boolean send = false; + boolean done; synchronized (stateLock) { + done = state == State.DONE; if (state == State.IDLE) { state = State.DOING; - executor.execute(() -> { - sendEntries(); - }); - } else if (state == State.DONE) { - sendNoDataEntries(); + send = true; } } + if (send) { + submit(this::sendEntries); + } else if (done) { + sendNoDataEntries(); + } } break; case CANCEL_REQUEST: // close stream closeQuery(); break; default: - sender.onError( - HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); + closeQuery(HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); } } @Override public void onError(Throwable t) { log.error("onError ", t); - closeQuery(); + closeQuery(t); } @Override @@ -138,18 +155,43 @@ public void onCompleted() { * * @param request */ + private void submit(Runnable task) { + if (this.cancelled.get()) { + return; + } + try { + this.lifecycle.execute(this.executor, task); + } catch (io.grpc.StatusRuntimeException e) { + closeQuery(e); + } + } + private void startQuery(String graphName, ScanQueryRequest request) { - this.query = request; - this.limit = request.getLimit(); - this.count = 0; - this.iterator = getParallelIterator(graphName, request, this.wrapper, executor); - synchronized (stateLock) { - if (state == State.IDLE) { + this.iteratorLock.lock(); + try { + if (this.cancelled.get() || this.query != null) { + return; + } + this.query = request; + this.limit = request.getLimit(); + this.count = 0; + this.iterator = getParallelIterator(graphName, request, this.wrapper, executor, this.lifecycle); + if (this.cancelled.get()) { + closeIter(); + return; + } + synchronized (stateLock) { + if (state != State.IDLE) { + return; + } state = State.DOING; - executor.execute(() -> { - sendEntries(); - }); } + submit(this::sendEntries); + } catch (RuntimeException e) { + closeQuery(e); + log.warn("Failed to start scan", e); + } finally { + this.iteratorLock.unlock(); } } @@ -157,25 +199,51 @@ private void startQuery(String graphName, ScanQueryRequest request) { * Generate iterator */ private void closeQuery() { + closeQuery(null); + } + + private void closeQuery(Throwable error) { + if (error instanceof java.util.concurrent.RejectedExecutionException) { + error = ScanLifecycle.rejected(this.executor, error); + } + if (!this.cancelled.compareAndSet(false, true)) { + return; + } setStateDone(); + ScanIterator current = this.iterator; + if (current instanceof ParallelScanIterator) { + ((ParallelScanIterator) current).requestStop(); + } try { closeIter(); - this.sender.onCompleted(); - } catch (Exception e) { - log.error("exception ", e); + if (this.lifecycle.cleanupFailure() != null) { + error = Status.INTERNAL.withDescription("Failed to close batch scan iterator") + .withCause(this.lifecycle.cleanupFailure()).asRuntimeException(); + } + if (error == null) { + this.sender.onCompleted(); + } else { + this.sender.onError(error); + } + } catch (RuntimeException | Error e) { + log.error("Failed to terminate batch scan response", e); + } finally { + int active = ScanBatchResponseFactory.getInstance().removeStreamObserver(this); + log.info("ScanBatchResponse closeQuery, active count is {}", active); } - int active = ScanBatchResponseFactory.getInstance().removeStreamObserver(this); - log.info("ScanBatchResponse closeQuery, active count is {}", active); } private void closeIter() { + this.iteratorLock.lock(); try { - if (this.iterator != null) { - this.iterator.close(); - this.iterator = null; + if (this.iterator != null && !this.iteratorClosed) { + this.iteratorClosed = true; + if (this.lifecycle.close(this.iterator)) { + this.iterator = null; + } } - } catch (Exception e) { - + } finally { + this.iteratorLock.unlock(); } } @@ -209,8 +277,12 @@ private void sendEntries() { this.sender.onNext(dataBuilder.build()); this.activeTime = System.currentTimeMillis(); } - if (!iterator.hasNext() || this.count >= limit || state == State.DONE) { + if (state == State.DONE || this.count >= limit || !iterator.hasNext()) { closeIter(); + if (this.lifecycle.cleanupFailure() != null) { + closeQuery(this.lifecycle.cleanupFailure()); + return; + } this.sender.onNext(KvStream.newBuilder().setOver(true).build()); setStateDone(); } else { @@ -219,14 +291,7 @@ private void sendEntries() { } catch (Throwable e) { if (this.state != State.DONE) { log.error(" send data exception: ", e); - setStateIdle(); - if (this.sender != null) { - try { - this.sender.onError(e); - } catch (Exception ex) { - log.warn("Error when call sender.onError {}", e.getMessage()); - } - } + closeQuery(e); } } finally { iteratorLock.unlock(); @@ -234,9 +299,13 @@ private void sendEntries() { } private void sendNoDataEntries() { + this.iteratorLock.lock(); try { - this.sender.onNext(KvStream.newBuilder().setOver(true).build()); - } catch (Exception e) { + if (!this.cancelled.get()) { + this.sender.onNext(KvStream.newBuilder().setOver(true).build()); + } + } finally { + this.iteratorLock.unlock(); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java index d0c52b372e..619b701991 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponse3.java @@ -61,7 +61,13 @@ public static StreamObserver of(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { HgAssert.isArgumentNotNull(responseObserver, "responseObserver"); HgAssert.isArgumentNotNull(wrapper, "wrapper"); - return new Broker(responseObserver, wrapper, executor); + return of(responseObserver, wrapper, executor, new ScanLifecycle()); + } + + static StreamObserver of(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + return new Broker(lifecycle.response(responseObserver), wrapper, executor, lifecycle); } private enum OrderState { @@ -81,10 +87,12 @@ private static class Broker implements StreamObserver { private final HgStoreWrapperEx wrapper; private final ThreadPoolExecutor executor; private final OrderManager manager = new OrderManager(); + private final ScanLifecycle lifecycle; private String graph; Broker(StreamObserver responseObserver, HgStoreWrapperEx wrapper, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; this.responseObserver = responseObserver; this.wrapper = wrapper; this.executor = executor; @@ -101,11 +109,10 @@ public void onNext(ScanStreamBatchReq request) { this.manager.receipt(request.getReceiptRequest().getTimes()); break; case CANCEL_REQUEST: - this.manager.finished(); + this.onCompleted(); break; default: - responseObserver.onError( - HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); + this.onError(HgGrpc.toErr("Unsupported sub-request: [ " + request + " ]")); } } @@ -113,11 +120,13 @@ public void onNext(ScanStreamBatchReq request) { public void onError(Throwable t) { log.warn(t.getMessage()); this.manager.breakdown(); + this.responseObserver.onError(t); } @Override public void onCompleted() { this.manager.finished(); + this.responseObserver.onCompleted(); } private void handleHeader(ScanStreamBatchReq request) { @@ -145,9 +154,9 @@ private void makeADeal(ScanQueryRequest request) { OrderWorker worker = new OrderWorker( request.getLimit(), request.getPageSize(), - ScanUtil.getIterator(this.graph, request, this.wrapper), + ScanUtil.getIterator(this.graph, request, this.wrapper, this.lifecycle::failedCleanup), deliverer, - this.executor); + this.executor, this.lifecycle); this.manager.deal(worker, deliverer); } @@ -160,16 +169,19 @@ private static class OrderManager { OrderState state = OrderState.NEW; OrderWorker worker; OrderDeliverer deliverer; + boolean cancelled; synchronized void deal(OrderWorker worker, OrderDeliverer deliverer) { if (log.isDebugEnabled()) { log.debug("Receiving query request."); } - if (this.state == OrderState.NEW) { + if (!this.cancelled && this.state == OrderState.NEW) { this.worker = worker; this.deliverer = deliverer; this.worker.hereWeGo(); this.state = OrderState.WORKING; + } else { + worker.closeIterator(); } } @@ -177,7 +189,9 @@ synchronized void receipt(int receiptTimes) { if (log.isDebugEnabled()) { log.debug("Receiving receipt request."); } - this.worker.setReceipt(receiptTimes); + if (this.worker != null) { + this.worker.setReceipt(receiptTimes); + } } synchronized void finished() { @@ -188,6 +202,7 @@ synchronized void finished() { } synchronized void breakdown() { + this.cancelled = true; if (this.worker != null) { this.worker.breakdown(); } @@ -240,9 +255,12 @@ void error(String msg) { } void error(String msg, Throwable t) { + this.error(Status.INTERNAL, msg, t); + } + + void error(Status status, String msg, Throwable t) { if (!finishFlag.getAndSet(true)) { - this.responseObserver.onError(HgGrpc.toErr(Status.INTERNAL, - msg, t)); + this.responseObserver.onError(HgGrpc.toErr(status, msg, t)); } } } @@ -261,11 +279,14 @@ private static class OrderWorker { private final AtomicInteger curTimes = new AtomicInteger(); private final ThreadPoolExecutor executor; private final long limit; + private final ScanLifecycle lifecycle; + private final AtomicBoolean iteratorClosed = new AtomicBoolean(); private long packageSize; private long counter; OrderWorker(long limit, long packageSize, ScanIterator iterator, OrderDeliverer deliverer, - ThreadPoolExecutor executor) { + ThreadPoolExecutor executor, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; this.limit = limit; this.packageSize = packageSize; this.iterator = iterator; @@ -297,7 +318,13 @@ void hereWeGo() { return; } - executor.execute(() -> working()); + try { + this.lifecycle.execute(executor, this::working); + } catch (io.grpc.StatusRuntimeException e) { + this.completeFlag.set(true); + closeIterator(); + this.deliverer.error(e.getStatus(), "Store scan task rejected", e); + } Thread.yield(); } @@ -325,6 +352,13 @@ void breakdown() { } } + private void closeIterator() { + if (this.iteratorClosed.compareAndSet(false, true) && + !this.lifecycle.close(this.iterator)) { + this.deliverer.error("Failed to close scan iterator", this.lifecycle.cleanupFailure()); + } + } + private void working() { if (this.isWorking.getAndSet(true)) { return; @@ -338,7 +372,8 @@ private void working() { Kv.Builder kvBuilder = Kv.newBuilder(); long packageCount = 0; - while (iterator.hasNext()) { + while (!this.breakdown.get() && !Thread.currentThread().isInterrupted() && + iterator.hasNext()) { if (++this.counter > limit) { this.completeFlag.set(true); break; @@ -353,7 +388,7 @@ private void working() { deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), false); Thread.yield(); - if (!this.checkContinue()) { + if (!this.breakdown.get() && !this.checkContinue()) { long start = System.currentTimeMillis(); iterator.wait( HgStoreConst.SCAN_WAIT_CLIENT_TAKING_TIME_OUT_SECONDS * @@ -390,15 +425,22 @@ private void working() { } + if (!this.breakdown.get() && Thread.currentThread().isInterrupted()) { + throw new InterruptedException("Scanning interrupted"); + } this.completeFlag.set(true); - deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), true); + closeIterator(); + if (this.lifecycle.cleanupFailure() == null) { + deliverer.deliver(dataBuilder, curTimes.incrementAndGet(), true); + } } } catch (InterruptedException e) { log.error("Interrupted waiting of iterator, canceled while.", e); - this.deliverer.error("Failed to finish scanning, cause by InterruptedException."); + this.deliverer.error(Status.CANCELLED, + "Failed to finish scanning, cause by InterruptedException.", e); } catch (TimeoutException t) { log.info(t.getMessage()); this.deliverer.error("Sever waiting exceeded [" @@ -409,7 +451,7 @@ private void working() { this.deliverer.error("Failed to finish scanning ", t); } finally { this.workingLock.unlock(); - this.iterator.close(); + closeIterator(); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java index 43abfd97ce..a0ab7f2920 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanBatchResponseFactory.java @@ -37,7 +37,13 @@ public static ScanBatchResponseFactory getInstance() { public static StreamObserver of(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - StreamObserver observer = new ScanBatchResponse(responseObserver, wrapper, executor); + return of(responseObserver, wrapper, executor, new ScanLifecycle()); + } + + static StreamObserver of(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + StreamObserver observer = new ScanBatchResponse(responseObserver, wrapper, executor, lifecycle); getInstance().addStreamObserver(observer); getInstance().checkStreamActive(); return observer; diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java new file mode 100644 index 0000000000..6cce65ad02 --- /dev/null +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanLifecycle.java @@ -0,0 +1,192 @@ +/* + * 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; + +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hugegraph.rocksdb.access.ScanIterator; + +import io.grpc.Status; +import io.grpc.stub.StreamObserver; + +/** Per-call response and cleanup accounting; a failed release remains registered. */ +final class ScanLifecycle { + + // 0 = open, 1 = cancellation won, 2 = response terminated. + private final AtomicInteger terminal = new AtomicInteger(); + private final Object responseLock = new Object(); + private int activities; + private volatile Throwable cleanupFailure; + private Runnable finished = () -> { }; + + synchronized void onFinished(Runnable action) { + this.finished = action; + } + + synchronized void enter() { + this.activities++; + } + + synchronized boolean tryEnter() { + if (this.terminal.get() != 0) { + return false; + } + this.activities++; + return true; + } + + synchronized boolean tryEnterCancellation() { + if (!tryCancel()) { + return false; + } + this.activities++; + return true; + } + + synchronized void leave() { + this.activities--; + checkFinished(); + } + + synchronized void failedCleanup(Throwable failure) { + if (this.cleanupFailure == null) { + this.cleanupFailure = failure; + } + } + + Throwable cleanupFailure() { + return this.cleanupFailure; + } + + private synchronized void checkFinished() { + if (this.terminal.get() == 2 && this.activities == 0 && this.cleanupFailure == null) { + this.finished.run(); + } + } + + boolean tryCancel() { + return this.terminal.compareAndSet(0, 1); + } + + boolean isCancelled() { + return this.terminal.get() == 1; + } + + void finishWithoutResponse() { + this.terminal.set(2); + checkFinished(); + } + + boolean close(ScanIterator iterator) { + if (iterator == null) { + return true; + } + try { + iterator.close(); + return true; + } catch (RuntimeException | Error failure) { + failedCleanup(failure); + return false; + } + } + + void execute(ThreadPoolExecutor executor, Runnable task) { + enter(); + AtomicBoolean released = new AtomicBoolean(); + Runnable release = () -> { + if (released.compareAndSet(false, true)) { + leave(); + } + }; + try { + executor.execute(() -> { + try { + task.run(); + } finally { + release.run(); + } + }); + } catch (RuntimeException | Error failure) { + release.run(); + if (failure instanceof RejectedExecutionException) { + throw rejected(executor, failure); + } + throw failure; + } + } + + static RuntimeException rejected(ThreadPoolExecutor executor, Throwable failure) { + Status status = executor.isShutdown() ? Status.UNAVAILABLE : Status.RESOURCE_EXHAUSTED; + return status.withDescription("Store scan task rejected").withCause(failure) + .asRuntimeException(); + } + + StreamObserver response(StreamObserver delegate) { + return new StreamObserver() { + @Override + public void onNext(T value) { + enter(); + try { + synchronized (responseLock) { + if (terminal.get() == 0) { + delegate.onNext(value); + } + } + } finally { + leave(); + } + } + + @Override + public void onError(Throwable failure) { + enter(); + try { + synchronized (responseLock) { + if (terminal.getAndSet(2) != 2) { + delegate.onError(failure); + } + } + } finally { + leave(); + } + } + + @Override + public void onCompleted() { + enter(); + try { + synchronized (responseLock) { + int previous = terminal.getAndSet(2); + if (previous == 0) { + delegate.onCompleted(); + } else if (previous == 1) { + delegate.onError(Status.CANCELLED + .withDescription("Store scan cancelled before completion") + .asRuntimeException()); + } + } + } finally { + leave(); + } + } + }; + } +} diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java index 5b77a4449a..e800cc1166 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java @@ -27,6 +27,7 @@ import com.google.protobuf.ByteString; +import io.grpc.Context; import io.grpc.Status; import io.grpc.stub.StreamObserver; import lombok.extern.slf4j.Slf4j; @@ -48,6 +49,12 @@ public class ScanOneShotResponse { public static void scanOneShot(ScanStreamReq request, StreamObserver responseObserver, HgStoreWrapperEx wrapper) { + scanOneShot(request, responseObserver, wrapper, new ScanLifecycle()); + } + + static void scanOneShot(ScanStreamReq request, StreamObserver output, + HgStoreWrapperEx wrapper, ScanLifecycle lifecycle) { + StreamObserver responseObserver = lifecycle.response(output); KvPageRes.Builder resBuilder = KvPageRes.newBuilder(); Kv.Builder kvBuilder = Kv.newBuilder(); @@ -57,12 +64,17 @@ public static void scanOneShot(ScanStreamReq request, responseObserver.onError(HgGrpc.toErr("limit<=0, please to invoke stream scan.")); return; } - ScanIterator iterator = ScanUtil.getIterator(request, wrapper); + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + ScanIterator iterator = ScanUtil.getIterator(request, wrapper, lifecycle::failedCleanup); int count = 0; + boolean cleanupAttempted = false; try { - while (iterator.hasNext()) { + while (!(Context.current().isCancelled() || lifecycle.isCancelled()) && + !Thread.currentThread().isInterrupted() && iterator.hasNext()) { if (++count > limit) { break; @@ -78,6 +90,19 @@ public static void scanOneShot(ScanStreamReq request, } + if ((Context.current().isCancelled() || lifecycle.isCancelled())) { + return; + } + if (Thread.currentThread().isInterrupted()) { + responseObserver.onError(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } + cleanupAttempted = true; + if (!lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + return; + } responseObserver.onNext( resBuilder.setVersion(ScanUtil.responseVersion(request)) .build()); @@ -87,7 +112,10 @@ public static void scanOneShot(ScanStreamReq request, String msg = "an exception occurred during data scanning"; responseObserver.onError(HgGrpc.toErr(Status.INTERNAL, msg, t)); } finally { - iterator.close(); + if (!cleanupAttempted && !lifecycle.close(iterator)) { + responseObserver.onError(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(lifecycle.cleanupFailure()).asRuntimeException()); + } } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java index e572862823..8628e0c92b 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java @@ -61,6 +61,21 @@ public class ScanStreamResponse implements StreamObserver { private final int waitTime; private final HgChannel channel; private ScanIterator iterator; + private final Object cancellationLock = new Object(); + private final Object responseLock = new Object(); + private Thread worker; + private final ScanLifecycle lifecycle; + + private void cancel() { + this.isStop.set(true); + this.channel.close(); + Thread current = Thread.currentThread(); + synchronized (this.cancellationLock) { + if (this.worker != null && this.worker != current) { + this.worker.interrupt(); + } + } + } private long limit = 0; private int times = 0; private long pageSize = 0; @@ -72,7 +87,14 @@ public class ScanStreamResponse implements StreamObserver { ScanStreamResponse(StreamObserver responseObserver, HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, AppConfig appConfig) { - this.responseObserver = responseObserver; + this(responseObserver, wrapper, executor, appConfig, new ScanLifecycle()); + } + + ScanStreamResponse(StreamObserver responseObserver, + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + AppConfig appConfig, ScanLifecycle lifecycle) { + this.lifecycle = lifecycle; + this.responseObserver = lifecycle.response(responseObserver); this.wrapper = wrapper; this.executor = executor; this.config = appConfig; @@ -91,6 +113,9 @@ public static ScanStreamResponse of(StreamObserver responseObserver, @Override public void onNext(ScanStreamReq request) { + if (this.isStop.get()) { + return; + } try { this.responseVersion = Math.max( this.responseVersion, ScanUtil.responseVersion(request)); @@ -100,22 +125,20 @@ public void onNext(ScanStreamReq request) { next(request); } } catch (Exception e) { - responseObserver.onError(e); + this.failServer(e); } } @Override public void onError(Throwable t) { - this.isStop.set(true); - this.finishServer(); + this.cancelServer(false); log.warn("onError from client [ graph: {} , table: {}]; Reason: {}]", graph, table, t.getMessage()); } @Override public void onCompleted() { - this.isStop.set(true); - this.finishServer(); + this.cancelServer(false); } private void initIterator(ScanStreamReq request) { @@ -123,7 +146,6 @@ private void initIterator(ScanStreamReq request) { if (this.isStarted.getAndSet(true)) { return; } - this.iterator = getIterator(request, this.wrapper); this.graph = request.getHeader().getGraph(); this.table = request.getTable(); this.limit = request.getLimit(); @@ -143,7 +165,15 @@ private void initIterator(ScanStreamReq request) { Kv.Builder kvBuilder = Kv.newBuilder(); int pageCount = 0; try { - while (iterator.hasNext()) { + synchronized (this.cancellationLock) { + if (this.isStop.get()) { + return; + } + this.worker = Thread.currentThread(); + } + this.iterator = getIterator(request, this.wrapper, this.lifecycle::failedCleanup); + while (!this.isStop.get() && !Thread.currentThread().isInterrupted() && + iterator.hasNext()) { if (limit > 0 && ++this.total > limit) { break; } @@ -164,32 +194,45 @@ private void initIterator(ScanStreamReq request) { } dataBuilder.addData(toKv(kvBuilder, iterator.next(), iterator.position())); } + if (this.isStop.get()) { + return; + } + if (Thread.currentThread().isInterrupted()) { + this.failServer(HgGrpc.toErr(Status.Code.CANCELLED, "Scanning interrupted")); + return; + } this.channel.send(dataBuilder); } catch (Throwable t) { + if (this.isStop.get()) { + return; + } String msg = "an exception occurred while scanning data:"; - StatusRuntimeException ex = - HgGrpc.toErr(Status.INTERNAL, msg + t.getMessage(), t); - responseObserver.onError(ex); + Status status = t instanceof InterruptedException || + Thread.currentThread().isInterrupted() ? + Status.CANCELLED : Status.INTERNAL; + StatusRuntimeException ex = HgGrpc.toErr(status, msg + t.getMessage(), t); + this.failServer(ex); } finally { try { - this.iterator.close(); + if (!this.lifecycle.close(this.iterator)) { + this.failServer(Status.INTERNAL.withDescription("Failed to close scan iterator") + .withCause(this.lifecycle.cleanupFailure()) + .asRuntimeException()); + } + } finally { this.channel.close(); - } catch (Exception e) { - + synchronized (this.cancellationLock) { + this.worker = null; + } } } }; - this.executor.execute(scanning); + this.lifecycle.execute(this.executor, scanning); } catch (Exception e) { - StatusRuntimeException ex = HgGrpc.toErr(Status.INTERNAL, null, e); - responseObserver.onError(ex); - try { - this.iterator.close(); - this.channel.close(); - } catch (Exception exception) { - - } + StatusRuntimeException ex = e instanceof StatusRuntimeException ? + (StatusRuntimeException) e : HgGrpc.toErr(Status.INTERNAL, null, e); + this.failServer(ex); } /*** Scanning loop end ***/ @@ -205,19 +248,7 @@ private Kv toKv(Kv.Builder kvBuilder, RocksDBSession.BackendColumn col, } private void close() { - this.isStop.set(true); - this.channel.close(); - if (!this.finishFlag.get()) { - responseObserver.onNext(KvPageRes.newBuilder() - .addAllData(Collections.EMPTY_LIST) - .setOver(true) - .setTimes(++times) - .setVersion(this.responseVersion) - .build() - ); - } - - this.finishServer(); + this.cancelServer(true); } private void next(ScanStreamReq request) { @@ -225,12 +256,19 @@ private void next(ScanStreamReq request) { KvPageRes.Builder resBuilder; try { + if (this.isStop.get()) { + return; + } resBuilder = this.channel.receive(); times++; } catch (Exception e) { + if (this.isStop.get()) { + return; + } String msg = "failed to poll a page of data, cause by:"; + Status status = Thread.currentThread().isInterrupted() ? Status.CANCELLED : Status.INTERNAL; log.error(msg, e); - responseObserver.onError(HgGrpc.toErr(msg + e.getMessage())); + this.failServer(HgGrpc.toErr(status, msg + e.getMessage(), e)); return; } boolean isOver = false; @@ -239,31 +277,65 @@ private void next(ScanStreamReq request) { isOver = true; resBuilder = KvPageRes.newBuilder().addAllData(Collections.EMPTY_LIST); } - if (!this.finishFlag.get()) { - responseObserver.onNext(resBuilder.setOver(isOver) - .setTimes(times) - .setVersion(this.responseVersion) - .build()); - } + this.sendPage(resBuilder.setOver(isOver) + .setTimes(times) + .setVersion(this.responseVersion) + .build()); if (isOver) { this.finishServer(); } } + private void cancelServer(boolean sendFinalPage) { + // Cancellation can make the worker throw; establish the normal terminal first. + boolean complete = this.finishFlag.compareAndSet(false, true); + this.cancel(); + if (complete) { + synchronized (this.responseLock) { + if (sendFinalPage) { + this.responseObserver.onNext(KvPageRes.newBuilder() + .setOver(true) + .setTimes(++times) + .setVersion(this.responseVersion) + .build()); + } + this.responseObserver.onCompleted(); + } + } + } + + private void sendPage(KvPageRes page) { + synchronized (this.responseLock) { + if (!this.finishFlag.get()) { + this.responseObserver.onNext(page); + } + } + } + + private void failServer(Throwable failure) { + // Claim the terminal signal before waking a receiver on the closed channel. + boolean reportFailure = this.finishFlag.compareAndSet(false, true); + this.cancel(); + if (reportFailure) { + synchronized (this.responseLock) { + this.responseObserver.onError(failure); + } + } + } + private void finishServer() { if (!this.finishFlag.getAndSet(true)) { - responseObserver.onCompleted(); + synchronized (this.responseLock) { + this.responseObserver.onCompleted(); + } } } private void timeoutSever() { - if (!this.finishFlag.getAndSet(true)) { - String msg = "server wait time exceeds the threshold[" + waitTime + - "] seconds."; - responseObserver.onError( - HgGrpc.toErr(Status.Code.DEADLINE_EXCEEDED, msg)); - } + String msg = "server wait time exceeds the threshold[" + waitTime + + "] seconds."; + this.failServer(HgGrpc.toErr(Status.Code.DEADLINE_EXCEEDED, msg)); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java index b7678f91d9..8f873cbaeb 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java @@ -29,6 +29,7 @@ import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Supplier; +import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.concurrent.NotThreadSafe; @@ -58,6 +59,11 @@ class ScanUtil { private final static Map tableKeyMap = new HashMap<>(); static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) { + return getIterator(request, wrapper, ignored -> { }); + } + + static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { boolean ordered = isOrdered(request); if (ordered && !request.getPosition().isEmpty()) { throw new IllegalArgumentException( @@ -109,16 +115,19 @@ static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) throw new IllegalArgumentException("Property selection requires a vertex or edge table: " + table); } iter = new SelectIterator(iter, properties, BusinessHandlerImpl.getGraphSupplier(graph), isVertex); - } catch (RuntimeException failure) { + } catch (RuntimeException | Error failure) { try { iter.close(); - } catch (RuntimeException closeFailure) { - failure.addSuppressed(closeFailure); + } catch (RuntimeException | Error closeFailure) { + cleanupFailure.accept(closeFailure); + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } } throw failure; } } - iter.seek(request.getPosition().toByteArray()); + seek(iter, request.getPosition().toByteArray(), cleanupFailure); return iter; } @@ -137,6 +146,11 @@ static int responseVersion(ScanStreamReq request) { } static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper) { + return getIterator(sq, wrapper, ignored -> { }); + } + + private static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { if (log.isDebugEnabled()) { log.debug("{}", sq); } @@ -161,12 +175,29 @@ static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper) { iter = new EmptyIterator(); } - iter.seek(sq.position); + seek(iter, sq.position, cleanupFailure); return iter; } + private static void seek(ScanIterator iterator, byte[] position, + Consumer cleanupFailure) { + try { + iterator.seek(position); + } catch (RuntimeException | Error failure) { + try { + iterator.close(); + } catch (RuntimeException | Error closing) { + cleanupFailure.accept(closing); + if (failure != closing) { + failure.addSuppressed(closing); + } + } + throw failure; + } + } + static ScanQuery toSq(ScanStreamReq request) { ScanQuery res = ScanQuery.of(); @@ -192,7 +223,12 @@ static ScanQuery toSq(ScanStreamReq request) { static ScanIterator getIterator(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper) { - ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper); + return getIterator(graph, request, wrapper, ignored -> { }); + } + + static ScanIterator getIterator(String graph, ScanQueryRequest request, + HgStoreWrapperEx wrapper, Consumer cleanupFailure) { + ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper, cleanupFailure); return BatchScanIterator.of(supplier, supplier.getLimitSupplier()); } @@ -200,10 +236,11 @@ static ScanIterator getIterator(String graph, ScanQueryRequest request, * Support for multi-iterators with parallel reading */ static ScanIterator getParallelIterator(String graph, ScanQueryRequest request, - HgStoreWrapperEx wrapper, ThreadPoolExecutor executor) { - ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper); + HgStoreWrapperEx wrapper, ThreadPoolExecutor executor, + ScanLifecycle lifecycle) { + ScanIteratorSupplier supplier = new ScanIteratorSupplier(graph, request, wrapper, lifecycle::failedCleanup); return ParallelScanIterator.of(supplier, supplier.getLimitSupplier(), - request, executor); + request, executor, lifecycle::failedCleanup); } @NotThreadSafe @@ -216,6 +253,7 @@ private static class ScanIteratorSupplier implements private final long perKeyMax; private final long skipDegree; private final HgStoreWrapperEx wrapper; + private final Consumer cleanupFailure; private long perKeyLimit; private List sqs = new LinkedList<>(); private Iterator sqIterator; @@ -223,7 +261,9 @@ private static class ScanIteratorSupplier implements private ScanQueryProducer scanQueryProducer; private Iterator scanQueryIterator; - ScanIteratorSupplier(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper) { + ScanIteratorSupplier(String graph, ScanQueryRequest request, HgStoreWrapperEx wrapper, + Consumer cleanupFailure) { + this.cleanupFailure = cleanupFailure; this.graph = graph; this.perKeyLimit = request.getPerKeyLimit(); this.perKeyMax = request.getPerKeyMax(); @@ -311,7 +351,7 @@ public KVPair get1() { ScanQuery query = null; if (this.sqIterator != null && this.sqIterator.hasNext()) { query = this.sqIterator.next(); - iterator = getIterator(query, this.wrapper); + iterator = getIterator(query, this.wrapper, this.cleanupFailure); } else { this.sqs.clear(); this.sqIterator = null; @@ -369,7 +409,7 @@ public ScanIterator get() { if (index + 1 > queries.length) { return null; } - return getIterator(queries[index++], wrapper); + return getIterator(queries[index++], wrapper, cleanupFailure); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java index 0ba569cb92..f8edd1d707 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryObserver.java @@ -21,13 +21,17 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import org.apache.hugegraph.backend.BackendColumn; import org.apache.hugegraph.rocksdb.access.RocksDBSession; @@ -65,18 +69,47 @@ public class AggregativeQueryObserver implements StreamObserver { private volatile ScanIterator iterator = null; private QueryPlan plan = null; private String queryId; + private final Consumer completion; + private final Set workers = new HashSet<>(); + private int pendingTasks; + private boolean finished; + private volatile boolean requestCompleted; + private final Object responseLock = new Object(); + private boolean responseFinished; + private boolean completeResponse = true; + private QueryResponse finalResponse; + private boolean errorReported; + private final AtomicReference cleanupFailure = new AtomicReference<>(); public AggregativeQueryObserver(StreamObserver sender, ExecutorService threadPool, long timeout, int batchSize) { + this(sender, threadPool, timeout, batchSize, ignored -> { }); + } + + AggregativeQueryObserver(StreamObserver sender, + ExecutorService threadPool, long timeout, int batchSize, + Consumer completion) { this.sender = sender; this.threadPool = threadPool; this.batchSize = batchSize; this.timeout = timeout; + this.completion = completion; + } + + ScanIterator getIterator(QueryRequest request) { + return QueryUtil.getIterator(request); + } + + QueryPlan buildPlan(QueryRequest request) { + return QueryUtil.buildPlan(request); } @Override - public void onNext(QueryRequest request) { + public synchronized void onNext(QueryRequest request) { + if (this.clientCanceled.get() || this.finished || this.requestCompleted) { + return; + } if (this.queryId == null) { log.debug("got request: {}", request); this.queryId = request.getQueryId(); @@ -85,9 +118,23 @@ public void onNext(QueryRequest request) { // the first request, start the sending thread if (iterator == null) { long current = System.nanoTime(); - iterator = QueryUtil.getIterator(request); - plan = QueryUtil.buildPlan(request); - threadPool.submit(this::sendData); + this.pendingTasks = 1; + try { + iterator = getIterator(request); + plan = buildPlan(request); + threadPool.execute(this::sendData); + } catch (RuntimeException | Error e) { + // The framework must report the original synchronous failure, not normal completion. + synchronized (this.responseLock) { + this.completeResponse = false; + } + taskFinished(); + Throwable failure = this.cleanupFailure.get(); + if (failure != null && failure != e) { + e.addSuppressed(failure); + } + throw e; + } log.debug("query id: {}, init data cost: {} ms", queryId, (System.nanoTime() - current) * 1.0 / 1000000); } else { @@ -98,18 +145,132 @@ public void onNext(QueryRequest request) { @Override public void onError(Throwable t) { - // Stop calculating when channel got error - this.clientCanceled.set(true); + onTransportCancel(); log.error("AggregativeQueryService, query id: {}, got error", this.queryId, t); } + void onTransportCancel() { + // Both transport cancellation and request onError can report the same event. + synchronized (this.responseLock) { + if (this.responseFinished) { + return; + } + this.completeResponse = false; + this.responseFinished = true; + } + cancel(); + } + @Override - public void onCompleted() { - // client my be cancelled earlier + public synchronized void onCompleted() { + // Half-close ends feedback, not work already supported by its existing credit. + this.requestCompleted = true; + if (this.pendingTasks == 0 && !this.finished) { + this.finished = true; + finishResponse(); + } + } + + public void cancel() { this.clientCanceled.set(true); + synchronized (this) { + for (Thread worker : this.workers) { + worker.interrupt(); + } + // An RPC with no first request has no iterator or cleanup task. + if (this.pendingTasks == 0 && !this.finished) { + this.finished = true; + finishResponse(); + } + } + } + + private synchronized void taskStarted() { + this.workers.add(Thread.currentThread()); + } + + private void workerFinished() { + synchronized (this) { + this.workers.remove(Thread.currentThread()); + } + taskFinished(); + } + + private void taskFinished() { + boolean cleanup; + synchronized (this) { + cleanup = --this.pendingTasks == 0; + if (cleanup) { + this.finished = true; + } + } + if (cleanup) { + if (this.plan != null) { + cleanup(this.plan::clear); + } + if (this.iterator != null) { + cleanup(this.iterator::close); + } + finishResponse(); + } + } + + private void cleanup(Runnable release) { + try { + release.run(); + } catch (RuntimeException | Error failure) { + synchronized (this.cleanupFailure) { + Throwable first = this.cleanupFailure.get(); + if (first == null) { + this.cleanupFailure.set(failure); + } else if (first != failure) { + first.addSuppressed(failure); + } + } + log.error("Aggregate query {} failed to release resources; Store shutdown remains blocked", + this.queryId, failure); + } + } + + private void finishResponse() { + try { + synchronized (this.responseLock) { + if (this.completeResponse && !this.responseFinished) { + this.responseFinished = true; + Throwable failure = this.cleanupFailure.get(); + if (failure != null && !this.errorReported) { + this.sender.onNext(errorResponse(getBuilder(), this.queryId, failure)); + } else if (failure == null && this.finalResponse != null && + !this.clientCanceled.get()) { + this.sender.onNext(this.finalResponse); + } + this.sender.onCompleted(); + } + } + } finally { + // A terminal RPC does not prove that its native resources were released. + if (this.cleanupFailure.get() == null) { + this.completion.accept(this); + } + } + } + + private void sendResponse(QueryResponse response) { + synchronized (this.responseLock) { + if (!this.responseFinished && !this.clientCanceled.get()) { + if (response.getIsOk() && response.getIsFinished()) { + // The client treats this batch as success; publish it only after cleanup. + this.finalResponse = response; + } else { + this.errorReported |= !response.getIsOk(); + this.sender.onNext(response); + } + } + } } public void sendData() { + taskStarted(); try { long lastSend = System.currentTimeMillis(); var responseBuilder = getBuilder(); @@ -118,12 +279,18 @@ public void sendData() { while (!this.clientCanceled.get()) { // produces more result than consumer, just waiting if (sendCount.get() - consumeCount.get() >= RESULT_COUNT) { + if (this.requestCompleted) { + sendResponse(errorResponse(getBuilder(), queryId, + new IllegalStateException( + "Request completed without enough feedback"))); + cancel(); + return; + } // read timeout, takes long time not to read data if (System.currentTimeMillis() - lastSend > timeout) { - this.sender.onNext(errorResponse(getBuilder(), queryId, - new RuntimeException( - "sending-timeout, server closed"))); - this.sender.onCompleted(); + sendResponse(errorResponse(getBuilder(), queryId, + new RuntimeException("sending-timeout, server closed"))); + cancel(); return; } @@ -141,23 +308,33 @@ public void sendData() { } else { try { builder.setQueryId(queryId); - sender.onNext(builder.build()); + sendResponse(builder.build()); this.sendCount.incrementAndGet(); lastSend = System.currentTimeMillis(); } catch (Exception e) { log.error("send data got error: ", e); + cancel(); break; } } - if (builder.getIsFinished() || !builder.getIsOk()) { + if (!builder.getIsOk()) { + // Report the internal error before cancelling remaining partition work. + cancel(); + break; + } + if (builder.getIsFinished()) { break; } } + } catch (Exception e) { + try { + sendResponse(errorResponse(getBuilder(), queryId, e)); + } finally { + cancel(); + } } finally { - this.plan.clear(); - this.iterator.close(); - this.sender.onCompleted(); + workerFinished(); } } @@ -188,32 +365,64 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, return builder; } else if (this.plan.hasIteratorResult()) { checkIterator = false; - AtomicReference exception = new AtomicReference<>(); + AtomicReference exception = new AtomicReference<>(); if (this.iterator instanceof MultiPartitionIterator) { var iterators = ((MultiPartitionIterator) this.iterator).getIterators(); CountDownLatch latch = new CountDownLatch(iterators.size()); - for (var itr2 : iterators) { - threadPool.execute(() -> { + synchronized (this) { + this.pendingTasks += iterators.size(); + } + int submitted = 0; + try { + for (var itr2 : iterators) { + threadPool.execute(() -> { + taskStarted(); + try { + execute(itr2); + } catch (RuntimeException | Error e) { + exception.compareAndSet(null, e); + } finally { + try { + cleanup(itr2::close); + Throwable failure = this.cleanupFailure.get(); + if (failure != null) { + exception.compareAndSet(null, new RuntimeException( + "partition iterator cleanup failed", failure)); + } + } finally { + latch.countDown(); + workerFinished(); + } + } + }); + submitted++; + } + } catch (RuntimeException | Error failure) { + // Rejection must also release iterators whose tasks were not accepted. + for (int i = submitted; i < iterators.size(); i++) { try { - execute(itr2); - } catch (Exception e) { - exception.set(e); + cleanup(iterators.get(i)::close); } finally { - // MultiPartitionIterator close() not working - itr2.close(); latch.countDown(); + taskFinished(); } - }); + } + throw failure; + } + if (!latch.await(timeout, TimeUnit.MILLISECONDS)) { + throw new TimeoutException("partition query timed out"); } - latch.await(timeout, TimeUnit.MILLISECONDS); if (exception.get() != null) { - throw exception.get(); + throw new RuntimeException("partition query failed", exception.get()); } } else { // can't be parallel, but has agg like stage execute(this.iterator); } + if (this.clientCanceled.get()) { + return builder.setIsOk(false).setIsFinished(false); + } try { // last empty element itr = (ScanIterator) plan.execute(PipelineResult.EMPTY); @@ -227,7 +436,7 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, builder.clear(); List batchResult = new ArrayList<>(); - while (itr.hasNext() && !this.clientCanceled.get()) { + while (!this.clientCanceled.get() && itr.hasNext()) { if (count >= batchSize) { break; } @@ -265,6 +474,9 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, builder.addAllData(batchResult); } catch (Exception e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } log.error("readBatchData got error: ", e); return builder.setIsOk(false).setIsFinished(false).setMessage("Store Server Error: " + Arrays.toString( @@ -273,7 +485,7 @@ private QueryResponse.Builder readBatchData(QueryResponse.Builder builder, if (checkIterator) { // check the iterator - finish = !itr.hasNext(); + finish = this.clientCanceled.get() || !itr.hasNext(); } log.debug("query id: {}, finished batch, with size :{}, finish:{}, cost: {} ms", queryId, count, @@ -320,7 +532,7 @@ public void close() { private void execute(ScanIterator itr) { long recordCount = 0; long current = System.nanoTime(); - while (itr.hasNext() && !this.clientCanceled.get()) { + while (!this.clientCanceled.get() && itr.hasNext()) { try { recordCount++; executePipeline(itr.next()); diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java index 3eb81be851..36e5500680 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryService.java @@ -18,7 +18,9 @@ package org.apache.hugegraph.store.node.grpc.query; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.atomic.AtomicLong; @@ -35,6 +37,8 @@ import com.google.protobuf.ByteString; +import io.grpc.Status; +import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -43,6 +47,9 @@ @GRpcService public class AggregativeQueryService extends QueryServiceGrpc.QueryServiceImplBase { + private final Set queries = new HashSet<>(); + private boolean closing; + private final int batchSize; private final Long timeout; @@ -62,6 +69,65 @@ public AggregativeQueryService() { 10000, true); } + AggregativeQueryService(ThreadPoolExecutor threadPool, long timeout, int batchSize) { + this.threadPool = threadPool; + this.timeout = timeout; + this.batchSize = batchSize; + } + + public synchronized void stopAcceptingQueries() { + this.closing = true; + } + + /** Cancel queries, but let every queued parent and partition release its iterator. */ + public void shutdownQueries() { + AggregativeQueryObserver[] active; + synchronized (this) { + this.closing = true; + active = this.queries.toArray(new AggregativeQueryObserver[0]); + } + for (AggregativeQueryObserver query : active) { + try { + query.cancel(); + } catch (RuntimeException | Error failure) { + // Response callbacks cannot skip cancellation or cleanup waits for other queries. + log.error("Failed to cancel aggregate query response; continuing shutdown", failure); + } + } + boolean interrupted = false; + synchronized (this) { + try { + while (!this.queries.isEmpty()) { + try { + this.wait(5000); + if (!this.queries.isEmpty()) { + log.warn("Waiting for {} aggregate queries to release iterators", + this.queries.size()); + } + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + // Parents can still enqueue partition cleanup until the last query is finished. + this.threadPool.shutdown(); + } + + private synchronized void finished(AggregativeQueryObserver observer) { + this.queries.remove(observer); + this.notifyAll(); + } + + AggregativeQueryObserver newObserver(StreamObserver sender) { + return new AggregativeQueryObserver(sender, this.threadPool, this.timeout, + this.batchSize, this::finished); + } + /** * Generate error response. * @@ -79,8 +145,17 @@ public static QueryResponse errorResponse(QueryResponse.Builder builder, String } @Override - public StreamObserver query(StreamObserver observer) { - return new AggregativeQueryObserver(observer, threadPool, timeout, batchSize); + public synchronized StreamObserver query(StreamObserver observer) { + if (this.closing) { + throw Status.UNAVAILABLE.withDescription("Store queries are stopping").asRuntimeException(); + } + AggregativeQueryObserver query = newObserver(observer); + this.queries.add(query); + if (observer instanceof ServerCallStreamObserver) { + ((ServerCallStreamObserver) observer).setOnCancelHandler( + query::onTransportCancel); + } + return query; } @Override diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java index 228850d65c..d3289c2492 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/query/model/QueryPlan.java @@ -118,10 +118,25 @@ public String toString() { } public void clear() { + Throwable failure = null; for (var stage : stages) { - stage.close(); + try { + stage.close(); + } catch (RuntimeException | Error e) { + if (failure == null) { + failure = e; + } else if (failure != e) { + failure.addSuppressed(e); + } + } } this.stages.clear(); + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw (RuntimeException) failure; + } } public boolean isEmpty() { diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java index 6583e59745..31bd69c40f 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/GraphStoreImpl.java @@ -77,7 +77,7 @@ public ThreadPoolExecutor getExecutor() { @Override public StreamObserver scanPartition( StreamObserver ro) { - return new ScanResponseObserver(ro, getHandler(), getExecutor()); + return this.storeStream.scanGraphPartition(ro, getHandler()); } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java index b4d10f7c43..a0b087acf0 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/scan/ScanResponseObserver.java @@ -18,249 +18,232 @@ package org.apache.hugegraph.store.node.grpc.scan; import java.util.ArrayList; -import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import org.apache.hugegraph.store.business.BusinessHandler; import org.apache.hugegraph.store.business.GraphStoreIterator; -import org.apache.hugegraph.store.grpc.Graphpb.Error; +import org.apache.hugegraph.store.grpc.Graphpb; import org.apache.hugegraph.store.grpc.Graphpb.ErrorType; import org.apache.hugegraph.store.grpc.Graphpb.ResponseHeader; import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; -import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest.Request; import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest.ScanType; import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; import com.google.protobuf.Descriptors; +import io.grpc.Status; import io.grpc.stub.StreamObserver; -import lombok.extern.slf4j.Slf4j; -@Slf4j -public class ScanResponseObserver implements - StreamObserver { +public class ScanResponseObserver implements StreamObserver { private static final int BATCH_SIZE = 100000; - private static final int MAX_PAGE = 8; // - private static final Error ok = Error.newBuilder().setType(ErrorType.OK).build(); - private static final ResponseHeader okHeader = - ResponseHeader.newBuilder().setError(ok).build(); + private static final int MAX_PAGE = 8; + private static final ResponseHeader OK = ResponseHeader.newBuilder() + .setError(Graphpb.Error.newBuilder().setType(ErrorType.OK)).build(); private final BusinessHandler handler; - private final AtomicInteger nextSeqNo = new AtomicInteger(0); - private final AtomicInteger cltSeqNo = new AtomicInteger(0); - private final ThreadPoolExecutor executor; - private final AtomicBoolean readOver = new AtomicBoolean(false); - private final LinkedBlockingQueue packages = - new LinkedBlockingQueue(MAX_PAGE * 2); - private final Descriptors.FieldDescriptor vertexField = - ScanResponse.getDescriptor().findFieldByNumber(3); - private final Descriptors.FieldDescriptor edgeField = - ScanResponse.getDescriptor().findFieldByNumber(4); - private final ReentrantLock readLock = new ReentrantLock(); - private final ReentrantLock sendLock = new ReentrantLock(); - private StreamObserver sender; + private final StreamObserver sender; + private final Consumer execute; + private final Consumer failedCleanup; + private final AtomicInteger nextSeqNo = new AtomicInteger(); + private final AtomicInteger cltSeqNo = new AtomicInteger(); + private final AtomicBoolean readOver = new AtomicBoolean(); + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicBoolean reading = new AtomicBoolean(); + private final AtomicBoolean sending = new AtomicBoolean(); + private final LinkedBlockingQueue packages = new LinkedBlockingQueue<>(MAX_PAGE * 2); + private final Object iteratorLock = new Object(); + private GraphStoreIterator iter; private ScanPartitionRequest scanReq; - private GraphStoreIterator iter; - private volatile long leftCount; - private volatile Future sendTask; - private volatile Future readTask; + private boolean iteratorClosed; + private long leftCount; - /* - * November 1, 2022 - * 1. onNext needs to be processed asynchronously to prevent the grpc call from being blocked. - * 2. Do not read iterators or send data do not produce thread waiting. - * 3. Before sending, try to prepare the data to be sent as much as possible. - * */ - - /* - * November 2, 2022 - * 1. Read the thread of rocksdb iterator read - * 2. Perform data conversion and send to the blocking queue thread offer - * 3. Thread for reading data from the blocking queue and sending, including waking up the - * reading and sending threads when no data is read - * */ + public ScanResponseObserver(StreamObserver sender, + BusinessHandler handler, ThreadPoolExecutor executor) { + this(sender, handler, executor::execute, failure -> { }); + } public ScanResponseObserver(StreamObserver sender, BusinessHandler handler, - ThreadPoolExecutor executor) { + Consumer execute, Consumer failedCleanup) { this.sender = sender; this.handler = handler; - this.executor = executor; + this.execute = execute; + this.failedCleanup = failedCleanup; } - private boolean readCondition() { - return packages.remainingCapacity() != 0 && !readOver.get(); + @Override + public void onNext(ScanPartitionRequest request) { + try { + if (request.hasScanRequest() && !request.hasReplyRequest()) { + synchronized (this.iteratorLock) { + if (this.closed.get()) { + return; + } + if (this.iter != null) { + throw Status.INVALID_ARGUMENT.withDescription("Scan request was already initialized") + .asRuntimeException(); + } + this.scanReq = request; + long limit = request.getScanRequest().getLimit(); + this.leftCount = limit > 0 ? limit : Long.MAX_VALUE; + this.iter = this.handler.scan(request); + } + startRead(); + } else { + this.cltSeqNo.incrementAndGet(); + startSend(); + } + } catch (RuntimeException | Error failure) { + terminate(failure); + } } - private boolean readTaskCondition() { - return readCondition() && (readTask == null || readTask.isDone()); + @Override + public void onError(Throwable failure) { + terminate(failure); } - private boolean sendCondition() { - return nextSeqNo.get() - cltSeqNo.get() < MAX_PAGE; + @Override + public void onCompleted() { + terminate(Status.CANCELLED.withDescription("Client closed the partition scan").asRuntimeException()); } - private boolean sendTaskCondition() { - return sendCondition() && (sendTask == null || sendTask.isDone()); + private boolean readCondition() { + return !this.closed.get() && !this.readOver.get() && this.packages.remainingCapacity() > 0; } - private void offer(Iterable data, boolean isVertex) { - ScanResponse.Builder builder = ScanResponse.newBuilder(); - builder.setHeader(okHeader).setSeqNo(nextSeqNo.get()); - if (isVertex) { - builder = builder.setField(vertexField, data); - } else { - builder = builder.setField(edgeField, data); - } - ScanResponse response = builder.build(); - packages.offer(response); - startSend(); + private boolean sendCondition() { + return !this.closed.get() && this.nextSeqNo.get() - this.cltSeqNo.get() < MAX_PAGE; } private void startRead() { - if (readTaskCondition()) { - if (readLock.tryLock()) { - if (readTaskCondition()) { - readTask = executor.submit(rr); + if (!readCondition() || !this.reading.compareAndSet(false, true)) { + return; + } + try { + this.execute.accept(() -> { + try { + read(); + } catch (RuntimeException | Error failure) { + terminate(failure); + } finally { + this.reading.set(false); + if (readCondition()) { + startRead(); + } } - readLock.unlock(); - } + }); + } catch (RuntimeException | Error failure) { + this.reading.set(false); + terminate(failure); } } - private void startSend() { - if (sendTaskCondition()) { - if (sendLock.tryLock()) { - if (sendTaskCondition()) { - sendTask = executor.submit(sr); + private void read() { + synchronized (this.iteratorLock) { + while (readCondition()) { + ArrayList data = new ArrayList<>(BATCH_SIZE); + while (!this.closed.get() && this.leftCount > 0 && + data.size() < BATCH_SIZE && this.iter.hasNext()) { + data.add(this.iter.next()); + this.leftCount--; + } + if (this.closed.get()) { + return; + } + boolean ended = this.leftCount == 0 || !this.iter.hasNext(); + if (!data.isEmpty()) { + Descriptors.FieldDescriptor field = ScanResponse.getDescriptor().findFieldByNumber( + this.scanReq.getScanRequest().getScanType() == ScanType.SCAN_VERTEX ? 3 : 4); + this.packages.add(ScanResponse.newBuilder().setHeader(OK) + .setSeqNo(this.nextSeqNo.get()).setField(field, data).build()); + } + if (ended) { + closeIterator(); + this.readOver.set(true); } - sendLock.unlock(); + startSend(); } } } - @Override - public void onNext(ScanPartitionRequest scanReq) { - if (scanReq.hasScanRequest() && !scanReq.hasReplyRequest()) { - this.scanReq = scanReq; - Request request = scanReq.getScanRequest(); - long rl = request.getLimit(); - leftCount = rl > 0 ? rl : Long.MAX_VALUE; - iter = handler.scan(scanReq); - if (!iter.hasNext()) { - close(); - sender.onCompleted(); - } else { - readTask = executor.submit(rr); - } - } else { - cltSeqNo.getAndIncrement(); - startSend(); + private void startSend() { + if (!sendCondition() || !this.sending.compareAndSet(false, true)) { + return; } - } - - @Override - public void onError(Throwable t) { - close(); - log.warn("receive client error:", t); - } - - @Override - public void onCompleted() { - close(); - } - - private void close() { try { - nextSeqNo.set(0); - if (sendTask != null) { - sendTask.cancel(true); - } - if (readTask != null) { - readTask.cancel(true); - } - readOver.set(true); - iter.close(); - } catch (Exception e) { - log.warn("on Complete with error:", e); - } - } - - Runnable rr = new Runnable() { - @Override - public void run() { - try { - if (readCondition()) { - synchronized (iter) { - while (readCondition()) { - Request r = scanReq.getScanRequest(); - ScanType t = r.getScanType(); - boolean isVertex = t.equals(ScanType.SCAN_VERTEX); - ArrayList data = new ArrayList<>(BATCH_SIZE); - int count = 0; - while (iter.hasNext() && leftCount > -1) { - count++; - leftCount--; - T next = (T) iter.next(); - data.add(next); - if (count >= BATCH_SIZE) { - offer(data, isVertex); - // data.clear(); - break; - } - } - if (!(iter.hasNext() && leftCount > -1)) { - if (data.size() > 0 && - data.size() < BATCH_SIZE) { - offer(data, isVertex); - } - readOver.set(true); - data = null; - //log.warn("scan complete , count: {},time: {}", - // sum, System.currentTimeMillis()); - return; - } + this.execute.accept(() -> { + try { + while (sendCondition()) { + ScanResponse response = this.packages.poll(10, TimeUnit.MILLISECONDS); + if (response != null) { + this.sender.onNext(response.toBuilder().setSeqNo(this.nextSeqNo.get()).build()); + this.nextSeqNo.incrementAndGet(); + startRead(); + } else if (this.readOver.get()) { + terminate(null); + return; + } else { + return; } } + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + terminate(failure); + } catch (RuntimeException | Error failure) { + terminate(failure); + } finally { + this.sending.set(false); + if (sendCondition() && (!this.packages.isEmpty() || this.readOver.get())) { + startSend(); + } } - } catch (Exception e) { - log.warn("read data with error: ", e); - sender.onError(e); - } + }); + } catch (RuntimeException | Error failure) { + this.sending.set(false); + terminate(failure); } - }; + } - Runnable sr = () -> { - while (sendCondition()) { - ScanResponse response; + private void closeIterator() { + synchronized (this.iteratorLock) { + if (this.iter == null || this.iteratorClosed) { + return; + } + this.iteratorClosed = true; try { - if (readOver.get()) { - if ((response = packages.poll()) == null) { - sender.onCompleted(); - } else { - sender.onNext(response); - nextSeqNo.incrementAndGet(); - } - } else { - response = packages.poll(10, - TimeUnit.MILLISECONDS); - if (response != null) { - sender.onNext(response); - nextSeqNo.incrementAndGet(); - startRead(); - } else { - break; - } - } + this.iter.close(); + } catch (RuntimeException | Error failure) { + this.failedCleanup.accept(failure); + throw failure; + } + } + } - } catch (InterruptedException e) { - break; + private void terminate(Throwable failure) { + if (!this.closed.compareAndSet(false, true)) { + return; + } + this.readOver.set(true); + this.packages.clear(); + try { + closeIterator(); + } catch (RuntimeException | Error cleanup) { + if (failure == null) { + failure = cleanup; + } else if (failure != cleanup) { + failure.addSuppressed(cleanup); } } - }; + if (failure == null) { + this.sender.onCompleted(); + } else { + this.sender.onError(failure); + } + } } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java index 285ec3b968..25a07a773e 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/listener/ContextClosedListener.java @@ -17,64 +17,125 @@ package org.apache.hugegraph.store.node.listener; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ThreadPoolExecutor; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryService; import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.lognet.springboot.grpc.context.GRpcServerInitializedEvent; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationListener; import org.springframework.context.event.ContextClosedEvent; +import org.springframework.context.event.EventListener; import org.springframework.stereotype.Service; +import io.grpc.Server; import lombok.extern.slf4j.Slf4j; @Slf4j @Service public class ContextClosedListener implements ApplicationListener { + private final List grpcServers = new CopyOnWriteArrayList<>(); + @Autowired HgStoreStreamImpl storeStream; @Autowired + AggregativeQueryService queryService; + @Autowired TTLCleaner cleaner; + @Autowired + GrpcShutdownBarrier grpcBarrier; + + @EventListener + public void onServerInitialized(GRpcServerInitializedEvent event) { + this.grpcServers.add(event.getServer()); + } @Override public void onApplicationEvent(ContextClosedEvent event) { + // Spring invokes HgStoreNodeService.destroy() after this event. Raft and + // its databases must remain available until local workers have stopped. + this.grpcBarrier.stopAcceptingCalls(); + if (storeStream != null) { + storeStream.stopAcceptingScans(); + } + if (queryService != null) { + queryService.stopAcceptingQueries(); + } + this.grpcServers.forEach(Server::shutdownNow); + if (cleaner != null) { + // The scheduler can create the worker pool while a job is starting. + stopAndWait(cleaner.getScheduler(), "TTL scheduler"); + stopAndWait(cleaner.getExecutor(), "TTL workers"); + } + if (storeStream != null) { + // Cancelled queued scans must run their finally blocks to release iterators. + storeStream.shutdownScans(); + awaitWorkers(storeStream.getRealExecutor(), "scan workers"); + } + if (queryService != null) { + queryService.shutdownQueries(); + awaitWorkers(queryService.getThreadPool(), "aggregate query workers"); + } + boolean interrupted = false; try { - log.info("closing scan threads...."); - if (storeStream != null) { - ThreadPoolExecutor executor = storeStream.getRealExecutor(); - if (executor != null) { + for (Server server : this.grpcServers) { + while (!server.isTerminated()) { try { - executor.shutdownNow(); - } catch (Exception e) { + if (!server.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn("Still waiting for gRPC callbacks before closing databases"); + } + } catch (InterruptedException e) { + interrupted = true; } } } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + this.grpcBarrier.awaitCallbacks(); + if (storeStream != null) { + storeStream.awaitScanCleanup(); + } + log.info("closed gRPC callbacks, scan, aggregate query and TTL workers"); + } - if (cleaner != null) { - ThreadPoolExecutor cleanerExecutor = cleaner.getExecutor(); - if (cleanerExecutor != null) { - try { - cleanerExecutor.shutdownNow(); - } catch (Exception e) { - - } - } - ScheduledExecutorService scheduler = cleaner.getScheduler(); - if (scheduler != null) { - try { - scheduler.shutdownNow(); - } catch (Exception e) { + private static void stopAndWait(ExecutorService executor, String name) { + if (executor == null) { + return; + } + executor.shutdownNow(); + awaitWorkers(executor, name); + } + private static void awaitWorkers(ExecutorService executor, String name) { + if (executor == null) { + return; + } + boolean interrupted = false; + try { + while (!executor.isTerminated()) { + try { + if (!executor.awaitTermination(5, TimeUnit.SECONDS)) { + log.warn("Still waiting for {} to stop before closing databases", name); } + } catch (InterruptedException e) { + // An interrupted shutdown thread must not close databases underneath + // a worker that still owns a native iterator. + interrupted = true; } } - } catch (Exception e) { - log.error("ContextClosedListener: ", e); } finally { - log.info("closed scan threads"); + if (interrupted) { + Thread.currentThread().interrupt(); + } } } - } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/task/TTLCleaner.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/task/TTLCleaner.java index 6e40dc8ba2..ef5ad36a8e 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/task/TTLCleaner.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/task/TTLCleaner.java @@ -301,7 +301,7 @@ private Runnable getTask( LinkedList all = new LinkedList<>(); AtomicBoolean state = new AtomicBoolean(true); AtomicLong partitionCounter = pc.get(id); - while (filter.hasNext() && state.get()) { + while (!Thread.currentThread().isInterrupted() && state.get() && filter.hasNext()) { RocksDBSession.BackendColumn current = filter.next(); byte[] realKey = Arrays.copyOfRange(current.name, 0, current.name.length - Short.BYTES); @@ -315,7 +315,7 @@ private Runnable getTask( all = new LinkedList<>(); } } - if (all.size() > 0 && state.get()) { + if (!Thread.currentThread().isInterrupted() && all.size() > 0 && state.get()) { submitter.submitClean(id, graph, table, all, state, tableCounter, partitionCounter); } diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java index c5b81fe5ca..debca4bea3 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/util/HgChannel.java @@ -37,6 +37,8 @@ @Slf4j public final class HgChannel { + // Recheck close without interrupting a borrowed caller thread. + private static final long CLOSE_POLL_NANOS = TimeUnit.MILLISECONDS.toNanos(100); private final BlockingQueue> queue; private final AtomicBoolean closed = new AtomicBoolean(false); private final long timeoutSeconds; @@ -69,14 +71,27 @@ public boolean send(T t) { if (this.closed.get()) { return false; } - boolean flag; + long timeoutNanos = Math.max(0L, TimeUnit.SECONDS.toNanos(this.timeoutSeconds)); + long started = System.nanoTime(); + long remaining = timeoutNanos; + Supplier supplier = () -> t; try { - flag = this.queue.offer(() -> t, timeoutSeconds, TimeUnit.SECONDS); + do { + if (this.closed.get()) { + return false; + } + if (this.queue.offer(supplier, Math.max(0L, Math.min(remaining, CLOSE_POLL_NANOS)), + TimeUnit.NANOSECONDS)) { + return true; + } + remaining = timeoutNanos - (System.nanoTime() - started); + } while (remaining > 0L); + return false; } catch (InterruptedException e) { + Thread.currentThread().interrupt(); log.error("failed to send a item to chan. cause by: ", t); throw new RuntimeException(e); } - return flag; } } @@ -105,11 +120,29 @@ public T receive(Consumer timeoutCallBack) { if (this.closed.get()) { s = this.queue.poll(); } else { + long timeoutNanos = Math.max(0L, TimeUnit.SECONDS.toNanos(this.timeoutSeconds)); + long started = System.nanoTime(); + long remaining = timeoutNanos; + s = null; try { - s = this.queue.poll(timeoutSeconds, TimeUnit.SECONDS); - } catch (Throwable t) { - log.error("Failed to receive a item from chan. cause by: ", t); - throw new RuntimeException(t); + do { + if (this.closed.get()) { + return null; + } + s = this.queue.poll(Math.max(0L, Math.min(remaining, CLOSE_POLL_NANOS)), + TimeUnit.NANOSECONDS); + if (s != null) { + break; + } + remaining = timeoutNanos - (System.nanoTime() - started); + } while (remaining > 0L); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.error("Failed to receive a item from chan. cause by: ", e); + throw new RuntimeException(e); + } + if (s == null && this.closed.get()) { + return null; } if (s == null) { if (timeoutCallBack == null) { @@ -136,14 +169,7 @@ public boolean isClosed() { * @throws RuntimeException when fail to close the chan */ public void close() { - if (this.closed.get()) { - return; - } this.closed.set(true); - this.queue.offer(() -> null); - Thread.yield(); - this.queue.poll(); - } } diff --git a/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java b/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java index dca8179308..7e23370cf6 100644 --- a/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java +++ b/hugegraph-store/hg-store-rocksdb/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIterator.java @@ -41,6 +41,7 @@ public class RocksDBScanIterator implements ScanIterator { private final AtomicBoolean closed = new AtomicBoolean(false); private final RocksDBSession.RefCounter iterReference; private final Consumer closeOp; + private Throwable closeFailure; private byte[] key; private boolean matched; @@ -225,13 +226,26 @@ private boolean filter(byte[] v) { } @Override - public void close() { + public synchronized void close() { if (!this.closed.getAndSet(true)) { - if (this.rawIt.isOwningHandle()) { - this.rawIt.close(); + try { + if (this.rawIt.isOwningHandle()) { + this.rawIt.close(); + } + this.closeOp.accept(true); + this.iterReference.release(); + } catch (RuntimeException | Error failure) { + this.closeFailure = failure; + throw failure; } - this.closeOp.accept(true); - this.iterReference.release(); + } + // Automatic exhaustion can initiate close before its owner enters finally. + // A repeated close must wait for that release and retain any failure, without retrying JNI. + if (this.closeFailure instanceof RuntimeException) { + throw (RuntimeException) this.closeFailure; + } + if (this.closeFailure instanceof Error) { + throw (Error) this.closeFailure; } } diff --git a/hugegraph-store/hg-store-test/pom.xml b/hugegraph-store/hg-store-test/pom.xml index 9804435aba..bf6ee0bb70 100644 --- a/hugegraph-store/hg-store-test/pom.xml +++ b/hugegraph-store/hg-store-test/pom.xml @@ -178,6 +178,22 @@ + + org.apache.hugegraph + hg-store-node + ${revision} + compile + + + org.springframework + spring-web + + + org.springframework + spring-webmvc + + + diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java new file mode 100644 index 0000000000..c66c8532b4 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/rocksdb/access/RocksDBScanIteratorCloseTest.java @@ -0,0 +1,444 @@ +/* + * 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.rocksdb.access; + +import static org.junit.Assert.assertArrayEquals; +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 static org.junit.Assert.fail; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; +import java.util.function.Consumer; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; +import org.apache.hugegraph.store.business.InnerKeyFilter; +import org.apache.hugegraph.store.business.MultiPartitionIterator; +import org.apache.hugegraph.store.term.Bits; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.InOrder; +import org.rocksdb.Options; +import org.rocksdb.RocksDB; +import org.rocksdb.RocksIterator; + +public class RocksDBScanIteratorCloseTest { + + @Rule + public TemporaryFolder directory = new TemporaryFolder(); + + @Test + @SuppressWarnings("unchecked") + public void testRealNativeEmptyIteratorRetainsOneTimeCallbackFailure() throws Exception { + RocksDB.loadLibrary(); + try (Options options = new Options().setCreateIfMissing(true); + RocksDB database = RocksDB.open(options, this.directory.newFolder("empty").getAbsolutePath())) { + RocksIterator raw = database.newIterator(); + RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + Consumer closeOp = mock(Consumer.class); + IllegalStateException failure = new IllegalStateException("close callback failed once"); + doThrow(failure).doNothing().when(closeOp).accept(true); + RocksDBScanIterator iterator = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, reference, closeOp); + + assertSame(failure, assertThrows(IllegalStateException.class, iterator::hasNext)); + assertFalse("JNI must release the native iterator before callback failure", raw.isOwningHandle()); + assertSame(failure, assertThrows(IllegalStateException.class, iterator::close)); + assertSame(failure, assertThrows(IllegalStateException.class, iterator::close)); + verify(closeOp).accept(true); + verify(reference, never()).release(); + } + } + + @Test + @SuppressWarnings("unchecked") + public void testRealNativeExhaustedIteratorRetainsOneTimeReferenceReleaseError() throws Exception { + RocksDB.loadLibrary(); + try (Options options = new Options().setCreateIfMissing(true); + RocksDB database = RocksDB.open(options, this.directory.newFolder("exhausted").getAbsolutePath())) { + database.put(new byte[]{1}, new byte[]{2}); + RocksIterator raw = database.newIterator(); + RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + Consumer closeOp = mock(Consumer.class); + AssertionError failure = new AssertionError("reference release failed once"); + doThrow(failure).doNothing().when(reference).release(); + RocksDBScanIterator iterator = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, reference, closeOp); + + assertTrue(iterator.hasNext()); + iterator.next(); + assertSame(failure, assertThrows(AssertionError.class, iterator::hasNext)); + assertFalse("JNI must release the native iterator before reference failure", raw.isOwningHandle()); + assertSame(failure, assertThrows(AssertionError.class, iterator::close)); + assertSame(failure, assertThrows(AssertionError.class, iterator::close)); + InOrder order = inOrder(closeOp, reference); + order.verify(closeOp).accept(true); + order.verify(reference).release(); + } + } + + @Test + public void testEmptyIteratorRetainsOneTimeNativeCloseFailure() { + Fixture fixture = new Fixture(); + IllegalStateException failure = new IllegalStateException("native close failed once"); + doThrow(failure).doNothing().when(fixture.raw).close(); + + assertAutomaticFailureIsRetained(fixture, failure); + verify(fixture.closeOp, never()).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testExhaustedIteratorRetainsOneTimeNativeCloseError() { + Fixture fixture = new Fixture(); + when(fixture.raw.isValid()).thenReturn(true, false); + when(fixture.raw.key()).thenReturn(new byte[]{1}); + when(fixture.raw.value()).thenReturn(new byte[]{2}); + AssertionError failure = new AssertionError("native close failed once"); + doThrow(failure).doNothing().when(fixture.raw).close(); + + assertTrue(fixture.iterator.hasNext()); + fixture.iterator.next(); + assertAutomaticFailureIsRetained(fixture, failure); + verify(fixture.closeOp, never()).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testAutomaticCloseRetainsOneTimeCallbackFailure() { + Fixture fixture = new Fixture(); + IllegalStateException failure = new IllegalStateException("close callback failed once"); + doThrow(failure).doNothing().when(fixture.closeOp).accept(true); + + assertAutomaticFailureIsRetained(fixture, failure); + InOrder order = inOrder(fixture.raw, fixture.closeOp); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + verify(fixture.reference, never()).release(); + } + + @Test + public void testAutomaticCloseRetainsOneTimeReferenceReleaseError() { + Fixture fixture = new Fixture(); + AssertionError failure = new AssertionError("reference release failed once"); + doThrow(failure).doNothing().when(fixture.reference).release(); + + assertAutomaticFailureIsRetained(fixture, failure); + InOrder order = inOrder(fixture.raw, fixture.closeOp, fixture.reference); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + order.verify(fixture.reference).release(); + } + + @Test + public void testSuccessfulCloseIsIdempotentAndPreservesReleaseOrder() { + Fixture fixture = new Fixture(); + + fixture.iterator.close(); + fixture.iterator.close(); + assertFalse(fixture.iterator.hasNext()); + InOrder order = inOrder(fixture.raw, fixture.closeOp, fixture.reference); + order.verify(fixture.raw).close(); + order.verify(fixture.closeOp).accept(true); + order.verify(fixture.reference).release(); + verify(fixture.raw).close(); + verify(fixture.closeOp).accept(true); + verify(fixture.reference).release(); + } + + @Test(timeout = 5000) + public void testConcurrentCloseWaitsForSuccessfulRelease() throws Exception { + assertConcurrentCloseWaits(null); + } + + @Test(timeout = 5000) + public void testConcurrentCloseWaitsAndReportsTheSameFailure() throws Exception { + assertConcurrentCloseWaits(new IllegalStateException("close callback failed once")); + } + + @Test + public void testAllInnerKeyConstructorsDeferReadingUntilOwnershipTransfer() { + for (int constructor = 0; constructor < 3; constructor++) { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + InnerKeyFilter filter = filter(raw, constructor); + + verify(raw, never()).hasNext(); + verify(raw, never()).next(); + assertTrue(filter.hasNext()); + verify(raw).hasNext(); + verify(raw).next(); + filter.close(); + verify(raw).close(); + } + } + + @Test + public void testPlainInnerKeyFilterStripsGraphAndCodeAndPreservesPrefetch() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, true, false); + when(raw.next()).thenReturn(column(1, 5), column(2, 19)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertTrue(filter.hasNext()); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{1}, filter.next().name); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{2}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw, times(2)).next(); + verify(raw, times(3)).hasNext(); + filter.close(); + } + + @Test + public void testCodeFilterKeepsHashSuffix() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(3, 15)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw, true); + + assertArrayEquals(new byte[]{3, 0, 15}, filter.next().name); + assertFalse(filter.hasNext()); + filter.close(); + } + + @Test + public void testCodeRangeFilterPreservesInclusiveLowerAndExclusiveUpperBounds() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, true, true, true, false); + when(raw.next()).thenReturn(column(1, 5), column(2, 10), column(3, 20), column(4, 19)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw, 10, 20); + + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{2, 0, 10}, filter.next().name); + assertTrue(filter.hasNext()); + assertArrayEquals(new byte[]{4, 0, 19}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw, times(4)).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterNextBeforeHasNextPreservesTheFirstColumn() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(7, 15)); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertArrayEquals(new byte[]{7}, filter.next().name); + assertFalse(filter.hasNext()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterIsValidUsesTheExistingPrefetchedDelegatePosition() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + when(raw.isValid()).thenAnswer(invocation -> { + verify(raw).next(); + return false; + }); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertFalse(filter.isValid()); + assertTrue(filter.hasNext()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterCountKeepsTheExistingPrefetchedDelegatePosition() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true); + when(raw.next()).thenReturn(column(1, 15)); + when(raw.count()).thenAnswer(invocation -> { + verify(raw).next(); + return 7L; + }); + InnerKeyFilter filter = new InnerKeyFilter<>(raw); + + assertEquals(7L, filter.count()); + verify(raw).next(); + filter.close(); + } + + @Test + public void testInnerKeyFilterCloseBeforeUseDoesNotPrefetch() { + for (int constructor = 0; constructor < 3; constructor++) { + ScanIterator raw = mock(ScanIterator.class); + InnerKeyFilter filter = filter(raw, constructor); + + filter.close(); + assertFalse(filter.hasNext()); + verify(raw, never()).hasNext(); + verify(raw, never()).next(); + verify(raw).close(); + } + } + + @Test + public void testInnerKeyFilterPreservesPartitionResumeKey() { + ScanIterator raw = mock(ScanIterator.class); + when(raw.hasNext()).thenReturn(true, false); + when(raw.next()).thenReturn(column(9, 15)); + AtomicInteger supplied = new AtomicInteger(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, position) -> { + assertEquals(Integer.valueOf(2), id); + assertArrayEquals(new byte[]{9}, position); + supplied.incrementAndGet(); + return new InnerKeyFilter<>(raw); + }); + iterator.seek(ByteBuffer.allocate(Integer.BYTES + 1).putInt(2).put((byte) 9).array()); + + assertTrue(iterator.hasNext()); + BackendColumn result = iterator.next(); + assertArrayEquals(new byte[]{9}, result.name); + assertArrayEquals(ByteBuffer.allocate(Integer.BYTES).putInt(2).array(), iterator.position()); + assertFalse(iterator.hasNext()); + assertEquals(1, supplied.get()); + verify(raw).close(); + } + + private static InnerKeyFilter filter(ScanIterator raw, int constructor) { + if (constructor == 0) { + return new InnerKeyFilter<>(raw); + } + if (constructor == 1) { + return new InnerKeyFilter<>(raw, true); + } + return new InnerKeyFilter<>(raw, 10, 20); + } + + private static BackendColumn column(int key, int code) { + byte[] name = new byte[Short.BYTES + 1 + Short.BYTES]; + Bits.putShort(name, 0, 1); + name[Short.BYTES] = (byte) key; + Bits.putShort(name, name.length - Short.BYTES, code); + return BackendColumn.of(name, new byte[]{(byte) key}); + } + + private static void assertAutomaticFailureIsRetained(Fixture fixture, Throwable failure) { + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::hasNext)); + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::close)); + assertSame(failure, assertThrows(Throwable.class, fixture.iterator::close)); + verify(fixture.raw).close(); + } + + private static void assertConcurrentCloseWaits(Throwable failure) throws Exception { + Fixture fixture = new Fixture(); + CountDownLatch nativeCloseStarted = new CountDownLatch(1); + CountDownLatch allowNativeClose = new CountDownLatch(1); + doAnswer(invocation -> { + nativeCloseStarted.countDown(); + allowNativeClose.await(); + return null; + }).when(fixture.raw).close(); + if (failure != null) { + doThrow(failure).doNothing().when(fixture.closeOp).accept(true); + } + FutureTask first = new FutureTask<>(() -> { + fixture.iterator.close(); + return null; + }); + FutureTask second = new FutureTask<>(() -> { + fixture.iterator.close(); + return null; + }); + Thread firstThread = new Thread(first, "test-first-native-close"); + Thread secondThread = new Thread(second, "test-second-native-close"); + firstThread.setDaemon(true); + secondThread.setDaemon(true); + try { + firstThread.start(); + assertTrue(nativeCloseStarted.await(1, TimeUnit.SECONDS)); + secondThread.start(); + awaitBlockedClose(secondThread); + assertFalse("second close cannot claim release before native close finishes", second.isDone()); + verify(fixture.reference, never()).release(); + allowNativeClose.countDown(); + if (failure == null) { + first.get(2, TimeUnit.SECONDS); + second.get(2, TimeUnit.SECONDS); + verify(fixture.reference).release(); + } else { + assertSame(failure, assertThrows(ExecutionException.class, + () -> first.get(2, TimeUnit.SECONDS)).getCause()); + assertSame(failure, assertThrows(ExecutionException.class, + () -> second.get(2, TimeUnit.SECONDS)).getCause()); + verify(fixture.reference, never()).release(); + } + verify(fixture.raw).close(); + verify(fixture.closeOp).accept(true); + } finally { + allowNativeClose.countDown(); + firstThread.join(1000); + secondThread.join(1000); + } + } + + private static void awaitBlockedClose(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (System.nanoTime() < deadline) { + if (thread.getState() == Thread.State.BLOCKED) { + return; + } + if (thread.getState() == Thread.State.TERMINATED) { + break; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("concurrent close must wait for the first resource release"); + } + + private static final class Fixture { + + private final RocksIterator raw = mock(RocksIterator.class); + private final RocksDBSession.RefCounter reference = mock(RocksDBSession.RefCounter.class); + @SuppressWarnings("unchecked") + private final Consumer closeOp = mock(Consumer.class); + private final RocksDBScanIterator iterator; + + private Fixture() { + when(this.raw.isOwningHandle()).thenReturn(true); + this.iterator = new RocksDBScanIterator<>(this.raw, null, null, + ScanIterator.Trait.SCAN_ANY, this.reference, this.closeOp); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java index 78445e5934..f47b571f2d 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java @@ -19,6 +19,8 @@ import org.apache.hugegraph.store.client.grpc.AbstractGrpcClientTest; import org.apache.hugegraph.store.client.grpc.KvPageScannerTest; +import org.apache.hugegraph.store.client.query.CommonKvStreamObserverTest; +import org.apache.hugegraph.store.client.query.MultiStreamIteratorTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -32,7 +34,9 @@ NodeTxExecutorTest.class, KvPageScannerTest.class, NodeTxSessionProxyTest.class, - OrderedKvIteratorTest.class + OrderedKvIteratorTest.class, + CommonKvStreamObserverTest.class, + MultiStreamIteratorTest.class }) public class ClientSuiteTest { } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java new file mode 100644 index 0000000000..633e330688 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/CommonKvStreamObserverTest.java @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.store.HgKvIterator; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class CommonKvStreamObserverTest { + + private static final Logger LOG = LoggerFactory.getLogger(CommonKvStreamObserverTest.class); + + @Test + public void testCompletedEmptyStreamSkipsTimedPoll() throws Exception { + RecordingQueue queue = new RecordingQueue(); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onCompleted(); + StreamKvIterator stream = stream(observer); + assertFalse(stream.hasNext()); + stream.close(); + assertEquals("completed empty stream must not enter timed poll", 0, queue.emptyTimedPolls.get()); + } + + @Test + public void testCompletedSingleStreamAndCloseSkipEmptyTimedPoll() throws Exception { + RecordingQueue queue = new RecordingQueue(); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onNext(Collections.singletonList(1)); + observer.onCompleted(); + StreamKvIterator stream = stream(observer); + assertTrue(stream.hasNext()); + assertEquals(Integer.valueOf(1), stream.next()); + assertFalse(stream.hasNext()); + assertFalse(stream.hasNext()); + stream.close(); + assertEquals("exhausted terminal stream must not enter empty timed poll", 0, queue.emptyTimedPolls.get()); + } + + @Test + public void testCompletedThreeStreamsMergeAndCloseSkipEmptyTimedPoll() throws Exception { + List queues = new ArrayList<>(); + MultiStreamIterator merged = merged(queues); + List actual = drain(merged); + merged.close(); + assertEquals(Arrays.asList(1, 2, 3), actual); + assertEquals("merged terminal streams and close must not enter empty timed poll", 0, + queues.stream().mapToInt(queue -> queue.emptyTimedPolls.get()).sum()); + } + + @Test(timeout = 5000) + public void testFinalBatchPublishedAtTerminalBoundaryIsNotLost() throws Exception { + CountDownLatch boundary = new CountDownLatch(1); + CountDownLatch published = new CountDownLatch(1); + GatedQueue queue = new GatedQueue(boundary, published); + CommonKvStreamObserver, Integer> observer = observer(queue); + StreamKvIterator stream = stream(observer); + AtomicReference failure = new AtomicReference<>(); + Thread producer = new Thread(() -> { + try { + assertTrue("consumer must reach terminal/queue boundary", boundary.await(2, TimeUnit.SECONDS)); + observer.onNext(Collections.singletonList(7)); + observer.onCompleted(); + } catch (Throwable error) { + failure.set(error); + } finally { + published.countDown(); + } + }, "publish-final-client-batch"); + producer.start(); + try { + assertTrue("final batch must remain visible when terminal state changes", stream.hasNext()); + assertEquals(Integer.valueOf(7), stream.next()); + assertFalse(stream.hasNext()); + assertEquals(null, failure.get()); + } finally { + boundary.countDown(); + published.countDown(); + producer.join(2000); + assertFalse("producer must finish", producer.isAlive()); + stream.close(); + } + } + + @Test(timeout = 20000) + public void testRealClientTerminalBenchmarks() throws Exception { + long[] single = new long[10]; + long[] merged = new long[10]; + long[] closed = new long[10]; + for (int i = 0; i < single.length; i++) { + CommonKvStreamObserver, Integer> observer = observer(new RecordingQueue()); + observer.onNext(Collections.singletonList(1)); + observer.onCompleted(); + StreamKvIterator stream = stream(observer); + assertTrue(stream.hasNext()); + assertEquals(Integer.valueOf(1), stream.next()); + long started = System.nanoTime(); + assertFalse(stream.hasNext()); + single[i] = System.nanoTime() - started; + stream.close(); + + MultiStreamIterator aggregate = merged(new ArrayList<>()); + started = System.nanoTime(); + assertEquals(Arrays.asList(1, 2, 3), drain(aggregate)); + merged[i] = System.nanoTime() - started; + started = System.nanoTime(); + aggregate.close(); + closed[i] = System.nanoTime() - started; + } + LOG.info("Real client terminal median micros: single={}, merged3={}, close={}", + medianMicros(single), medianMicros(merged), medianMicros(closed)); + } + + private static long medianMicros(long[] samples) { + Arrays.sort(samples); + return TimeUnit.NANOSECONDS.toMicros(samples[samples.length / 2]); + } + + private static List drain(HgKvIterator stream) { + List values = new ArrayList<>(); + while (stream.hasNext()) { + values.add(stream.next()); + } + return values; + } + + private static MultiStreamIterator merged(List queues) throws Exception { + List> streams = new ArrayList<>(); + for (int value : Arrays.asList(1, 2, 3)) { + RecordingQueue queue = new RecordingQueue(); + queues.add(queue); + CommonKvStreamObserver, Integer> observer = observer(queue); + observer.onNext(Collections.singletonList(value)); + observer.onCompleted(); + streams.add(stream(observer)); + } + return new MultiStreamIterator<>(streams); + } + + private static StreamKvIterator stream(CommonKvStreamObserver, Integer> observer) { + return new StreamKvIterator<>(ignored -> observer.clear(), observer::consume); + } + + private static CommonKvStreamObserver, Integer> observer(RecordingQueue queue) throws Exception { + CommonKvStreamObserver, Integer> observer = + new CommonKvStreamObserver<>(List::iterator, ignored -> ResultState.FINISHED); + observer.setRequestSender(ignored -> { }); + observer.setTransferComplete(ignored -> { }); + Field field = CommonKvStreamObserver.class.getDeclaredField("queue"); + field.setAccessible(true); + field.set(observer, queue); + return observer; + } + + private static class RecordingQueue extends LinkedBlockingQueue> { + + private static final long serialVersionUID = 1L; + final AtomicInteger emptyTimedPolls = new AtomicInteger(); + + @Override + public Iterator poll(long timeout, TimeUnit unit) throws InterruptedException { + if (super.isEmpty()) { + this.emptyTimedPolls.incrementAndGet(); + } + return super.poll(timeout, unit); + } + } + + private static final class GatedQueue extends RecordingQueue { + + private static final long serialVersionUID = 1L; + private final CountDownLatch boundary; + private final CountDownLatch published; + private final AtomicBoolean armed = new AtomicBoolean(true); + + private GatedQueue(CountDownLatch boundary, CountDownLatch published) { + this.boundary = boundary; + this.published = published; + } + + private void publishAtBoundary() { + if (this.armed.compareAndSet(true, false)) { + this.boundary.countDown(); + try { + assertTrue("final batch must be published", this.published.await(2, TimeUnit.SECONDS)); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError(error); + } + } + } + + @Override + public boolean isEmpty() { + boolean empty = super.isEmpty(); + this.publishAtBoundary(); + return empty; + } + + @Override + public Iterator poll(long timeout, TimeUnit unit) throws InterruptedException { + this.publishAtBoundary(); + return super.poll(timeout, unit); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java new file mode 100644 index 0000000000..675c6cb9b9 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/query/MultiStreamIteratorTest.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; + +import java.util.List; + +import org.apache.hugegraph.store.HgKvIterator; +import org.junit.Test; + +public class MultiStreamIteratorTest { + + @Test + public void testCloseCancelsEveryChildBeforeIterationAndIsTerminal() { + HgKvIterator first = mock(HgKvIterator.class); + HgKvIterator second = mock(HgKvIterator.class); + MultiStreamIterator composite = new MultiStreamIterator<>(List.of(first, second, first)); + composite.close(); + composite.close(); + assertFalse(composite.hasNext()); + verify(first).close(); + verify(second).close(); + verifyNoMoreInteractions(first, second); + try { + composite.next(); + fail("Closed iterator must not resume a child stream"); + } catch (java.util.NoSuchElementException expected) { + // Closed is terminal. + } + } + + @Test + public void testCloseAttemptsAllChildrenAndPreservesPrimaryFailure() { + HgKvIterator first = mock(HgKvIterator.class); + HgKvIterator second = mock(HgKvIterator.class); + HgKvIterator third = mock(HgKvIterator.class); + RuntimeException primary = new IllegalStateException("first"); + RuntimeException later = new IllegalArgumentException("second"); + doThrow(primary).when(first).close(); + doThrow(later).when(second).close(); + MultiStreamIterator composite = new MultiStreamIterator<>(List.of(first, second, third)); + try { + composite.close(); + fail("Cleanup failure must be visible"); + } catch (RuntimeException error) { + assertSame(primary, error); + assertEquals(1, error.getSuppressed().length); + assertSame(later, error.getSuppressed()[0]); + } + verify(first).close(); + verify(second).close(); + verify(third).close(); + assertFalse(composite.hasNext()); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryShutdownTest.java new file mode 100644 index 0000000000..53343e06a3 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryShutdownTest.java @@ -0,0 +1,1541 @@ +/* + * 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.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; + +import org.apache.hugegraph.rocksdb.access.RocksDBScanIterator; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.store.business.InnerKeyFilter; +import org.apache.hugegraph.store.business.MultiPartitionIterator; +import org.apache.hugegraph.store.grpc.query.QueryRequest; +import org.apache.hugegraph.store.grpc.query.QueryResponse; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.query.model.PipelineResult; +import org.apache.hugegraph.store.node.grpc.query.model.QueryPlan; +import org.apache.hugegraph.store.node.listener.ContextClosedListener; +import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.junit.Test; +import org.rocksdb.RocksIterator; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; + +public class AggregativeQueryShutdownTest { + + @Test(timeout = 5000) + public void testContextCloseContinuesAfterCancellationResponseFailure() throws Exception { + assertContextCloseContinuesAfterCancellationFailure( + new IllegalStateException("response completion failed")); + } + + @Test(timeout = 5000) + public void testContextCloseContinuesAfterCancellationResponseError() throws Exception { + assertContextCloseContinuesAfterCancellationFailure(new AssertionError("response completion failed")); + } + + @Test(timeout = 5000) + public void testContextWaitsForIteratorFinallyAfterRpcCancellation() throws Exception { + ThreadPoolExecutor pool = pool(1); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch closingIterator = new CountDownLatch(1); + CountDownLatch allowClose = new CountDownLatch(1); + AtomicBoolean databaseClosed = new AtomicBoolean(); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException ignored) { + return false; + } + return false; + }); + doAnswer(invocation -> { + closingIterator.countDown(); + awaitUninterruptibly(allowClose); + return null; + }).when(iterator).close(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + context.getBeanFactory().registerSingleton("queryService", service); + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("cleaner", mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean( + "database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + StreamObserver request = service.query(sender()); + request.onNext(QueryRequest.getDefaultInstance()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + request.onError(Status.CANCELLED.asRuntimeException()); + assertTrue(closingIterator.await(1, TimeUnit.SECONDS)); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return Thread.currentThread().isInterrupted(); + }); + Thread closer = start(closing); + awaitShutdownWait(closer); + closer.interrupt(); + awaitShutdownWait(closer); + assertFalse("RPC cancellation must not release a still-closing iterator", closing.isDone()); + assertFalse(databaseClosed.get()); + allowClose.countDown(); + assertTrue("context close must preserve interruption", closing.get(2, TimeUnit.SECONDS)); + assertTrue(databaseClosed.get()); + verify(iterator).close(); + } finally { + allowClose.countDown(); + context.close(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testPartitionTimeoutDoesNotClearPlanOrFinishBeforeChildFinally() throws Exception { + ThreadPoolExecutor pool = pool(2); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch releaseChild = new CountDownLatch(1); + CountDownLatch parentReturned = new CountDownLatch(1); + CountDownLatch childClosed = new CountDownLatch(1); + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + awaitUninterruptibly(releaseChild); + return false; + }); + doAnswer(invocation -> { + childClosed.countDown(); + return null; + }).when(child).close(); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + AggregativeQueryService service = new AggregativeQueryService(pool, 50, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = super.newObserver(sender); + return fixture(observer, iterator, plan, parentReturned); + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("timeout-query").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + assertTrue("parent timeout must not wait on its own executor's queued children", + parentReturned.await(1, TimeUnit.SECONDS)); + assertTrue("a live client must receive the timeout error", response.received.await(1, TimeUnit.SECONDS)); + assertError(response, "timeout-query"); + assertEquals("response must not complete before child finally", 0, response.completed.get()); + verify(plan, never()).clear(); + verify(iterator, never()).close(); + FutureTask closing = close(service); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + releaseChild.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertTrue(childClosed.await(1, TimeUnit.SECONDS)); + verify(plan).clear(); + verify(iterator).close(); + assertTrue(pool.awaitTermination(1, TimeUnit.SECONDS)); + assertEquals(1, response.completed.get()); + } finally { + releaseChild.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testQueuedParentsAreDrainedAndNewQueriesAreRejected() throws Exception { + ThreadPoolExecutor pool = pool(1); + CountDownLatch blockerStarted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ScanIterator iterator = mock(ScanIterator.class); + AtomicInteger closed = new AtomicInteger(); + doAnswer(invocation -> { + closed.incrementAndGet(); + return null; + }).when(iterator).close(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + pool.execute(() -> { + blockerStarted.countDown(); + awaitUninterruptibly(release); + }); + try { + assertTrue(blockerStarted.await(1, TimeUnit.SECONDS)); + service.query(sender()).onNext(QueryRequest.getDefaultInstance()); + service.query(sender()).onNext(QueryRequest.getDefaultInstance()); + service.stopAcceptingQueries(); + try { + service.query(sender()); + fail("new queries must be rejected after admission closes"); + } catch (StatusRuntimeException expected) { + assertEquals(Status.Code.UNAVAILABLE, expected.getStatus().getCode()); + } + FutureTask closing = close(service); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + release.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertEquals(2, closed.get()); + verify(iterator, never()).hasNext(); + assertTrue(pool.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + release.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testQueuedPartitionRunsCleanupAfterParentTimeout() throws Exception { + CountDownLatch childQueued = new CountDownLatch(1); + CountDownLatch releaseSubmission = new CountDownLatch(1); + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + super.execute(task); + if (submissions.incrementAndGet() == 2) { + childQueued.countDown(); + awaitUninterruptibly(releaseSubmission); + } + } + }; + ScanIterator child = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + AggregativeQueryService service = service(pool, iterator, plan, 50); + try { + StreamObserver request = service.query(sender()); + request.onNext(QueryRequest.getDefaultInstance()); + assertTrue(childQueued.await(1, TimeUnit.SECONDS)); + request.onError(Status.CANCELLED.asRuntimeException()); + FutureTask closing = close(service); + start(closing); + releaseSubmission.countDown(); + closing.get(2, TimeUnit.SECONDS); + verify(child).close(); + verify(child, never()).hasNext(); + verify(plan).clear(); + verify(iterator).close(); + } finally { + releaseSubmission.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedPartitionsReleaseAllOpenedIterators() throws Exception { + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + if (submissions.incrementAndGet() > 1) { + throw new RejectedExecutionException("partition queue is full"); + } + super.execute(task); + } + }; + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(first, second)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + CountDownLatch parentReturned = new CountDownLatch(1); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, parentReturned); + } + }; + try { + service.query(sender()).onNext(QueryRequest.getDefaultInstance()); + assertTrue(parentReturned.await(1, TimeUnit.SECONDS)); + FutureTask closing = close(service); + start(closing); + closing.get(2, TimeUnit.SECONDS); + verify(first).close(); + verify(second).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testIdleRpcAndInitializationFailureDoNotBlockShutdown() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = spy(super.newObserver(sender)); + doReturn(iterator).when(observer).getIterator(org.mockito.ArgumentMatchers.any()); + org.mockito.Mockito.doThrow(new IllegalArgumentException("invalid query plan")) + .when(observer).buildPlan(org.mockito.ArgumentMatchers.any()); + return observer; + } + }; + ResponseRecorder failedResponse = new ResponseRecorder(); + try { + service.query(sender()); + try { + service.query(failedResponse).onNext(QueryRequest.getDefaultInstance()); + fail("invalid query must fail initialization"); + } catch (IllegalArgumentException expected) { + assertEquals("invalid query plan", expected.getMessage()); + } + assertEquals("framework must retain the synchronous error", 0, failedResponse.completed.get()); + assertTrue(failedResponse.responses.isEmpty()); + service.shutdownQueries(); + verify(iterator).close(); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testHasNextFailureSendsErrorBeforeSingleCompletion() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenThrow(new IllegalStateException("iterator read failed")); + QueryPlan plan = mock(QueryPlan.class); + when(plan.isEmpty()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("read-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "read-query"); + assertEquals(1, response.completed.get()); + assertEquals(0, response.errors.get()); + service.shutdownQueries(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testPartitionFailureSendsErrorBeforeSingleCompletion() throws Exception { + ThreadPoolExecutor pool = pool(2); + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenThrow(new IllegalStateException("partition read failed")); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("partition-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "partition-query"); + assertEquals(1, response.completed.get()); + service.shutdownQueries(); + verify(child).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testAcceptedChildAndRejectedRemainingChildrenRetainCleanupOwnership() throws Exception { + CountDownLatch childStarted = new CountDownLatch(1); + CountDownLatch releaseChild = new CountDownLatch(1); + CountDownLatch parentReturned = new CountDownLatch(1); + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(2, 2, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + int submitted = submissions.incrementAndGet(); + if (submitted >= 3) { + throw new RejectedExecutionException("second partition rejected"); + } + super.execute(task); + if (submitted == 2) { + awaitUninterruptibly(childStarted); + } + } + }; + ScanIterator accepted = mock(ScanIterator.class); + when(accepted.hasNext()).thenAnswer(invocation -> { + childStarted.countDown(); + awaitUninterruptibly(releaseChild); + return false; + }); + ScanIterator rejected = mock(ScanIterator.class); + ScanIterator notSubmitted = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(accepted, rejected, notSubmitted)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, parentReturned); + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("mixed-query").build()); + assertTrue(childStarted.await(1, TimeUnit.SECONDS)); + assertTrue(parentReturned.await(1, TimeUnit.SECONDS)); + assertError(response, "mixed-query"); + assertEquals(0, response.completed.get()); + verify(rejected).close(); + verify(notSubmitted).close(); + verify(accepted, never()).close(); + verify(iterator, never()).close(); + verify(plan, never()).clear(); + FutureTask closing = close(service); + awaitShutdownWait(start(closing)); + assertFalse("accepted child's finally still owns the query", closing.isDone()); + releaseChild.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertEquals(1, response.completed.get()); + verify(accepted).close(); + verify(iterator).close(); + verify(plan).clear(); + } finally { + releaseChild.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testNormalIdleHalfCloseCompletesOnceAndExternalErrorNeverCompletesNormally() { + ThreadPoolExecutor pool = pool(1); + AggregativeQueryService service = service(pool, mock(ScanIterator.class), new QueryPlan(), 500); + try { + ResponseRecorder normal = new ResponseRecorder(); + StreamObserver normalRequest = service.query(normal); + normalRequest.onCompleted(); + normalRequest.onCompleted(); + assertEquals(1, normal.completed.get()); + ResponseRecorder failed = new ResponseRecorder(); + StreamObserver failedRequest = service.query(failed); + failedRequest.onError(Status.CANCELLED.asRuntimeException()); + failedRequest.onCompleted(); + assertEquals(0, failed.completed.get()); + service.shutdownQueries(); + assertEquals(1, normal.completed.get()); + assertEquals(0, failed.completed.get()); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testActiveNormalHalfCloseCompletesAfterWorkerCleanupOnce() throws Exception { + ThreadPoolExecutor pool = pool(1); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ScanIterator iterator = mock(ScanIterator.class); + AtomicInteger reads = new AtomicInteger(); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[]{1}, new byte[]{2})); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return reads.incrementAndGet() == 1; + }); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + try { + StreamObserver request = service.query(response); + request.onNext(QueryRequest.newBuilder().setQueryId("half-close-query").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + request.onCompleted(); + assertFalse(response.finished.await(50, TimeUnit.MILLISECONDS)); + release.countDown(); + assertTrue("normal half-close must complete after cleanup", response.finished.await(1, TimeUnit.SECONDS)); + request.onCompleted(); + service.shutdownQueries(); + assertEquals(1, response.completed.get()); + assertEquals(0, response.errors.get()); + assertEquals(1, response.responses.size()); + assertTrue(response.responses.get(0).getIsFinished()); + assertTrue(response.responses.get(0).getIsOk()); + assertEquals(1, response.responses.get(0).getDataCount()); + assertEquals(com.google.protobuf.ByteString.copyFrom(new byte[]{2}), + response.responses.get(0).getData(0).getValue()); + verify(iterator).close(); + } finally { + release.countDown(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testHalfCloseWithoutEnoughFeedbackReportsError() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[]{1}, new byte[]{2})); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 5000); + try { + StreamObserver request = service.query(response); + request.onNext(QueryRequest.newBuilder().setQueryId("feedback-half-close").build()); + request.onCompleted(); + assertTrue(response.finished.await(2, TimeUnit.SECONDS)); + assertEquals(17, response.responses.size()); + QueryResponse error = response.responses.get(16); + assertFalse(error.getIsOk()); + assertTrue(error.getMessage().contains("without enough feedback")); + assertEquals(1, response.completed.get()); + verify(iterator).close(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 15000) + public void testTransportCancellationAfterHalfCloseReleasesWorker() throws Exception { + assertTransportCancellationAfterHalfClose(false); + } + + @Test(timeout = 15000) + public void testDeadlineAfterHalfCloseReleasesWorker() throws Exception { + assertTransportCancellationAfterHalfClose(true); + } + + private static void assertTransportCancellationAfterHalfClose(boolean deadline) throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator source = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + CountDownLatch fallbackRelease = new CountDownLatch(1); + when(source.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + fallbackRelease.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }); + doAnswer(invocation -> { + released.countDown(); + return null; + }).when(source).close(); + AggregativeQueryService service = new AggregativeQueryService(pool, 5000, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), source, new QueryPlan(), null); + } + }; + String name = io.grpc.inprocess.InProcessServerBuilder.generateName(); + io.grpc.Server server = io.grpc.inprocess.InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + io.grpc.ManagedChannel channel = io.grpc.inprocess.InProcessChannelBuilder.forName(name) + .directExecutor().build(); + io.grpc.CallOptions options = deadline ? io.grpc.CallOptions.DEFAULT.withDeadlineAfter(2, TimeUnit.SECONDS) : + io.grpc.CallOptions.DEFAULT; + io.grpc.ClientCall call = channel.newCall( + org.apache.hugegraph.store.grpc.query.QueryServiceGrpc.getQueryMethod(), options); + CountDownLatch terminated = new CountDownLatch(1); + java.util.concurrent.atomic.AtomicReference terminal = + new java.util.concurrent.atomic.AtomicReference<>(); + try { + call.start(new io.grpc.ClientCall.Listener() { + @Override + public void onClose(Status status, io.grpc.Metadata trailers) { + terminal.set(status); + terminated.countDown(); + } + }, new io.grpc.Metadata()); + call.request(1); + call.sendMessage(QueryRequest.newBuilder().setQueryId("half-close-cancel").build()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + call.halfClose(); + if (!deadline) { + call.cancel("cancel after half-close", null); + } + assertTrue(terminated.await(3, TimeUnit.SECONDS)); + assertEquals(deadline ? Status.Code.DEADLINE_EXCEEDED : Status.Code.CANCELLED, + terminal.get().getCode()); + assertTrue("transport cancellation must interrupt the worker after half-close", + interrupted.await(1, TimeUnit.SECONDS)); + assertTrue("transport cancellation must release the source without manual progress", + released.await(1, TimeUnit.SECONDS)); + verify(source).close(); + } finally { + fallbackRelease.countDown(); + call.cancel("test cleanup", null); + channel.shutdownNow(); + server.shutdownNow(); + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testRealQueryClientCloseCancelsEarlyAndPreservesNormalCompletion() throws Exception { + assertClientIteratorClose(true); + assertClientIteratorClose(false); + } + + private static void assertClientIteratorClose(boolean early) throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator source = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + CountDownLatch completed = new CountDownLatch(1); + if (early) { + when(source.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException expected) { + Thread.currentThread().interrupt(); + } + return false; + }); + } + doAnswer(invocation -> { + released.countDown(); + return null; + }).when(source).close(); + java.util.concurrent.atomic.AtomicReference observed = + new java.util.concurrent.atomic.AtomicReference<>(); + AggregativeQueryService service = new AggregativeQueryService(pool, 5000, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + StreamObserver tracked = new StreamObserver() { + @Override + public void onNext(QueryResponse value) { + sender.onNext(value); + } + + @Override + public void onError(Throwable failure) { + sender.onError(failure); + } + + @Override + public void onCompleted() { + sender.onCompleted(); + completed.countDown(); + } + }; + AggregativeQueryObserver observer = fixture(super.newObserver(tracked), source, new QueryPlan(), null); + observed.set(observer); + return observer; + } + }; + String name = io.grpc.inprocess.InProcessServerBuilder.generateName(); + io.grpc.Server server = io.grpc.inprocess.InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + io.grpc.ManagedChannel channel = io.grpc.inprocess.InProcessChannelBuilder.forName(name) + .directExecutor().build(); + org.apache.hugegraph.store.client.query.QueryV2Client.setTestChannel(channel); + try { + org.apache.hugegraph.store.client.query.QueryExecutor executor = + new org.apache.hugegraph.store.client.query.QueryExecutor(null, null, 5000L); + java.lang.reflect.Method getIterator = executor.getClass().getDeclaredMethod( + "getIterator", String.class, QueryRequest.class); + getIterator.setAccessible(true); + org.apache.hugegraph.store.HgKvIterator iterator = + (org.apache.hugegraph.store.HgKvIterator) getIterator.invoke( + executor, "in-process-query-close", QueryRequest.newBuilder() + .setQueryId("client-close").setTable("test").build()); + if (early) { + assertTrue(reading.await(2, TimeUnit.SECONDS)); + FutureTask first = new FutureTask<>(() -> { + iterator.close(); + return null; + }); + Thread closer = start(first); + iterator.close(); + first.get(2, TimeUnit.SECONDS); + closer.join(1000); + } else { + assertFalse(iterator.hasNext()); + iterator.close(); + // A finished batch arrives before transport completion; don't cancel that RPC in teardown. + assertTrue("normal RPC must finish before channel teardown", completed.await(2, TimeUnit.SECONDS)); + } + assertTrue("iterator close must release resources before service shutdown", + released.await(2, TimeUnit.SECONDS)); + if (early) { + org.mockito.ArgumentCaptor failure = + org.mockito.ArgumentCaptor.forClass(Throwable.class); + verify(observed.get()).onError(failure.capture()); + verify(observed.get()).cancel(); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode()); + } else { + verify(observed.get(), never()).onError(org.mockito.ArgumentMatchers.any()); + verify(observed.get(), never()).cancel(); + } + verify(observed.get()).onNext(org.mockito.ArgumentMatchers.any()); + verify(source).close(); + } finally { + channel.shutdownNow(); + server.shutdownNow(); + service.shutdownQueries(); + pool.shutdownNow(); + org.apache.hugegraph.store.client.query.QueryV2Client.setTestChannel(null); + } + } + + @Test(timeout = 5000) + public void testIteratorInitializationFailureKeepsOriginalThrowWithoutNormalCompletion() { + ThreadPoolExecutor pool = pool(1); + IllegalStateException failure = new IllegalStateException("iterator initialization failed"); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = spy(super.newObserver(sender)); + doThrow(failure).when(observer).getIterator(org.mockito.ArgumentMatchers.any()); + return observer; + } + }; + ResponseRecorder response = new ResponseRecorder(); + try { + try { + service.query(response).onNext(QueryRequest.getDefaultInstance()); + fail("synchronous iterator initialization must fail"); + } catch (IllegalStateException expected) { + org.junit.Assert.assertSame(failure, expected); + } + assertEquals(0, response.completed.get()); + assertTrue(response.responses.isEmpty()); + service.shutdownQueries(); + } finally { + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParentSubmissionRejectionKeepsOriginalThrowWithoutNormalCompletion() { + ThreadPoolExecutor pool = pool(1); + pool.shutdown(); + ScanIterator iterator = mock(ScanIterator.class); + QueryPlan plan = mock(QueryPlan.class); + AggregativeQueryService service = service(pool, iterator, plan, 500); + ResponseRecorder response = new ResponseRecorder(); + try { + service.query(response).onNext(QueryRequest.getDefaultInstance()); + fail("synchronous parent submission must be rejected"); + } catch (RejectedExecutionException expected) { + assertEquals(0, response.completed.get()); + assertTrue(response.responses.isEmpty()); + verify(iterator).close(); + verify(plan).clear(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testFinalSuccessIsSentOnlyAfterPlanAndIteratorCleanup() throws Exception { + ThreadPoolExecutor pool = pool(2); + ScanIterator child = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + ScanIterator result = mock(ScanIterator.class); + when(result.hasNext()).thenReturn(true, false); + when(result.next()).thenReturn(PipelineResult.EMPTY); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + when(plan.execute(PipelineResult.EMPTY)).thenReturn(result); + AtomicBoolean cleared = new AtomicBoolean(); + AtomicBoolean closed = new AtomicBoolean(); + doAnswer(invocation -> { + cleared.set(true); + return null; + }).when(plan).clear(); + doAnswer(invocation -> { + closed.set(true); + return null; + }).when(iterator).close(); + ResponseRecorder response = new ResponseRecorder() { + @Override + public void onNext(QueryResponse batch) { + assertTrue("success cannot precede plan cleanup", cleared.get()); + assertTrue("success cannot precede iterator cleanup", closed.get()); + super.onNext(batch); + } + }; + AggregativeQueryService service = service(pool, iterator, plan, 500); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("success-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertEquals(1, response.responses.size()); + assertTrue(response.responses.get(0).getIsOk()); + assertTrue(response.responses.get(0).getIsFinished()); + service.shutdownQueries(); + verify(child).close(); + } finally { + service.shutdownQueries(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParentCleanupFailuresReportErrorAndKeepShutdownBlocked() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + QueryPlan plan = mock(QueryPlan.class); + when(plan.onlyStopStage()).thenReturn(true); + IllegalStateException planFailure = new IllegalStateException("plan cleanup failed"); + IllegalStateException iteratorFailure = new IllegalStateException("iterator cleanup failed"); + doThrow(planFailure).when(plan).clear(); + doThrow(iteratorFailure).when(iterator).close(); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("cleanup-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "cleanup-query"); + assertEquals("plan cleanup failed", response.responses.get(0).getMessage()); + assertEquals(1, planFailure.getSuppressed().length); + org.junit.Assert.assertSame(iteratorFailure, planFailure.getSuppressed()[0]); + verify(plan).clear(); + verify(iterator).close(); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + assertFalse("failed cleanup cannot release the executor", pool.isShutdown()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testPartitionCloseFailurePreventsSuccessfulAggregation() throws Exception { + ThreadPoolExecutor pool = pool(2); + ScanIterator child = mock(ScanIterator.class); + doThrow(new IllegalStateException("child cleanup failed")).when(child).close(); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(child)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("child-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "child-query"); + verify(plan, never()).execute(PipelineResult.EMPTY); + verify(plan).clear(); + verify(iterator).close(); + verify(child).close(); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedPartitionCloseFailureStillClosesRemainingChildren() throws Exception { + AtomicInteger submissions = new AtomicInteger(); + ThreadPoolExecutor pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.SECONDS, + new LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable task) { + if (submissions.incrementAndGet() > 1) { + throw new RejectedExecutionException("partition rejected"); + } + super.execute(task); + } + }; + ScanIterator failed = mock(ScanIterator.class); + doThrow(new IllegalStateException("rejected child cleanup failed")).when(failed).close(); + ScanIterator remaining = mock(ScanIterator.class); + MultiPartitionIterator iterator = mock(MultiPartitionIterator.class); + when(iterator.getIterators()).thenReturn(Arrays.asList(failed, remaining)); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("reject-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "reject-query"); + verify(failed).close(); + verify(remaining).close(); + verify(plan).clear(); + verify(iterator).close(); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testSpringDoesNotDestroyDatabaseWhenIteratorCloseFails() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + doThrow(new IllegalStateException("native iterator close failed")).when(iterator).close(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + AtomicBoolean databaseClosed = new AtomicBoolean(); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + context.getBeanFactory().registerSingleton("queryService", service); + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("cleaner", mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean( + "database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + ResponseRecorder response = new ResponseRecorder(); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return null; + }); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("spring-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "spring-query"); + Thread closer = start(closing); + awaitShutdownWait(closer); + closer.interrupt(); + awaitShutdownWait(closer); + assertFalse(closing.isDone()); + assertFalse("Spring must retain the database after cleanup failure", databaseClosed.get()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + context.close(); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testSynchronousFailureRetainsCleanupErrorAndShutdownBarrier() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("plan initialization failed"); + IllegalStateException cleanup = new IllegalStateException("iterator cleanup failed"); + doThrow(cleanup).when(iterator).close(); + AggregativeQueryService service = new AggregativeQueryService(pool, 500, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + AggregativeQueryObserver observer = spy(super.newObserver(sender)); + doReturn(iterator).when(observer).getIterator(org.mockito.ArgumentMatchers.any()); + doThrow(original).when(observer).buildPlan(org.mockito.ArgumentMatchers.any()); + return observer; + } + }; + ResponseRecorder response = new ResponseRecorder(); + FutureTask closing = close(service); + try { + IllegalStateException thrown = org.junit.Assert.assertThrows(IllegalStateException.class, + () -> service.query(response).onNext(QueryRequest.getDefaultInstance())); + org.junit.Assert.assertSame(original, thrown); + assertEquals(1, thrown.getSuppressed().length); + org.junit.Assert.assertSame(cleanup, thrown.getSuppressed()[0]); + assertEquals(0, response.completed.get()); + assertTrue(response.responses.isEmpty()); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testPartialPartitionInitializationCleanupFailureRetainsQueryBarrier() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator first = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + doThrow(new IllegalStateException("opened partition cleanup failed")).when(first).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, key) -> { + if (id == 1) { + return first; + } + throw new IllegalStateException("second partition initialization failed"); + }); + QueryPlan plan = mock(QueryPlan.class); + when(plan.hasIteratorResult()).thenReturn(true); + ResponseRecorder response = new ResponseRecorder(); + AggregativeQueryService service = service(pool, iterator, plan, 500); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("init-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "init-query"); + verify(first).close(); + verify(plan).clear(); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test + public void testPartitionInitializationFailureClosesEveryOpenedChild() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenThrow(new IllegalStateException("child initialization failed")); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::getIterators); + verify(first).close(); + verify(second).close(); + iterator.close(); + } + + @Test + public void testPartitionInitializationCleanupFailureIsRetainedByParent() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenThrow(new IllegalStateException("child initialization failed")); + doThrow(new IllegalStateException("first cleanup failed")).when(first).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::getIterators); + verify(first).close(); + verify(second).close(); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(first).close(); + } + + @Test + public void testEmptyPartitionIteratorsAreReleasedBeforeOwnershipTransfer() { + ScanIterator empty = mock(ScanIterator.class); + ScanIterator nonempty = mock(ScanIterator.class); + when(nonempty.hasNext()).thenReturn(true); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2, 3), + (id, key) -> id == 1 ? empty : + id == 2 ? null : nonempty); + assertEquals(Arrays.asList(nonempty), iterator.getIterators()); + verify(empty).close(); + verify(nonempty, never()).close(); + nonempty.close(); + } + + @Test + public void testSequentialScanReleasesEveryEmptyPartitionBeforeAdvancing() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + ScanIterator nonempty = mock(ScanIterator.class); + when(nonempty.hasNext()).thenReturn(true); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2, 3, 4), + (id, key) -> id == 1 ? first : id == 2 ? second : id == 3 ? null : nonempty); + + assertTrue(iterator.hasNext()); + verify(first).close(); + verify(second).close(); + verify(nonempty, never()).close(); + iterator.close(); + iterator.close(); + verify(nonempty).close(); + } + + @Test + public void testSequentialInitializationFailureReleasesItsCreatedIterator() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition initialization failed"); + when(child.hasNext()).thenThrow(original); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::hasNext)); + verify(child).close(); + iterator.close(); + } + + @Test + public void testSequentialInitializationCleanupFailureRetainsOriginalErrorAndBarrier() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition initialization failed"); + IllegalStateException cleanup = new IllegalStateException("partition cleanup failed"); + when(child.hasNext()).thenThrow(original); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::hasNext)); + assertEquals(1, original.getSuppressed().length); + org.junit.Assert.assertSame(cleanup, original.getSuppressed()[0]); + IllegalStateException retained = org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + org.junit.Assert.assertSame(cleanup, retained.getCause()); + verify(child).close(); + } + + @Test + public void testSequentialCountIncludesAlreadyOpenedPartitionAndClosesAllSources() { + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(second.hasNext()).thenReturn(true); + when(first.count()).thenReturn(7L); + when(second.count()).thenReturn(11L); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), + (id, key) -> id == 1 ? first : second); + + assertTrue(iterator.hasNext()); + assertEquals(18L, iterator.count()); + assertFalse(iterator.hasNext()); + iterator.close(); + verify(first).close(); + verify(second).close(); + } + + @Test + public void testSequentialCountFailureReleasesCurrentSource() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition count failed"); + when(child.hasNext()).thenReturn(true); + when(child.count()).thenThrow(original); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::count)); + iterator.close(); + verify(child).close(); + } + + @Test + public void testSequentialCountCleanupFailureRetainsOriginalErrorAndBarrier() { + ScanIterator child = mock(ScanIterator.class); + IllegalStateException original = new IllegalStateException("partition count failed"); + IllegalStateException cleanup = new IllegalStateException("partition cleanup failed"); + when(child.hasNext()).thenReturn(true); + when(child.count()).thenThrow(original); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + org.junit.Assert.assertSame(original, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::count)); + assertEquals(1, original.getSuppressed().length); + org.junit.Assert.assertSame(cleanup, original.getSuppressed()[0]); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(child).close(); + } + + @Test + public void testSequentialExhaustionCleanupFailureRemainsVisibleToParent() { + ScanIterator child = mock(ScanIterator.class); + when(child.hasNext()).thenReturn(true, false); + IllegalStateException cleanup = new IllegalStateException("exhausted partition cleanup failed"); + doThrow(cleanup).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + + assertTrue(iterator.hasNext()); + org.junit.Assert.assertSame(cleanup, + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::next)); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(child).close(); + } + + @Test(timeout = 5000) + public void testSequentialEmptyPartitionCleanupFailureRetainsQueryBarrier() throws Exception { + ThreadPoolExecutor pool = pool(1); + ScanIterator child = mock(ScanIterator.class); + doThrow(new IllegalStateException("empty partition cleanup failed")).when(child).close(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + ResponseRecorder response = new ResponseRecorder(); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("empty-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "empty-query"); + verify(child).close(); + awaitShutdownWait(start(closing)); + assertFalse(closing.isDone()); + assertFalse("failed sequential cleanup cannot release the executor", pool.isShutdown()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testNativeAutomaticCloseFailureRetainsSequentialQueryBarrier() throws Exception { + ThreadPoolExecutor pool = pool(1); + RocksIterator raw = mock(RocksIterator.class); + when(raw.isOwningHandle()).thenReturn(true); + IllegalStateException failure = new IllegalStateException("native close failed once"); + doThrow(failure).doNothing().when(raw).close(); + // Native close fails before the callback or reference release can be reached. + RocksDBScanIterator child = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, null, ignored -> { }); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> child); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + ResponseRecorder response = new ResponseRecorder(); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("native-close-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "native-close-query"); + org.junit.Assert.assertSame(failure, + org.junit.Assert.assertThrows(IllegalStateException.class, child::close)); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(raw).close(); + awaitShutdownWait(start(closing)); + assertFalse("automatic native-close failure must retain the query registry", closing.isDone()); + assertFalse(pool.isShutdown()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testSupplierInnerKeyFilterCloseFailureRetainsQueryBarrier() throws Exception { + for (int constructor = 0; constructor < 3; constructor++) { + assertSupplierFilterCloseFailureRetainsQueryBarrier(constructor); + } + } + + private static void assertSupplierFilterCloseFailureRetainsQueryBarrier(int constructor) throws Exception { + ThreadPoolExecutor pool = pool(1); + RocksIterator raw = mock(RocksIterator.class); + when(raw.isOwningHandle()).thenReturn(true); + IllegalStateException failure = new IllegalStateException("supplier iterator close failed once"); + doThrow(failure).doNothing().when(raw).close(); + RocksDBScanIterator child = new RocksDBScanIterator<>(raw, null, null, + ScanIterator.Trait.SCAN_ANY, null, ignored -> { }); + AtomicInteger supplierReturned = new AtomicInteger(); + MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1), (id, key) -> { + InnerKeyFilter filter; + if (constructor == 0) { + filter = new InnerKeyFilter<>(child); + } else if (constructor == 1) { + filter = new InnerKeyFilter<>(child, true); + } else { + filter = new InnerKeyFilter<>(child, 10, 20); + } + supplierReturned.incrementAndGet(); + return filter; + }); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + ResponseRecorder response = new ResponseRecorder(); + FutureTask closing = close(service); + try { + service.query(response).onNext(QueryRequest.newBuilder().setQueryId("supplier-close-query").build()); + assertTrue(response.finished.await(1, TimeUnit.SECONDS)); + assertError(response, "supplier-close-query"); + assertEquals("factory must transfer ownership before any failing prefetch", 1, supplierReturned.get()); + org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close); + verify(raw).close(); + awaitShutdownWait(start(closing)); + assertFalse("factory prefetch failure must retain the aggregate query barrier", closing.isDone()); + assertFalse(pool.isShutdown()); + } finally { + releaseMockQueries(service); + closing.run(); + closing.get(1, TimeUnit.SECONDS); + pool.shutdownNow(); + } + } + + @Test + public void testPlanCleanupContinuesAfterStageFailure() { + QueryStage failed = mock(QueryStage.class); + QueryStage remaining = mock(QueryStage.class); + doThrow(new IllegalStateException("stage cleanup failed")).when(failed).close(); + QueryPlan plan = new QueryPlan(); + plan.addStage(failed); + plan.addStage(remaining); + org.junit.Assert.assertThrows(IllegalStateException.class, plan::clear); + verify(failed).close(); + verify(remaining).close(); + assertTrue(plan.isEmpty()); + } + + private static void awaitShutdownWait(Thread closer) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (System.nanoTime() < deadline) { + if (closer.getState() == Thread.State.TIMED_WAITING && + Arrays.stream(closer.getStackTrace()).anyMatch(frame -> + frame.getClassName().equals(AggregativeQueryService.class.getName()) && + frame.getMethodName().equals("shutdownQueries"))) { + return; + } + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); + } + fail("shutdown must enter the aggregate-query cleanup wait"); + } + + private static void assertContextCloseContinuesAfterCancellationFailure(Throwable failure) throws Exception { + ThreadPoolExecutor pool = pool(1); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch closingIterator = new CountDownLatch(1); + CountDownLatch allowClose = new CountDownLatch(1); + AtomicBoolean iteratorClosed = new AtomicBoolean(); + AtomicBoolean databaseClosed = new AtomicBoolean(); + AtomicBoolean prematureDestruction = new AtomicBoolean(); + ScanIterator iterator = mock(ScanIterator.class); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + new CountDownLatch(1).await(); + } catch (InterruptedException ignored) { + return false; + } + return false; + }); + doAnswer(invocation -> { + closingIterator.countDown(); + awaitUninterruptibly(allowClose); + iteratorClosed.set(true); + return null; + }).when(iterator).close(); + AggregativeQueryService service = service(pool, iterator, new QueryPlan(), 500); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + context.getBeanFactory().registerSingleton("queryService", service); + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("cleaner", mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean("database", () -> { + prematureDestruction.set(!iteratorClosed.get() || !pool.isTerminated()); + databaseClosed.set(true); + }); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + StreamObserver firstSender = sender(); + StreamObserver secondSender = sender(); + StreamObserver first = service.query(firstSender); + StreamObserver second = service.query(secondSender); + Field field = AggregativeQueryService.class.getDeclaredField("queries"); + field.setAccessible(true); + // Fail the first cancellation in the actual HashSet snapshot, regardless of identity hashes. + boolean firstCancelsFirst = ((Set) field.get(service)).iterator().next() == first; + StreamObserver broken = firstCancelsFirst ? firstSender : secondSender; + StreamObserver active = firstCancelsFirst ? second : first; + doThrow(failure).when(broken).onCompleted(); + active.onNext(QueryRequest.getDefaultInstance()); + assertTrue(reading.await(1, TimeUnit.SECONDS)); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return null; + }); + Thread closer = start(closing); + assertTrue("other queries must still be cancelled", closingIterator.await(1, TimeUnit.SECONDS)); + awaitShutdownWait(closer); + assertFalse(closing.isDone()); + assertFalse(databaseClosed.get()); + assertFalse(pool.isShutdown()); + allowClose.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertTrue(databaseClosed.get()); + assertFalse("Spring must destroy databases only after query cleanup and worker termination", + prematureDestruction.get()); + verify(broken).onCompleted(); + verify(iterator).close(); + } finally { + allowClose.countDown(); + service.shutdownQueries(); + context.close(); + pool.shutdownNow(); + } + } + + private static void releaseMockQueries(AggregativeQueryService service) throws Exception { + // Only test mocks are discarded here; production deliberately has no bypass for this barrier. + Field field = AggregativeQueryService.class.getDeclaredField("queries"); + field.setAccessible(true); + synchronized (service) { + ((Set) field.get(service)).clear(); + service.notifyAll(); + } + } + + private static void assertError(ResponseRecorder response, String queryId) { + assertEquals(1, response.responses.size()); + QueryResponse error = response.responses.get(0); + assertEquals(queryId, error.getQueryId()); + assertFalse(error.getIsOk()); + assertFalse(error.getIsFinished()); + assertFalse(error.getMessage().isEmpty()); + } + + private static class ResponseRecorder implements StreamObserver { + + private final List responses = new CopyOnWriteArrayList<>(); + private final AtomicInteger completed = new AtomicInteger(); + private final AtomicInteger errors = new AtomicInteger(); + private final CountDownLatch received = new CountDownLatch(1); + private final CountDownLatch finished = new CountDownLatch(1); + + @Override + public void onNext(QueryResponse response) { + this.responses.add(response); + this.received.countDown(); + } + + @Override + public void onError(Throwable error) { + this.errors.incrementAndGet(); + this.finished.countDown(); + } + + @Override + public void onCompleted() { + this.completed.incrementAndGet(); + this.finished.countDown(); + } + } + + private static ThreadPoolExecutor pool(int size) { + return (ThreadPoolExecutor) Executors.newFixedThreadPool(size); + } + + private static AggregativeQueryService service(ThreadPoolExecutor pool, ScanIterator iterator, + QueryPlan plan, long timeout) { + return new AggregativeQueryService(pool, timeout, 10) { + @Override + AggregativeQueryObserver newObserver(StreamObserver sender) { + return fixture(super.newObserver(sender), iterator, plan, null); + } + }; + } + + private static AggregativeQueryObserver fixture(AggregativeQueryObserver observer, + ScanIterator iterator, + QueryPlan plan, CountDownLatch parentReturned) { + AggregativeQueryObserver managed = spy(observer); + doReturn(iterator).when(managed).getIterator(org.mockito.ArgumentMatchers.any()); + doReturn(plan).when(managed).buildPlan(org.mockito.ArgumentMatchers.any()); + if (parentReturned != null) { + doAnswer(invocation -> { + try { + invocation.callRealMethod(); + return null; + } finally { + parentReturned.countDown(); + } + }).when(managed).sendData(); + } + return managed; + } + + @SuppressWarnings("unchecked") + private static StreamObserver sender() { + return mock(StreamObserver.class); + } + + private static FutureTask close(AggregativeQueryService service) { + return new FutureTask<>(() -> { + service.shutdownQueries(); + return null; + }); + } + + private static Thread start(FutureTask task) { + Thread thread = new Thread(task, "test-aggregate-query-close"); + thread.setDaemon(true); + thread.start(); + return thread; + } + + private static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + try { + for (;;) { + try { + latch.await(); + return; + } catch (InterruptedException e) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ContextClosedListenerTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ContextClosedListenerTest.java new file mode 100644 index 0000000000..b351c2fe5c --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ContextClosedListenerTest.java @@ -0,0 +1,166 @@ +/* + * 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.service; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryService; +import org.apache.hugegraph.store.node.listener.ContextClosedListener; +import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.junit.Test; +import org.lognet.springboot.grpc.context.GRpcServerInitializedEvent; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + +import io.grpc.Server; + +public class ContextClosedListenerTest { + + @Test(timeout = 5000) + public void testContextCloseStopsActiveWorkersWithoutPartitionDelay() throws Exception { + ThreadPoolExecutor scan = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ThreadPoolExecutor ttl = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + CountDownLatch started = new CountDownLatch(3); + CountDownLatch interrupted = new CountDownLatch(3); + CountDownLatch cleanupStarted = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch cleanupAllowed = new CountDownLatch(1); + Runnable work = () -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + interrupted.countDown(); + cleanupStarted.countDown(); + boolean done = false; + while (!done) { + try { + cleanupAllowed.await(); + done = true; + } catch (InterruptedException ignored) { + // Model resource cleanup which must finish before DB destruction. + } + } + Thread.currentThread().interrupt(); + } + }; + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + HgStoreStreamImpl stream = mock(HgStoreStreamImpl.class); + TTLCleaner cleaner = mock(TTLCleaner.class); + when(stream.getRealExecutor()).thenReturn(scan); + doAnswer(invocation -> { + scan.shutdownNow(); + return null; + }).when(stream).shutdownScans(); + when(cleaner.getExecutor()).thenReturn(ttl); + when(cleaner.getScheduler()).thenReturn(scheduler); + context.getBeanFactory().registerSingleton("storeStream", stream); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", cleaner); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + scan.execute(work); + ttl.execute(work); + scheduler.execute(work); + assertTrue("all shutdown participants must be active", started.await(1, TimeUnit.SECONDS)); + + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return null; + }); + Thread closeThread = new Thread(closing, "test-context-close"); + closeThread.setDaemon(true); + closeThread.start(); + // Scheduler shutdown must not let DB destruction overtake resource cleanup. + assertTrue(cleanupStarted.await(1, TimeUnit.SECONDS)); + assertFalse("context close must wait for worker cleanup", closing.isDone()); + cleanupAllowed.countDown(); + closing.get(2, TimeUnit.SECONDS); + assertTrue("close must interrupt active workers", interrupted.await(1, TimeUnit.SECONDS)); + assertTrue(scan.awaitTermination(1, TimeUnit.SECONDS)); + assertTrue(ttl.awaitTermination(1, TimeUnit.SECONDS)); + assertTrue(scheduler.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + cleanupAllowed.countDown(); + release.countDown(); + scan.shutdownNow(); + ttl.shutdownNow(); + scheduler.shutdownNow(); + context.close(); + } + } + + @Test(timeout = 5000) + public void testGrpcCallbacksFinishBeforeBeanDestructionEvenWhenInterrupted() throws Exception { + CountDownLatch awaitingCallbacks = new CountDownLatch(1); + CountDownLatch releaseCallbacks = new CountDownLatch(1); + AtomicBoolean terminated = new AtomicBoolean(); + AtomicBoolean databaseClosed = new AtomicBoolean(); + Server server = mock(Server.class); + when(server.shutdownNow()).thenReturn(server); + when(server.isTerminated()).thenAnswer(invocation -> terminated.get()); + when(server.awaitTermination(5, TimeUnit.SECONDS)).thenAnswer(invocation -> { + awaitingCallbacks.countDown(); + releaseCallbacks.await(); + terminated.set(true); + return true; + }); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean("database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + context.publishEvent(new GRpcServerInitializedEvent(context, server)); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return Thread.currentThread().isInterrupted(); + }); + Thread closeThread = new Thread(closing, "test-grpc-context-close"); + closeThread.setDaemon(true); + closeThread.start(); + assertTrue(awaitingCallbacks.await(1, TimeUnit.SECONDS)); + closeThread.interrupt(); + assertFalse("DB must remain open while callbacks own its resources", databaseClosed.get()); + releaseCallbacks.countDown(); + assertTrue("shutdown must preserve interruption", closing.get(2, TimeUnit.SECONDS)); + assertTrue(terminated.get()); + assertTrue(databaseClosed.get()); + } finally { + releaseCallbacks.countDown(); + context.close(); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanShutdownTest.java new file mode 100644 index 0000000000..be0fcbbe78 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanShutdownTest.java @@ -0,0 +1,397 @@ +/* + * 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.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.business.BusinessHandler; +import org.apache.hugegraph.store.business.GraphStoreIterator; +import org.apache.hugegraph.store.grpc.Graphpb; +import org.apache.hugegraph.store.grpc.Graphpb.ScanPartitionRequest; +import org.apache.hugegraph.store.grpc.Graphpb.ScanResponse; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.junit.Test; + +import io.grpc.stub.StreamObserver; + +public class GraphPartitionScanShutdownTest { + + @Test + public void testNaturalEndAndLimitCloseIteratorBeforeCompleting() throws Exception { + for (int limit : new int[]{0, 1}) { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + AtomicInteger consumed = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(call -> consumed.get() < 2); + when(iterator.next()).thenAnswer(call -> { + consumed.incrementAndGet(); + return Graphpb.Vertex.getDefaultInstance(); + }); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + AtomicInteger delivered = new AtomicInteger(); + doAnswer(call -> { + delivered.addAndGet(((ScanResponse) call.getArgument(0)).getVertexCount()); + return null; + }).when(response).onNext(any()); + StreamObserver request = service.scanGraphPartition(response, handler); + request.onNext(request(limit)); + verify(response, timeout(2000)).onCompleted(); + service.awaitScanCleanup(); + verify(iterator).close(); + assertEquals(limit == 0 ? 2 : 1, consumed.get()); + assertEquals(consumed.get(), delivered.get()); + assertTrue(registry(service).isEmpty()); + request.onError(new IllegalStateException("late cancellation")); + verify(iterator).close(); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + } + + @Test + public void testEmptyScanClosesOnceAndReleasesBarrier() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onCompleted(); + service.awaitScanCleanup(); + verify(iterator).close(); + assertTrue(registry(service).isEmpty()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testRejectedWorkerClosesIteratorAndTerminatesResponse() throws Exception { + ThreadPoolExecutor executor = executor(); + executor.shutdown(); + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response).onError(any()); + verify(iterator).close(); + assertTrue(registry(service).isEmpty()); + } + + @Test + public void testReadFailureClosesIteratorAndPreservesPrimaryError() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + RuntimeException failure = new IllegalStateException("native read failed"); + when(iterator.hasNext()).thenThrow(failure); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onError(failure); + service.awaitScanCleanup(); + verify(iterator).close(); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testFailedNativeCloseRemainsInShutdownBarrier() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + RuntimeException failure = new IllegalStateException("native close failed"); + doThrow(failure).when(iterator).close(); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onError(failure); + assertEquals(1, registry(service).size()); + Object lifecycle = registry(service).keySet().iterator().next(); + java.lang.reflect.Method cleanup = lifecycle.getClass().getDeclaredMethod("cleanupFailure"); + cleanup.setAccessible(true); + assertSame(failure, cleanup.invoke(lifecycle)); + verify(iterator).close(); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testShutdownWaitsForActiveReadBeforeClosingNativeIterator() throws Exception { + ThreadPoolExecutor executor = executor(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch stopped = new CountDownLatch(1); + CountDownLatch stoppingStarted = new CountDownLatch(1); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + when(iterator.hasNext()).thenAnswer(call -> { + entered.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return false; + }); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + assertTrue(entered.await(2, TimeUnit.SECONDS)); + Thread stopping = new Thread(() -> { + service.stopAcceptingScans(); + stoppingStarted.countDown(); + service.shutdownScans(); + service.awaitScanCleanup(); + stopped.countDown(); + }); + stopping.start(); + assertTrue(stoppingStarted.await(2, TimeUnit.SECONDS)); + assertFalse(stopped.await(100, TimeUnit.MILLISECONDS)); + release.countDown(); + assertTrue(stopped.await(2, TimeUnit.SECONDS)); + stopping.join(2000); + verify(iterator).close(); + verify(response).onError(any()); + assertTrue(registry(service).isEmpty()); + } finally { + release.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testClientCompletionBeforeInitializationReleasesBarrier() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + StreamObserver response = mock(StreamObserver.class); + StreamObserver request = + service.scanGraphPartition(response, mock(BusinessHandler.class)); + request.onCompleted(); + request.onCompleted(); + verify(response).onError(any()); + assertTrue(registry(service).isEmpty()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testReadFailureKeepsCloseFailureSuppressedAndSticky() throws Exception { + ThreadPoolExecutor executor = executor(); + try { + HgStoreStreamImpl service = service(executor); + GraphStoreIterator iterator = mock(GraphStoreIterator.class); + RuntimeException primary = new IllegalStateException("read"); + RuntimeException cleanup = new IllegalArgumentException("close"); + when(iterator.hasNext()).thenThrow(primary); + doThrow(cleanup).when(iterator).close(); + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver response = mock(StreamObserver.class); + service.scanGraphPartition(response, handler).onNext(request(0)); + verify(response, timeout(2000)).onError(primary); + assertEquals(1, primary.getSuppressed().length); + assertSame(cleanup, primary.getSuppressed()[0]); + verify(iterator).close(); + assertEquals(1, registry(service).size()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } + + @Test + public void testFullCreditWindowResumesWithOrderedPagesAndKeepsFinalPartialBatch() throws Exception { + ThreadPoolExecutor executor = executor(); + HgStoreStreamImpl service = service(executor); + final int batchSize = 100000; + final int window = 8; + final int total = batchSize * window + 7; + AtomicInteger closes = new AtomicInteger(); + Field suppliers = org.apache.hugegraph.store.business.BusinessHandlerImpl.class + .getDeclaredField("GRAPH_SUPPLIER_CACHE"); + suppliers.setAccessible(true); + @SuppressWarnings("unchecked") + Map graphCache = + (Map) suppliers.get(null); + String graphName = "TEST/credit-window"; + org.apache.hugegraph.HugeGraphSupplier previous = + graphCache.put(graphName, mock(org.apache.hugegraph.HugeGraphSupplier.class)); + GraphStoreIterator iterator; + try { + iterator = new GraphStoreIterator( + mock(ScanIterator.class), request(0)) { + private int row; + + @Override + public boolean hasNext() { + return this.row < total; + } + + @Override + public Graphpb.Vertex next() { + return Graphpb.Vertex.newBuilder().setId(Graphpb.Variant.newBuilder() + .setType(Graphpb.VariantType.VT_LONG).setValueInt64(this.row++)).build(); + } + + @Override + public void close() { + closes.incrementAndGet(); + } + }; + } finally { + if (previous == null) { + graphCache.remove(graphName); + } else { + graphCache.put(graphName, previous); + } + } + LinkedBlockingQueue responses = new LinkedBlockingQueue<>(); + AtomicInteger terminals = new AtomicInteger(); + AtomicReference error = new AtomicReference<>(); + CountDownLatch finished = new CountDownLatch(1); + StreamObserver response = new StreamObserver() { + @Override + public void onNext(ScanResponse value) { + responses.add(value); + } + + @Override + public void onError(Throwable failure) { + error.set(failure); + terminals.incrementAndGet(); + finished.countDown(); + } + + @Override + public void onCompleted() { + terminals.incrementAndGet(); + finished.countDown(); + } + }; + try { + BusinessHandler handler = mock(BusinessHandler.class); + doReturn(iterator).when(handler).scan(any()); + StreamObserver request = service.scanGraphPartition(response, handler); + request.onNext(request(0)); + long expectedId = 0; + for (int page = 0; page < window; page++) { + ScanResponse value = responses.poll(20, TimeUnit.SECONDS); + assertTrue("Expected page " + page, value != null); + assertEquals(page, value.getSeqNo()); + assertEquals(batchSize, value.getVertexCount()); + for (Graphpb.Vertex vertex : value.getVertexList()) { + assertEquals(expectedId++, vertex.getId().getValueInt64()); + } + } + assertNull("No ninth page is allowed without credit", responses.poll(100, TimeUnit.MILLISECONDS)); + assertEquals(0, terminals.get()); + for (int page = 0; page < window; page++) { + request.onNext(ScanPartitionRequest.newBuilder().setReplyRequest( + ScanPartitionRequest.Reply.newBuilder().setSeqNo(page)).build()); + } + ScanResponse last = responses.poll(20, TimeUnit.SECONDS); + assertTrue("The last partial batch must be delivered", last != null); + assertEquals(window, last.getSeqNo()); + assertEquals(7, last.getVertexCount()); + for (Graphpb.Vertex vertex : last.getVertexList()) { + assertEquals(expectedId++, vertex.getId().getValueInt64()); + } + assertEquals(total, expectedId); + assertTrue(finished.await(20, TimeUnit.SECONDS)); + assertNull(error.get()); + service.awaitScanCleanup(); + request.onCompleted(); + assertEquals(1, terminals.get()); + assertEquals(1, closes.get()); + assertTrue(responses.isEmpty()); + assertTrue(registry(service).isEmpty()); + } finally { + service.shutdownScans(); + executor.shutdown(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static ScanPartitionRequest request(long limit) { + return ScanPartitionRequest.newBuilder().setScanRequest(Graphpb.ScanPartitionRequest.Request.newBuilder() + .setGraphName("TEST/credit-window") + .setScanType(Graphpb.ScanPartitionRequest.ScanType.SCAN_VERTEX).setLimit(limit)).build(); + } + + private static ThreadPoolExecutor executor() { + return new ThreadPoolExecutor(2, 2, 0, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>()); + } + + private static HgStoreStreamImpl service(ThreadPoolExecutor executor) throws Exception { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + Field field = HgStoreStreamImpl.class.getDeclaredField("executor"); + field.setAccessible(true); + field.set(service, executor); + return service; + } + + private static Map registry(HgStoreStreamImpl service) throws Exception { + Field field = HgStoreStreamImpl.class.getDeclaredField("scans"); + field.setAccessible(true); + return (Map) field.get(service); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java new file mode 100644 index 0000000000..c32bcb0393 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GrpcShutdownBarrierTest.java @@ -0,0 +1,131 @@ +/* + * 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.service; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.io.ByteArrayInputStream; +import java.io.InputStream; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.junit.Test; + +import io.grpc.CallOptions; +import io.grpc.ManagedChannel; +import io.grpc.MethodDescriptor; +import io.grpc.Server; +import io.grpc.ServerServiceDefinition; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.ClientCalls; +import io.grpc.stub.ServerCalls; +import io.grpc.stub.StreamObserver; + +public class GrpcShutdownBarrierTest { + + @Test(timeout = 7000) + public void testTransportTerminationDoesNotReleaseActiveCallback() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AtomicBoolean finished = new AtomicBoolean(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + MethodDescriptor.Marshaller marshaller = new MethodDescriptor.Marshaller() { + @Override + public InputStream stream(byte[] value) { + return new ByteArrayInputStream(value); + } + @Override + public byte[] parse(InputStream stream) { + return new byte[0]; + } + }; + MethodDescriptor method = MethodDescriptor.newBuilder() + .setType(MethodDescriptor.MethodType.UNARY).setFullMethodName("probe/hold") + .setRequestMarshaller(marshaller).setResponseMarshaller(marshaller).build(); + String name = InProcessServerBuilder.generateName(); + GrpcShutdownBarrier barrier = new GrpcShutdownBarrier(); + Server server = InProcessServerBuilder.forName(name).executor(executor).intercept(barrier) + .addService(ServerServiceDefinition.builder("probe").addMethod(method, + ServerCalls.asyncUnaryCall((request, response) -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + finished.set(true); + } + response.onNext(new byte[0]); + response.onCompleted(); + })).build()).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + try { + ClientCalls.asyncUnaryCall(channel.newCall(method, CallOptions.DEFAULT), new byte[0], + new StreamObserver() { + @Override + public void onNext(byte[] value) { + } + @Override + public void onError(Throwable error) { + } + @Override + public void onCompleted() { + } + }); + assertTrue("handler must start", started.await(2, TimeUnit.SECONDS)); + barrier.stopAcceptingCalls(); + server.shutdownNow(); + assertTrue(server.awaitTermination(1, TimeUnit.SECONDS)); + assertFalse("transport termination must not be mistaken for callback exit", finished.get()); + CountDownLatch waiting = new CountDownLatch(1); + FutureTask drained = new FutureTask<>(() -> { + waiting.countDown(); + barrier.awaitCallbacks(); + return null; + }); + Thread waiter = new Thread(drained, "test-rpc-drain"); + waiter.setDaemon(true); + waiter.start(); + assertTrue(waiting.await(1, TimeUnit.SECONDS)); + try { + drained.get(100, TimeUnit.MILLISECONDS); + fail("barrier passed while a handler still owns resources"); + } catch (TimeoutException expected) { + // The active handler must prevent database destruction. + } + release.countDown(); + drained.get(2, TimeUnit.SECONDS); + assertTrue(finished.get()); + } finally { + release.countDown(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdown(); + assertTrue("cleanup must finish", executor.awaitTermination(2, TimeUnit.SECONDS)); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanShutdownTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanShutdownTest.java new file mode 100644 index 0000000000..d7f382e708 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanShutdownTest.java @@ -0,0 +1,1730 @@ +/* + * 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.service; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; + +import org.apache.hugegraph.pd.common.KVPair; +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; +import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.KvStream; +import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; +import org.apache.hugegraph.store.grpc.stream.ScanStreamBatchReq; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; +import org.apache.hugegraph.store.node.AppConfig; +import org.apache.hugegraph.store.node.grpc.GrpcShutdownBarrier; +import org.apache.hugegraph.store.node.grpc.HgStoreStreamImpl; +import org.apache.hugegraph.store.node.grpc.HgStoreWrapperEx; +import org.apache.hugegraph.store.node.grpc.ParallelScanIterator; +import org.apache.hugegraph.store.node.grpc.QueryCondition; +import org.apache.hugegraph.store.node.grpc.ScanBatchOneShotResponse; +import org.apache.hugegraph.store.node.grpc.ScanBatchResponse; +import org.apache.hugegraph.store.node.grpc.ScanBatchResponse3; +import org.apache.hugegraph.store.node.grpc.ScanOneShotResponse; +import org.apache.hugegraph.store.node.grpc.ScanStreamResponse; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryService; +import org.apache.hugegraph.store.node.listener.ContextClosedListener; +import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.apache.hugegraph.store.node.util.HgChannel; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + +import io.grpc.Context; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; + +public class ScanShutdownTest { + + @Test + public void testFailedSelectionCleanupBlocksLifecycleCompletion() throws Exception { + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException closing = new IllegalStateException("native release failed"); + doThrow(closing).when(iterator).close(); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + Class lifecycleType = Class.forName("org.apache.hugegraph.store.node.grpc.ScanLifecycle"); + java.lang.reflect.Constructor constructor = lifecycleType.getDeclaredConstructor(); + constructor.setAccessible(true); + Object lifecycle = constructor.newInstance(); + Method failedCleanup = lifecycleType.getDeclaredMethod("failedCleanup", Throwable.class); + failedCleanup.setAccessible(true); + AtomicBoolean finished = new AtomicBoolean(); + Method onFinished = lifecycleType.getDeclaredMethod("onFinished", Runnable.class); + onFinished.setAccessible(true); + onFinished.invoke(lifecycle, (Runnable) () -> finished.set(true)); + java.util.function.Consumer cleanup = failure -> { + try { + failedCleanup.invoke(lifecycle, failure); + } catch (ReflectiveOperationException error) { + throw new AssertionError(error); + } + }; + Class selector = Class.forName("org.apache.hugegraph.store.node.grpc.ScanUtil"); + Method getIterator = selector.getDeclaredMethod("getIterator", ScanStreamReq.class, + HgStoreWrapperEx.class, + java.util.function.Consumer.class); + getIterator.setAccessible(true); + ScanStreamReq request = ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setHeader(Header.newBuilder().setGraph("selection-cleanup")) + .setTable("unsupported-property-table") + .setSelects(org.apache.hugegraph.store.grpc.stream.SelectParam.newBuilder().addProperties(1)).build(); + try { + getIterator.invoke(null, request, wrapper, cleanup); + fail("Invalid property selection must fail"); + } catch (java.lang.reflect.InvocationTargetException error) { + assertTrue(error.getCause() instanceof IllegalArgumentException); + assertEquals(Collections.singletonList(closing), Arrays.asList(error.getCause().getSuppressed())); + } + Method finish = lifecycleType.getDeclaredMethod("finishWithoutResponse"); + finish.setAccessible(true); + finish.invoke(lifecycle); + assertFalse("Failed native release must remain a shutdown blocker", finished.get()); + verify(iterator).close(); + } + + @Test(timeout = 5000) + public void testInFlightReceiptDoesNotPublishOverAfterErrorCancellation() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + Field lockField = ScanBatchResponse.class.getDeclaredField("stateLock"); + lockField.setAccessible(true); + Field state = ScanBatchResponse.class.getDeclaredField("state"); + state.setAccessible(true); + Field cancelled = ScanBatchResponse.class.getDeclaredField("cancelled"); + cancelled.setAccessible(true); + AtomicBoolean cancelling = (AtomicBoolean) cancelled.get(response); + Object done = Arrays.stream(state.getType().getEnumConstants()) + .filter(value -> value.toString().equals("DONE")).findFirst().get(); + FutureTask receipt = new FutureTask<>(() -> { + response.onNext(ScanStreamBatchReq.newBuilder().setReceiptRequest( + org.apache.hugegraph.store.grpc.stream.ScanReceiptRequest.newBuilder().setTimes(1)).build()); + return null; + }); + FutureTask failing = new FutureTask<>(() -> { + response.onError(Status.INTERNAL.asRuntimeException()); + return null; + }); + Thread reader = new Thread(receipt, "receipt-at-terminal-transition"); + Thread closer = new Thread(failing, "error-at-terminal-transition"); + try { + synchronized (lockField.get(response)) { + reader.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (reader.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, reader.getState()); + closer.start(); + deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (!cancelling.get() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue(cancelling.get()); + state.set(response, done); + } + receipt.get(2, TimeUnit.SECONDS); + failing.get(2, TimeUnit.SECONDS); + verify(output, never()).onNext(any(KvStream.class)); + verify(output).onError(any(Throwable.class)); + } finally { + executor.shutdownNow(); + reader.join(1000); + closer.join(1000); + } + } + + @Test(timeout = 10000) + public void testOrderedVertexRemainsConsecutiveAcrossBackpressureWithTwoScanners() throws Exception { + AtomicBoolean defer = new AtomicBoolean(true); + List initial = new ArrayList<>(); + ThreadPoolExecutor executor = new ThreadPoolExecutor(2, 2, 0, TimeUnit.SECONDS, + new java.util.concurrent.LinkedBlockingQueue<>()) { + @Override + public void execute(Runnable command) { + if (defer.get()) { + initial.add(command); + } else { + super.execute(command); + } + } + }; + ScanIterator first = mock(ScanIterator.class); + ScanIterator second = mock(ScanIterator.class); + AtomicInteger firstRows = new AtomicInteger(); + AtomicInteger secondRows = new AtomicInteger(); + CountDownLatch firstBackpressure = new CountDownLatch(1); + CountDownLatch secondReading = new CountDownLatch(1); + when(first.hasNext()).thenAnswer(ignored -> firstRows.get() < 8); + when(first.next()).thenAnswer(ignored -> { + if (firstRows.incrementAndGet() == 5) { + firstBackpressure.countDown(); + } + return RocksDBSession.BackendColumn.of(new byte[]{1, 0, 0, 0}, new byte[16]); + }); + when(second.hasNext()).thenAnswer(ignored -> secondRows.get() < 2); + when(second.next()).thenAnswer(ignored -> { + secondRows.incrementAndGet(); + secondReading.countDown(); + return RocksDBSession.BackendColumn.of(new byte[]{2, 0, 0, 0}, new byte[16]); + }); + AtomicInteger supplies = new AtomicInteger(); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int previous = bodySize.getInt(null); + bodySize.setInt(null, 16); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of(() -> { + int index = supplies.incrementAndGet(); + return new KVPair<>(mock(QueryCondition.class), index == 1 ? first : index == 2 ? second : null); + }, () -> Long.MAX_VALUE, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), executor); + // Force two owned scanners independently of the container's CPU-derived default. + Field registry = ParallelScanIterator.class.getDeclaredField("scanners"); + registry.setAccessible(true); + java.util.Queue scanners = (java.util.Queue) registry.get(scan); + Class scannerType = scanners.peek().getClass(); + java.lang.reflect.Constructor constructor = + scannerType.getDeclaredConstructor(ParallelScanIterator.class); + constructor.setAccessible(true); + Object other = constructor.newInstance(scan); + scanners.add(other); + Method run = scannerType.getDeclaredMethod("scanKV"); + run.setAccessible(true); + defer.set(false); + executor.execute(initial.get(0)); + assertTrue(firstBackpressure.await(2, TimeUnit.SECONDS)); + executor.execute(() -> { + try { + run.invoke(other); + } catch (ReflectiveOperationException e) { + throw new AssertionError(e); + } + }); + assertTrue(secondReading.await(2, TimeUnit.SECONDS)); + List vertices = new ArrayList<>(); + while (scan.hasNext()) { + for (Object row : scan.next()) { + Field key = row.getClass().getField("key"); + key.setAccessible(true); + vertices.add((int) ((byte[]) key.get(row))[0]); + } + } + assertEquals(Arrays.asList(1, 1, 1, 1, 1, 1, 1, 1, 2, 2), vertices); + verify(first).close(); + verify(second).close(); + } finally { + if (scan != null) { + scan.close(); + } + executor.shutdownNow(); + bodySize.setInt(null, previous); + } + } + + @Test(timeout = 20000) + public void testOrdinaryScanCleanupFailureBlocksSpringDestruction() throws Exception { + for (int mode = 0; mode < 5; mode++) { + assertCleanupBlocksDestruction(mode); + } + } + + private static void assertCleanupBlocksDestruction(int mode) throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(4); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator broken = mock(ScanIterator.class); + ScanIterator healthy = mock(ScanIterator.class); + IllegalStateException failure = new IllegalStateException("injected native release failure"); + doThrow(failure).when(broken).close(); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(broken, healthy); + HgStoreStreamImpl service = scanService(executor, wrapper); + AnnotationConfigApplicationContext context = + new AnnotationConfigApplicationContext(); + AtomicBoolean destroyed = new AtomicBoolean(); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return null; + }); + Thread closer = new Thread(closing, "scan-failed-cleanup-context-close"); + try { + context.getBeanFactory().registerSingleton("storeStream", service); + context.getBeanFactory().registerSingleton("queryService", + mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", + mock(TTLCleaner.class)); + context.getDefaultListableBeanFactory().registerDisposableBean("database", () -> destroyed.set(true)); + context.register(ContextClosedListener.class, + GrpcShutdownBarrier.class); + context.refresh(); + for (int i = 0; i < 2; i++) { + ScanStreamReq request = ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build(); + switch (mode) { + case 0: + service.scan(mock(StreamObserver.class)).onNext(request); + break; + case 1: + service.scanBatch2(mock(StreamObserver.class)).onNext(batchRequest()); + break; + case 2: + service.scanBatch(mock(StreamObserver.class)).onNext(batchRequest()); + break; + case 3: + service.scanOneShot(request, mock(StreamObserver.class)); + break; + default: + service.scanBatchOneShot(batchRequest(), mock(StreamObserver.class)); + } + } + verify(broken, timeout(2000)).close(); + verify(healthy, timeout(2000)).close(); + closer.start(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + Thread.State observed = closer.getState(); + while (observed != Thread.State.TIMED_WAITING && System.nanoTime() < deadline) { + Thread.yield(); + observed = closer.getState(); + } + // Assert the captured wait observation: a second state read can see a + // notification wakeup even though database destruction remains blocked. + assertEquals(Thread.State.TIMED_WAITING, observed); + assertFalse("terminated workers do not confirm native cleanup", closing.isDone()); + assertFalse(destroyed.get()); + verify(broken).close(); + verify(healthy).close(); + Map pending = scanRegistry(service); + assertEquals(1, pending.size()); + Object retained = pending.keySet().iterator().next(); + Field cleanupFailure = retained.getClass().getDeclaredField("cleanupFailure"); + cleanupFailure.setAccessible(true); + assertSame(failure, cleanupFailure.get(retained)); + } finally { + // Test-only teardown: production never clears a failed native release. + Map pending = scanRegistry(service); + synchronized (pending) { + pending.clear(); + pending.notifyAll(); + } + executor.shutdownNow(); + if (closer.isAlive()) { + closing.get(2, TimeUnit.SECONDS); + } + context.close(); + } + } + + @Test(timeout = 5000) + public void testUnsupportedPausedBatchClosesOwnedIteratorAndRetainsFailure() throws Exception { + for (boolean fails : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreStreamImpl service = scanService(executor, mock(HgStoreWrapperEx.class)); + StreamObserver output = mock(StreamObserver.class); + doThrow(new AssertionError("injected response callback failure")) + .when(output).onError(any(Throwable.class)); + StreamObserver input = service.scanBatch2(output); + ScanBatchResponse response = null; + for (Field field : input.getClass().getDeclaredFields()) { + field.setAccessible(true); + Object captured = field.get(input); + if (captured instanceof ScanBatchResponse) { + response = (ScanBatchResponse) captured; + } + } + org.junit.Assert.assertNotNull(response); + ScanIterator iterator = mock(ScanIterator.class); + if (fails) { + doThrow(new IllegalStateException("paused iterator close failed")).when(iterator).close(); + } + Field resource = ScanBatchResponse.class.getDeclaredField("iterator"); + resource.setAccessible(true); + resource.set(response, iterator); + try { + // An idle response can retain a native iterator while waiting for receipts. + input.onNext(ScanStreamBatchReq.getDefaultInstance()); + verify(iterator).close(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + assertEquals(fails ? 1 : 0, scanRegistry(service).size()); + service.shutdownScans(); + input.onCompleted(); + verify(iterator).close(); + } finally { + scanRegistry(service).clear(); + executor.shutdownNow(); + } + } + } + + @Test(timeout = 5000) + public void testFailedCleanupAndResponseCallbackStillClearWorkerOwnership() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-close-callback-failure"); + thread.setUncaughtExceptionHandler((ignored, failure) -> { }); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + StreamObserver output = mock(StreamObserver.class); + CountDownLatch sending = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + // Exercise cleanup failure while the response callback is already in flight. + assertTrue(sending.await(2, TimeUnit.SECONDS)); + throw new IllegalStateException("native close failed"); + }).when(iterator).close(); + doAnswer(invocation -> { + sending.countDown(); + assertTrue(release.await(2, TimeUnit.SECONDS)); + return null; + }).when(output).onNext(any(KvPageRes.class)); + doThrow(new AssertionError("error callback failed")).when(output).onError(any(Throwable.class)); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(5); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-close-callback-request"); + try { + caller.start(); + assertTrue(sending.await(1, TimeUnit.SECONDS)); + Field finished = ScanStreamResponse.class.getDeclaredField("finishFlag"); + finished.setAccessible(true); + AtomicBoolean terminal = (AtomicBoolean) finished.get(response); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (!terminal.get() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue("cleanup failure must claim error before normal completion", terminal.get()); + release.countDown(); + request.get(2, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + Field worker = ScanStreamResponse.class.getDeclaredField("worker"); + worker.setAccessible(true); + assertNull(worker.get(response)); + verify(iterator).close(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + } finally { + release.countDown(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testTerminalAdmissionRejectsRequestWaitingForLifecycleLock() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver input = service.scanBatch(mock(StreamObserver.class)); + Object lifecycle = scanRegistry(service).keySet().iterator().next(); + Method finished = lifecycle.getClass().getDeclaredMethod("finishWithoutResponse"); + finished.setAccessible(true); + FutureTask request = new FutureTask<>(() -> { + input.onNext(batchRequest()); + return null; + }); + Thread late = new Thread(request, "late-scan-admission"); + try { + synchronized (lifecycle) { + late.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (late.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, late.getState()); + // Exercise the atomic admission invariant directly, without claiming a native fault. + finished.invoke(lifecycle); + assertTrue(scanRegistry(service).isEmpty()); + } + request.get(1, TimeUnit.SECONDS); + service.shutdownScans(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verifyNoInteractions(wrapper); + assertEquals(0, executor.getTaskCount()); + } finally { + executor.shutdownNow(); + late.join(1000); + } + } + + @Test(timeout = 10000) + public void testInProcessNormalCompletionDoesNotCancelCompletedScan() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + HgStoreStreamImpl service = scanService(executor, wrapper); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name) + .directExecutor().addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name) + .directExecutor().build(); + org.apache.logging.log4j.core.Logger logger = (org.apache.logging.log4j.core.Logger) + org.apache.logging.log4j.LogManager.getLogger(ScanStreamResponse.class); + List messages = new java.util.concurrent.CopyOnWriteArrayList<>(); + org.apache.logging.log4j.core.appender.AbstractAppender logs = + new org.apache.logging.log4j.core.appender.AbstractAppender( + "scan-normal-completion", null, null, false, + org.apache.logging.log4j.core.config.Property.EMPTY_ARRAY) { + @Override + public void append(org.apache.logging.log4j.core.LogEvent event) { + messages.add(event.getMessage().getFormattedMessage()); + } + }; + logs.start(); + logger.addAppender(logs); + CountDownLatch completed = new CountDownLatch(1); + AtomicInteger errors = new AtomicInteger(); + try { + StreamObserver request = + HgStoreStreamGrpc.newStub(channel) + .scan(new StreamObserver() { + @Override + public void onNext(KvPageRes value) { } + + @Override + public void onError(Throwable failure) { + errors.incrementAndGet(); + completed.countDown(); + } + + @Override + public void onCompleted() { + completed.countDown(); + } + }); + request.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + assertTrue(completed.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertEquals(0, errors.get()); + assertTrue(messages.stream().noneMatch(message -> message.contains("onError from client"))); + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + logger.removeAppender(logs); + logs.stop(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testInProcessClientCancellationDrainsIterator() throws Exception { + assertInProcessCancellation(false); + assertInProcessCancellation(true); + } + + private static void assertInProcessCancellation(boolean completionRace) throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(2); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + CountDownLatch reading = new CountDownLatch(1); + CountDownLatch releaseRead = new CountDownLatch(1); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenAnswer(invocation -> { + reading.countDown(); + try { + releaseRead.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return false; + }); + HgStoreStreamImpl service = scanService(executor, wrapper); + // gRPC 1.39 dispatches context cancellation through the application executor too. + // Keep a slot for cancellation while scan.onNext waits for its worker's first page. + ExecutorService callbacks = Executors.newFixedThreadPool(2); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name) + .executor(callbacks).addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name) + .directExecutor().build(); + Context.CancellableContext context = Context.current().withCancellation(); + CountDownLatch ended = new CountDownLatch(1); + AtomicInteger terminals = new AtomicInteger(); + AtomicReference failure = new AtomicReference<>(); + try { + context.run(() -> { + StreamObserver request = + HgStoreStreamGrpc.newStub(channel) + .scan(new StreamObserver() { + @Override + public void onNext(KvPageRes value) { } + + @Override + public void onError(Throwable error) { + failure.set(error); + terminals.incrementAndGet(); + ended.countDown(); + } + + @Override + public void onCompleted() { + terminals.incrementAndGet(); + ended.countDown(); + } + }); + request.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + }); + assertTrue(reading.await(2, TimeUnit.SECONDS)); + if (completionRace) { + CountDownLatch start = new CountDownLatch(1); + Thread cancel = new Thread(() -> { + try { + start.await(); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + context.cancel(null); + }, "scan-completion-cancel-race"); + cancel.start(); + start.countDown(); + releaseRead.countDown(); + cancel.join(2000); + assertFalse(cancel.isAlive()); + } else { + context.cancel(null); + } + assertTrue(ended.await(2, TimeUnit.SECONDS)); + verify(iterator, timeout(2000)).close(); + service.shutdownScans(); + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertEquals(1, terminals.get()); + if (!completionRace || failure.get() != null) { + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.get()).getCode()); + } + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + } finally { + releaseRead.countDown(); + context.cancel(null); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + callbacks.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testBatchRejectionReportsStatusAndClosesUnstartedIterator() throws Exception { + for (boolean stopped : new boolean[]{false, true}) { + ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 1, TimeUnit.SECONDS, + new SynchronousQueue<>()); + CountDownLatch release = new CountDownLatch(1); + if (stopped) { + executor.shutdown(); + } else { + executor.execute(() -> { + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + } + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + HgStoreStreamImpl service = scanService(executor, wrapper); + StreamObserver output = mock(StreamObserver.class); + try { + service.scanBatch(output).onNext(batchRequest()); + ArgumentCaptor error = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(error.capture()); + assertEquals(stopped ? Status.Code.UNAVAILABLE : Status.Code.RESOURCE_EXHAUSTED, + Status.fromThrowable(error.getValue()).getCode()); + verify(output, never()).onCompleted(); + assertTrue(scanRegistry(service).isEmpty()); + verifyNoInteractions(wrapper); + } finally { + release.countDown(); + executor.shutdownNow(); + } + } + } + + private static HgStoreStreamImpl scanService(ThreadPoolExecutor executor, + HgStoreWrapperEx wrapper) throws Exception { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(5); + for (String name : new String[]{"executor", "wrapper", "appConfig"}) { + Field field = HgStoreStreamImpl.class.getDeclaredField(name); + field.setAccessible(true); + field.set(service, name.equals("executor") ? executor : name.equals("wrapper") ? wrapper : config); + } + return service; + } + + private static Map scanRegistry(HgStoreStreamImpl service) throws Exception { + Field scans = HgStoreStreamImpl.class.getDeclaredField("scans"); + scans.setAccessible(true); + return (Map) scans.get(service); + } + + @Test + public void testClosingPreventsLazyExecutorAndStateCreation() { + HgStoreStreamImpl service = new HgStoreStreamImpl(); + service.stopAcceptingScans(); + assertUnavailable(service::getExecutor); + assertUnavailable(service::getState); + assertUnavailable(() -> service.scan(mock(StreamObserver.class))); + service.shutdownScans(); + assertNull(service.getRealExecutor()); + } + + @Test(timeout = 5000) + public void testQueuedStreamCancellationDrainsWithoutOpeningIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch occupied = new CountDownLatch(1); + executor.execute(() -> { + occupied.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-shutdown-test"); + try { + assertTrue(occupied.await(1, TimeUnit.SECONDS)); + caller.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (executor.getQueue().isEmpty() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(1, executor.getQueue().size()); + response.onError(Status.CANCELLED.asRuntimeException()); + executor.shutdown(); + release.countDown(); + request.get(1, TimeUnit.SECONDS); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verifyNoInteractions(wrapper); + assertEquals(2L, executor.getCompletedTaskCount()); + } finally { + response.onCompleted(); + release.countDown(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 10000) + public void testInProcessOneShotCancellationDuringCloseReportsCancelled() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(false); + CountDownLatch closing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + closing.countDown(); + assertTrue(release.await(5, TimeUnit.SECONDS)); + return null; + }).when(iterator).close(); + HgStoreStreamImpl service = scanService(executor, wrapper); + String name = InProcessServerBuilder.generateName(); + Server server = InProcessServerBuilder.forName(name).directExecutor() + .addService(service).build().start(); + ManagedChannel channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + FutureTask call = new FutureTask<>(() -> { + try { + HgStoreStreamGrpc.HgStoreStreamBlockingStub stub = + HgStoreStreamGrpc.newBlockingStub(channel).withDeadlineAfter(3, TimeUnit.SECONDS); + if (batch) { + stub.scanBatchOneShot(batchRequest()); + } else { + stub.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build()); + } + fail("cancelled unary scan reported success"); + return Status.Code.OK; + } catch (StatusRuntimeException error) { + return error.getStatus().getCode(); + } + }); + Thread caller = new Thread(call, "one-shot-cancellation-client"); + try { + caller.start(); + assertTrue(closing.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertEquals(1, scanRegistry(service).size()); + assertFalse(call.isDone()); + release.countDown(); + assertEquals(Status.Code.CANCELLED, call.get(2, TimeUnit.SECONDS)); + service.awaitScanCleanup(); + assertTrue(scanRegistry(service).isEmpty()); + verify(iterator).close(); + } finally { + release.countDown(); + channel.shutdownNow(); + server.shutdownNow(); + executor.shutdownNow(); + caller.join(2000); + assertFalse(caller.isAlive()); + assertTrue(channel.awaitTermination(2, TimeUnit.SECONDS)); + assertTrue(server.awaitTermination(2, TimeUnit.SECONDS)); + } + } + } + + @Test(timeout = 10000) + public void testOneShotCancelledDuringCloseDoesNotCompleteSuccessfully() throws Exception { + for (boolean batch : new boolean[]{false, true}) { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(false); + CountDownLatch closing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + doAnswer(invocation -> { + closing.countDown(); + assertTrue(release.await(5, TimeUnit.SECONDS)); + return null; + }).when(iterator).close(); + HgStoreStreamImpl service = scanService(executor, wrapper); + FutureTask scanning = new FutureTask<>(() -> { + if (batch) { + service.scanBatchOneShot(batchRequest(), output); + } else { + service.scanOneShot(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setLimit(10).build(), output); + } + return null; + }); + Thread caller = new Thread(scanning, "one-shot-close-cancellation"); + try { + caller.start(); + assertTrue(closing.await(2, TimeUnit.SECONDS)); + service.shutdownScans(); + assertEquals(1, scanRegistry(service).size()); + assertFalse(scanning.isDone()); + verifyNoInteractions(output); + release.countDown(); + scanning.get(2, TimeUnit.SECONDS); + ArgumentCaptor error = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(error.capture()); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(error.getValue()).getCode()); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + assertTrue(scanRegistry(service).isEmpty()); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + release.countDown(); + caller.join(2000); + executor.shutdownNow(); + assertFalse(caller.isAlive()); + } + } + } + + @Test(timeout = 5000) + public void testOneShotObservesContextCancellationAndClosesIterator() { + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + Context.CancellableContext context = Context.current().withCancellation(); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenAnswer(invocation -> { + context.cancel(null); + return true; + }); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + try { + context.run(() -> ScanOneShotResponse.scanOneShot( + ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL).setLimit(100).build(), + output, wrapper)); + verify(iterator).close(); + verifyNoInteractions(output); + } finally { + context.cancel(null); + } + } + + @Test(timeout = 5000) + public void testBatchDoneRejectsSecondQueryBeforeAllocatingIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator source = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(source); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, wrapper, executor); + Field state = ScanBatchResponse.class.getDeclaredField("state"); + state.setAccessible(true); + Field iterator = ScanBatchResponse.class.getDeclaredField("iterator"); + iterator.setAccessible(true); + try { + response.onNext(batchRequest()); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2); + while (!"DONE".equals(state.get(response).toString()) && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals("DONE", state.get(response).toString()); + verify(source).close(); + assertNull(iterator.get(response)); + response.onNext(batchRequest()); + // The single worker processes the queued request before this barrier. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + assertNull("a completed query must reject another request before allocating", iterator.get(response)); + response.onCompleted(); + response.onCompleted(); + verify(wrapper).scanAll(anyString(), anyString(), any(byte[].class)); + verify(source).close(); + verify(output).onCompleted(); + verify(output, never()).onError(any(Throwable.class)); + } finally { + // Preserve cleanup even when the old implementation allocated a second producer. + Object remaining = iterator.get(response); + response.onCompleted(); + if (remaining instanceof ScanIterator) { + ((ScanIterator) remaining).close(); + } + executor.shutdownNow(); + } + } + + @Test + public void testPausedBatchClosesIteratorAndRejectsLaterQuery() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, wrapper, executor); + Field field = ScanBatchResponse.class.getDeclaredField("iterator"); + field.setAccessible(true); + field.set(response, iterator); + try { + response.onCompleted(); + response.onNext(ScanStreamBatchReq.newBuilder() + .setQueryRequest(ScanQueryRequest.getDefaultInstance()) + .build()); + response.onCompleted(); + verify(iterator).close(); + verifyNoInteractions(wrapper); + assertNull(field.get(response)); + assertFalse(executor.isShutdown()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testRejectedStreamReportsOnlyFailure() { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + executor.shutdown(); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + ScanStreamReq request = ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build(); + + response.onNext(request); + response.onNext(request); + response.onCompleted(); + + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verifyNoInteractions(wrapper); + } + + @Test(timeout = 5000) + public void testStreamIteratorFailureReportsOnlyFailure() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))) + .thenThrow(new IllegalStateException("iterator failed")); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-failure-test"); + try { + caller.start(); + request.get(1, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + response.onCompleted(); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testReceiveFailureDoesNotInterruptItsCallbackThread() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + CountDownLatch occupied = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + executor.execute(() -> { + occupied.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(0); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-receive-failure-test"); + try { + assertTrue(occupied.await(1, TimeUnit.SECONDS)); + caller.start(); + assertFalse(request.get(1, TimeUnit.SECONDS)); + release.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verifyNoInteractions(wrapper); + } finally { + release.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 10000) + public void testStreamFailureWaitsForPageCallback() throws Exception { + AtomicReference scanWorker = new AtomicReference<>(); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-page-failure-worker"); + scanWorker.set(thread); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + CountDownLatch pageEntered = new CountDownLatch(1); + CountDownLatch releasePage = new CountDownLatch(1); + CountDownLatch iteratorFailed = new CountDownLatch(1); + CountDownLatch errorReceived = new CountDownLatch(1); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + if (advances.incrementAndGet() <= 2) { + return true; + } + assertTrue(pageEntered.await(1, TimeUnit.SECONDS)); + iteratorFailed.countDown(); + throw new IllegalStateException("failure after first page"); + }); + StreamObserver output = mock(StreamObserver.class); + List signals = Collections.synchronizedList(new ArrayList<>()); + doAnswer(invocation -> { + signals.add("page-start"); + pageEntered.countDown(); + assertTrue(releasePage.await(2, TimeUnit.SECONDS)); + signals.add("page-end"); + return null; + }).when(output).onNext(any(KvPageRes.class)); + doAnswer(invocation -> { + signals.add("error"); + errorReceived.countDown(); + return null; + }).when(output).onError(any(Throwable.class)); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return null; + }); + Thread caller = new Thread(request, "scan-page-failure-test"); + try { + caller.start(); + assertTrue(pageEntered.await(1, TimeUnit.SECONDS)); + assertTrue(iteratorFailed.await(1, TimeUnit.SECONDS)); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (scanWorker.get().getState() != Thread.State.BLOCKED && + errorReceived.getCount() != 0 && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals("worker must reach the response serialization point", + Thread.State.BLOCKED, scanWorker.get().getState()); + assertEquals(1L, errorReceived.getCount()); + releasePage.countDown(); + request.get(1, TimeUnit.SECONDS); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + assertEquals(Arrays.asList("page-start", "page-end", "error"), signals); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + releasePage.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 10000) + public void testNormalStreamCancellationClaimsTerminalBeforeInterruptingWorker() throws Exception { + for (boolean closeRequest : new boolean[]{false, true}) { + CountDownLatch workerEntered = new CountDownLatch(1); + CountDownLatch iteratorClosed = new CountDownLatch(1); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + if (advances.incrementAndGet() <= 2) { + return true; + } + workerEntered.countDown(); + new CountDownLatch(1).await(); + return false; + }); + doAnswer(invocation -> { + iteratorClosed.countDown(); + return null; + }).when(iterator).close(); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + assertTrue(workerEntered.await(1, TimeUnit.SECONDS)); + if (closeRequest) { + response.onNext(ScanStreamReq.newBuilder().setCloseFlag(1).build()); + } else { + response.onCompleted(); + } + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output, never()).onError(any(Throwable.class)); + verify(output).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + } + + @Test + public void testRejectedBatchReportsFailureWithoutSuccessfulCompletion() { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + executor.shutdown(); + StreamObserver output = mock(StreamObserver.class); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + response.onNext(ScanStreamBatchReq.newBuilder() + .setQueryRequest(ScanQueryRequest.getDefaultInstance()).build()); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onCompleted(); + } + + @Test + public void testBatchIteratorFailureClosesResourcesWithoutSuccessfulCompletion() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + StreamObserver output = mock(StreamObserver.class); + ScanIterator iterator = mock(ScanIterator.class); + IllegalStateException failure = new IllegalStateException("iterator failed"); + when(iterator.hasNext()).thenThrow(failure); + ScanBatchResponse response = new ScanBatchResponse(output, mock(HgStoreWrapperEx.class), executor); + Field field = ScanBatchResponse.class.getDeclaredField("iterator"); + field.setAccessible(true); + field.set(response, iterator); + Method send = ScanBatchResponse.class.getDeclaredMethod("sendEntries"); + send.setAccessible(true); + try { + send.invoke(response); + verify(iterator).close(); + verify(output).onError(failure); + verify(output, never()).onCompleted(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testActiveStreamCancellationReleasesBlockedProducer() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + // Consume just one page, then leave the producer waiting for its consumer. + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(1000).build()); + response.onError(Status.CANCELLED.asRuntimeException()); + executor.shutdown(); + assertTrue("cancel must release a producer without its 60-second timeout", + executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelPausedScannerReleasesIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator source = mock(ScanIterator.class); + when(source.hasNext()).thenReturn(true); + when(source.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int originalBodySize = bodySize.getInt(null); + bodySize.setInt(null, 2); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of( + () -> new KVPair<>(mock(QueryCondition.class), source), () -> Long.MAX_VALUE, + ScanQueryRequest.getDefaultInstance(), executor); + // One worker fills its output allowance and pauses while retaining its iterator. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + verify(source, atLeastOnce()).next(); + verify(source, never()).close(); + scan.close(); + scan.close(); + verify(source).close(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + if (scan != null) { + scan.close(); + } + bodySize.setInt(null, originalBodySize); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelOrderedScannerReleasesQueueLockOnEmptyIterator() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator first = mock(ScanIterator.class); + ScanIterator empty = mock(ScanIterator.class); + when(first.hasNext()).thenReturn(true); + when(first.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(empty.hasNext()).thenReturn(false); + + AtomicInteger supplies = new AtomicInteger(); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + int originalBodySize = bodySize.getInt(null); + bodySize.setInt(null, 2); + ParallelScanIterator scan = null; + try { + scan = ParallelScanIterator.of( + () -> { + int supply = supplies.incrementAndGet(); + if (supply == 1) { + return new KVPair<>(mock(QueryCondition.class), first); + } + if (supply == 2) { + return new KVPair<>(mock(QueryCondition.class), empty); + } + return new KVPair<>(mock(QueryCondition.class), null); + }, + () -> 1L, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), + executor); + + // Wait for the scanner task to finish before checking the lock owner. + executor.submit(() -> { }).get(2, TimeUnit.SECONDS); + Field queueLock = ParallelScanIterator.class.getDeclaredField("queueLock"); + queueLock.setAccessible(true); + assertFalse(((ReentrantLock) queueLock.get(scan)).isLocked()); + } finally { + if (scan != null) { + scan.close(); + } + bodySize.setInt(null, originalBodySize); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testParallelOrderedProducerCancelsWithFullOutputQueue() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + ScanIterator source = mock(ScanIterator.class); + Field bodySize = ParallelScanIterator.class.getDeclaredField("maxBodySize"); + bodySize.setAccessible(true); + byte[] value = new byte[bodySize.getInt(null)]; + CountDownLatch full = new CountDownLatch(1); + AtomicInteger rows = new AtomicInteger(); + when(source.hasNext()).thenReturn(true); + when(source.next()).thenAnswer(invocation -> { + if (rows.incrementAndGet() == 5) { + full.countDown(); + } + return RocksDBSession.BackendColumn.of(new byte[4], value); + }); + ParallelScanIterator scan = ParallelScanIterator.of( + () -> new KVPair<>(mock(QueryCondition.class), source), () -> Long.MAX_VALUE, + ScanQueryRequest.newBuilder().setOrderType(ScanOrderType.ORDER_WITHIN_VERTEX).build(), + executor); + try { + assertTrue("four queued batches must fill the single-scanner queue", + full.await(2, TimeUnit.SECONDS)); + scan.close(); + executor.shutdown(); + assertTrue("cancellation must release the blocked ordered producer", + executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(source).close(); + } finally { + scan.close(); + executor.shutdownNow(); + } + } + + @Test + public void testInterruptedOneShotDoesNotSendPartialSuccess() { + for (boolean batch : new boolean[]{false, true}) { + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + try { + if (batch) { + ScanBatchOneShotResponse.scanOneShot(batchRequest(), output, wrapper); + } else { + ScanOneShotResponse.scanOneShot( + ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL).setLimit(10).build(), + output, wrapper); + } + verify(iterator).next(); + assertTrue("scan must preserve an external interrupt", + Thread.currentThread().isInterrupted()); + verify(output).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + Thread.interrupted(); + } + } + } + + @Test(timeout = 5000) + public void testInterruptedBatchWorkerDoesNotCompletePartialPage() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + StreamObserver output = mock(StreamObserver.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + StreamObserver response = ScanBatchResponse3.of(output, wrapper, executor); + try { + response.onNext(batchRequest()); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).next(); + assertCancelled(output); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void testCancellationDoesNotInterruptReceiverAfterHandoff() throws Exception { + assertReceiverInterruptAfterHandoff(false); + } + + @Test(timeout = 10000) + public void testCancellationPreservesExternalReceiverInterruptAfterHandoff() throws Exception { + assertReceiverInterruptAfterHandoff(true); + } + + @Test(timeout = 5000) + public void testReceivePreservesExternalInterrupt() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + FutureTask request = new FutureTask<>(() -> { + Thread.currentThread().interrupt(); + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-external-interrupt-test"); + try { + caller.start(); + assertTrue(request.get(1, TimeUnit.SECONDS)); + assertCancelled(output); + verify(output, never()).onCompleted(); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + } finally { + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + private static void assertReceiverInterruptAfterHandoff(boolean externalInterrupt) throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + CountDownLatch beforePage = new CountDownLatch(1); + CountDownLatch releasePage = new CountDownLatch(1); + AtomicInteger advances = new AtomicInteger(); + when(iterator.hasNext()).thenAnswer(invocation -> { + int advance = advances.incrementAndGet(); + if (advance == 2) { + beforePage.countDown(); + assertTrue(releasePage.await(2, TimeUnit.SECONDS)); + } + return advance <= 2; + }); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + Field responseLock = ScanStreamResponse.class.getDeclaredField("responseLock"); + responseLock.setAccessible(true); + FutureTask request = new FutureTask<>(() -> { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(request, "scan-late-cancellation-test"); + try { + caller.start(); + assertTrue(beforePage.await(1, TimeUnit.SECONDS)); + synchronized (responseLock.get(response)) { + // The real channel hands off a page before its callback reaches responseLock. + releasePage.countDown(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + while (caller.getState() != Thread.State.BLOCKED && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(Thread.State.BLOCKED, caller.getState()); + if (externalInterrupt) { + caller.interrupt(); + } + response.onCompleted(); + } + assertEquals(externalInterrupt, request.get(1, TimeUnit.SECONDS)); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(output).onCompleted(); + verify(output, never()).onError(any(Throwable.class)); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(iterator).close(); + } finally { + releasePage.countDown(); + response.onCompleted(); + executor.shutdownNow(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testInterruptedStreamWorkerDoesNotCompletePartialPage() throws Exception { + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.position()).thenReturn(new byte[4]); + when(iterator.next()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return RocksDBSession.BackendColumn.of(new byte[4], new byte[0]); + }); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + verify(iterator).next(); + assertCancelled(output); + verify(output, never()).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testInterruptedStreamSenderReportsCancellation() throws Exception { + AtomicReference worker = new AtomicReference<>(); + ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1, task -> { + Thread thread = new Thread(task, "scan-interrupted-sender-test"); + worker.set(thread); + return thread; + }); + HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class); + ScanIterator iterator = mock(ScanIterator.class); + when(wrapper.scanAll(anyString(), anyString(), any(byte[].class))).thenReturn(iterator); + when(iterator.hasNext()).thenReturn(true); + when(iterator.next()).thenReturn(RocksDBSession.BackendColumn.of(new byte[4], new byte[0])); + when(iterator.position()).thenReturn(new byte[4]); + AppConfig config = mock(AppConfig.class); + when(config.getServerWaitTime()).thenReturn(60); + StreamObserver output = mock(StreamObserver.class); + CountDownLatch errorReceived = new CountDownLatch(1); + doAnswer(invocation -> { + errorReceived.countDown(); + return null; + }).when(output).onError(any(Throwable.class)); + ScanStreamResponse response = ScanStreamResponse.of(output, wrapper, executor, config); + try { + // Consume one page, then interrupt the actual worker waiting to hand off its next page. + response.onNext(ScanStreamReq.newBuilder().setMethod(ScanMethod.ALL) + .setPageSize(1).setLimit(10).build()); + awaitTimedWaiting(worker.get()); + worker.get().interrupt(); + assertTrue(errorReceived.await(1, TimeUnit.SECONDS)); + executor.shutdown(); + assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS)); + assertCancelled(output); + verify(output).onNext(any(KvPageRes.class)); + verify(output, never()).onCompleted(); + verify(iterator).close(); + } finally { + response.onCompleted(); + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testChannelCloseReleasesWaitingReceiverWithoutInterrupt() throws Exception { + HgChannel channel = HgChannel.of(60); + FutureTask receive = new FutureTask<>(() -> { + assertNull(channel.receive()); + return Thread.currentThread().isInterrupted(); + }); + Thread caller = new Thread(receive, "channel-close-receiver-test"); + try { + caller.start(); + awaitTimedWaiting(caller); + channel.close(); + assertFalse(receive.get(1, TimeUnit.SECONDS)); + assertNull(channel.receive()); + assertFalse(channel.send("after-close")); + } finally { + channel.close(); + caller.join(1000); + } + } + + @Test(timeout = 5000) + public void testChannelCloseReleasesWaitingProducer() throws Exception { + HgChannel channel = HgChannel.of(60); + FutureTask send = new FutureTask<>(() -> channel.send("page")); + Thread producer = new Thread(send, "channel-close-producer-test"); + try { + producer.start(); + awaitTimedWaiting(producer); + channel.close(); + assertFalse(send.get(1, TimeUnit.SECONDS)); + } finally { + channel.close(); + producer.join(1000); + } + } + + @Test(timeout = 5000) + public void testChannelRetainsTimeoutAndNormalHandoff() throws Exception { + HgChannel channel = HgChannel.of(1); + AtomicInteger timeouts = new AtomicInteger(); + long started = System.nanoTime(); + assertNull(channel.receive(timeout -> { + assertEquals(Long.valueOf(1L), timeout); + timeouts.incrementAndGet(); + })); + assertEquals(1, timeouts.get()); + assertTrue("short close polling must not shorten the configured timeout", + System.nanoTime() - started >= TimeUnit.MILLISECONDS.toNanos(900)); + FutureTask send = new FutureTask<>(() -> channel.send("page")); + Thread producer = new Thread(send, "channel-handoff-test"); + try { + producer.start(); + assertEquals("page", channel.receive()); + assertTrue(send.get(1, TimeUnit.SECONDS)); + } finally { + channel.close(); + producer.join(1000); + } + } + + private static void awaitTimedWaiting(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + Thread.State observed = thread.getState(); + while (observed != Thread.State.TIMED_WAITING && System.nanoTime() < deadline) { + Thread.yield(); + observed = thread.getState(); + } + assertEquals(Thread.State.TIMED_WAITING, observed); + } + + private static ScanStreamBatchReq batchRequest() { + return ScanStreamBatchReq.newBuilder().setHeader(Header.newBuilder().setGraph("g")) + .setQueryRequest(ScanQueryRequest.newBuilder().setMethod(ScanMethod.ALL) + .setTable("t").setLimit(10) + .setPerKeyMax(Long.MAX_VALUE) + .setPageSize(1)).build(); + } + + private static void assertCancelled(StreamObserver output) { + ArgumentCaptor failure = ArgumentCaptor.forClass(Throwable.class); + verify(output).onError(failure.capture()); + assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode()); + } + + private static void assertUnavailable(Runnable action) { + try { + action.run(); + fail("Scan admission must be closed"); + } catch (StatusRuntimeException e) { + assertEquals(Status.Code.UNAVAILABLE, e.getStatus().getCode()); + } + } +} 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 2d115e78ed..dd9ccf4add 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 @@ -17,7 +17,9 @@ package org.apache.hugegraph.store.service; +import org.apache.hugegraph.rocksdb.access.RocksDBScanIteratorCloseTest; import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryShutdownTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -25,7 +27,13 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ - OrderedMultiPartitionIteratorTest.class + RocksDBScanIteratorCloseTest.class, + OrderedMultiPartitionIteratorTest.class, + ContextClosedListenerTest.class, + AggregativeQueryShutdownTest.class, + GrpcShutdownBarrierTest.class, + ScanShutdownTest.class, + GraphPartitionScanShutdownTest.class }) @Slf4j