diff --git a/.github/workflows/pd-store-ci.yml b/.github/workflows/pd-store-ci.yml index 47a8abaefc..734af925a4 100644 --- a/.github/workflows/pd-store-ci.yml +++ b/.github/workflows/pd-store-ci.yml @@ -8,6 +8,19 @@ 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: bash hugegraph-server/hugegraph-dist/src/assembly/travis/test-stop-hugegraph-store.sh + rocksdb-compatibility: permissions: contents: read @@ -431,7 +444,7 @@ jobs: mvn test -pl hugegraph-store/hg-store-test -am \ -P store-core-test -Djacoco.sessionId=store-core-test - - name: Run server lifecycle test + - name: Run server lifecycle and shutdown tests run: | mvn test -pl hugegraph-store/hg-store-test -am \ -P store-server-test -Djacoco.sessionId=store-server-test diff --git a/docs/storage-lifecycle.md b/docs/storage-lifecycle.md index 374a4572b3..f225d3c9c1 100644 --- a/docs/storage-lifecycle.md +++ b/docs/storage-lifecycle.md @@ -12,8 +12,8 @@ atomic restore across all databases in a graph. - [RPC and Scan ownership](#rpc-and-scan-ownership): final batches, half-close, cancellation, iterator cleanup and callback dispatch. -- [Shutdown coordination](#shutdown-coordination): admission, worker drain and - application callbacks before database teardown. +- [Shutdown coordination](#shutdown-coordination): admission, worker drain, TTL cleanup + and [Store node shutdown](#store-node-shutdown). - [Standalone snapshot recovery](#standalone-snapshot-recovery): pending records, checkpoint identity, WAL paths and safe retry. - [Storage provider selection and ToplingDB](#storage-provider-selection-and-toplingdb): @@ -114,6 +114,38 @@ releases graph database and schema resources. Keep this ordering when modifying database teardown. Request drain and database release have different responsibilities; reuse their existing owners instead of adding a parallel registry. +### Store node shutdown + +Spring coordinates Store shutdown: new RPCs are refused, active requests are cancelled, +and their application callbacks and scan workers are drained before database teardown. +The shutdown coordinator continues cancellation and independent cleanup even if one +callback fails. It reuses the existing per-call owners. + +TTL cleanup closes the scan before returning its borrowed RocksDB session lease. Both +releases are attempted when either fails. The first cleanup failure is retained, later +failures are attached, and completion is signalled after these release attempts. +Executor termination alone does not establish that native cleanup succeeded. + +Shutdown waits for TTL cleanup as well as ordinary request callbacks. Failed native +releases remain diagnostic drain blockers. The engine stops heartbeat producers and +joins partition Raft services before releasing native databases. Do not add a concurrent +JVM shutdown hook that closes those databases independently of this drain. + +Cancellation can fail in-flight writes. Before planned maintenance, stop new writes and +check outstanding write outcomes. + +For an unpacked distribution, invoke its normal stop script: + +```bash +bash bin/stop-hugegraph-store.sh +``` + +The script 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 data. +A stuck callback or failed iterator close can keep shutdown pending. Inspect Store logs +and a JVM thread dump to identify the remaining owner before retrying. Preserve the data, +WAL and Raft directories, and do not bypass the resource drain to make the process exit. + ## Standalone snapshot recovery ### Installation and pending metadata 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-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-store/README.md b/hugegraph-store/README.md index d038c9d52d..5747d66ba6 100644 --- a/hugegraph-store/README.md +++ b/hugegraph-store/README.md @@ -397,10 +397,15 @@ For Docker and Kubernetes deployment details, see [Deployment Guide](docs/deploy --- +## Stopping a Store node + +See [Store-wide shutdown](../docs/storage-lifecycle.md#store-node-shutdown) for stop commands, +timeout handling, diagnostics and resource drain behavior. + ## Documentation -See [RPC and Scan lifecycle](../docs/storage-lifecycle.md#rpc-and-scan-ownership) for request ownership, -cancellation and feedback semantics. +See the [storage lifecycle guide](../docs/storage-lifecycle.md) for request ownership, +shutdown, recovery and provider configuration. Comprehensive documentation for HugeGraph Store: 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/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 a46a4a9b2c..03b6443f87 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 @@ -69,13 +69,9 @@ public void onApplicationEvent(ContextClosedEvent event) { } this.grpcServers.forEach(Server::shutdownNow); if (cleaner != null) { - // TTL failure ownership and worker drain are added by the native shutdown PR. - if (cleaner.getExecutor() != null) { - cleaner.getExecutor().shutdownNow(); - } - if (cleaner.getScheduler() != null) { - cleaner.getScheduler().shutdownNow(); - } + // 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. @@ -108,7 +104,18 @@ public void onApplicationEvent(ContextClosedEvent event) { if (storeStream != null) { storeStream.awaitScanCleanup(); } - log.info("closed gRPC callbacks, scan and aggregate query workers"); + if (cleaner != null) { + cleaner.awaitCleanup(); + } + log.info("closed gRPC callbacks, scan, aggregate query and TTL workers"); + } + + 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) { 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..ee25a0e1c6 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 @@ -33,6 +33,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.apache.commons.lang3.ArrayUtils; @@ -97,6 +98,7 @@ public class TTLCleaner implements Runnable { @Autowired private HgStoreNodeService service; private final AtomicBoolean running = new AtomicBoolean(false); + private final AtomicReference cleanupFailure = new AtomicReference<>(); private final AppConfig appConfig; private final AppConfig.JobConfig jobConfig; @@ -285,11 +287,12 @@ private Runnable getTask( String table = t.getRight(); TaskInfo taskInfo = counter.get(graph); ScanIterator scan = null; + RocksDBSession session = null; try { Map graphCounter = taskInfo.getTableCounter(); TaskSubmitter submitter = taskInfo.getTaskSubmitter(); AtomicLong tableCounter = graphCounter.get(table); - RocksDBSession session = handler.getSession(id); + session = handler.getSession(id); InnerKeyCreator keyCreator = handler.getKeyCreator(); SessionOperator op = session.sessionOp(); BiFunction judge = getJudge(graph, table); @@ -301,7 +304,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 +318,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); } @@ -326,14 +329,59 @@ private Runnable getTask( String msg = String.format(s, id, graph, table); log.error(msg, e); } finally { - latch.countDown(); - if (scan != null) { - scan.close(); + try { + if (scan != null) { + scan.close(); + } + } catch (RuntimeException | Error failure) { + recordCleanupFailure(failure); + } finally { + try { + if (session != null) { + session.close(); + } + } catch (RuntimeException | Error failure) { + recordCleanupFailure(failure); + } finally { + latch.countDown(); + } } } }; } + private void recordCleanupFailure(Throwable failure) { + this.cleanupFailure.compareAndSet(null, failure); + Throwable first = this.cleanupFailure.get(); + if (first != failure) { + first.addSuppressed(failure); + } + log.error("TTL native cleanup failed; database close stays blocked", failure); + } + + /** Call after workers terminate: a failed native release must remain a shutdown blocker. */ + public void awaitCleanup() { + boolean interrupted = false; + long nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + synchronized (this.cleanupFailure) { + while (this.cleanupFailure.get() != null) { + try { + this.cleanupFailure.wait(5000L); + } catch (InterruptedException e) { + interrupted = true; + } + if (System.nanoTime() - nextLog >= 0) { + log.warn("TTL scan cleanup failed; database close stays blocked", + this.cleanupFailure.get()); + nextLog = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + public ScheduledFuture getFuture() { return future; } 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..6712961320 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryShutdownTest.java @@ -0,0 +1,653 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.node.grpc.query; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +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.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +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 org.apache.hugegraph.rocksdb.access.RocksDBScanIterator; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +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; + +public class AggregativeQueryShutdownTest extends AggregativeQueryTestSupport { + + @Test(timeout = 5000) + public void testContextCloseContinuesAfterCancellationResponseFailure() throws Exception { + assertContextCloseContinuesAfterCancellationFailure( + new IllegalStateException("shutdown response failed")); + } + + @Test(timeout = 5000) + public void testContextCloseContinuesAfterCancellationResponseError() throws Exception { + assertContextCloseContinuesAfterCancellationFailure(new AssertionError("shutdown response 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 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 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 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 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(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); + } + } + + @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(); + } + } + + 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(); + } + } + + 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).onError(org.mockito.ArgumentMatchers.any(Throwable.class)); + 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).onError(org.mockito.ArgumentMatchers.any(Throwable.class)); + verify(broken, never()).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(); + } + } +} 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..d640e56cfe --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ContextClosedListenerTest.java @@ -0,0 +1,102 @@ +/* + * 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 testTtlCleanupBlocksBeanDestructionAfterWorkersTerminate() throws Exception { + CountDownLatch awaitingCleanup = new CountDownLatch(1); + CountDownLatch interruptedCleanup = new CountDownLatch(1); + CountDownLatch releaseCleanup = new CountDownLatch(1); + AtomicBoolean databaseClosed = new AtomicBoolean(); + TTLCleaner cleaner = mock(TTLCleaner.class); + doAnswer(invocation -> { + boolean interrupted = false; + awaitingCleanup.countDown(); + boolean released = false; + while (!released) { + try { + releaseCleanup.await(); + released = true; + } catch (InterruptedException e) { + interrupted = true; + interruptedCleanup.countDown(); + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + return null; + }).when(cleaner).awaitCleanup(); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + try { + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", cleaner); + context.getDefaultListableBeanFactory().registerDisposableBean("database", () -> databaseClosed.set(true)); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return Thread.currentThread().isInterrupted(); + }); + Thread closeThread = new Thread(closing, "test-ttl-context-close"); + closeThread.setDaemon(true); + closeThread.start(); + assertTrue(awaitingCleanup.await(1, TimeUnit.SECONDS)); + closeThread.interrupt(); + assertTrue(interruptedCleanup.await(1, TimeUnit.SECONDS)); + assertFalse("TTL native cleanup must finish before DB destruction", databaseClosed.get()); + assertFalse(closing.isDone()); + releaseCleanup.countDown(); + assertTrue("shutdown must preserve interruption", closing.get(2, TimeUnit.SECONDS)); + assertTrue(databaseClosed.get()); + } finally { + releaseCleanup.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..076a981e58 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/GraphPartitionScanShutdownTest.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.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.assertTrue; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doReturn; +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.when; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import org.apache.hugegraph.store.business.BusinessHandler; +import org.apache.hugegraph.store.business.GraphStoreIterator; +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 extends GraphPartitionScanTestSupport { + + @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(); + boolean interrupted = false; + try { + // Model native work that cannot exit until the read itself finishes. + while (release.getCount() != 0) { + try { + release.await(); + } catch (InterruptedException cancellation) { + interrupted = true; + } + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + 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)); + verify(iterator, never()).close(); + 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)); + } + } +} 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..0a1003ceb0 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanShutdownTest.java @@ -0,0 +1,159 @@ +/* + * 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.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +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.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +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.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.junit.Test; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + +import io.grpc.Status; +import io.grpc.stub.StreamObserver; + +public class ScanShutdownTest extends ScanTestSupport { + + @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(); + } + } + + + + +} 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 9804cde7b0..082f9d30ba 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 @@ -20,6 +20,7 @@ import org.apache.hugegraph.rocksdb.access.RocksDBScanIteratorCloseTest; import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest; import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryLifecycleTest; +import org.apache.hugegraph.store.node.grpc.query.AggregativeQueryShutdownTest; import org.apache.hugegraph.store.node.grpc.query.UnaryQueryLifecycleTest; import org.apache.hugegraph.store.node.grpc.BatchScanHalfCloseTest; import org.junit.runner.RunWith; @@ -32,13 +33,18 @@ RocksDBScanIteratorCloseTest.class, OrderedMultiPartitionIteratorTest.class, ScanContextCleanupTest.class, + ContextClosedListenerTest.class, + TTLCleanerCleanupTest.class, GrpcShutdownBarrierTest.class, PartitionEngineShutdownTest.class, AggregativeQueryLifecycleTest.class, + AggregativeQueryShutdownTest.class, UnaryQueryLifecycleTest.class, BatchScanHalfCloseTest.class, ScanLifecycleTest.class, - GraphPartitionScanLifecycleTest.class + ScanShutdownTest.class, + GraphPartitionScanLifecycleTest.class, + GraphPartitionScanShutdownTest.class }) @Slf4j diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/TTLCleanerCleanupTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/TTLCleanerCleanupTest.java new file mode 100644 index 0000000000..489eff98e9 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/TTLCleanerCleanupTest.java @@ -0,0 +1,374 @@ +/* + * 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.assertArrayEquals; +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.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Future; +import java.util.concurrent.FutureTask; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; + +import org.apache.commons.configuration2.MapConfiguration; +import org.apache.commons.lang3.tuple.Triple; +import org.apache.hugegraph.config.HugeConfig; +import org.apache.hugegraph.config.OptionSpace; +import org.apache.hugegraph.rocksdb.access.RocksDBOptions; +import org.apache.hugegraph.rocksdb.access.RocksDBSession; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.rocksdb.access.SessionOperator; +import org.apache.hugegraph.store.business.BusinessHandlerImpl; +import org.apache.hugegraph.store.business.InnerKeyCreator; +import org.apache.hugegraph.store.constant.HugeServerTables; +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.query.AggregativeQueryService; +import org.apache.hugegraph.store.node.listener.ContextClosedListener; +import org.apache.hugegraph.store.node.task.TTLCleaner; +import org.apache.hugegraph.store.node.task.ttl.TaskInfo; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.InOrder; +import org.rocksdb.RocksDB; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; + +public class TTLCleanerCleanupTest { + + @Rule + public TemporaryFolder directory = new TemporaryFolder(); + + @Test(timeout = 5000) + public void testCompletionWaitsForNativeClose() throws Exception { + TTLCleaner cleaner = newCleaner(); + CountDownLatch completed = new CountDownLatch(1); + CountDownLatch closing = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ScanIterator scan = mock(ScanIterator.class); + doAnswer(invocation -> { + closing.countDown(); + release.await(); + return null; + }).when(scan).close(); + RocksDBSession session = mock(RocksDBSession.class); + FutureTask task = new FutureTask<>(task(cleaner, scan, session, completed), null); + Thread worker = new Thread(task, "test-ttl-close"); + worker.setDaemon(true); + try { + worker.start(); + assertTrue(closing.await(1, TimeUnit.SECONDS)); + assertEquals("TTL must not compact before native close finishes", 1L, + completed.getCount()); + release.countDown(); + task.get(1, TimeUnit.SECONDS); + assertEquals(0L, completed.getCount()); + InOrder order = inOrder(scan, session); + order.verify(scan).close(); + order.verify(session).close(); + cleaner.awaitCleanup(); + } finally { + release.countDown(); + cleaner.getScheduler().shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testFailedCloseStaysBlockedAfterInterruption() throws Exception { + assertFailedCloseStaysBlocked(true, false); + assertFailedCloseStaysBlocked(false, true); + assertFailedCloseStaysBlocked(true, true); + } + + private static void assertFailedCloseStaysBlocked(boolean scanFails, boolean sessionFails) throws Exception { + TTLCleaner cleaner = newCleaner(); + RuntimeException failure = new IllegalStateException("native release failed"); + ScanIterator scan = mock(ScanIterator.class); + RocksDBSession session = mock(RocksDBSession.class); + RuntimeException sessionFailure = scanFails ? new IllegalStateException("session release failed") : failure; + if (scanFails) { + doAnswer(invocation -> { throw failure; }).when(scan).close(); + } + if (sessionFails) { + doAnswer(invocation -> { throw sessionFailure; }).when(session).close(); + } + CountDownLatch completed = new CountDownLatch(1); + CountDownLatch awaiting = new CountDownLatch(1); + FutureTask cleanup = new FutureTask<>(() -> { + awaiting.countDown(); + cleaner.awaitCleanup(); + return Thread.currentThread().isInterrupted(); + }); + Thread shutdown = new Thread(cleanup, "test-ttl-cleanup-failure"); + shutdown.setDaemon(true); + Field field = TTLCleaner.class.getDeclaredField("cleanupFailure"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + AtomicReference retained = (AtomicReference) field.get(cleaner); + try { + task(cleaner, scan, session, completed).run(); + InOrder order = inOrder(scan, session); + order.verify(scan).close(); + order.verify(session).close(); + if (scanFails && sessionFails) { + assertEquals(1, failure.getSuppressed().length); + assertSame(sessionFailure, failure.getSuppressed()[0]); + } + assertEquals(0L, completed.getCount()); + assertSame(failure, retained.get()); + shutdown.start(); + assertTrue(awaiting.await(1, TimeUnit.SECONDS)); + shutdown.interrupt(); + assertFalse("interruption cannot authorize native storage destruction", + cleanup.isDone()); + assertSame(failure, retained.get()); + } finally { + // Production never clears this failure. Release only this test's daemon waiter. + synchronized (retained) { + retained.set(null); + retained.notifyAll(); + } + shutdown.interrupt(); + assertTrue(cleanup.get(1, TimeUnit.SECONDS)); + cleaner.getScheduler().shutdownNow(); + } + } + + @Test(timeout = 5000) + public void testNativeSessionCloneAndScanReleaseAllReferences() throws Exception { + RocksDB.loadLibrary(); + OptionSpace.register("rocksdb", "org.apache.hugegraph.rocksdb.access.RocksDBOptions"); + RocksDBOptions.instance(); + HugeConfig config = new HugeConfig(new MapConfiguration( + Collections.singletonMap("rocksdb.write_buffer_size", "1048576"))); + TTLCleaner cleaner = newCleaner(); + RocksDBSession owner = new RocksDBSession(config, directory.newFolder("ttl-native").getAbsolutePath(), + "ttl-native", 0L); + RocksDBSession lease = owner.clone(); + try { + owner.checkTable(HugeServerTables.VERTEX_TABLE); + assertEquals(2, owner.getRefCount()); + BusinessHandlerImpl handler = mock(BusinessHandlerImpl.class); + when(handler.getSession(1)).thenReturn(lease); + InnerKeyCreator keys = mock(InnerKeyCreator.class); + when(keys.getStartKey(1, "graph")).thenReturn(new byte[]{0}); + when(keys.getEndKey(1, "graph")).thenReturn(new byte[]{1}); + when(handler.getKeyCreator()).thenReturn(keys); + CountDownLatch completed = new CountDownLatch(1); + task(cleaner, handler, completed).run(); + assertEquals(0L, completed.getCount()); + cleaner.awaitCleanup(); + assertEquals("TTL must release both the iterator lease and cloned session", 1, owner.getRefCount()); + RocksDB database = owner.getDB(); + assertTrue(database.isOwningHandle()); + owner.close(); + assertEquals(0, owner.getRefCount()); + assertNull(owner.getDB()); + assertFalse(database.isOwningHandle()); + } finally { + lease.close(); + owner.close(); + cleaner.getScheduler().shutdownNow(); + } + } + + @Test(timeout = 30000) + public void testActiveNativeTtlScanDrainsBeforeInterruptedSpringClose() throws Exception { + RocksDB.loadLibrary(); + OptionSpace.register("rocksdb", "org.apache.hugegraph.rocksdb.access.RocksDBOptions"); + RocksDBOptions.instance(); + HugeConfig config = new HugeConfig(new MapConfiguration( + Collections.singletonMap("rocksdb.write_buffer_size", "1048576"))); + RocksDBSession owner = new RocksDBSession(config, directory.newFolder("ttl-spring").getAbsolutePath(), + "ttl-spring", 0L); + RocksDBSession lease = owner.clone(); + RocksDB database = owner.getDB(); + TTLCleaner cleaner = spy(newCleaner()); + CountDownLatch judging = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + CountDownLatch workerInterrupted = new CountDownLatch(1); + CountDownLatch closeInterrupted = new CountDownLatch(1); + ThreadPoolExecutor workers = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>()) { + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + try { + return super.awaitTermination(timeout, unit); + } catch (InterruptedException e) { + closeInterrupted.countDown(); + throw e; + } + } + }; + Field executor = TTLCleaner.class.getDeclaredField("executor"); + executor.setAccessible(true); + executor.set(cleaner, workers); + Field registry = RocksDBSession.class.getDeclaredField("iteratorMap"); + registry.setAccessible(true); + Map iterators = (Map) registry.get(owner); + BiFunction judge = (key, value) -> { + assertArrayEquals(new byte[]{10}, key); + assertArrayEquals(new byte[]{42}, value); + judging.countDown(); + boolean interrupted = false; + while (release.getCount() != 0) { + try { + release.await(); + } catch (InterruptedException e) { + interrupted = true; + workerInterrupted.countDown(); + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + return false; + }; + doReturn(judge).when(cleaner).getJudge("graph", HugeServerTables.VERTEX_TABLE); + BusinessHandlerImpl handler = mock(BusinessHandlerImpl.class); + when(handler.getSession(1)).thenReturn(lease); + InnerKeyCreator keys = mock(InnerKeyCreator.class); + when(keys.getStartKey(1, "graph")).thenReturn(new byte[]{0}); + when(keys.getEndKey(1, "graph")).thenReturn(new byte[]{1}); + when(handler.getKeyCreator()).thenReturn(keys); + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(); + CountDownLatch completed = new CountDownLatch(1); + Future worker = null; + FutureTask destruction = new FutureTask<>(() -> { + assertEquals(0L, completed.getCount()); + assertTrue(iterators.isEmpty()); + assertEquals(1, owner.getRefCount()); + assertTrue(database.isOwningHandle()); + owner.close(); + assertEquals(0, owner.getRefCount()); + assertFalse(database.isOwningHandle()); + return null; + }); + FutureTask closing = new FutureTask<>(() -> { + context.close(); + return Thread.currentThread().isInterrupted(); + }); + try { + owner.checkTable(HugeServerTables.VERTEX_TABLE); + SessionOperator operator = owner.sessionOp(); + operator.prepare(); + operator.put(HugeServerTables.VERTEX_TABLE, new byte[]{0, 0, 10, 0, 1}, new byte[]{42}); + operator.put(HugeServerTables.VERTEX_TABLE, new byte[]{0, 0, 11, 0, 1}, new byte[]{43}); + operator.commit(); + context.getBeanFactory().registerSingleton("storeStream", mock(HgStoreStreamImpl.class)); + context.getBeanFactory().registerSingleton("queryService", mock(AggregativeQueryService.class)); + context.getBeanFactory().registerSingleton("cleaner", cleaner); + context.getDefaultListableBeanFactory().registerDisposableBean("database", destruction::run); + context.register(ContextClosedListener.class, GrpcShutdownBarrier.class); + context.refresh(); + worker = workers.submit(task(cleaner, handler, completed)); + assertTrue(judging.await(2, TimeUnit.SECONDS)); + assertFalse("the nonempty scan must still own its native iterator", iterators.isEmpty()); + Thread closeThread = new Thread(closing, "test-active-ttl-spring-close"); + closeThread.start(); + assertTrue(workerInterrupted.await(2, TimeUnit.SECONDS)); + closeThread.interrupt(); + assertTrue(closeInterrupted.await(2, TimeUnit.SECONDS)); + assertFalse(closing.isDone()); + assertFalse(destruction.isDone()); + assertEquals(1L, completed.getCount()); + assertFalse(iterators.isEmpty()); + assertTrue(database.isOwningHandle()); + release.countDown(); + worker.get(2, TimeUnit.SECONDS); + assertTrue("Spring close must preserve interruption", closing.get(2, TimeUnit.SECONDS)); + destruction.get(2, TimeUnit.SECONDS); + assertTrue(workers.isTerminated()); + assertNull(owner.getDB()); + } finally { + release.countDown(); + workers.shutdownNow(); + try { + if (worker != null) { + worker.get(2, TimeUnit.SECONDS); + } + } finally { + try { + context.close(); + } finally { + lease.close(); + owner.close(); + cleaner.getScheduler().shutdownNow(); + } + } + } + } + + private static TTLCleaner newCleaner() { + AppConfig config = mock(AppConfig.class); + AppConfig.JobConfig job = mock(AppConfig.JobConfig.class); + when(config.getJobConfig()).thenReturn(job); + when(job.getStartTime()).thenReturn(19); + when(job.getBatchSize()).thenReturn(128); + return new TTLCleaner(config); + } + + private static Runnable task(TTLCleaner cleaner, ScanIterator scan, RocksDBSession session, + CountDownLatch completed) throws Exception { + BusinessHandlerImpl handler = mock(BusinessHandlerImpl.class); + SessionOperator operator = mock(SessionOperator.class); + when(handler.getSession(1)).thenReturn(session); + when(handler.getKeyCreator()).thenReturn(mock(InnerKeyCreator.class)); + when(session.sessionOp()).thenReturn(operator); + when(operator.scan(anyString(), any(), any(), anyInt())).thenReturn(scan); + return task(cleaner, handler, completed); + } + + private static Runnable task(TTLCleaner cleaner, BusinessHandlerImpl handler, + CountDownLatch completed) throws Exception { + TaskInfo info = mock(TaskInfo.class); + when(info.getTableCounter()).thenReturn(new ConcurrentHashMap<>( + Collections.singletonMap(HugeServerTables.VERTEX_TABLE, new AtomicLong()))); + Method method = TTLCleaner.class.getDeclaredMethod("getTask", BusinessHandlerImpl.class, + CountDownLatch.class, Triple.class, Map.class, Map.class); + method.setAccessible(true); + return (Runnable) method.invoke(cleaner, handler, completed, + Triple.of(1, "graph", HugeServerTables.VERTEX_TABLE), + Collections.singletonMap("graph", info), + Collections.singletonMap(1, new AtomicLong())); + } +}