You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hugegraph.apache.org by ji...@apache.org on 2026/10/10 09:50:06 UTC
(hugegraph) 08/11: fix(store): stop cancelled batch readers
This is an automated email from the ASF dual-hosted git repository.
imbajin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hugegraph.git
commit bd584058f21f15f14d53eb17000d94f80f73991a
Author: imbajin <ji...@apache.org>
AuthorDate: Sat Oct 10 06:07:16 2026 +0800
fix(store): stop cancelled batch readers
- interrupt owned batch readers before iterator locks
- keep closed partition iterators terminal
- refresh the timeout when accepting a response
- cover blocked shutdown and parsing boundaries
---
.../store/client/query/CommonKvStreamObserver.java | 1 +
.../store/business/MultiPartitionIterator.java | 7 +-
.../store/node/grpc/ParallelScanIterator.java | 37 +++++++
.../store/node/grpc/ScanBatchResponse3.java | 33 +++++-
.../client/query/CommonKvStreamObserverTest.java | 49 +++++++++
.../grpc/query/AggregativeQueryLifecycleTest.java | 60 +++++++++++
.../hugegraph/store/service/ScanLifecycleTest.java | 114 +++++++++++++++++++++
7 files changed, 298 insertions(+), 3 deletions(-)
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 3f83c1f6c..21da9a39f 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
@@ -107,6 +107,7 @@ public class CommonKvStreamObserver<R, T> implements StreamObserver<R> {
return;
}
this.parsing++;
+ this.current = System.nanoTime();
}
Iterator<T> iterator = null;
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 030e72407..54fd80be7 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
@@ -46,6 +46,7 @@ public class MultiPartitionIterator implements ScanIterator {
private Integer positionPartitionId;
private byte[] positionKey;
private RuntimeException cleanupFailure;
+ private boolean closed;
private MultiPartitionIterator(List<Integer> partitionIds,
BiFunction<Integer, byte[], ScanIterator> supplier) {
@@ -100,7 +101,7 @@ public class MultiPartitionIterator implements ScanIterator {
}
private void init() {
- if (this.iterator == null) {
+ if (!this.closed && this.iterator == null) {
this.iterator = this.getIterator();
}
}
@@ -173,6 +174,7 @@ public class MultiPartitionIterator implements ScanIterator {
@Override
public void close() {
+ this.closed = true;
try {
closeCurrentIterator();
} catch (RuntimeException | Error failure) {
@@ -231,6 +233,9 @@ public class MultiPartitionIterator implements ScanIterator {
*/
public List<ScanIterator> getIterators() {
List<ScanIterator> opened = new ArrayList<>();
+ if (this.closed) {
+ return opened;
+ }
try {
for (int id : this.partitions) {
ScanIterator child = this.supplier.apply(id, getPositionKey(id));
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 02a05292a..65823da8a 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
@@ -190,6 +190,10 @@ public class ParallelScanIterator implements ScanIterator {
synchronized (scanners) {
pending = new ArrayList<>(scanners);
}
+ // Stop active scanner reads before close() waits for their iterator locks.
+ for (KVScanner scanner : pending) {
+ scanner.interruptReader();
+ }
// Never hold a registry lock while waiting for a scanner's iterator lock.
for (KVScanner scanner : pending) {
scanner.close();
@@ -363,12 +367,40 @@ public class ParallelScanIterator implements ScanIterator {
class KVScanner {
private final ReentrantLock iteratorLock = new ReentrantLock();
+ private final Object readerLock = new Object();
+ private Thread reader;
private ScanIterator iterator = null;
private QueryCondition query = null;
private long limit;
private long counter;
private volatile boolean closed = false;
+ private void interruptReader() {
+ synchronized (this.readerLock) {
+ if (this.reader != null && this.reader != Thread.currentThread()) {
+ this.reader.interrupt();
+ }
+ }
+ }
+
+ private boolean startReader() {
+ synchronized (this.readerLock) {
+ if (this.reader != null) {
+ return false;
+ }
+ this.reader = Thread.currentThread();
+ return true;
+ }
+ }
+
+ private void finishReader() {
+ synchronized (this.readerLock) {
+ if (this.reader == Thread.currentThread()) {
+ this.reader = null;
+ }
+ }
+ }
+
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) {
@@ -387,6 +419,9 @@ public class ParallelScanIterator implements ScanIterator {
boolean canNext = true;
ArrayList<KV> dataList = new ArrayList<>(batchSize);
dataList.ensureCapacity(batchSize);
+ if (!startReader()) {
+ return;
+ }
iteratorLock.lock();
try {
long entriesSize = 0, bodySize = 0;
@@ -436,6 +471,8 @@ public class ParallelScanIterator implements ScanIterator {
queueLock.unlock();
}
iteratorLock.unlock();
+ // Clear ownership before a paused scanner can be scheduled again.
+ finishReader();
if (iterator != null && counter < limit && !closed && !finished) {
suspendScanner(this);
} else {
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 07ffcf366..3c39ac40f 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
@@ -288,6 +288,8 @@ public class ScanBatchResponse3 {
private final long limit;
private final ScanLifecycle lifecycle;
private final AtomicBoolean iteratorClosed = new AtomicBoolean();
+ private final Object readerLock = new Object();
+ private Thread reader;
private long packageSize;
private long counter;
private volatile boolean inputCompleted;
@@ -362,6 +364,12 @@ public class ScanBatchResponse3 {
void breakdown() {
this.breakdown.set(true);
+ // Interrupt before waiting for the iterator monitor held during hasNext().
+ synchronized (this.readerLock) {
+ if (this.reader != null && this.reader != Thread.currentThread()) {
+ this.reader.interrupt();
+ }
+ }
synchronized (this.iterator) {
this.iterator.notify();
}
@@ -379,6 +387,19 @@ public class ScanBatchResponse3 {
return;
}
+ boolean startReading;
+ synchronized (this.readerLock) {
+ startReading = !this.breakdown.get() && !this.completeFlag.get();
+ if (startReading) {
+ this.reader = Thread.currentThread();
+ }
+ }
+ if (!startReading) {
+ this.completeFlag.set(true);
+ closeIterator();
+ return;
+ }
+
this.workingLock.lock();
try {
@@ -468,8 +489,16 @@ public class ScanBatchResponse3 {
log.error("Failed to do while for scanning, cause by:", t);
this.deliverer.error("Failed to finish scanning ", t);
} finally {
- this.workingLock.unlock();
- closeIterator();
+ try {
+ this.workingLock.unlock();
+ closeIterator();
+ } finally {
+ synchronized (this.readerLock) {
+ if (this.reader == Thread.currentThread()) {
+ this.reader = null;
+ }
+ }
+ }
}
}
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
index 181c82d4c..71e88f7fa 100644
--- 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
@@ -110,6 +110,55 @@ public class CommonKvStreamObserverTest {
}
}
+ @Test(timeout = 5000)
+ public void testAcceptedResponseRefreshesTimeoutBeforeParsing() throws Exception {
+ CountDownLatch parsing = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicInteger polls = new AtomicInteger();
+ RecordingQueue queue = new RecordingQueue() {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public Iterator<Integer> poll(long timeout, TimeUnit unit) throws InterruptedException {
+ if (polls.getAndIncrement() == 0) {
+ // Exercise the timeout check while the accepted response is still being parsed.
+ return null;
+ }
+ release.countDown();
+ return super.poll(timeout, unit);
+ }
+ };
+ CommonKvStreamObserver<List<Integer>, Integer> observer = parsingObserver(parsing, release);
+ observer.setTimeout(1000);
+ Field queueField = CommonKvStreamObserver.class.getDeclaredField("queue");
+ queueField.setAccessible(true);
+ queueField.set(observer, queue);
+ Field current = CommonKvStreamObserver.class.getDeclaredField("current");
+ current.setAccessible(true);
+ current.setLong(observer, System.nanoTime() - TimeUnit.SECONDS.toNanos(2));
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+ Thread producer = new Thread(() -> {
+ try {
+ observer.onNext(Collections.singletonList(7));
+ } catch (Throwable error) {
+ failure.set(error);
+ }
+ });
+ producer.start();
+ try {
+ assertTrue(parsing.await(2, TimeUnit.SECONDS));
+ assertEquals(Integer.valueOf(7), observer.consume().next());
+ producer.join(2000);
+ assertFalse(producer.isAlive());
+ assertEquals(null, failure.get());
+ assertEquals(null, observer.consume());
+ } finally {
+ release.countDown();
+ producer.join(2000);
+ observer.clear();
+ }
+ }
+
@Test(timeout = 5000)
public void testCompletionWaitsForResponseBeingParsed() throws Exception {
CountDownLatch parsing = new CountDownLatch(1);
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java
index a0affeca5..688537639 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/query/AggregativeQueryLifecycleTest.java
@@ -702,6 +702,66 @@ public class AggregativeQueryLifecycleTest extends AggregativeQueryTestSupport {
verify(nonempty).close();
}
+ @Test
+ public void testClosedSequentialIteratorDoesNotOpenRemainingPartitions() {
+ for (boolean initialized : new boolean[]{false, true}) {
+ ScanIterator first = mock(ScanIterator.class);
+ ScanIterator second = mock(ScanIterator.class);
+ when(first.hasNext()).thenReturn(true);
+ when(second.hasNext()).thenReturn(true);
+ AtomicInteger opened = new AtomicInteger();
+ MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, key) -> {
+ opened.incrementAndGet();
+ return id == 1 ? first : second;
+ });
+ if (initialized) {
+ assertTrue(iterator.hasNext());
+ }
+ iterator.close();
+ iterator.close();
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.isValid());
+ assertEquals(0L, iterator.count());
+ assertTrue(iterator.getIterators().isEmpty());
+ org.junit.Assert.assertThrows(java.util.NoSuchElementException.class, iterator::next);
+ assertEquals(initialized ? 1 : 0, opened.get());
+ if (initialized) {
+ verify(first).close();
+ } else {
+ verify(first, never()).close();
+ }
+ verify(second, never()).close();
+ }
+ }
+
+ @Test
+ public void testSequentialCloseFailureDoesNotReopenRemainingPartitions() {
+ ScanIterator first = mock(ScanIterator.class);
+ ScanIterator second = mock(ScanIterator.class);
+ when(first.hasNext()).thenReturn(true);
+ when(second.hasNext()).thenReturn(true);
+ IllegalStateException cleanup = new IllegalStateException("partition cleanup failed");
+ doThrow(cleanup).when(first).close();
+ AtomicInteger opened = new AtomicInteger();
+ MultiPartitionIterator iterator = MultiPartitionIterator.of(Arrays.asList(1, 2), (id, key) -> {
+ opened.incrementAndGet();
+ return id == 1 ? first : second;
+ });
+ assertTrue(iterator.hasNext());
+ IllegalStateException retained = org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close);
+ org.junit.Assert.assertSame(cleanup, retained.getCause());
+ assertFalse(iterator.hasNext());
+ assertFalse(iterator.isValid());
+ assertEquals(0L, iterator.count());
+ assertTrue(iterator.getIterators().isEmpty());
+ org.junit.Assert.assertThrows(java.util.NoSuchElementException.class, iterator::next);
+ org.junit.Assert.assertSame(retained,
+ org.junit.Assert.assertThrows(IllegalStateException.class, iterator::close));
+ assertEquals(1, opened.get());
+ verify(first).close();
+ verify(second, never()).close();
+ }
+
@Test
public void testSequentialInitializationFailureReleasesItsCreatedIterator() {
ScanIterator child = mock(ScanIterator.class);
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java
index e6de90080..6e9f791f1 100644
--- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ScanLifecycleTest.java
@@ -731,6 +731,120 @@ public class ScanLifecycleTest extends ScanTestSupport {
}
}
+ @Test(timeout = 5000)
+ public void testBatchShutdownInterruptsBlockedReaderBeforeIteratorClose() 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);
+ CountDownLatch reading = new CountDownLatch(1);
+ CountDownLatch interrupted = new CountDownLatch(1);
+ CountDownLatch readExited = new CountDownLatch(1);
+ AtomicBoolean closedAfterRead = new AtomicBoolean();
+ when(iterator.hasNext()).thenAnswer(invocation -> {
+ reading.countDown();
+ try {
+ new CountDownLatch(1).await();
+ return true;
+ } catch (InterruptedException expected) {
+ interrupted.countDown();
+ return false;
+ } finally {
+ readExited.countDown();
+ }
+ });
+ doAnswer(invocation -> {
+ closedAfterRead.set(readExited.getCount() == 0);
+ return null;
+ }).when(iterator).close();
+
+ HgStoreStreamImpl service = scanService(executor, wrapper);
+ StreamObserver<KvPageRes> output = mock(StreamObserver.class);
+ StreamObserver<ScanStreamBatchReq> input = service.scanBatch(output);
+ FutureTask<Void> shutdown = new FutureTask<>(() -> {
+ service.shutdownScans();
+ return null;
+ });
+ Thread shutdownThread = new Thread(shutdown, "batch-scan-shutdown-test");
+ try {
+ input.onNext(batchRequest());
+ assertTrue("batch reader must enter its blocking iterator read",
+ reading.await(1, TimeUnit.SECONDS));
+ shutdownThread.start();
+ shutdown.get(2, TimeUnit.SECONDS);
+ assertTrue(interrupted.await(1, TimeUnit.SECONDS));
+ service.awaitScanCleanup();
+ assertTrue(closedAfterRead.get());
+ assertTrue(scanRegistry(service).isEmpty());
+ assertCancelled(output);
+ verify(output, never()).onCompleted();
+ verify(iterator).close();
+ assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS));
+ } finally {
+ executor.shutdownNow();
+ shutdownThread.interrupt();
+ shutdownThread.join(1000);
+ }
+ }
+
+ @Test(timeout = 5000)
+ public void testEncodedBatchShutdownInterruptsBlockedScannerBeforeIteratorClose() 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);
+ CountDownLatch reading = new CountDownLatch(1);
+ CountDownLatch interrupted = new CountDownLatch(1);
+ CountDownLatch readExited = new CountDownLatch(1);
+ AtomicBoolean closedAfterRead = new AtomicBoolean();
+ when(iterator.hasNext()).thenAnswer(invocation -> {
+ reading.countDown();
+ try {
+ new CountDownLatch(1).await();
+ return true;
+ } catch (InterruptedException expected) {
+ interrupted.countDown();
+ return false;
+ } finally {
+ readExited.countDown();
+ }
+ });
+ doAnswer(invocation -> {
+ closedAfterRead.set(readExited.getCount() == 0);
+ return null;
+ }).when(iterator).close();
+
+ HgStoreStreamImpl service = scanService(executor, wrapper);
+ StreamObserver<KvStream> output = mock(StreamObserver.class);
+ StreamObserver<ScanStreamBatchReq> input = service.scanBatch2(output);
+ FutureTask<Void> shutdown = new FutureTask<>(() -> {
+ service.shutdownScans();
+ return null;
+ });
+ Thread shutdownThread = new Thread(shutdown, "encoded-batch-scan-shutdown-test");
+ try {
+ input.onNext(batchRequest());
+ assertTrue("parallel scanner must enter its blocking iterator read",
+ reading.await(1, TimeUnit.SECONDS));
+ shutdownThread.start();
+ shutdown.get(2, TimeUnit.SECONDS);
+ assertTrue(interrupted.await(1, TimeUnit.SECONDS));
+ service.awaitScanCleanup();
+ assertTrue(closedAfterRead.get());
+ assertTrue(scanRegistry(service).isEmpty());
+ ArgumentCaptor<Throwable> failure = ArgumentCaptor.forClass(Throwable.class);
+ verify(output).onError(failure.capture());
+ assertEquals(Status.Code.CANCELLED, Status.fromThrowable(failure.getValue()).getCode());
+ verify(output, never()).onCompleted();
+ verify(iterator).close();
+ assertTrue(executor.awaitTermination(1, TimeUnit.SECONDS));
+ } finally {
+ executor.shutdownNow();
+ shutdownThread.interrupt();
+ shutdownThread.join(1000);
+ }
+ }
+
@Test(timeout = 5000)
public void testOneShotObservesContextCancellationAndClosesIterator() {
HgStoreWrapperEx wrapper = mock(HgStoreWrapperEx.class);