From 9dc3fa07d18e6ee19af023854cc027ec8357157a Mon Sep 17 00:00:00 2001 From: contrueCT Date: Thu, 6 Aug 2026 17:29:42 +0800 Subject: [PATCH 1/6] fix(hstore): preserve range index order across partitions --- .../hugegraph/backend/page/IdHolder.java | 25 +- .../hugegraph/backend/page/QueryList.java | 9 +- .../hugegraph/backend/query/QueryResults.java | 19 +- .../backend/tx/GraphIndexTransaction.java | 16 +- .../backend/store/hstore/HstoreFeatures.java | 3 +- .../backend/store/hstore/HstoreSessions.java | 5 + .../store/hstore/HstoreSessionsImpl.java | 25 ++ .../backend/store/hstore/HstoreTable.java | 44 +- .../store/hstore/HstoreSessionsImplTest.java | 111 +++++ .../backend/store/hstore/HstoreTableTest.java | 180 ++++++++ .../backend/tx/GraphIndexTransactionTest.java | 49 +++ .../apache/hugegraph/unit/UnitTestSuite.java | 6 + .../hugegraph/unit/core/IdHolderTest.java | 105 +++++ .../hugegraph/unit/core/QueryResultsTest.java | 81 ++++ .../org/apache/hugegraph/store/HgKvStore.java | 13 + .../hugegraph/store/client/NodeTkv.java | 23 +- .../store/client/NodeTxSessionProxy.java | 93 +++++ .../store/client/OrderedKvIterator.java | 193 +++++++++ .../store/client/ClientSuiteTest.java | 4 +- .../store/client/NodeTxSessionProxyTest.java | 393 ++++++++++++++++++ .../store/client/OrderedKvIteratorTest.java | 185 +++++++++ 21 files changed, 1545 insertions(+), 37 deletions(-) create mode 100644 hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreTableTest.java create mode 100644 hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransactionTest.java create mode 100644 hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/IdHolderTest.java create mode 100644 hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java create mode 100644 hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/IdHolder.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/IdHolder.java index b420648767..ab64959a01 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/IdHolder.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/IdHolder.java @@ -35,11 +35,17 @@ public abstract class IdHolder { protected final Query query; + private final boolean keepOrder; protected boolean exhausted; public IdHolder(Query query) { + this(query, false); + } + + public IdHolder(Query query, boolean keepOrder) { E.checkNotNull(query, "query"); this.query = query; + this.keepOrder = keepOrder; this.exhausted = false; } @@ -48,7 +54,7 @@ public Query query() { } public boolean keepOrder() { - return false; + return this.keepOrder; } @Override @@ -97,7 +103,13 @@ public static class PagingIdHolder extends IdHolder { public PagingIdHolder(ConditionQuery query, Function fetcher) { - super(query.copy()); + this(query, fetcher, false); + } + + public PagingIdHolder(ConditionQuery query, + Function fetcher, + boolean keepOrder) { + super(query.copy(), keepOrder); E.checkArgument(query.paging(), "Query '%s' must include page info", query); this.fetcher = fetcher; @@ -142,7 +154,14 @@ public static class BatchIdHolder extends IdHolder public BatchIdHolder(ConditionQuery query, Iterator entries, Function> fetcher) { - super(query); + this(query, entries, fetcher, false); + } + + public BatchIdHolder(ConditionQuery query, + Iterator entries, + Function> fetcher, + boolean keepOrder) { + super(query, keepOrder); this.entries = entries; this.fetcher = fetcher; this.count = 0L; diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java index d1e11e9220..b581a6e714 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java @@ -258,7 +258,7 @@ private QueryResults each(IdHolder holder) { return null; } - return this.queryByIndexIds(ids); + return this.queryByIndexIds(ids, holder.keepOrder()); }); } @@ -275,7 +275,8 @@ public PageResults iterator(int index, String page, long pageSize) { return PageResults.emptyIterator(); } - QueryResults results = this.queryByIndexIds(pageIds.ids()); + QueryResults results = this.queryByIndexIds(pageIds.ids(), + holder.keepOrder()); return new PageResults<>(results, pageIds.pageState()); } @@ -311,10 +312,6 @@ private void updateResultsFilter(Query query) { } } - private QueryResults queryByIndexIds(Set ids) { - return this.queryByIndexIds(ids, false); - } - private QueryResults queryByIndexIds(Set ids, boolean inOrder) { IdQuery query = new IdQuery(parent(), ids); query.mustSortByInput(inOrder); diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java index a03e5c9aee..aa7fb70263 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java @@ -102,12 +102,12 @@ public Iterator keepInputOrderIfNeeded( return origin; } Collection ids; - if (!this.mustSortByInputIds() || this.paging() || + if (!this.mustSortByInputIds() || (ids = this.queryIds()).size() <= 1) { /* - * Return the original iterator if it's paging query or if the - * query input is less than one id, or don't have to do sort. - * NOTE: queryIds() only return the first batch of index query + * Return the original iterator if the query input is less than one + * id, or don't have to do sort. + * NOTE: queryIds() only return the first batch of index query. */ return origin; } @@ -142,17 +142,6 @@ private boolean mustSortByInputIds() { return false; } - private boolean paging() { - assert !this.queries.isEmpty(); - for (Query query : this.queries) { - Query origin = query.originQuery(); - if (query.paging() || origin != null && origin.paging()) { - return true; - } - } - return false; - } - @SuppressWarnings("unused") private boolean bigCapacity() { assert !this.queries.isEmpty(); diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransaction.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransaction.java index 7388425167..6faace9671 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransaction.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransaction.java @@ -653,10 +653,22 @@ private IdHolder doIndexQuery(IndexLabel indexLabel, ConditionQuery query) { } else { return new PagingIdHolder(query, q -> { return this.doIndexQueryOnce(indexLabel, q); - }); + }, this.keepBackendIndexOrder(indexLabel, query)); } } + private boolean keepBackendIndexOrder(IndexLabel indexLabel, + ConditionQuery query) { + return keepBackendIndexOrder(this.store().provider().isHstore(), + indexLabel.indexType(), query); + } + + static boolean keepBackendIndexOrder(boolean hstore, IndexType indexType, + ConditionQuery query) { + return hstore && indexType.isRange() && + (query.paging() || !query.noLimitAndOffset()); + } + @Watched(prefix = "index") private IdHolder doIndexQueryBatch(IndexLabel indexLabel, ConditionQuery query) { @@ -693,7 +705,7 @@ private IdHolder doIndexQueryBatch(IndexLabel indexLabel, } finally { locks.unlock(); } - }); + }, this.keepBackendIndexOrder(indexLabel, query)); } private void recordIndexValue(ConditionQuery query, HugeIndex index) { diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreFeatures.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreFeatures.java index e5aa4b943c..437f3f24ea 100644 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreFeatures.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreFeatures.java @@ -58,7 +58,8 @@ public boolean supportsQueryWithRangeCondition() { @Override public boolean supportsQuerySortByInputIds() { - return true; + // Multi-node batch scans group input keys by Store and lose global order + return false; } @Override diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java index 0abb6458b9..958e1bf0c1 100755 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java @@ -159,6 +159,11 @@ public abstract BackendColumnIterator scan(String table, byte[] query, byte[] position); + public abstract BackendColumnIterator scanOrdered( + String table, byte[] ownerKeyFrom, byte[] ownerKeyTo, + byte[] keyFrom, byte[] keyTo, int scanType, byte[] query, + long limit); + public abstract BackendColumnIterator scan(String table, int codeFrom, int codeTo, diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java index e619aff4fb..d172651393 100755 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImpl.java @@ -728,6 +728,31 @@ public BackendColumnIterator scan(String table, byte[] ownerKeyFrom, scanType); } + @Override + public BackendColumnIterator scanOrdered(String table, + byte[] ownerKeyFrom, + byte[] ownerKeyTo, + byte[] keyFrom, + byte[] keyTo, + int scanType, + byte[] query, + long limit) { + assert !this.hasChanges(); + HgKvIterator result = this.graph.scanIteratorOrdered( + table, HgOwnerKey.of(ownerKeyFrom, keyFrom), + HgOwnerKey.of(ownerKeyTo, keyTo), toHstoreLimit(limit), + scanType, query); + return new ColumnIterator<>(table, result, keyFrom, keyTo, + scanType); + } + + private long toHstoreLimit(long limit) { + if (limit <= 0L || limit == Query.NO_LIMIT) { + return HgStoreClientConst.NO_LIMIT; + } + return limit; + } + @Override public BackendColumnIterator scan(String table, int codeFrom, int codeTo, int scanType, diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreTable.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreTable.java index e1830111c3..c6e0823288 100755 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreTable.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreTable.java @@ -626,13 +626,20 @@ protected BackendColumnIterator queryByRange(Session session, ConditionQuery cq; Query origin = query.originQuery(); byte[] position = null; - if (query.paging() && !query.page().isEmpty()) { - position = PageState.fromString(query.page()).position(); - } byte[] ownerStart = this.ownerByQueryDelegate.apply(query.resultType(), query.start()); byte[] ownerEnd = this.ownerByQueryDelegate.apply(query.resultType(), query.end()); + if (shouldUseOrderedRangeScan(query)) { + start = rangeIndexScanStart(query, start); + type = rangeIndexScanType(query, type); + return session.scanOrdered(this.table(), ownerStart, ownerEnd, + start, end, type, null, + rangeScanBudget(query)); + } + if (query.paging() && !query.page().isEmpty()) { + position = PageState.fromString(query.page()).position(); + } if (origin instanceof ConditionQuery && (query.resultType().isEdge() || query.resultType().isVertex())) { cq = (ConditionQuery) query.originQuery(); @@ -650,6 +657,37 @@ protected BackendColumnIterator queryByRange(Session session, ownerEnd, start, end, type, null, position); } + static boolean shouldUseOrderedRangeScan(IdRangeQuery query) { + return query.resultType().isRangeIndex() && + (query.paging() || !query.noLimitAndOffset()); + } + + static byte[] rangeIndexScanStart(IdRangeQuery query, byte[] start) { + if (query.paging() && !query.page().isEmpty()) { + return PageState.fromString(query.page()).position(); + } + return start; + } + + static int rangeIndexScanType(IdRangeQuery query, int scanType) { + if (query.paging() && !query.page().isEmpty()) { + scanType &= ~Session.SCAN_GTE_BEGIN; + scanType |= Session.SCAN_GTE_BEGIN; + } + return scanType; + } + + static long rangeScanBudget(IdRangeQuery query) { + if (query.noLimit()) { + return HgStoreClientConst.NO_LIMIT; + } + long total = query.total(); + if (total < 0L || total == Long.MAX_VALUE) { + return HgStoreClientConst.NO_LIMIT; + } + return total + 1L; + } + protected BackendColumnIterator queryByCond(Session session, ConditionQuery query) { if (query.containsScanCondition()) { diff --git a/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java index 69f2591e25..b962ae0b0d 100644 --- a/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java +++ b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreSessionsImplTest.java @@ -19,8 +19,16 @@ import java.io.IOException; import java.io.InputStream; +import java.lang.reflect.Constructor; import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.hugegraph.backend.store.BackendEntry.BackendColumn; +import org.apache.hugegraph.backend.store.BackendEntry.BackendColumnIterator; +import org.apache.hugegraph.store.HgKvEntry; +import org.apache.hugegraph.store.HgKvIterator; import org.junit.Assert; import org.junit.Test; @@ -36,4 +44,107 @@ public void testProductionClassDoesNotReferenceTestAssert() throws IOException { "org/apache/hugegraph/testutil/Assert")); } } + + @Test + public void testColumnIteratorPositionTracksNextUnreadKey() + throws Exception { + BackendColumnIterator iterator = newColumnIterator( + new TestIterator(1, 2)); + + Assert.assertTrue(iterator.hasNext()); + BackendColumn first = iterator.next(); + Assert.assertArrayEquals(keyBytes(1), first.name); + + Assert.assertTrue(iterator.hasNext()); + Assert.assertArrayEquals(keyBytes(2), iterator.position()); + } + + @Test + public void testColumnIteratorClearsPositionWhenFullyExhausted() + throws Exception { + BackendColumnIterator iterator = newColumnIterator( + new TestIterator(1)); + + Assert.assertTrue(iterator.hasNext()); + iterator.next(); + Assert.assertFalse(iterator.hasNext()); + Assert.assertNull(iterator.position()); + } + + private static BackendColumnIterator newColumnIterator( + HgKvIterator iterator) throws Exception { + Class clazz = Class.forName(HstoreSessionsImpl.class.getName() + + "$ColumnIterator"); + Constructor constructor = clazz.getDeclaredConstructor( + String.class, HgKvIterator.class, byte[].class, byte[].class, + int.class); + constructor.setAccessible(true); + return (BackendColumnIterator) constructor.newInstance( + "test", iterator, null, null, 0); + } + + private static byte[] keyBytes(int key) { + return new byte[]{(byte) key}; + } + + private static final class TestIterator implements HgKvIterator { + + private final List keys; + private int offset; + private HgKvEntry current; + + private TestIterator(Integer... keys) { + this.keys = Arrays.asList(keys); + this.offset = 0; + this.current = null; + } + + @Override + public boolean hasNext() { + return this.offset < this.keys.size(); + } + + @Override + public HgKvEntry next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + this.current = new TestEntry(keyBytes(this.keys.get(this.offset++))); + return this.current; + } + + @Override + public byte[] key() { + return this.current == null ? null : this.current.key(); + } + + @Override + public byte[] value() { + return this.current == null ? null : this.current.value(); + } + + @Override + public byte[] position() { + return this.key(); + } + } + + private static final class TestEntry implements HgKvEntry { + + private final byte[] key; + + private TestEntry(byte[] key) { + this.key = key; + } + + @Override + public byte[] key() { + return this.key; + } + + @Override + public byte[] value() { + return this.key; + } + } } diff --git a/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreTableTest.java b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreTableTest.java new file mode 100644 index 0000000000..b61b65a768 --- /dev/null +++ b/hugegraph-server/hugegraph-hstore/src/test/java/org/apache/hugegraph/backend/store/hstore/HstoreTableTest.java @@ -0,0 +1,180 @@ +/* + * 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.backend.store.hstore; + +import java.util.Arrays; +import java.util.List; +import java.util.NoSuchElementException; + +import org.apache.hugegraph.backend.id.Id.IdType; +import org.apache.hugegraph.backend.id.IdGenerator; +import org.apache.hugegraph.backend.page.PageInfo; +import org.apache.hugegraph.backend.page.PageState; +import org.apache.hugegraph.backend.query.IdRangeQuery; +import org.apache.hugegraph.backend.query.Query; +import org.apache.hugegraph.backend.store.BackendEntry; +import org.apache.hugegraph.backend.store.BackendEntry.BackendColumn; +import org.apache.hugegraph.backend.store.BackendEntry.BackendColumnIterator; +import org.apache.hugegraph.backend.store.BackendEntryIterator; +import org.apache.hugegraph.store.client.util.HgStoreClientConst; +import org.apache.hugegraph.type.HugeType; +import org.junit.Assert; +import org.junit.Test; + +public class HstoreTableTest { + + @Test + public void testHstoreDoesNotAdvertiseInputIdOrdering() { + Assert.assertFalse(new HstoreFeatures() + .supportsQuerySortByInputIds()); + } + + @Test + public void testRangeIndexPageStateUsesNextUnreadPhysicalKey() { + Query query = new Query(HugeType.RANGE_INT_INDEX); + query.page(""); + query.limit(1L); + + BackendEntryIterator iterator = HstoreTable.newEntryIterator( + new TestColumnIterator(1, 2), query); + + Assert.assertTrue(iterator.hasNext()); + BackendEntry entry = iterator.next(); + Assert.assertArrayEquals(keyBytes(1), entry.id().asBytes()); + + PageState pageState = PageInfo.pageState(iterator); + Assert.assertArrayEquals(keyBytes(2), pageState.position()); + Assert.assertEquals(1L, pageState.total()); + } + + @Test + public void testRangeIndexPagingUsesPagePositionAsInclusiveScanStart() { + byte[] originalStart = keyBytes(1); + byte[] pagePosition = keyBytes(2); + IdRangeQuery query = rangeIndexQuery(); + + query.page(""); + Assert.assertArrayEquals(originalStart, + HstoreTable.rangeIndexScanStart( + query, originalStart)); + + query.page(new PageState(pagePosition, 0, 1).toString()); + Assert.assertArrayEquals(pagePosition, + HstoreTable.rangeIndexScanStart( + query, originalStart)); + int type = HstoreTable.rangeIndexScanType( + query, HstoreSessions.Session.SCAN_GT_BEGIN | + HstoreSessions.Session.SCAN_LT_END); + Assert.assertTrue(HstoreSessions.Session.matchScanType( + HstoreSessions.Session.SCAN_GTE_BEGIN, type)); + Assert.assertTrue(HstoreSessions.Session.matchScanType( + HstoreSessions.Session.SCAN_LT_END, type)); + } + + @Test + public void testOrderedRangeScanIsScopedToOrderSensitiveIndexes() { + IdRangeQuery query = rangeIndexQuery(); + Assert.assertFalse(HstoreTable.shouldUseOrderedRangeScan(query)); + + query.limit(10L); + Assert.assertTrue(HstoreTable.shouldUseOrderedRangeScan(query)); + + query = rangeIndexQuery(); + query.offset(1L); + Assert.assertTrue(HstoreTable.shouldUseOrderedRangeScan(query)); + + query = rangeIndexQuery(); + query.page(""); + Assert.assertTrue(HstoreTable.shouldUseOrderedRangeScan(query)); + + query = new IdRangeQuery(HugeType.VERTEX, null, + IdGenerator.of(keyBytes(1), IdType.STRING), + true, + IdGenerator.of(keyBytes(9), IdType.STRING), + false); + query.limit(10L); + Assert.assertFalse(HstoreTable.shouldUseOrderedRangeScan(query)); + } + + @Test + public void testRangeScanBudgetIncludesOneLookaheadRecord() { + IdRangeQuery query = rangeIndexQuery(); + Assert.assertEquals(HgStoreClientConst.NO_LIMIT, + HstoreTable.rangeScanBudget(query)); + + query.limit(10L); + Assert.assertEquals(11L, HstoreTable.rangeScanBudget(query)); + + query.offset(3L); + Assert.assertEquals(14L, HstoreTable.rangeScanBudget(query)); + } + + private static IdRangeQuery rangeIndexQuery() { + return new IdRangeQuery(HugeType.RANGE_INT_INDEX, null, + IdGenerator.of(keyBytes(1), IdType.STRING), + true, + IdGenerator.of(keyBytes(9), IdType.STRING), + false); + } + + private static byte[] keyBytes(int key) { + byte[] bytes = new byte[9]; + bytes[0] = HugeType.RANGE_INT_INDEX.code(); + bytes[8] = (byte) key; + return bytes; + } + + private static final class TestColumnIterator + implements BackendColumnIterator { + + private final List keys; + private int offset; + private byte[] position; + + private TestColumnIterator(Integer... keys) { + this.keys = Arrays.asList(keys); + this.offset = 0; + this.position = null; + } + + @Override + public boolean hasNext() { + return this.offset < this.keys.size(); + } + + @Override + public BackendColumn next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + byte[] key = keyBytes(this.keys.get(this.offset++)); + this.position = key; + return BackendColumn.of(key, key); + } + + @Override + public void close() { + // pass + } + + @Override + public byte[] position() { + return this.position; + } + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransactionTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransactionTest.java new file mode 100644 index 0000000000..65940d3e3a --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/backend/tx/GraphIndexTransactionTest.java @@ -0,0 +1,49 @@ +/* + * 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.backend.tx; + +import org.apache.hugegraph.backend.query.ConditionQuery; +import org.apache.hugegraph.backend.query.Query; +import org.apache.hugegraph.testutil.Assert; +import org.apache.hugegraph.type.HugeType; +import org.apache.hugegraph.type.define.IndexType; +import org.junit.Test; + +public class GraphIndexTransactionTest { + + @Test + public void testKeepBackendIndexOrderOnlyForOrderedHstoreRangeQuery() { + ConditionQuery query = new ConditionQuery(HugeType.RANGE_INT_INDEX); + + Assert.assertFalse(GraphIndexTransaction.keepBackendIndexOrder( + true, IndexType.RANGE_INT, query)); + + query.limit(10L); + Assert.assertTrue(GraphIndexTransaction.keepBackendIndexOrder( + true, IndexType.RANGE_INT, query)); + Assert.assertFalse(GraphIndexTransaction.keepBackendIndexOrder( + false, IndexType.RANGE_INT, query)); + Assert.assertFalse(GraphIndexTransaction.keepBackendIndexOrder( + true, IndexType.SECONDARY, query)); + + query.limit(Query.NO_LIMIT); + query.page(""); + Assert.assertTrue(GraphIndexTransaction.keepBackendIndexOrder( + true, IndexType.RANGE_INT, query)); + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java index 1733680e3f..b629cf2326 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java @@ -17,6 +17,7 @@ package org.apache.hugegraph.unit; +import org.apache.hugegraph.backend.tx.GraphIndexTransactionTest; import org.apache.hugegraph.api.auth.GraphSpaceAuthPayloadTest; import org.apache.hugegraph.api.auth.GraphSpaceGroupAPITest; import org.apache.hugegraph.auth.StandardAuthManagerV2Test; @@ -48,8 +49,10 @@ import org.apache.hugegraph.unit.core.ExceptionTest; import org.apache.hugegraph.unit.core.GraphManagerAdminInitTest; import org.apache.hugegraph.unit.core.GraphManagerConfigTest; +import org.apache.hugegraph.unit.core.IdHolderTest; import org.apache.hugegraph.unit.core.LocksTableTest; import org.apache.hugegraph.unit.core.PageStateTest; +import org.apache.hugegraph.unit.core.QueryResultsTest; import org.apache.hugegraph.unit.core.QueryTest; import org.apache.hugegraph.unit.core.RangeTest; import org.apache.hugegraph.unit.core.RolePermissionTest; @@ -140,7 +143,9 @@ BackendMutationTest.class, ConditionTest.class, ConditionQueryFlattenTest.class, + GraphIndexTransactionTest.class, QueryTest.class, + QueryResultsTest.class, RangeTest.class, SecurityManagerTest.class, RolePermissionTest.class, @@ -150,6 +155,7 @@ BackendStoreInfoTest.class, TraversalUtilTest.class, TraversalUtilOptimizeTest.class, + IdHolderTest.class, PageStateTest.class, SystemSchemaStoreTest.class, ServerInfoManagerTest.class, diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/IdHolderTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/IdHolderTest.java new file mode 100644 index 0000000000..c6bdcc32d5 --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/IdHolderTest.java @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.unit.core; + +import java.util.Collections; +import java.util.Iterator; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.backend.id.Id; +import org.apache.hugegraph.backend.id.IdGenerator; +import org.apache.hugegraph.backend.page.IdHolder.BatchIdHolder; +import org.apache.hugegraph.backend.page.IdHolder.PagingIdHolder; +import org.apache.hugegraph.backend.page.IdHolderList; +import org.apache.hugegraph.backend.page.PageIds; +import org.apache.hugegraph.backend.page.PageState; +import org.apache.hugegraph.backend.page.QueryList; +import org.apache.hugegraph.backend.query.ConditionQuery; +import org.apache.hugegraph.backend.query.IdQuery; +import org.apache.hugegraph.backend.query.Query; +import org.apache.hugegraph.backend.query.QueryResults; +import org.apache.hugegraph.backend.store.BackendEntry; +import org.apache.hugegraph.testutil.Assert; +import org.apache.hugegraph.type.HugeType; +import org.apache.hugegraph.util.InsertionOrderUtil; +import org.junit.Test; + +public class IdHolderTest { + + @Test + public void testBatchIndexQueryKeepsHolderOrder() { + ConditionQuery parent = new ConditionQuery(HugeType.VERTEX); + ConditionQuery indexQuery = + new ConditionQuery(HugeType.RANGE_INT_INDEX); + Set ids = ids(IdGenerator.of(2L), IdGenerator.of(1L)); + Iterator entries = + Collections.singletonList(null).iterator(); + BatchIdHolder holder = new BatchIdHolder(indexQuery, entries, + batch -> ids, true); + + IdQuery idQuery = fetchIdQuery(parent, holder, 2); + + Assert.assertTrue(holder.keepOrder()); + Assert.assertTrue(idQuery.mustSortByInput()); + } + + @Test + public void testPagingIndexQueryKeepsHolderOrder() { + ConditionQuery parent = new ConditionQuery(HugeType.VERTEX); + parent.page(""); + ConditionQuery indexQuery = + new ConditionQuery(HugeType.RANGE_INT_INDEX); + indexQuery.page(""); + Set ids = ids(IdGenerator.of(2L), IdGenerator.of(1L)); + PagingIdHolder holder = new PagingIdHolder( + indexQuery, + query -> new PageIds(ids, PageState.EMPTY), true); + + IdQuery idQuery = fetchIdQuery(parent, holder, 2); + + Assert.assertTrue(holder.keepOrder()); + Assert.assertTrue(idQuery.mustSortByInput()); + } + + private static IdQuery fetchIdQuery(ConditionQuery parent, + org.apache.hugegraph.backend.page.IdHolder holder, + int pageSize) { + AtomicReference captured = new AtomicReference<>(); + QueryList queries = new QueryList<>(parent, query -> { + Assert.assertTrue(query instanceof IdQuery); + IdQuery idQuery = (IdQuery) query; + captured.set(idQuery); + return new QueryResults<>(idQuery.ids().iterator(), idQuery); + }); + IdHolderList holders = new IdHolderList(holder.paging()); + holders.add(holder); + queries.add(holders, Query.QUERY_BATCH); + + Iterator results = queries.fetch(pageSize).iterator(); + Assert.assertTrue(results.hasNext()); + Assert.assertNotNull(captured.get()); + return captured.get(); + } + + private static Set ids(Id... ids) { + Set result = InsertionOrderUtil.newSet(); + Collections.addAll(result, ids); + return result; + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java new file mode 100644 index 0000000000..2728de0015 --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.unit.core; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; + +import org.apache.hugegraph.backend.id.Id; +import org.apache.hugegraph.backend.id.IdGenerator; +import org.apache.hugegraph.backend.query.IdQuery; +import org.apache.hugegraph.backend.query.Query; +import org.apache.hugegraph.backend.query.QueryResults; +import org.apache.hugegraph.testutil.Assert; +import org.apache.hugegraph.type.HugeType; +import org.apache.hugegraph.type.Idfiable; +import org.apache.hugegraph.util.InsertionOrderUtil; +import org.junit.Test; + +import com.google.common.collect.ImmutableList; + +public class QueryResultsTest { + + @Test + public void testKeepInputOrderForPagingIdQuery() { + Id id1 = IdGenerator.of(1L); + Id id2 = IdGenerator.of(2L); + Query pagingQuery = new Query(HugeType.VERTEX); + pagingQuery.page("page-1"); + pagingQuery.limit(2L); + + Set ids = InsertionOrderUtil.newSet(); + ids.add(id2); + ids.add(id1); + + IdQuery idQuery = new IdQuery(pagingQuery, ids); + idQuery.mustSortByInput(true); + QueryResults results = new QueryResults<>( + Arrays.asList(new TestIdfiable(id1), + new TestIdfiable(id2)).iterator(), + idQuery); + + List orderedIds = new ArrayList<>(); + results.keepInputOrderIfNeeded( + Arrays.asList(new TestIdfiable(id1), + new TestIdfiable(id2)).iterator()) + .forEachRemaining(item -> orderedIds.add(item.id())); + + Assert.assertEquals(ImmutableList.of(id2, id1), orderedIds); + } + + private static final class TestIdfiable implements Idfiable { + + private final Id id; + + private TestIdfiable(Id id) { + this.id = id; + } + + @Override + public Id id() { + return this.id; + } + } +} diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgKvStore.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgKvStore.java index dcce95ba1e..3332fc897a 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgKvStore.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgKvStore.java @@ -95,6 +95,19 @@ HgKvIterator scanIterator(String table, HgOwnerKey startKey, HgOwnerK HgKvIterator scanIterator(String table, HgOwnerKey startKey, HgOwnerKey endKey, long limit, int scanType, byte[] query); + /** + * Scan a range in global key order across all partitions. + * + * Low-level node sessions don't provide cross-partition ordering. The + * multi-node session proxy overrides this capability. + */ + default HgKvIterator scanIteratorOrdered( + String table, HgOwnerKey startKey, HgOwnerKey endKey, long limit, + int scanType, byte[] query) { + throw new UnsupportedOperationException( + "Global ordered scan is not supported"); + } + HgKvIterator scanIterator(String table, int codeFrom, int codeTo, int scanType, byte[] query); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTkv.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTkv.java index e78ced4c10..b9c28bcd9d 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTkv.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTkv.java @@ -17,6 +17,7 @@ package org.apache.hugegraph.store.client; +import java.util.Arrays; import java.util.Objects; import javax.annotation.concurrent.ThreadSafe; @@ -39,16 +40,16 @@ class NodeTkv { NodeTkv(HgNodePartition nodePartition, String table, HgOwnerKey key) { this.nodePartition = nodePartition; this.table = table; - this.key = key; - this.endKey = key; + this.key = copyOf(key); + this.endKey = this.key; this.key.setKeyCode(this.nodePartition.getKeyCode()); } NodeTkv(HgNodePartition nodePartition, String table, HgOwnerKey key, int keyCode) { this.nodePartition = nodePartition; this.table = table; - this.key = key; - this.endKey = key; + this.key = copyOf(key); + this.endKey = this.key; this.key.setKeyCode(keyCode); } @@ -57,8 +58,8 @@ class NodeTkv { HgOwnerKey endKey) { this.nodePartition = nodePartition; this.table = table; - this.key = startKey; - this.endKey = endKey; + this.key = copyOf(startKey); + this.endKey = copyOf(endKey); this.key.setKeyCode(nodePartition.getStartKey()); this.endKey.setKeyCode(nodePartition.getEndKey()); } @@ -122,4 +123,14 @@ public HgStoreSession getSession() { public void setSession(HgStoreSession session) { this.session = session; } + + private static HgOwnerKey copyOf(HgOwnerKey key) { + HgOwnerKey copy = HgOwnerKey.of(Arrays.copyOf(key.getOwner(), + key.getOwner().length), + Arrays.copyOf(key.getKey(), + key.getKey().length)); + copy.setKeyCode(key.getKeyCode()); + copy.setSerialNo(key.getSerialNo()); + return copy; + } } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java index 65e3b76ec0..9802abc639 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java @@ -41,6 +41,7 @@ import org.apache.hugegraph.HugeGraphSupplier; import org.apache.hugegraph.pd.common.PDException; +import org.apache.hugegraph.pd.common.PartitionUtils; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; import org.apache.hugegraph.store.HgKvOrderedIterator; @@ -54,12 +55,17 @@ import org.apache.hugegraph.store.client.util.HgAssert; import org.apache.hugegraph.store.client.util.HgStoreClientConst; import org.apache.hugegraph.store.client.util.HgStoreClientUtil; +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder; import org.apache.hugegraph.store.query.StoreQueryParam; import org.apache.hugegraph.store.term.HgPair; import org.apache.hugegraph.store.term.HgTriple; import org.apache.hugegraph.structure.BaseElement; +import com.google.protobuf.ByteString; + import lombok.extern.slf4j.Slf4j; /** @@ -485,6 +491,40 @@ public HgKvIterator scanIterator(String table, HgOwnerKey startKey, H } + @Override + public HgKvIterator scanIteratorOrdered(String table, + HgOwnerKey startKey, + HgOwnerKey endKey, + long limit, + int scanType, + byte[] query) { + HgAssert.isFalse(HgAssert.isInvalid(table), + "The argument is invalid: table"); + HgAssert.isFalse(startKey == null, + "The argument is invalid: startKey"); + HgAssert.isFalse(endKey == null, + "The argument is invalid: endKey"); + + List nodeTkvs = + this.toOrderedRangeNodeTkvList(table, startKey, endKey); + List> iterators = + new ArrayList<>(nodeTkvs.size()); + try { + for (NodeTkv nodeTkv : nodeTkvs) { + HgKvIterator iterator = + this.getStoreNode(nodeTkv.getNodeId()) + .openSession(this.graphName) + .scanIterator(this.orderedRangeScanBuilder( + nodeTkv, limit, scanType, query)); + iterators.add(iterator); + } + } catch (RuntimeException | Error e) { + closeIteratorsAfterFailure(iterators, e); + throw e; + } + return mergeOrderedRangeScanIterators(iterators, limit); + } + @Override public HgKvIterator scanIterator(String table, int codeFrom, int codeTo, int scanType, byte[] query) { @@ -642,6 +682,23 @@ private BiFunction, HgScanQuery.ScanBuilder> toScanQueryFu } /*-- common --*/ + static HgKvIterator mergeOrderedRangeScanIterators( + List> iteratorList, + long limit) { + return new OrderedKvIterator(iteratorList, limit); + } + + private static void closeIteratorsAfterFailure( + List> iterators, Throwable failure) { + for (HgKvIterator iterator : iterators) { + try { + iterator.close(); + } catch (RuntimeException | Error closeFailure) { + failure.addSuppressed(closeFailure); + } + } + } + private HgKvIterator toHgKvIteratorProxy(List iteratorList, long limit) { boolean isAllOrderedLimiter = iteratorList.stream() .allMatch( @@ -748,6 +805,42 @@ private List toNodeTkvList(String table, HgOwnerKey startKey, HgOwnerKe return nodeTkvs; } + private List toOrderedRangeNodeTkvList(String table, + HgOwnerKey startKey, + HgOwnerKey endKey) { + // One Store may host multiple partitions with independent key order + Collection partitions = + this.doPartition(table, 0, PartitionUtils.MAX_VALUE); + List nodeTkvs = new ArrayList<>(partitions.size()); + for (HgNodePartition partition : partitions) { + nodeTkvs.add(new NodeTkv(partition, table, startKey, endKey)); + } + return nodeTkvs; + } + + private Builder orderedRangeScanBuilder(NodeTkv nodeTkv, long limit, + int scanType, byte[] query) { + long scanLimit = limit <= HgStoreClientConst.NO_LIMIT ? + Integer.MAX_VALUE : limit; + return ScanStreamReq.newBuilder() + .setHeader(Header.newBuilder() + .setGraph(this.graphName) + .build()) + .setMethod(ScanMethod.RANGE) + .setTable(nodeTkv.getTable()) + .setStart(toByteString(nodeTkv.getKey().getKey())) + .setEnd(toByteString(nodeTkv.getEndKey().getKey())) + .setLimit(scanLimit) + .setCode(nodeTkv.getKey().getKeyCode()) + .setScanType(scanType) + .setQuery(toByteString(query)); + } + + private static ByteString toByteString(byte[] bytes) { + return ByteString.copyFrom(bytes != null ? bytes : + HgStoreClientConst.EMPTY_BYTES); + } + private List toNodeTkvList(String table, int startCode, int endCode) { Collection partitions = this.doPartition(table, startCode, endCode); ArrayList nodeTkvs = new ArrayList<>(partitions.size()); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java new file mode 100644 index 0000000000..c8a1eca53f --- /dev/null +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java @@ -0,0 +1,193 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client; + +import java.util.Arrays; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.PriorityQueue; + +import org.apache.hugegraph.store.HgKvEntry; +import org.apache.hugegraph.store.HgKvIterator; +import org.apache.hugegraph.store.client.util.HgStoreClientConst; + +final class OrderedKvIterator implements HgKvIterator { + + private final List> iterators; + private final PriorityQueue queue; + private final boolean[] sourceClosed; + private final long limit; + + private boolean initialized; + private boolean closed; + private long count; + private HgKvEntry current; + private byte[] position; + + OrderedKvIterator(List> iterators, + long limit) { + this.iterators = iterators; + this.queue = new PriorityQueue<>((left, right) -> { + int result = Arrays.compareUnsigned(left.entry.key(), + right.entry.key()); + if (result != 0) { + return result; + } + return Integer.compare(left.source, right.source); + }); + this.sourceClosed = new boolean[iterators.size()]; + this.limit = limit <= HgStoreClientConst.NO_LIMIT ? Long.MAX_VALUE : + limit; + this.initialized = false; + this.closed = false; + this.count = 0L; + this.current = null; + this.position = HgStoreClientConst.EMPTY_BYTES; + } + + @Override + public boolean hasNext() { + if (this.closed) { + return false; + } + this.initialize(); + boolean hasNext = this.count < this.limit && !this.queue.isEmpty(); + if (!hasNext) { + this.close(); + } + return hasNext; + } + + @Override + public HgKvEntry next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + + SourceEntry sourceEntry = this.queue.poll(); + this.current = sourceEntry.entry; + this.position = this.current.key(); + this.count++; + + try { + if (this.count < this.limit) { + this.addNext(sourceEntry.source); + } else { + this.close(); + } + } catch (RuntimeException | Error e) { + this.closeAfterFailure(e); + throw e; + } + return this.current; + } + + @Override + public byte[] key() { + return this.current == null ? null : this.current.key(); + } + + @Override + public byte[] value() { + return this.current == null ? null : this.current.value(); + } + + @Override + public byte[] position() { + return this.position; + } + + @Override + public void close() { + if (this.closed) { + return; + } + this.closed = true; + Throwable failure = null; + for (int i = 0; i < this.iterators.size(); i++) { + try { + this.closeSource(i); + } catch (RuntimeException | Error e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + this.queue.clear(); + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure != null) { + throw (Error) failure; + } + } + + private void initialize() { + if (this.initialized) { + return; + } + this.initialized = true; + try { + for (int i = 0; i < this.iterators.size(); i++) { + this.addNext(i); + } + } catch (RuntimeException | Error e) { + this.closeAfterFailure(e); + throw e; + } + } + + private void addNext(int source) { + HgKvIterator iterator = + this.iterators.get(source); + if (iterator.hasNext()) { + this.queue.add(new SourceEntry(source, iterator.next())); + } else { + this.closeSource(source); + } + } + + private void closeSource(int source) { + if (this.sourceClosed[source]) { + return; + } + this.sourceClosed[source] = true; + this.iterators.get(source).close(); + } + + private void closeAfterFailure(Throwable failure) { + try { + this.close(); + } catch (RuntimeException | Error closeFailure) { + failure.addSuppressed(closeFailure); + } + } + + private static final class SourceEntry { + + private final int source; + private final HgKvEntry entry; + + private SourceEntry(int source, HgKvEntry entry) { + this.source = source; + this.entry = entry; + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java index 4217a4c1de..ebbe27ac41 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java @@ -27,7 +27,9 @@ */ @RunWith(Suite.class) @Suite.SuiteClasses({ - AbstractGrpcClientTest.class + AbstractGrpcClientTest.class, + NodeTxSessionProxyTest.class, + OrderedKvIteratorTest.class }) public class ClientSuiteTest { } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java new file mode 100644 index 0000000000..d5f38cb6e6 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java @@ -0,0 +1,393 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client; + +import java.lang.reflect.Field; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import org.apache.hugegraph.pd.common.PartitionUtils; +import org.apache.hugegraph.store.HgKvEntry; +import org.apache.hugegraph.store.HgKvIterator; +import org.apache.hugegraph.store.HgOwnerKey; +import org.apache.hugegraph.store.HgStoreSession; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder; +import org.junit.Assert; +import org.junit.Test; + +public class NodeTxSessionProxyTest { + + @Test + public void testNodeTkvDoesNotMutateSharedOwnerKeys() { + HgOwnerKey start = HgOwnerKey.of(keyBytes(9), keyBytes(1)); + HgOwnerKey end = HgOwnerKey.of(keyBytes(9), keyBytes(5)); + start.setSerialNo(7); + end.setSerialNo(8); + + NodeTkv first = new NodeTkv(HgNodePartition.of(1L, 0, 10, 20), + "g+index", start, end); + NodeTkv second = new NodeTkv(HgNodePartition.of(2L, 0, 30, 40), + "g+index", start, end); + + Assert.assertEquals(0, start.getKeyCode()); + Assert.assertEquals(0, end.getKeyCode()); + Assert.assertEquals(10, first.getKey().getKeyCode()); + Assert.assertEquals(20, first.getEndKey().getKeyCode()); + Assert.assertEquals(30, second.getKey().getKeyCode()); + Assert.assertEquals(40, second.getEndKey().getKeyCode()); + Assert.assertEquals(7, first.getKey().getSerialNo()); + Assert.assertEquals(8, first.getEndKey().getSerialNo()); + } + + @Test + public void testScanIteratorOrderedUsesPerPartitionBuildersLazily() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + RecordingPartitioner partitioner = + new RecordingPartitioner(firstNodeId, secondNodeId); + TestIterator firstIterator = new TestIterator(3, 4); + TestIterator secondIterator = new TestIterator(1, 2); + RecordingSession firstSession = new RecordingSession(firstIterator); + RecordingSession secondSession = new RecordingSession(secondIterator); + String graph = "graph-" + firstNodeId; + manager.addNode(graph, new RecordingStoreNode(firstNodeId, + firstSession.proxy())); + manager.addNode(graph, new RecordingStoreNode(secondNodeId, + secondSession.proxy())); + manager.setNodePartitioner(partitioner); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + HgKvIterator iterator = proxy.scanIteratorOrdered( + "table", HgOwnerKey.of(keyBytes(9), keyBytes(1)), + HgOwnerKey.of(keyBytes(9), keyBytes(5)), 5L, 123, + keyBytes(7)); + + Assert.assertEquals(1, firstSession.builders.size()); + Assert.assertEquals(1, secondSession.builders.size()); + Assert.assertEquals(0, partitioner.ownerRangeCalls); + Assert.assertEquals(1, partitioner.codeRangeCalls); + Assert.assertEquals(0, partitioner.startCode); + Assert.assertEquals(PartitionUtils.MAX_VALUE, + partitioner.endCode); + Assert.assertEquals(0, firstSession.rangeScanCalls); + Assert.assertEquals(0, secondSession.rangeScanCalls); + Assert.assertEquals(0, firstIterator.nextCalls); + Assert.assertEquals(0, secondIterator.nextCalls); + assertOrderedRangeBuilder(firstSession.builders.get(0), 5L, + 123, 10, keyBytes(7)); + assertOrderedRangeBuilder(secondSession.builders.get(0), 5L, + 123, 30, keyBytes(7)); + + Assert.assertEquals(1, key(iterator.next())); + Assert.assertEquals(2, key(iterator.next())); + Assert.assertEquals(3, key(iterator.next())); + Assert.assertEquals(4, key(iterator.next())); + Assert.assertFalse(iterator.hasNext()); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + + @Test + public void testScanIteratorOrderedClosesOpenedIteratorsOnOpenFailure() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + RecordingPartitioner partitioner = + new RecordingPartitioner(firstNodeId, secondNodeId); + CountDownLatch firstOpened = new CountDownLatch(1); + TestIterator firstIterator = new TestIterator(1); + RecordingSession firstSession = new RecordingSession( + firstIterator, firstOpened, null, false); + RecordingSession secondSession = new RecordingSession( + new TestIterator(2), null, firstOpened, true); + String graph = "graph-open-failure-" + firstNodeId; + manager.addNode(graph, new RecordingStoreNode(firstNodeId, + firstSession.proxy())); + manager.addNode(graph, new RecordingStoreNode(secondNodeId, + secondSession.proxy())); + manager.setNodePartitioner(partitioner); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + + Assert.assertThrows(RuntimeException.class, + () -> proxy.scanIteratorOrdered( + "table", + HgOwnerKey.of(keyBytes(9), keyBytes(1)), + HgOwnerKey.of(keyBytes(9), keyBytes(5)), + 5L, 123, keyBytes(7))); + Assert.assertTrue(firstIterator.closed); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + + private static void assertOrderedRangeBuilder(Builder builder, long limit, + int scanType, int code, + byte[] query) { + Assert.assertEquals(ScanMethod.RANGE, builder.getMethod()); + Assert.assertEquals("table", builder.getTable()); + Assert.assertEquals(limit, builder.getLimit()); + Assert.assertEquals(scanType, builder.getScanType()); + Assert.assertEquals(code, builder.getCode()); + Assert.assertArrayEquals(keyBytes(1), builder.getStart().toByteArray()); + Assert.assertArrayEquals(keyBytes(5), builder.getEnd().toByteArray()); + Assert.assertArrayEquals(query, builder.getQuery().toByteArray()); + } + + private static void restoreNodePartitioner(HgStoreNodeManager manager, + HgStoreNodePartitioner old) + throws Exception { + Field field = HgStoreNodeManager.class.getDeclaredField( + "nodePartitioner"); + field.setAccessible(true); + field.set(manager, old); + } + + private static int key(HgKvEntry entry) { + return entry.key()[0] & 0xff; + } + + private static byte[] keyBytes(int key) { + return new byte[]{(byte) key}; + } + + private static final class TestIterator implements HgKvIterator { + + private final List keys; + private int offset; + private int nextCalls; + private HgKvEntry current; + private boolean closed; + + private TestIterator(Integer... keys) { + this.keys = Arrays.asList(keys); + this.offset = 0; + this.nextCalls = 0; + this.current = null; + this.closed = false; + } + + @Override + public boolean hasNext() { + return this.offset < this.keys.size(); + } + + @Override + public HgKvEntry next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + this.nextCalls++; + this.current = new TestEntry(keyBytes(this.keys.get(this.offset++))); + return this.current; + } + + @Override + public byte[] key() { + return this.current == null ? null : this.current.key(); + } + + @Override + public byte[] value() { + return this.current == null ? null : this.current.value(); + } + + @Override + public void close() { + this.closed = true; + } + } + + private static final class TestEntry implements HgKvEntry { + + private final byte[] key; + + private TestEntry(byte[] key) { + this.key = key; + } + + @Override + public byte[] key() { + return this.key; + } + + @Override + public byte[] value() { + return this.key; + } + } + + private static final class RecordingPartitioner + implements HgStoreNodePartitioner { + + private final long firstNodeId; + private final long secondNodeId; + private int ownerRangeCalls; + private int codeRangeCalls; + private int startCode; + private int endCode; + + private RecordingPartitioner(long firstNodeId, long secondNodeId) { + this.firstNodeId = firstNodeId; + this.secondNodeId = secondNodeId; + this.ownerRangeCalls = 0; + this.codeRangeCalls = 0; + this.startCode = -1; + this.endCode = -1; + } + + @Override + public int partition(HgNodePartitionerBuilder builder, + String graphName, byte[] startKey, + byte[] endKey) { + this.ownerRangeCalls++; + return this.setPartitions(builder); + } + + @Override + public int partition(HgNodePartitionerBuilder builder, + String graphName, int startCode, + int endCode) { + this.codeRangeCalls++; + this.startCode = startCode; + this.endCode = endCode; + return this.setPartitions(builder); + } + + private int setPartitions(HgNodePartitionerBuilder builder) { + Set partitions = new LinkedHashSet<>(); + partitions.add(HgNodePartition.of(this.firstNodeId, 10, 10, 20)); + partitions.add(HgNodePartition.of(this.secondNodeId, 30, 30, 40)); + builder.setPartitions(partitions); + return 0; + } + } + + private static final class RecordingStoreNode implements HgStoreNode { + + private final Long nodeId; + private final HgStoreSession session; + + private RecordingStoreNode(Long nodeId, HgStoreSession session) { + this.nodeId = nodeId; + this.session = session; + } + + @Override + public Long getNodeId() { + return this.nodeId; + } + + @Override + public String getAddress() { + return "127.0.0.1:" + this.nodeId; + } + + @Override + public HgStoreSession openSession(String graphName) { + return this.session; + } + } + + private static final class RecordingSession implements InvocationHandler { + + private final List builders; + private final TestIterator iterator; + private final CountDownLatch opened; + private final CountDownLatch waitBeforeFailure; + private final boolean failOnBuilder; + private int rangeScanCalls; + + private RecordingSession(TestIterator iterator) { + this(iterator, null, null, false); + } + + private RecordingSession(TestIterator iterator, + CountDownLatch opened, + CountDownLatch waitBeforeFailure, + boolean failOnBuilder) { + this.builders = Collections.synchronizedList(new ArrayList<>()); + this.iterator = iterator; + this.opened = opened; + this.waitBeforeFailure = waitBeforeFailure; + this.failOnBuilder = failOnBuilder; + this.rangeScanCalls = 0; + } + + private HgStoreSession proxy() { + return (HgStoreSession) Proxy.newProxyInstance( + HgStoreSession.class.getClassLoader(), + new Class[]{HgStoreSession.class}, this); + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + if ("scanIterator".equals(method.getName())) { + if (args != null && args.length == 1 && + args[0] instanceof Builder) { + if (this.waitBeforeFailure != null) { + try { + if (!this.waitBeforeFailure.await(5L, + TimeUnit.SECONDS)) { + throw new IllegalStateException( + "Timed out waiting for first iterator"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } + if (this.failOnBuilder) { + throw new IllegalStateException("injected failure"); + } + this.builders.add(((Builder) args[0]).clone()); + if (this.opened != null) { + this.opened.countDown(); + } + return this.iterator; + } + this.rangeScanCalls++; + return this.iterator; + } + if ("isTx".equals(method.getName())) { + return false; + } + if ("toString".equals(method.getName())) { + return "RecordingSession"; + } + throw new UnsupportedOperationException(method.toString()); + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java new file mode 100644 index 0000000000..b7c7310e35 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java @@ -0,0 +1,185 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.NoSuchElementException; + +import org.apache.hugegraph.store.HgKvEntry; +import org.apache.hugegraph.store.HgKvIterator; +import org.junit.Assert; +import org.junit.Test; + +public class OrderedKvIteratorTest { + + @Test + public void testMergeInterleavedSourcesByUnsignedKey() { + TestIterator first = new TestIterator(1, 4); + TestIterator second = new TestIterator(2, 3); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L); + + Assert.assertEquals(Arrays.asList(1, 2, 3, 4), keys(iterator)); + Assert.assertArrayEquals(keyBytes(4), iterator.position()); + Assert.assertTrue(first.closed); + Assert.assertTrue(second.closed); + } + + @Test + public void testMergeIsLazyAndStopsAtRawLimit() { + TestIterator first = new TestIterator(1, 4, 5); + TestIterator second = new TestIterator(2, 3, 6); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 3L); + + Assert.assertEquals(0, first.nextCalls); + Assert.assertEquals(0, second.nextCalls); + + Assert.assertEquals(Arrays.asList(1, 2, 3), keys(iterator)); + Assert.assertEquals(4, first.nextCalls + second.nextCalls); + Assert.assertTrue(first.closed); + Assert.assertTrue(second.closed); + } + + @Test + public void testMergeComparesKeysAsUnsignedBytes() { + TestIterator first = new TestIterator(0x80); + TestIterator second = new TestIterator(0x7f); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L); + + Assert.assertEquals(Arrays.asList(0x7f, 0x80), keys(iterator)); + } + + @Test + public void testMergeUsesStableSourceOrderForEqualKeys() { + TestIterator first = new TestIterator(1); + TestIterator second = new TestIterator(1); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L); + + Assert.assertSame(first.entries.get(0), iterator.next()); + Assert.assertSame(second.entries.get(0), iterator.next()); + Assert.assertFalse(iterator.hasNext()); + } + + @Test + public void testMergeClosesAllSourcesWhenAdvanceFails() { + TestIterator first = new TestIterator(1, 3); + TestIterator second = new TestIterator(2, 4); + first.failOnHasNextAfter(1); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L); + + Assert.assertThrows(IllegalStateException.class, iterator::next); + Assert.assertTrue(first.closed); + Assert.assertTrue(second.closed); + } + + private static List keys(HgKvIterator iterator) { + List keys = new ArrayList<>(); + while (iterator.hasNext()) { + keys.add(iterator.next().key()[0] & 0xff); + } + return keys; + } + + private static byte[] keyBytes(int key) { + return new byte[]{(byte) key}; + } + + private static final class TestIterator implements HgKvIterator { + + private final List entries; + private int offset; + private int nextCalls; + private HgKvEntry current; + private boolean closed; + private int failOnHasNextAfter; + + private TestIterator(Integer... keys) { + this.entries = new ArrayList<>(keys.length); + for (int key : keys) { + this.entries.add(new TestEntry(keyBytes(key))); + } + this.offset = 0; + this.nextCalls = 0; + this.current = null; + this.closed = false; + this.failOnHasNextAfter = -1; + } + + private void failOnHasNextAfter(int nextCalls) { + this.failOnHasNextAfter = nextCalls; + } + + @Override + public boolean hasNext() { + if (this.nextCalls == this.failOnHasNextAfter) { + throw new IllegalStateException("injected failure"); + } + return this.offset < this.entries.size(); + } + + @Override + public HgKvEntry next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + this.nextCalls++; + this.current = this.entries.get(this.offset++); + return this.current; + } + + @Override + public byte[] key() { + return this.current == null ? null : this.current.key(); + } + + @Override + public byte[] value() { + return this.current == null ? null : this.current.value(); + } + + @Override + public void close() { + this.closed = true; + } + } + + private static final class TestEntry implements HgKvEntry { + + private final byte[] key; + + private TestEntry(byte[] key) { + this.key = key; + } + + @Override + public byte[] key() { + return this.key; + } + + @Override + public byte[] value() { + return this.key; + } + } +} From 5fbcefef8fc7f2d50a69f89dfc295ed98ce55f67 Mon Sep 17 00:00:00 2001 From: contrueCT Date: Thu, 6 Aug 2026 19:28:30 +0800 Subject: [PATCH 2/6] fix(hstore): merge ordered scans within store nodes --- .../store/client/NodeTxSessionProxy.java | 10 +- .../store/client/OrderedKvIterator.java | 65 +++++- .../store/client/grpc/KvPageScanner.java | 21 +- .../store/business/BusinessHandler.java | 3 + .../store/business/BusinessHandlerImpl.java | 23 ++ .../OrderedMultiPartitionIterator.java | 215 ++++++++++++++++++ .../src/main/proto/store_common.proto | 4 +- .../src/main/proto/store_stream_meta.proto | 1 + .../store/node/grpc/HgStoreWrapperEx.java | 7 + .../hugegraph/store/node/grpc/ScanUtil.java | 12 +- .../OrderedMultiPartitionIteratorTest.java | 154 +++++++++++++ .../store/client/ClientSuiteTest.java | 2 + .../store/client/NodeTxSessionProxyTest.java | 29 ++- .../store/client/OrderedKvIteratorTest.java | 64 ++++++ .../store/client/grpc/KvPageScannerTest.java | 125 ++++++++++ .../store/node/grpc/ScanUtilTest.java | 140 ++++++++++++ .../store/service/ServerSuiteTest.java | 4 + 17 files changed, 854 insertions(+), 25 deletions(-) create mode 100644 hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java index 9802abc639..6b92b2028f 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java @@ -41,7 +41,6 @@ import org.apache.hugegraph.HugeGraphSupplier; import org.apache.hugegraph.pd.common.PDException; -import org.apache.hugegraph.pd.common.PartitionUtils; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; import org.apache.hugegraph.store.HgKvOrderedIterator; @@ -57,6 +56,7 @@ import org.apache.hugegraph.store.client.util.HgStoreClientUtil; import org.apache.hugegraph.store.grpc.common.Header; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder; import org.apache.hugegraph.store.query.StoreQueryParam; @@ -77,6 +77,8 @@ @NotThreadSafe public class NodeTxSessionProxy implements HgStoreSession { + private static final int ORDERED_SCAN_PAGE_SIZE = 64; + private final HgSessionConfig sessionConfig; private final HgStoreNodeManager nodeManager; private final HgStoreNodePartitioner nodePartitioner; @@ -808,9 +810,9 @@ private List toNodeTkvList(String table, HgOwnerKey startKey, HgOwnerKe private List toOrderedRangeNodeTkvList(String table, HgOwnerKey startKey, HgOwnerKey endKey) { - // One Store may host multiple partitions with independent key order + byte[] allOwner = HgStoreClientConst.ALL_PARTITION_OWNER; Collection partitions = - this.doPartition(table, 0, PartitionUtils.MAX_VALUE); + this.doPartition(table, allOwner, allOwner); List nodeTkvs = new ArrayList<>(partitions.size()); for (HgNodePartition partition : partitions) { nodeTkvs.add(new NodeTkv(partition, table, startKey, endKey)); @@ -833,6 +835,8 @@ private Builder orderedRangeScanBuilder(NodeTkv nodeTkv, long limit, .setLimit(scanLimit) .setCode(nodeTkv.getKey().getKeyCode()) .setScanType(scanType) + .setPageSize(ORDERED_SCAN_PAGE_SIZE) + .setOrderType(ScanOrderType.ORDER_BY_KEY) .setQuery(toByteString(query)); } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java index c8a1eca53f..3848496d16 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java @@ -17,21 +17,36 @@ package org.apache.hugegraph.store.client; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.NoSuchElementException; +import java.util.Objects; import java.util.PriorityQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; +import org.apache.hugegraph.store.client.util.ExecutorPool; import org.apache.hugegraph.store.client.util.HgStoreClientConst; final class OrderedKvIterator implements HgKvIterator { + private static final int INITIALIZE_THREADS = 8; + private static final ExecutorService INITIALIZER = + Executors.newFixedThreadPool( + INITIALIZE_THREADS, + ExecutorPool.newThreadFactory("ordered-scan-init")); + private final List> iterators; private final PriorityQueue queue; private final boolean[] sourceClosed; private final long limit; + private final ExecutorService initializer; private boolean initialized; private boolean closed; @@ -41,6 +56,11 @@ final class OrderedKvIterator implements HgKvIterator { OrderedKvIterator(List> iterators, long limit) { + this(iterators, limit, INITIALIZER); + } + + OrderedKvIterator(List> iterators, + long limit, ExecutorService initializer) { this.iterators = iterators; this.queue = new PriorityQueue<>((left, right) -> { int result = Arrays.compareUnsigned(left.entry.key(), @@ -53,6 +73,7 @@ final class OrderedKvIterator implements HgKvIterator { this.sourceClosed = new boolean[iterators.size()]; this.limit = limit <= HgStoreClientConst.NO_LIMIT ? Long.MAX_VALUE : limit; + this.initializer = Objects.requireNonNull(initializer); this.initialized = false; this.closed = false; this.count = 0L; @@ -144,16 +165,45 @@ private void initialize() { return; } this.initialized = true; + List> tasks = + new ArrayList<>(this.iterators.size()); + for (int i = 0; i < this.iterators.size(); i++) { + int source = i; + tasks.add(() -> this.firstEntry(source)); + } try { - for (int i = 0; i < this.iterators.size(); i++) { - this.addNext(i); + List> futures = + this.initializer.invokeAll(tasks); + for (int i = 0; i < futures.size(); i++) { + SourceEntry entry = futures.get(i).get(); + if (entry == null) { + this.closeSource(i); + } else { + this.queue.add(entry); + } } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw this.initializationFailure( + new IllegalStateException( + "Interrupted while initializing ordered scan", e)); + } catch (ExecutionException e) { + throw this.initializationFailure(e.getCause()); } catch (RuntimeException | Error e) { this.closeAfterFailure(e); throw e; } } + private SourceEntry firstEntry(int source) { + HgKvIterator iterator = + this.iterators.get(source); + if (!iterator.hasNext()) { + return null; + } + return new SourceEntry(source, iterator.next()); + } + private void addNext(int source) { HgKvIterator iterator = this.iterators.get(source); @@ -180,6 +230,17 @@ private void closeAfterFailure(Throwable failure) { } } + private RuntimeException initializationFailure(Throwable failure) { + this.closeAfterFailure(failure); + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure instanceof RuntimeException) { + return (RuntimeException) failure; + } + return new IllegalStateException(failure); + } + private static final class SourceEntry { private final int source; diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java index e0a5389fa3..fb6a08a9d9 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java @@ -24,6 +24,7 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; import javax.annotation.concurrent.NotThreadSafe; @@ -60,7 +61,8 @@ class KvPageScanner implements KvCloseableIterator, HgPageSize, HgSeekAble { private static final HgStoreClientConfig clientConfig = HgStoreClientConfig.of(); private static final int nextTimeout = clientConfig.getNetKvScannerHaveNextTimeout(); private final HgStoreNodeSession session; - private final HgStoreStreamStub stub; + private final Function, + StreamObserver> streamFactory; private final AtomicBoolean completed = new AtomicBoolean(false); private final SelectParam.Builder selectBuilder = SelectParam.newBuilder(); private final BlockingQueue reqQueue = new LinkedBlockingQueue<>(); @@ -78,7 +80,7 @@ private KvPageScanner(ScanMethod scanMethod, HgStoreNodeSession session, HgStore int partition, int scanType, byte[] query) { this.session = session; - this.stub = stub; + this.streamFactory = stub::scan; this.pageSize = clientConfig.getNetKvScannerPageSize(); this.reqBuilder.setHeader(this.getHeader(this.session)) .setMethod(scanMethod) @@ -97,9 +99,18 @@ private KvPageScanner(ScanMethod scanMethod, HgStoreNodeSession session, HgStore public KvPageScanner(HgStoreNodeSession session, HgStoreStreamStub stub, ScanStreamReq.Builder reqBuilder) { + this(session, reqBuilder, stub::scan); + } + + KvPageScanner(HgStoreNodeSession session, ScanStreamReq.Builder reqBuilder, + Function, + StreamObserver> streamFactory) { this.session = session; - this.stub = stub; - reqBuilder.setPageSize(pageSize); + this.streamFactory = streamFactory; + this.pageSize = reqBuilder.getPageSize() > 0 ? + reqBuilder.getPageSize() : + clientConfig.getNetKvScannerPageSize(); + reqBuilder.setPageSize(this.pageSize); reqBuilder.setPosition(toBs(this.nodePosition)); this.reqBuilder = reqBuilder; this.init(); @@ -149,7 +160,7 @@ private ScanStreamReq createStopReq() { private void init() { this.proxy = HgBufferProxy.of(() -> this.serverScan()); - this.observer = this.stub.scan(new ServeObserverImpl()); + this.observer = this.streamFactory.apply(new ServeObserverImpl()); } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java index 8133654387..8f390db871 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java @@ -69,6 +69,9 @@ void doPut(String graph, int code, String table, byte[] key, byte[] value) throw ScanIterator scan(String graph, int code, String table, byte[] start, byte[] end, int scanType) throws HgStoreException; + ScanIterator scanOrdered(String graph, String table, byte[] start, + byte[] end, int scanType) throws HgStoreException; + /** * primary index scan */ diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java index 9287bfe267..542e82fe35 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java @@ -344,6 +344,29 @@ public ScanIterator scan(String graph, int code, String table, byte[] start, byt return MultiPartitionIterator.of(ids, function); } + @Override + public ScanIterator scanOrdered(String graph, String table, byte[] start, + byte[] end, int scanType) throws HgStoreException { + List ids = this.getLeaderPartitionIds(graph); + Function function = id -> { + byte[] endKey; + int type; + if (ArrayUtils.isEmpty(end)) { + endKey = keyCreator.getEndKey(id, graph); + type = ScanIterator.Trait.SCAN_LT_END; + } else { + endKey = keyCreator.getEndKey(id, graph, end); + type = scanType; + } + try (RocksDBSession dbSession = getSession(graph, table, id)) { + return new InnerKeyFilter(dbSession.sessionOp().scan( + table, keyCreator.getStartKey(id, graph, start), + endKey, type)); + } + }; + return OrderedMultiPartitionIterator.of(ids, function); + } + /** * Merge ID scans into a single list, and invoke the scan function for others * diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java new file mode 100644 index 0000000000..a2ff8c99db --- /dev/null +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java @@ -0,0 +1,215 @@ +/* + * 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.business; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.PriorityQueue; +import java.util.function.Function; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; +import org.apache.hugegraph.rocksdb.access.ScanIterator; + +public final class OrderedMultiPartitionIterator implements ScanIterator { + + private static final byte[] EMPTY_BYTES = new byte[0]; + + private final List partitionIds; + private final Function supplier; + private final List sources; + private final PriorityQueue queue; + + private boolean initialized; + private boolean closed; + private Integer currentPartitionId; + + private OrderedMultiPartitionIterator(List partitionIds, + Function supplier) { + this.partitionIds = new ArrayList<>(Objects.requireNonNull(partitionIds)); + Collections.sort(this.partitionIds); + this.supplier = Objects.requireNonNull(supplier); + this.sources = new ArrayList<>(this.partitionIds.size()); + this.queue = new PriorityQueue<>((left, right) -> { + int result = Arrays.compareUnsigned(left.entry.name, + right.entry.name); + if (result != 0) { + return result; + } + return Integer.compare(left.partitionId, right.partitionId); + }); + this.initialized = false; + this.closed = false; + this.currentPartitionId = null; + } + + public static OrderedMultiPartitionIterator of( + List partitionIds, + Function supplier) { + return new OrderedMultiPartitionIterator(partitionIds, supplier); + } + + @Override + public boolean hasNext() { + if (this.closed) { + return false; + } + this.initialize(); + if (this.queue.isEmpty()) { + this.close(); + return false; + } + return true; + } + + @Override + public boolean isValid() { + return this.hasNext(); + } + + @Override + @SuppressWarnings("unchecked") + public T next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + + SourceEntry source = this.queue.poll(); + BackendColumn current = source.entry; + this.currentPartitionId = source.partitionId; + try { + if (source.iterator.hasNext()) { + source.entry = source.iterator.next(); + this.queue.add(source); + } else { + this.closeSource(source); + } + } catch (RuntimeException | Error e) { + this.closeAfterFailure(e); + throw e; + } + return (T) current; + } + + @Override + public byte[] position() { + if (this.currentPartitionId == null) { + return EMPTY_BYTES; + } + return ByteBuffer.allocate(Integer.BYTES) + .putInt(this.currentPartitionId) + .array(); + } + + @Override + public void seek(byte[] position) { + if (position == null || position.length == 0) { + return; + } + throw new UnsupportedOperationException( + "Ordered scans resume from their physical start key"); + } + + @Override + public void close() { + if (this.closed) { + return; + } + this.closed = true; + Throwable failure = null; + for (SourceEntry source : this.sources) { + try { + this.closeSource(source); + } catch (RuntimeException | Error e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + this.queue.clear(); + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure != null) { + throw (Error) failure; + } + } + + private void initialize() { + if (this.initialized) { + return; + } + this.initialized = true; + try { + for (int partitionId : this.partitionIds) { + ScanIterator iterator = this.supplier.apply(partitionId); + if (iterator == null) { + continue; + } + SourceEntry source = new SourceEntry(partitionId, iterator); + this.sources.add(source); + if (iterator.hasNext()) { + source.entry = iterator.next(); + this.queue.add(source); + } else { + this.closeSource(source); + } + } + } catch (RuntimeException | Error e) { + this.closeAfterFailure(e); + throw e; + } + } + + private void closeSource(SourceEntry source) { + if (source.closed) { + return; + } + source.closed = true; + source.iterator.close(); + } + + private void closeAfterFailure(Throwable failure) { + try { + this.close(); + } catch (RuntimeException | Error closeFailure) { + failure.addSuppressed(closeFailure); + } + } + + private static final class SourceEntry { + + private final int partitionId; + private final ScanIterator iterator; + private BackendColumn entry; + private boolean closed; + + private SourceEntry(int partitionId, ScanIterator iterator) { + this.partitionId = partitionId; + this.iterator = iterator; + this.entry = null; + this.closed = false; + } + } +} diff --git a/hugegraph-store/hg-store-grpc/src/main/proto/store_common.proto b/hugegraph-store/hg-store-grpc/src/main/proto/store_common.proto index 06d161c70f..4005cbafe4 100644 --- a/hugegraph-store/hg-store-grpc/src/main/proto/store_common.proto +++ b/hugegraph-store/hg-store-grpc/src/main/proto/store_common.proto @@ -81,11 +81,11 @@ enum ScanMethod { RANGE = 3; } -enum ScanOrderType{ - // Under batch interface, the requirement for return order +enum ScanOrderType { ORDER_NONE = 0; // Allow unordered ORDER_WITHIN_VERTEX = 1; // Edges within a vertex will not be broken, but the order between different vertices is unordered. ORDER_STRICT = 2; // Ensure the original input point order + ORDER_BY_KEY = 3; // Sort scan results by their complete key } enum OpType { diff --git a/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto b/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto index 0a08114b3b..018b917e86 100644 --- a/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto +++ b/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto @@ -83,6 +83,7 @@ message ScanStreamReq { bytes position = 12; uint32 closeFlag = 13; SelectParam selects = 14; + ScanOrderType orderType = 15; } message SelectParam { diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreWrapperEx.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreWrapperEx.java index 26e7f2357f..eae9534e89 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreWrapperEx.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/HgStoreWrapperEx.java @@ -63,6 +63,13 @@ public ScanIterator scan(String graph, int partId, String table, byte[] start, b return FilterIterator.of(scanIterator, query); } + public ScanIterator scanOrdered(String graph, String table, byte[] start, + byte[] end, int scanType, byte[] query) { + ScanIterator scanIterator = + this.handler.scanOrdered(graph, table, start, end, scanType); + return FilterIterator.of(scanIterator, query); + } + public void batchGet(String graph, String table, Supplier> s, Consumer> c) { this.handler.batchGet(graph, table, s, (pair -> { diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java index e2fcff42e5..0cf53c2fb8 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java @@ -17,6 +17,8 @@ package org.apache.hugegraph.store.node.grpc; +import static org.apache.hugegraph.store.node.util.HgStoreConst.SCAN_ALL_PARTITIONS_ID; + import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -35,6 +37,7 @@ import org.apache.hugegraph.rocksdb.access.ScanIterator; import org.apache.hugegraph.store.business.SelectIterator; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.grpc.stream.SelectParam; @@ -71,7 +74,14 @@ static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) iter = wrapper.scanPrefix(graph, partition, table, prefix, scanType, query); break; case RANGE: - iter = wrapper.scan(graph, partition, table, start, end, scanType, query); + if (partition == SCAN_ALL_PARTITIONS_ID && + request.getOrderType() == ScanOrderType.ORDER_BY_KEY) { + iter = wrapper.scanOrdered(graph, table, start, end, + scanType, query); + } else { + iter = wrapper.scan(graph, partition, table, start, end, + scanType, query); + } break; } if (iter == null) { diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java new file mode 100644 index 0000000000..630febadae --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java @@ -0,0 +1,154 @@ +/* + * 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.business; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.junit.Assert; +import org.junit.Test; + +public class OrderedMultiPartitionIteratorTest { + + @Test + public void testMergeByUnsignedKeyAndTrackPartitionPosition() { + Map sources = new HashMap<>(); + sources.put(1, new TestIterator(1, 4)); + sources.put(3, new TestIterator(2, 3)); + AtomicInteger supplierCalls = new AtomicInteger(); + OrderedMultiPartitionIterator iterator = + OrderedMultiPartitionIterator.of(Arrays.asList(3, 1), id -> { + supplierCalls.incrementAndGet(); + return sources.get(id); + }); + + Assert.assertEquals(0, supplierCalls.get()); + List keys = new ArrayList<>(); + List positions = new ArrayList<>(); + while (iterator.hasNext()) { + BackendColumn column = iterator.next(); + keys.add(column.name[0] & 0xff); + positions.add(ByteBuffer.wrap(iterator.position()).getInt()); + } + + Assert.assertEquals(Arrays.asList(1, 2, 3, 4), keys); + Assert.assertEquals(Arrays.asList(1, 3, 3, 1), positions); + Assert.assertEquals(2, supplierCalls.get()); + Assert.assertTrue(sources.get(1).closed); + Assert.assertTrue(sources.get(3).closed); + } + + @Test + public void testMergeComparesKeysAsUnsignedBytes() { + Map sources = new HashMap<>(); + sources.put(1, new TestIterator(0x80)); + sources.put(2, new TestIterator(0x7f)); + OrderedMultiPartitionIterator iterator = + OrderedMultiPartitionIterator.of(Arrays.asList(1, 2), sources::get); + + List keys = new ArrayList<>(); + while (iterator.hasNext()) { + BackendColumn column = iterator.next(); + keys.add(column.name[0] & 0xff); + } + + Assert.assertEquals(Arrays.asList(0x7f, 0x80), keys); + } + + @Test + public void testInitializationFailureClosesOpenedIterators() { + TestIterator first = new TestIterator(1); + OrderedMultiPartitionIterator iterator = + OrderedMultiPartitionIterator.of(Arrays.asList(1, 2), id -> { + if (id == 1) { + return first; + } + throw new IllegalStateException("injected failure"); + }); + + Assert.assertThrows(IllegalStateException.class, iterator::hasNext); + Assert.assertTrue(first.closed); + } + + @Test + public void testCloseBeforeInitializationDoesNotOpenSources() { + AtomicInteger supplierCalls = new AtomicInteger(); + OrderedMultiPartitionIterator iterator = + OrderedMultiPartitionIterator.of(Arrays.asList(1, 2), id -> { + supplierCalls.incrementAndGet(); + return new TestIterator(id); + }); + + iterator.close(); + + Assert.assertEquals(0, supplierCalls.get()); + Assert.assertFalse(iterator.hasNext()); + } + + private static byte[] keyBytes(int key) { + return new byte[]{(byte) key}; + } + + private static final class TestIterator implements ScanIterator { + + private final List columns; + private int offset; + private boolean closed; + + private TestIterator(Integer... keys) { + this.columns = new ArrayList<>(keys.length); + for (int key : keys) { + this.columns.add(BackendColumn.of(keyBytes(key), keyBytes(key))); + } + this.offset = 0; + this.closed = false; + } + + @Override + public boolean hasNext() { + return !this.closed && this.offset < this.columns.size(); + } + + @Override + public boolean isValid() { + return this.hasNext(); + } + + @Override + @SuppressWarnings("unchecked") + public T next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + return (T) this.columns.get(this.offset++); + } + + @Override + public void close() { + this.closed = true; + } + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java index ebbe27ac41..ece66898d0 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/ClientSuiteTest.java @@ -18,6 +18,7 @@ package org.apache.hugegraph.store.client; import org.apache.hugegraph.store.client.grpc.AbstractGrpcClientTest; +import org.apache.hugegraph.store.client.grpc.KvPageScannerTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -28,6 +29,7 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ AbstractGrpcClientTest.class, + KvPageScannerTest.class, NodeTxSessionProxyTest.class, OrderedKvIteratorTest.class }) diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java index d5f38cb6e6..44ba5b804b 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java @@ -31,12 +31,12 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; -import org.apache.hugegraph.pd.common.PartitionUtils; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; import org.apache.hugegraph.store.HgOwnerKey; import org.apache.hugegraph.store.HgStoreSession; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder; import org.junit.Assert; import org.junit.Test; @@ -66,7 +66,7 @@ public void testNodeTkvDoesNotMutateSharedOwnerKeys() { } @Test - public void testScanIteratorOrderedUsesPerPartitionBuildersLazily() + public void testScanIteratorOrderedUsesOneStreamPerStoreLazily() throws Exception { HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); @@ -93,19 +93,16 @@ public void testScanIteratorOrderedUsesPerPartitionBuildersLazily() Assert.assertEquals(1, firstSession.builders.size()); Assert.assertEquals(1, secondSession.builders.size()); - Assert.assertEquals(0, partitioner.ownerRangeCalls); - Assert.assertEquals(1, partitioner.codeRangeCalls); - Assert.assertEquals(0, partitioner.startCode); - Assert.assertEquals(PartitionUtils.MAX_VALUE, - partitioner.endCode); + Assert.assertEquals(1, partitioner.ownerRangeCalls); + Assert.assertEquals(0, partitioner.codeRangeCalls); Assert.assertEquals(0, firstSession.rangeScanCalls); Assert.assertEquals(0, secondSession.rangeScanCalls); Assert.assertEquals(0, firstIterator.nextCalls); Assert.assertEquals(0, secondIterator.nextCalls); assertOrderedRangeBuilder(firstSession.builders.get(0), 5L, - 123, 10, keyBytes(7)); + 123, keyBytes(7)); assertOrderedRangeBuilder(secondSession.builders.get(0), 5L, - 123, 30, keyBytes(7)); + 123, keyBytes(7)); Assert.assertEquals(1, key(iterator.next())); Assert.assertEquals(2, key(iterator.next())); @@ -154,13 +151,16 @@ public void testScanIteratorOrderedClosesOpenedIteratorsOnOpenFailure() } private static void assertOrderedRangeBuilder(Builder builder, long limit, - int scanType, int code, + int scanType, byte[] query) { Assert.assertEquals(ScanMethod.RANGE, builder.getMethod()); Assert.assertEquals("table", builder.getTable()); Assert.assertEquals(limit, builder.getLimit()); Assert.assertEquals(scanType, builder.getScanType()); - Assert.assertEquals(code, builder.getCode()); + Assert.assertEquals(-1, builder.getCode()); + Assert.assertEquals(64, builder.getPageSize()); + Assert.assertEquals(ScanOrderType.ORDER_BY_KEY, + builder.getOrderType()); Assert.assertArrayEquals(keyBytes(1), builder.getStart().toByteArray()); Assert.assertArrayEquals(keyBytes(5), builder.getEnd().toByteArray()); Assert.assertArrayEquals(query, builder.getQuery().toByteArray()); @@ -273,7 +273,11 @@ public int partition(HgNodePartitionerBuilder builder, String graphName, byte[] startKey, byte[] endKey) { this.ownerRangeCalls++; - return this.setPartitions(builder); + Set stores = new LinkedHashSet<>(); + stores.add(HgNodePartition.of(this.firstNodeId, -1)); + stores.add(HgNodePartition.of(this.secondNodeId, -1)); + builder.setPartitions(stores); + return 0; } @Override @@ -289,6 +293,7 @@ public int partition(HgNodePartitionerBuilder builder, private int setPartitions(HgNodePartitionerBuilder builder) { Set partitions = new LinkedHashSet<>(); partitions.add(HgNodePartition.of(this.firstNodeId, 10, 10, 20)); + partitions.add(HgNodePartition.of(this.firstNodeId, 20, 20, 30)); partitions.add(HgNodePartition.of(this.secondNodeId, 30, 30, 40)); builder.setPartitions(partitions); return 0; diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java index b7c7310e35..104ee4530b 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java @@ -21,6 +21,10 @@ import java.util.Arrays; import java.util.List; import java.util.NoSuchElementException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; @@ -93,6 +97,43 @@ public void testMergeClosesAllSourcesWhenAdvanceFails() { Assert.assertTrue(second.closed); } + @Test + public void testMergePrimesSourcesConcurrently() { + ExecutorService executor = Executors.newFixedThreadPool(2); + CountDownLatch bothStarted = new CountDownLatch(2); + TestIterator first = new TestIterator(1); + TestIterator second = new TestIterator(2); + first.blockFirstHasNext(bothStarted); + second.blockFirstHasNext(bothStarted); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L, executor); + try { + Assert.assertTrue(iterator.hasNext()); + Assert.assertEquals(0L, bothStarted.getCount()); + } finally { + iterator.close(); + executor.shutdownNow(); + } + } + + @Test + public void testMergeClosesAllSourcesWhenConcurrentInitializeFails() { + ExecutorService executor = Executors.newFixedThreadPool(2); + TestIterator first = new TestIterator(1); + TestIterator second = new TestIterator(2); + first.failOnHasNextAfter(0); + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(first, second), 0L, executor); + try { + Assert.assertThrows(IllegalStateException.class, + iterator::hasNext); + Assert.assertTrue(first.closed); + Assert.assertTrue(second.closed); + } finally { + executor.shutdownNow(); + } + } + private static List keys(HgKvIterator iterator) { List keys = new ArrayList<>(); while (iterator.hasNext()) { @@ -113,6 +154,8 @@ private static final class TestIterator implements HgKvIterator { private HgKvEntry current; private boolean closed; private int failOnHasNextAfter; + private CountDownLatch firstHasNextBarrier; + private boolean firstHasNextBlocked; private TestIterator(Integer... keys) { this.entries = new ArrayList<>(keys.length); @@ -124,17 +167,38 @@ private TestIterator(Integer... keys) { this.current = null; this.closed = false; this.failOnHasNextAfter = -1; + this.firstHasNextBarrier = null; + this.firstHasNextBlocked = false; } private void failOnHasNextAfter(int nextCalls) { this.failOnHasNextAfter = nextCalls; } + private void blockFirstHasNext(CountDownLatch barrier) { + this.firstHasNextBarrier = barrier; + } + @Override public boolean hasNext() { if (this.nextCalls == this.failOnHasNextAfter) { throw new IllegalStateException("injected failure"); } + if (this.firstHasNextBarrier != null && + !this.firstHasNextBlocked) { + this.firstHasNextBlocked = true; + this.firstHasNextBarrier.countDown(); + try { + if (!this.firstHasNextBarrier.await(5L, + TimeUnit.SECONDS)) { + throw new IllegalStateException( + "Timed out waiting for concurrent source"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } return this.offset < this.entries.size(); } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java new file mode 100644 index 0000000000..89fc8f0662 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client.grpc; + +import java.lang.reflect.Proxy; +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; + +import org.apache.hugegraph.store.client.HgStoreNodeSession; +import org.apache.hugegraph.store.grpc.common.Kv; +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.junit.Assert; +import org.junit.Test; + +import com.google.protobuf.ByteString; + +import io.grpc.stub.StreamObserver; + +public class KvPageScannerTest { + + @Test + public void testKeepsExplicitPageSizeAndRequestsNextPageOnDemand() { + List requests = new ArrayList<>(); + Function, StreamObserver> + streamFactory = response -> requestObserver(requests, response); + ScanStreamReq.Builder builder = ScanStreamReq.newBuilder() + .setMethod(ScanMethod.RANGE) + .setTable("table") + .setPageSize(2); + KvPageScanner scanner = new KvPageScanner(session(), builder, + streamFactory); + + Assert.assertEquals(0, requests.size()); + Assert.assertTrue(scanner.hasNext()); + Assert.assertEquals(1, requests.size()); + Assert.assertEquals(2, requests.get(0).getPageSize()); + + Assert.assertEquals(1, scanner.next().getKey().byteAt(0) & 0xff); + Assert.assertTrue(scanner.hasNext()); + Assert.assertEquals(1, requests.size()); + Assert.assertEquals(2, scanner.next().getKey().byteAt(0) & 0xff); + + Assert.assertTrue(scanner.hasNext()); + Assert.assertEquals(2, requests.size()); + Assert.assertEquals(2, requests.get(1).getPageSize()); + Assert.assertEquals(3, scanner.next().getKey().byteAt(0) & 0xff); + Assert.assertFalse(scanner.hasNext()); + Assert.assertEquals(2, requests.size()); + } + + private static HgStoreNodeSession session() { + return (HgStoreNodeSession) Proxy.newProxyInstance( + HgStoreNodeSession.class.getClassLoader(), + new Class[]{HgStoreNodeSession.class}, + (proxy, method, args) -> { + if ("getGraphName".equals(method.getName())) { + return "graph"; + } + if ("isTx".equals(method.getName())) { + return false; + } + if ("toString".equals(method.getName())) { + return "TestNodeSession"; + } + throw new UnsupportedOperationException(method.toString()); + }); + } + + private static StreamObserver requestObserver( + List requests, + StreamObserver response) { + return new StreamObserver() { + @Override + public void onNext(ScanStreamReq request) { + if (request.getCloseFlag() != 0) { + return; + } + requests.add(request); + if (requests.size() == 1) { + response.onNext(page(false, 1, 2)); + } else if (requests.size() == 2) { + response.onNext(page(true, 3)); + } else { + Assert.fail("Unexpected page request"); + } + } + + @Override + public void onError(Throwable throwable) { + throw new AssertionError(throwable); + } + + @Override + public void onCompleted() { + } + }; + } + + private static KvPageRes page(boolean over, int... keys) { + KvPageRes.Builder page = KvPageRes.newBuilder().setOver(over); + for (int key : keys) { + ByteString bytes = ByteString.copyFrom(new byte[]{(byte) key}); + page.addData(Kv.newBuilder().setKey(bytes).setValue(bytes)); + } + return page.build(); + } +} diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java new file mode 100644 index 0000000000..13f32a1cab --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.node.grpc; + +import java.util.NoSuchElementException; + +import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; +import org.apache.hugegraph.rocksdb.access.ScanIterator; +import org.apache.hugegraph.store.grpc.common.Header; +import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; +import org.junit.Assert; +import org.junit.Test; + +public class ScanUtilTest { + + @Test + public void testOrderedAllPartitionRangeUsesOrderedScan() { + RecordingWrapper wrapper = new RecordingWrapper(); + ScanStreamReq request = rangeRequest(ScanOrderType.ORDER_BY_KEY); + + ScanIterator iterator = ScanUtil.getIterator(request, wrapper); + BackendColumn column = iterator.next(); + + Assert.assertEquals(1, column.name[0] & 0xff); + Assert.assertTrue(wrapper.orderedCalled); + Assert.assertFalse(wrapper.legacyCalled); + } + + @Test + public void testUnorderedAllPartitionRangeUsesLegacyScan() { + RecordingWrapper wrapper = new RecordingWrapper(); + ScanStreamReq request = rangeRequest(ScanOrderType.ORDER_NONE); + + ScanIterator iterator = ScanUtil.getIterator(request, wrapper); + BackendColumn column = iterator.next(); + + Assert.assertEquals(9, column.name[0] & 0xff); + Assert.assertFalse(wrapper.orderedCalled); + Assert.assertTrue(wrapper.legacyCalled); + } + + private static ScanStreamReq rangeRequest(ScanOrderType orderType) { + return ScanStreamReq.newBuilder() + .setHeader(Header.newBuilder().setGraph("graph")) + .setMethod(ScanMethod.RANGE) + .setTable("table") + .setCode(-1) + .setStart(bytes(1)) + .setEnd(bytes(5)) + .setScanType(ScanIterator.Trait.SCAN_LT_END) + .setOrderType(orderType) + .build(); + } + + private static com.google.protobuf.ByteString bytes(int value) { + return com.google.protobuf.ByteString.copyFrom(new byte[]{(byte) value}); + } + + private static final class RecordingWrapper extends HgStoreWrapperEx { + + private boolean orderedCalled; + private boolean legacyCalled; + + private RecordingWrapper() { + super(null); + this.orderedCalled = false; + this.legacyCalled = false; + } + + @Override + public ScanIterator scanOrdered(String graph, String table, + byte[] start, byte[] end, + int scanType, byte[] query) { + this.orderedCalled = true; + return new SingleColumnIterator(1); + } + + @Override + public ScanIterator scan(String graph, int partId, String table, + byte[] start, byte[] end, int scanType, + byte[] query) { + this.legacyCalled = true; + return new SingleColumnIterator(9); + } + } + + private static final class SingleColumnIterator implements ScanIterator { + + private final BackendColumn column; + private boolean consumed; + + private SingleColumnIterator(int key) { + byte[] bytes = new byte[]{(byte) key}; + this.column = BackendColumn.of(bytes, bytes); + this.consumed = false; + } + + @Override + public boolean hasNext() { + return !this.consumed; + } + + @Override + public boolean isValid() { + return this.hasNext(); + } + + @Override + @SuppressWarnings("unchecked") + public T next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + this.consumed = true; + return (T) this.column; + } + + @Override + public void close() { + this.consumed = true; + } + } +} 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 00a58d490e..b4328e0f2b 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/service/ServerSuiteTest.java @@ -17,6 +17,8 @@ package org.apache.hugegraph.store.service; +import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest; +import org.apache.hugegraph.store.node.grpc.ScanUtilTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -24,6 +26,8 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ + OrderedMultiPartitionIteratorTest.class, + ScanUtilTest.class }) @Slf4j From b45fdede3451c6f3115b46fe370f90c43e307f1c Mon Sep 17 00:00:00 2001 From: contrueCT Date: Thu, 6 Aug 2026 20:34:35 +0800 Subject: [PATCH 3/6] fix(store): relocate scan util test to node module --- .../org/apache/hugegraph/store/node/grpc/ScanUtilTest.java | 0 .../org/apache/hugegraph/store/service/ServerSuiteTest.java | 4 +--- 2 files changed, 1 insertion(+), 3 deletions(-) rename hugegraph-store/{hg-store-test/src/main => hg-store-node/src/test}/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java (100%) diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java b/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java similarity index 100% rename from hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java rename to hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java 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 b4328e0f2b..2d115e78ed 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 @@ -18,7 +18,6 @@ package org.apache.hugegraph.store.service; import org.apache.hugegraph.store.business.OrderedMultiPartitionIteratorTest; -import org.apache.hugegraph.store.node.grpc.ScanUtilTest; import org.junit.runner.RunWith; import org.junit.runners.Suite; @@ -26,8 +25,7 @@ @RunWith(Suite.class) @Suite.SuiteClasses({ - OrderedMultiPartitionIteratorTest.class, - ScanUtilTest.class + OrderedMultiPartitionIteratorTest.class }) @Slf4j From a05e86f0d1c2fcdd5a719e37b6d030265cf3bf82 Mon Sep 17 00:00:00 2001 From: contrueCT Date: Fri, 7 Aug 2026 13:00:00 +0800 Subject: [PATCH 4/6] fix(hstore): address ordered scan review feedback --- .../hugegraph/backend/page/QueryList.java | 3 +- .../hugegraph/backend/query/QueryResults.java | 187 ++++++++--- .../backend/store/hstore/HstoreSessions.java | 7 +- .../apache/hugegraph/unit/UnitTestSuite.java | 2 + .../unit/core/HstoreSessionsTest.java | 39 +++ .../hugegraph/unit/core/QueryResultsTest.java | 147 +++++++++ .../store/client/OrderedKvIterator.java | 66 +++- .../store/client/grpc/KvPageScanner.java | 16 +- .../store/business/BusinessHandler.java | 8 +- .../OrderedMultiPartitionIterator.java | 10 + .../src/main/proto/store_stream_meta.proto | 5 + .../store/node/grpc/ScanOneShotResponse.java | 6 +- .../store/node/grpc/ScanStreamResponse.java | 9 +- .../hugegraph/store/node/grpc/ScanUtil.java | 24 +- .../store/node/grpc/ScanUtilTest.java | 106 +++++++ .../OrderedMultiPartitionIteratorTest.java | 23 ++ .../store/client/NodeTxSessionProxyTest.java | 300 ++++++++++++++++++ .../store/client/OrderedKvIteratorTest.java | 115 +++++++ .../store/client/grpc/KvPageScannerTest.java | 62 ++++ .../client/grpc/KvPageScannerTestSupport.java | 43 +++ 20 files changed, 1113 insertions(+), 65 deletions(-) create mode 100644 hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/HstoreSessionsTest.java create mode 100644 hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTestSupport.java diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java index b581a6e714..327c4aab47 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/page/QueryList.java @@ -89,7 +89,8 @@ public QueryResults fetch(int pageSize) { PageEntryIterator iter = new PageEntryIterator<>(this, pageSize); /* * NOTE: PageEntryIterator query will change every fetch time. - * TODO: sort results by input ids in each page. + * QueryResults tracks this change and restores input-id order + * within each page without fetching later pages eagerly. */ return iter.results(); } diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java index aa7fb70263..01c2024e45 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/QueryResults.java @@ -17,6 +17,7 @@ package org.apache.hugegraph.backend.query; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.Iterator; @@ -32,6 +33,7 @@ import org.apache.hugegraph.iterator.FlatMapperIterator; import org.apache.hugegraph.iterator.ListIterator; import org.apache.hugegraph.iterator.MapperIterator; +import org.apache.hugegraph.iterator.WrappedIterator; import org.apache.hugegraph.perf.PerfUtil.Watched; import org.apache.hugegraph.type.Idfiable; import org.apache.hugegraph.util.E; @@ -47,6 +49,8 @@ public class QueryResults { private final Iterator results; private final List queries; + private List currentQueries; + private long queryVersion; public QueryResults(Iterator results, Query query) { this(results); @@ -56,6 +60,8 @@ public QueryResults(Iterator results, Query query) { private QueryResults(Iterator results) { this.results = results; this.queries = InsertionOrderUtil.newList(); + this.currentQueries = Collections.emptyList(); + this.queryVersion = 0L; } public void setQuery(Query query) { @@ -67,14 +73,17 @@ public void setQuery(Query query) { private void addQuery(Query query) { E.checkNotNull(query, "query"); - this.queries.add(query); + this.addQueries(Collections.singletonList(query)); } private void addQueries(List queries) { assert !queries.isEmpty(); for (Query query : queries) { - this.addQuery(query); + E.checkNotNull(query, "query"); + this.queries.add(query); } + this.currentQueries = new ArrayList<>(queries); + this.queryVersion++; } public Iterator iterator() { @@ -101,39 +110,15 @@ public Iterator keepInputOrderIfNeeded( // None result found return origin; } - Collection ids; - if (!this.mustSortByInputIds() || - (ids = this.queryIds()).size() <= 1) { - /* - * Return the original iterator if the query input is less than one - * id, or don't have to do sort. - * NOTE: queryIds() only return the first batch of index query. - */ + if (!mustSortByInputIds(this.currentQueries)) { return origin; } - - // Fill map with all elements - Map map = InsertionOrderUtil.newMap(); - QueryResults.fillMap(origin, map); - - if (map.size() > ids.size()) { - /* - * This means current query is part of QueryResults. For example, - * g.V().has('country', 'china').has('city', within('HK', 'BJ')) - * will be converted to - * g.V().has('country', 'china').has('city', 'HK') or - * g.V().has('country', 'china').has('city', 'BJ'), - * and ids is just first index subquery's id, not all. - */ - ids = map.keySet(); - } - - return new MapperIterator<>(ids.iterator(), map::get); + return new InputOrderIterator<>(this, origin); } - private boolean mustSortByInputIds() { - assert !this.queries.isEmpty() : this; - for (Query query : this.queries) { + private static boolean mustSortByInputIds(List queries) { + assert !queries.isEmpty() : queries; + for (Query query : queries) { if (query instanceof IdQuery && ((IdQuery) query).mustSortByInput()) { return true; @@ -153,14 +138,14 @@ private boolean bigCapacity() { return false; } - private Collection queryIds() { - assert !this.queries.isEmpty(); - if (this.queries.size() == 1) { - return this.queries.get(0).ids(); + private static Collection queryIds(List queries) { + assert !queries.isEmpty(); + if (queries.size() == 1) { + return queries.get(0).ids(); } Set ids = InsertionOrderUtil.newSet(); - for (Query query : this.queries) { + for (Query query : queries) { ids.addAll(query.ids()); } return ids; @@ -212,16 +197,19 @@ public static QueryResults flatMap( if (results == null || !results.iterator().hasNext()) { return null; } - /* - * NOTE: should call results.iterator().hasNext() before - * results.queries() to collect sub-query with index query - */ - qr[0].addQueries(results.queries()); - return results.iterator(); + return new QueryTrackingIterator<>(qr[0], results); })); return qr[0]; } + private long queryVersion() { + return this.queryVersion; + } + + private List currentQueries() { + return new ArrayList<>(this.currentQueries); + } + @Watched public static T one(Iterator iterator) { try { @@ -315,4 +303,119 @@ public void close() throws Exception { // pass } } + + private static class QueryTrackingIterator + extends WrappedIterator { + + private final QueryResults parent; + private final QueryResults child; + private long childQueryVersion; + + public QueryTrackingIterator(QueryResults parent, + QueryResults child) { + this.parent = parent; + this.child = child; + this.childQueryVersion = -1L; + } + + @Override + protected Iterator originIterator() { + return this.child.iterator(); + } + + @Override + protected boolean fetch() { + Iterator origin = this.child.iterator(); + if (!origin.hasNext()) { + return false; + } + R result = origin.next(); + long queryVersion = this.child.queryVersion(); + if (this.childQueryVersion != queryVersion) { + this.parent.addQueries(this.child.currentQueries()); + this.childQueryVersion = queryVersion; + } + assert this.current == none(); + this.current = result; + return true; + } + } + + private static class InputOrderIterator + extends WrappedIterator { + + private final QueryResults queryResults; + private final Iterator origin; + private Iterator currentBatch; + + public InputOrderIterator(QueryResults queryResults, + Iterator origin) { + this.queryResults = queryResults; + this.origin = origin; + this.currentBatch = Collections.emptyIterator(); + } + + @Override + protected Iterator originIterator() { + return this.origin; + } + + @Override + protected boolean fetch() { + while (true) { + if (this.currentBatch.hasNext()) { + assert this.current == none(); + this.current = this.currentBatch.next(); + return true; + } + if (!this.origin.hasNext()) { + return false; + } + this.currentBatch = this.fetchBatch(); + } + } + + private Iterator fetchBatch() { + long queryVersion = this.queryResults.queryVersion(); + List queries = this.queryResults.currentQueries(); + List results = InsertionOrderUtil.newList(); + do { + results.add(this.origin.next()); + Query.checkForceCapacity(results.size()); + } while (this.origin.hasNext() && + queryVersion == this.queryResults.queryVersion()); + + if (!mustSortByInputIds(queries)) { + return results.iterator(); + } + Collection ids = queryIds(queries); + if (ids.size() <= 1) { + return results.iterator(); + } + + Map byId = InsertionOrderUtil.newMap(); + for (T result : results) { + assert result.id() != null; + byId.put(result.id(), result); + } + if (byId.size() > ids.size()) { + /* + * The current query only describes part of this segment. + * Preserve backend order because it can't fully define the + * order of every returned result. + */ + return results.iterator(); + } + + List ordered = new ArrayList<>(results.size()); + for (Id id : ids) { + T result = byId.remove(id); + if (result != null) { + ordered.add(result); + } + } + ordered.addAll(byId.values()); + return ordered.iterator(); + } + } } diff --git a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java index 958e1bf0c1..85d664c587 100755 --- a/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java +++ b/hugegraph-server/hugegraph-hstore/src/main/java/org/apache/hugegraph/backend/store/hstore/HstoreSessions.java @@ -159,10 +159,13 @@ public abstract BackendColumnIterator scan(String table, byte[] query, byte[] position); - public abstract BackendColumnIterator scanOrdered( + public BackendColumnIterator scanOrdered( String table, byte[] ownerKeyFrom, byte[] ownerKeyTo, byte[] keyFrom, byte[] keyTo, int scanType, byte[] query, - long limit); + long limit) { + throw new UnsupportedOperationException( + "Ordered scan is not supported"); + } public abstract BackendColumnIterator scan(String table, int codeFrom, diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java index b629cf2326..c93ac70ee3 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/UnitTestSuite.java @@ -49,6 +49,7 @@ import org.apache.hugegraph.unit.core.ExceptionTest; import org.apache.hugegraph.unit.core.GraphManagerAdminInitTest; import org.apache.hugegraph.unit.core.GraphManagerConfigTest; +import org.apache.hugegraph.unit.core.HstoreSessionsTest; import org.apache.hugegraph.unit.core.IdHolderTest; import org.apache.hugegraph.unit.core.LocksTableTest; import org.apache.hugegraph.unit.core.PageStateTest; @@ -152,6 +153,7 @@ ExceptionTest.class, GraphManagerAdminInitTest.class, GraphManagerConfigTest.class, + HstoreSessionsTest.class, BackendStoreInfoTest.class, TraversalUtilTest.class, TraversalUtilOptimizeTest.class, diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/HstoreSessionsTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/HstoreSessionsTest.java new file mode 100644 index 0000000000..c7879beb9b --- /dev/null +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/HstoreSessionsTest.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.unit.core; + +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; + +import org.apache.hugegraph.backend.store.hstore.HstoreSessions; +import org.junit.Assert; +import org.junit.Test; + +public class HstoreSessionsTest { + + @Test + public void testOrderedScanDoesNotAddAbstractSubclassRequirement() + throws Exception { + Method method = HstoreSessions.Session.class.getDeclaredMethod( + "scanOrdered", String.class, byte[].class, byte[].class, + byte[].class, byte[].class, int.class, byte[].class, + long.class); + + Assert.assertFalse(Modifier.isAbstract(method.getModifiers())); + } +} diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java index 2728de0015..0ee2cf1678 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryResultsTest.java @@ -19,8 +19,11 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Iterator; import java.util.List; +import java.util.NoSuchElementException; import java.util.Set; +import java.util.function.Consumer; import org.apache.hugegraph.backend.id.Id; import org.apache.hugegraph.backend.id.IdGenerator; @@ -65,6 +68,102 @@ public void testKeepInputOrderForPagingIdQuery() { Assert.assertEquals(ImmutableList.of(id2, id1), orderedIds); } + @Test + public void testKeepInputOrderAcrossBatches() { + List firstInput = new ArrayList<>(); + List firstOutput = new ArrayList<>(); + for (long id = 0L; id < Query.QUERY_BATCH; id++) { + firstInput.add(Query.QUERY_BATCH - id - 1L); + firstOutput.add(id); + } + List secondInput = ImmutableList.of( + Query.QUERY_BATCH + 1L, Query.QUERY_BATCH); + List secondOutput = ImmutableList.of( + Query.QUERY_BATCH, Query.QUERY_BATCH + 1L); + QueryResults first = resultsOf( + firstInput, firstOutput); + QueryResults second = resultsOf( + secondInput, secondOutput); + QueryResults results = QueryResults.flatMap( + ImmutableList.of(first, second).iterator(), result -> result); + + List orderedIds = new ArrayList<>(); + results.keepInputOrderIfNeeded(results.iterator()) + .forEachRemaining(item -> orderedIds.add(item.id())); + + List expected = new ArrayList<>(); + firstInput.forEach(id -> expected.add(IdGenerator.of(id))); + secondInput.forEach(id -> expected.add(IdGenerator.of(id))); + Assert.assertTrue(orderedIds.size() > Query.QUERY_BATCH); + Assert.assertEquals(expected, orderedIds); + } + + @Test + public void testKeepBackendOrderWhenQueryOnlyDescribesPartOfResults() { + QueryResults results = resultsOf( + ImmutableList.of(2L, 3L), + ImmutableList.of(1L, 2L, 3L)); + + List orderedIds = new ArrayList<>(); + results.keepInputOrderIfNeeded(results.iterator()) + .forEachRemaining(item -> orderedIds.add(item.id())); + + Assert.assertEquals(ImmutableList.of(IdGenerator.of(1L), + IdGenerator.of(2L), + IdGenerator.of(3L)), + orderedIds); + } + + @Test + public void testKeepInputOrderDoesNotDrainFollowingPages() { + IdQuery firstQuery = queryOf(2L, 1L); + IdQuery secondQuery = queryOf(4L, 3L); + @SuppressWarnings("unchecked") + QueryResults[] holder = new QueryResults[1]; + PagingIterator origin = new PagingIterator( + ImmutableList.of(new TestIdfiable(IdGenerator.of(1L)), + new TestIdfiable(IdGenerator.of(2L)), + new TestIdfiable(IdGenerator.of(3L)), + new TestIdfiable(IdGenerator.of(4L))), + 2, + query -> holder[0].setQuery(query), + ImmutableList.of(firstQuery, secondQuery)); + holder[0] = new QueryResults<>(origin, firstQuery); + + Iterator ordered = + holder[0].keepInputOrderIfNeeded(holder[0].iterator()); + + Assert.assertEquals(IdGenerator.of(2L), ordered.next().id()); + Assert.assertEquals(IdGenerator.of(1L), ordered.next().id()); + Assert.assertEquals(2, origin.consumed()); + Assert.assertEquals(IdGenerator.of(4L), ordered.next().id()); + Assert.assertEquals(IdGenerator.of(3L), ordered.next().id()); + Assert.assertFalse(ordered.hasNext()); + } + + private static QueryResults resultsOf(List input, + List output) { + List results = new ArrayList<>(output.size()); + for (Long id : output) { + results.add(new TestIdfiable(IdGenerator.of(id))); + } + return new QueryResults<>(results.iterator(), queryOf(input)); + } + + private static IdQuery queryOf(Long... ids) { + return queryOf(Arrays.asList(ids)); + } + + private static IdQuery queryOf(List ids) { + Set queryIds = InsertionOrderUtil.newSet(); + for (Long id : ids) { + queryIds.add(IdGenerator.of(id)); + } + IdQuery query = new IdQuery(new Query(HugeType.VERTEX), queryIds); + query.mustSortByInput(true); + return query; + } + private static final class TestIdfiable implements Idfiable { private final Id id; @@ -78,4 +177,52 @@ public Id id() { return this.id; } } + + private static final class PagingIterator + implements Iterator { + + private final List results; + private final int pageSize; + private final Consumer pageListener; + private final List queries; + + private int current; + private int announcedPage; + + private PagingIterator(List results, int pageSize, + Consumer pageListener, + List queries) { + this.results = results; + this.pageSize = pageSize; + this.pageListener = pageListener; + this.queries = queries; + this.current = 0; + this.announcedPage = 0; + } + + @Override + public boolean hasNext() { + if (this.current >= this.results.size()) { + return false; + } + int page = this.current / this.pageSize; + if (page != this.announcedPage) { + this.pageListener.accept(this.queries.get(page)); + this.announcedPage = page; + } + return true; + } + + @Override + public TestIdfiable next() { + if (!this.hasNext()) { + throw new NoSuchElementException(); + } + return this.results.get(this.current++); + } + + private int consumed() { + return this.current; + } + } } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java index 3848496d16..b234706caa 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java @@ -24,8 +24,10 @@ import java.util.Objects; import java.util.PriorityQueue; import java.util.concurrent.Callable; +import java.util.concurrent.CompletionService; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; +import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -37,16 +39,13 @@ final class OrderedKvIterator implements HgKvIterator { private static final int INITIALIZE_THREADS = 8; - private static final ExecutorService INITIALIZER = - Executors.newFixedThreadPool( - INITIALIZE_THREADS, - ExecutorPool.newThreadFactory("ordered-scan-init")); private final List> iterators; private final PriorityQueue queue; private final boolean[] sourceClosed; private final long limit; private final ExecutorService initializer; + private final boolean ownsInitializer; private boolean initialized; private boolean closed; @@ -56,11 +55,18 @@ final class OrderedKvIterator implements HgKvIterator { OrderedKvIterator(List> iterators, long limit) { - this(iterators, limit, INITIALIZER); + this(iterators, limit, createInitializer(iterators), true); } OrderedKvIterator(List> iterators, long limit, ExecutorService initializer) { + this(iterators, limit, initializer, false); + } + + private OrderedKvIterator( + List> iterators, + long limit, ExecutorService initializer, + boolean ownsInitializer) { this.iterators = iterators; this.queue = new PriorityQueue<>((left, right) -> { int result = Arrays.compareUnsigned(left.entry.key(), @@ -74,6 +80,7 @@ final class OrderedKvIterator implements HgKvIterator { this.limit = limit <= HgStoreClientConst.NO_LIMIT ? Long.MAX_VALUE : limit; this.initializer = Objects.requireNonNull(initializer); + this.ownsInitializer = ownsInitializer; this.initialized = false; this.closed = false; this.count = 0L; @@ -152,6 +159,7 @@ public void close() { } } this.queue.clear(); + this.shutdownInitializer(); if (failure instanceof RuntimeException) { throw (RuntimeException) failure; } @@ -171,27 +179,37 @@ private void initialize() { int source = i; tasks.add(() -> this.firstEntry(source)); } + List> futures = + new ArrayList<>(this.iterators.size()); try { - List> futures = - this.initializer.invokeAll(tasks); - for (int i = 0; i < futures.size(); i++) { - SourceEntry entry = futures.get(i).get(); - if (entry == null) { - this.closeSource(i); + CompletionService completions = + new ExecutorCompletionService<>(this.initializer); + for (Callable task : tasks) { + futures.add(completions.submit(task)); + } + for (int i = 0; i < tasks.size(); i++) { + SourceEntry entry = completions.take().get(); + if (entry.entry == null) { + this.closeSource(entry.source); } else { this.queue.add(entry); } } } catch (InterruptedException e) { + this.cancel(futures); Thread.currentThread().interrupt(); throw this.initializationFailure( new IllegalStateException( "Interrupted while initializing ordered scan", e)); } catch (ExecutionException e) { + this.cancel(futures); throw this.initializationFailure(e.getCause()); } catch (RuntimeException | Error e) { + this.cancel(futures); this.closeAfterFailure(e); throw e; + } finally { + this.shutdownInitializer(); } } @@ -199,11 +217,35 @@ private SourceEntry firstEntry(int source) { HgKvIterator iterator = this.iterators.get(source); if (!iterator.hasNext()) { - return null; + return new SourceEntry(source, null); } return new SourceEntry(source, iterator.next()); } + private static ExecutorService createInitializer( + List> iterators) { + Objects.requireNonNull(iterators); + int threads = Math.max(1, Math.min(INITIALIZE_THREADS, + iterators.size())); + return Executors.newFixedThreadPool( + threads, + ExecutorPool.newThreadFactory("ordered-scan-init")); + } + + private void cancel(List> futures) { + for (Future future : futures) { + if (!future.isDone()) { + future.cancel(true); + } + } + } + + private void shutdownInitializer() { + if (this.ownsInitializer) { + this.initializer.shutdownNow(); + } + } + private void addNext(int source) { HgKvIterator iterator = this.iterators.get(source); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java index fb6a08a9d9..4bf938fa80 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScanner.java @@ -41,8 +41,10 @@ import org.apache.hugegraph.store.grpc.common.Header; import org.apache.hugegraph.store.grpc.common.Kv; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.HgStoreStreamGrpc.HgStoreStreamStub; import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.ScanStreamVersion; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.grpc.stream.SelectParam; @@ -63,6 +65,7 @@ class KvPageScanner implements KvCloseableIterator, HgPageSize, HgSeekAble { private final HgStoreNodeSession session; private final Function, StreamObserver> streamFactory; + private final boolean orderedScan; private final AtomicBoolean completed = new AtomicBoolean(false); private final SelectParam.Builder selectBuilder = SelectParam.newBuilder(); private final BlockingQueue reqQueue = new LinkedBlockingQueue<>(); @@ -81,6 +84,7 @@ private KvPageScanner(ScanMethod scanMethod, HgStoreNodeSession session, HgStore int scanType, byte[] query) { this.session = session; this.streamFactory = stub::scan; + this.orderedScan = false; this.pageSize = clientConfig.getNetKvScannerPageSize(); this.reqBuilder.setHeader(this.getHeader(this.session)) .setMethod(scanMethod) @@ -107,6 +111,8 @@ public KvPageScanner(HgStoreNodeSession session, HgStoreStreamStub stub, StreamObserver> streamFactory) { this.session = session; this.streamFactory = streamFactory; + this.orderedScan = reqBuilder.getOrderType() == + ScanOrderType.ORDER_BY_KEY; this.pageSize = reqBuilder.getPageSize() > 0 ? reqBuilder.getPageSize() : clientConfig.getNetKvScannerPageSize(); @@ -287,6 +293,14 @@ private class ServeObserverImpl implements StreamObserver { @Override public void onNext(KvPageRes value) { + if (orderedScan && value.getVersion() < + ScanStreamVersion + .SCAN_STREAM_VERSION_ORDERED_BY_KEY_VALUE) { + this.onError(new IllegalStateException( + "Store node doesn't support ordered scan; " + + "upgrade every Store node before using ORDER_BY_KEY")); + return; + } if (value.getOver()) { completed.set(true); observer.onCompleted(); @@ -305,8 +319,8 @@ public void onError(Throwable t) { } catch (Exception e) { log.warn("failed to invoke requestObserver.onCompleted(), reason:", e.getMessage()); } - proxy.close(); proxy.setError(t); + proxy.close(); log.error("failed to complete scan of session: " + session, t); } diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java index 8f390db871..b82317707d 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java @@ -69,8 +69,12 @@ void doPut(String graph, int code, String table, byte[] key, byte[] value) throw ScanIterator scan(String graph, int code, String table, byte[] start, byte[] end, int scanType) throws HgStoreException; - ScanIterator scanOrdered(String graph, String table, byte[] start, - byte[] end, int scanType) throws HgStoreException; + default ScanIterator scanOrdered(String graph, String table, byte[] start, + byte[] end, int scanType) + throws HgStoreException { + throw new UnsupportedOperationException( + "Ordered scan is not supported"); + } /** * primary index scan diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java index a2ff8c99db..f7f013b89f 100644 --- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java +++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIterator.java @@ -110,6 +110,16 @@ public T next() { return (T) current; } + @Override + public long count() { + long count = 0L; + while (this.hasNext()) { + this.next(); + count++; + } + return count; + } + @Override public byte[] position() { if (this.currentPartitionId == null) { diff --git a/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto b/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto index 018b917e86..21a7c9813a 100644 --- a/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto +++ b/hugegraph-store/hg-store-grpc/src/main/proto/store_stream_meta.proto @@ -100,6 +100,11 @@ message KvPageRes { bytes stream = 5; } +enum ScanStreamVersion { + SCAN_STREAM_VERSION_UNKNOWN = 0; + SCAN_STREAM_VERSION_ORDERED_BY_KEY = 1; +} + enum KvStreamType { STREAM_TYPE_NONE = 0; STREAM_TYPE_KV = 1; diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java index 203628f9e6..5b77a4449a 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanOneShotResponse.java @@ -50,7 +50,6 @@ public static void scanOneShot(ScanStreamReq request, HgStoreWrapperEx wrapper) { KvPageRes.Builder resBuilder = KvPageRes.newBuilder(); Kv.Builder kvBuilder = Kv.newBuilder(); - ScanIterator iterator = ScanUtil.getIterator(ScanUtil.toSq(request), wrapper); long limit = request.getLimit(); @@ -58,6 +57,7 @@ public static void scanOneShot(ScanStreamReq request, responseObserver.onError(HgGrpc.toErr("limit<=0, please to invoke stream scan.")); return; } + ScanIterator iterator = ScanUtil.getIterator(request, wrapper); int count = 0; @@ -78,7 +78,9 @@ public static void scanOneShot(ScanStreamReq request, } - responseObserver.onNext(resBuilder.build()); + responseObserver.onNext( + resBuilder.setVersion(ScanUtil.responseVersion(request)) + .build()); responseObserver.onCompleted(); } catch (Throwable t) { diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java index a4e7369f1d..e572862823 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanStreamResponse.java @@ -65,6 +65,7 @@ public class ScanStreamResponse implements StreamObserver { private int times = 0; private long pageSize = 0; private int total = 0; + private int responseVersion = 0; private String graph; private String table; @@ -91,6 +92,8 @@ public static ScanStreamResponse of(StreamObserver responseObserver, @Override public void onNext(ScanStreamReq request) { try { + this.responseVersion = Math.max( + this.responseVersion, ScanUtil.responseVersion(request)); if (request.getCloseFlag() == 1) { close(); } else { @@ -209,6 +212,7 @@ private void close() { .addAllData(Collections.EMPTY_LIST) .setOver(true) .setTimes(++times) + .setVersion(this.responseVersion) .build() ); } @@ -236,7 +240,10 @@ private void next(ScanStreamReq request) { resBuilder = KvPageRes.newBuilder().addAllData(Collections.EMPTY_LIST); } if (!this.finishFlag.get()) { - responseObserver.onNext(resBuilder.setOver(isOver).setTimes(times).build()); + responseObserver.onNext(resBuilder.setOver(isOver) + .setTimes(times) + .setVersion(this.responseVersion) + .build()); } if (isOver) { this.finishServer(); diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java index 0cf53c2fb8..a1a706edcc 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanUtil.java @@ -40,6 +40,7 @@ import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; +import org.apache.hugegraph.store.grpc.stream.ScanStreamVersion; import org.apache.hugegraph.store.grpc.stream.SelectParam; import lombok.extern.slf4j.Slf4j; @@ -55,6 +56,12 @@ class ScanUtil { private final static Map tableKeyMap = new HashMap<>(); static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) { + boolean ordered = isOrdered(request); + if (ordered && !request.getPosition().isEmpty()) { + throw new IllegalArgumentException( + "Ordered scan doesn't support an opaque position; " + + "continue with the physical key in start"); + } String graph = request.getHeader().getGraph(); String table = request.getTable(); ScanMethod method = request.getMethod(); @@ -74,8 +81,7 @@ static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) iter = wrapper.scanPrefix(graph, partition, table, prefix, scanType, query); break; case RANGE: - if (partition == SCAN_ALL_PARTITIONS_ID && - request.getOrderType() == ScanOrderType.ORDER_BY_KEY) { + if (ordered) { iter = wrapper.scanOrdered(graph, table, start, end, scanType, query); } else { @@ -98,6 +104,20 @@ static ScanIterator getIterator(ScanStreamReq request, HgStoreWrapperEx wrapper) return iter; } + static boolean isOrdered(ScanStreamReq request) { + return request.getMethod() == ScanMethod.RANGE && + request.getCode() == SCAN_ALL_PARTITIONS_ID && + request.getOrderType() == ScanOrderType.ORDER_BY_KEY; + } + + static int responseVersion(ScanStreamReq request) { + if (isOrdered(request)) { + return ScanStreamVersion + .SCAN_STREAM_VERSION_ORDERED_BY_KEY_VALUE; + } + return ScanStreamVersion.SCAN_STREAM_VERSION_UNKNOWN_VALUE; + } + static ScanIterator getIterator(ScanQuery sq, HgStoreWrapperEx wrapper) { if (log.isDebugEnabled()) { log.debug("{}", sq); diff --git a/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java b/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java index 13f32a1cab..9764d7264d 100644 --- a/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java +++ b/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java @@ -18,16 +18,23 @@ package org.apache.hugegraph.store.node.grpc; import java.util.NoSuchElementException; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicReference; import org.apache.hugegraph.rocksdb.access.RocksDBSession.BackendColumn; import org.apache.hugegraph.rocksdb.access.ScanIterator; import org.apache.hugegraph.store.grpc.common.Header; import org.apache.hugegraph.store.grpc.common.ScanMethod; import org.apache.hugegraph.store.grpc.common.ScanOrderType; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; +import org.apache.hugegraph.store.node.AppConfig; import org.junit.Assert; import org.junit.Test; +import io.grpc.stub.StreamObserver; + public class ScanUtilTest { @Test @@ -56,6 +63,71 @@ public void testUnorderedAllPartitionRangeUsesLegacyScan() { Assert.assertTrue(wrapper.legacyCalled); } + @Test + public void testOrderedScanRejectsOpaquePositionBeforeOpeningIterator() { + RecordingWrapper wrapper = new RecordingWrapper(); + ScanStreamReq request = rangeRequest(ScanOrderType.ORDER_BY_KEY) + .toBuilder() + .setPosition(bytes(3)) + .build(); + + try { + ScanUtil.getIterator(request, wrapper); + Assert.fail("Expected ordered cursor rejection"); + } catch (IllegalArgumentException e) { + Assert.assertTrue(e.getMessage().contains("position")); + } + Assert.assertFalse(wrapper.orderedCalled); + Assert.assertFalse(wrapper.legacyCalled); + } + + @Test + public void testOneShotOrderedRangeUsesOrderedScanAndAcknowledgesIt() { + RecordingWrapper wrapper = new RecordingWrapper(); + ScanStreamReq request = rangeRequest(ScanOrderType.ORDER_BY_KEY); + AtomicReference response = new AtomicReference<>(); + AtomicReference error = new AtomicReference<>(); + + ScanOneShotResponse.scanOneShot(request, + observer(response, error), wrapper); + + Assert.assertNull(error.get()); + Assert.assertNotNull(response.get()); + Assert.assertEquals(1, response.get().getVersion()); + Assert.assertTrue(wrapper.orderedCalled); + Assert.assertFalse(wrapper.legacyCalled); + } + + @Test + public void testOrderedStreamCloseAcknowledgesCapabilityBeforeFirstPage() { + RecordingWrapper wrapper = new RecordingWrapper(); + AtomicReference response = new AtomicReference<>(); + AtomicReference error = new AtomicReference<>(); + AppConfig config = new AppConfig(); + config.setServerWaitTime(1); + ThreadPoolExecutor executor = + (ThreadPoolExecutor) Executors.newFixedThreadPool(1); + try { + ScanStreamResponse stream = ScanStreamResponse.of( + observer(response, error), wrapper, executor, config); + ScanStreamReq close = rangeRequest(ScanOrderType.ORDER_BY_KEY) + .toBuilder() + .setCloseFlag(1) + .build(); + + stream.onNext(close); + + Assert.assertNull(error.get()); + Assert.assertNotNull(response.get()); + Assert.assertTrue(response.get().getOver()); + Assert.assertEquals(1, response.get().getVersion()); + Assert.assertFalse(wrapper.orderedCalled); + Assert.assertFalse(wrapper.legacyCalled); + } finally { + executor.shutdownNow(); + } + } + private static ScanStreamReq rangeRequest(ScanOrderType orderType) { return ScanStreamReq.newBuilder() .setHeader(Header.newBuilder().setGraph("graph")) @@ -65,10 +137,31 @@ private static ScanStreamReq rangeRequest(ScanOrderType orderType) { .setStart(bytes(1)) .setEnd(bytes(5)) .setScanType(ScanIterator.Trait.SCAN_LT_END) + .setLimit(1L) .setOrderType(orderType) .build(); } + private static StreamObserver observer( + AtomicReference response, + AtomicReference error) { + return new StreamObserver() { + @Override + public void onNext(KvPageRes value) { + response.set(value); + } + + @Override + public void onError(Throwable throwable) { + error.set(throwable); + } + + @Override + public void onCompleted() { + } + }; + } + private static com.google.protobuf.ByteString bytes(int value) { return com.google.protobuf.ByteString.copyFrom(new byte[]{(byte) value}); } @@ -136,5 +229,18 @@ public T next() { public void close() { this.consumed = true; } + + @Override + public byte[] position() { + return new byte[]{0, 0, 0, 1}; + } + + @Override + public void seek(byte[] position) { + if (position.length > 0) { + throw new IllegalArgumentException( + "Ordered scan position is unsupported"); + } + } } } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java index 630febadae..25b4f5fd8d 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/business/OrderedMultiPartitionIteratorTest.java @@ -33,6 +33,14 @@ public class OrderedMultiPartitionIteratorTest { + @Test + public void testOrderedScanIsBackwardCompatibleDefaultMethod() + throws Exception { + Assert.assertTrue(BusinessHandler.class.getMethod( + "scanOrdered", String.class, String.class, byte[].class, + byte[].class, int.class).isDefault()); + } + @Test public void testMergeByUnsignedKeyAndTrackPartitionPosition() { Map sources = new HashMap<>(); @@ -78,6 +86,21 @@ public void testMergeComparesKeysAsUnsignedBytes() { Assert.assertEquals(Arrays.asList(0x7f, 0x80), keys); } + @Test + public void testCountConsumesAllRemainingEntries() { + Map sources = new HashMap<>(); + sources.put(1, new TestIterator(1, 4)); + sources.put(2, new TestIterator(2, 3)); + OrderedMultiPartitionIterator iterator = + OrderedMultiPartitionIterator.of(Arrays.asList(1, 2), + sources::get); + + Assert.assertEquals(4L, iterator.count()); + Assert.assertFalse(iterator.hasNext()); + Assert.assertTrue(sources.get(1).closed); + Assert.assertTrue(sources.get(2).closed); + } + @Test public void testInitializationFailureClosesOpenedIterators() { TestIterator first = new TestIterator(1); diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java index 44ba5b804b..a35838f328 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java @@ -35,12 +35,21 @@ import org.apache.hugegraph.store.HgKvIterator; import org.apache.hugegraph.store.HgOwnerKey; import org.apache.hugegraph.store.HgStoreSession; +import org.apache.hugegraph.store.client.grpc.KvPageScannerTestSupport; +import org.apache.hugegraph.store.grpc.common.Kv; import org.apache.hugegraph.store.grpc.common.ScanMethod; import org.apache.hugegraph.store.grpc.common.ScanOrderType; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder; +import org.apache.hugegraph.store.grpc.stream.ScanStreamVersion; import org.junit.Assert; import org.junit.Test; +import com.google.protobuf.ByteString; + +import io.grpc.stub.StreamObserver; + public class NodeTxSessionProxyTest { @Test @@ -150,6 +159,141 @@ public void testScanIteratorOrderedClosesOpenedIteratorsOnOpenFailure() } } + @Test + public void testOrderedScanMergesPagedStoresWithLimitAndCursor() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + String graph = "graph-paged-" + firstNodeId; + PagedStream firstStream = new PagedStream(Arrays.asList( + range(0, 128, 2), Collections.singletonList(128)), -1); + PagedStream secondStream = new PagedStream(Arrays.asList( + range(1, 129, 2), Collections.singletonList(129)), -1); + addPagedNode(manager, graph, firstNodeId, firstStream); + addPagedNode(manager, graph, secondNodeId, secondStream); + manager.setNodePartitioner(new RecordingPartitioner(firstNodeId, + secondNodeId)); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + HgKvIterator iterator = proxy.scanIteratorOrdered( + "table", HgOwnerKey.of(keyBytes(9), keyBytes(0)), + HgOwnerKey.of(keyBytes(9), keyBytes(130)), 129L, 123, + keyBytes(7)); + + List actual = keys(iterator); + Assert.assertEquals(range(0, 129, 1), actual); + Assert.assertArrayEquals(keyBytes(128), iterator.position()); + Assert.assertEquals(2, firstStream.dataRequests); + Assert.assertEquals(2, secondStream.dataRequests); + assertOrderedPageRequests(firstStream.requests, 129L); + assertOrderedPageRequests(secondStream.requests, 129L); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + + @Test + public void testOrderedScanHandlesEmptyStoreAndEarlyClose() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + String graph = "graph-empty-" + firstNodeId; + PagedStream emptyStream = new PagedStream( + Collections.singletonList(Collections.emptyList()), -1); + PagedStream activeStream = new PagedStream(Arrays.asList( + Arrays.asList(1, 2), Collections.singletonList(3)), -1); + addPagedNode(manager, graph, firstNodeId, emptyStream); + addPagedNode(manager, graph, secondNodeId, activeStream); + manager.setNodePartitioner(new RecordingPartitioner(firstNodeId, + secondNodeId)); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + HgKvIterator iterator = proxy.scanIteratorOrdered( + "table", HgOwnerKey.of(keyBytes(9), keyBytes(0)), + HgOwnerKey.of(keyBytes(9), keyBytes(4)), 0L, 123, + keyBytes(7)); + + Assert.assertEquals(1, key(iterator.next())); + iterator.close(); + + Assert.assertEquals(1, emptyStream.dataRequests); + Assert.assertTrue(emptyStream.clientCompleted); + Assert.assertEquals(1, activeStream.dataRequests); + Assert.assertEquals(1, activeStream.closeRequests); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + + @Test + public void testOrderedScanClosesOtherStoreOnPagedFailure() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + String graph = "graph-page-failure-" + firstNodeId; + PagedStream failedStream = new PagedStream( + Collections.singletonList(Collections.singletonList(0)), 1); + PagedStream activeStream = new PagedStream(Arrays.asList( + Collections.singletonList(1), + Collections.singletonList(3)), -1); + addPagedNode(manager, graph, firstNodeId, failedStream); + addPagedNode(manager, graph, secondNodeId, activeStream); + manager.setNodePartitioner(new RecordingPartitioner(firstNodeId, + secondNodeId)); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + HgKvIterator iterator = proxy.scanIteratorOrdered( + "table", HgOwnerKey.of(keyBytes(9), keyBytes(0)), + HgOwnerKey.of(keyBytes(9), keyBytes(4)), 0L, 123, + keyBytes(7)); + + Assert.assertThrows(RuntimeException.class, iterator::next); + Assert.assertTrue(failedStream.clientCompleted); + Assert.assertEquals(1, activeStream.closeRequests); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + + @Test + public void testOrderedScanRejectsMixedStoreVersionsAndClosesAll() + throws Exception { + HgStoreNodeManager manager = HgStoreNodeManager.getInstance(); + HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner(); + long firstNodeId = System.nanoTime(); + long secondNodeId = firstNodeId + 1L; + String graph = "graph-mixed-version-" + firstNodeId; + PagedStream currentStream = new PagedStream(Arrays.asList( + Collections.singletonList(1), + Collections.singletonList(3)), -1); + PagedStream oldStream = new PagedStream( + Collections.singletonList(Collections.singletonList(2)), + -1, 0); + addPagedNode(manager, graph, firstNodeId, currentStream); + addPagedNode(manager, graph, secondNodeId, oldStream); + manager.setNodePartitioner(new RecordingPartitioner(firstNodeId, + secondNodeId)); + try { + HgStoreSession proxy = new NodeTxSessionProxy(graph, manager); + HgKvIterator iterator = proxy.scanIteratorOrdered( + "table", HgOwnerKey.of(keyBytes(9), keyBytes(0)), + HgOwnerKey.of(keyBytes(9), keyBytes(4)), 0L, 123, + keyBytes(7)); + + Assert.assertThrows(RuntimeException.class, iterator::hasNext); + Assert.assertTrue(oldStream.clientCompleted); + Assert.assertEquals(1, currentStream.closeRequests); + } finally { + restoreNodePartitioner(manager, oldPartitioner); + } + } + private static void assertOrderedRangeBuilder(Builder builder, long limit, int scanType, byte[] query) { @@ -166,6 +310,25 @@ private static void assertOrderedRangeBuilder(Builder builder, long limit, Assert.assertArrayEquals(query, builder.getQuery().toByteArray()); } + private static void assertOrderedPageRequests( + List requests, long limit) { + for (ScanStreamReq request : requests) { + Assert.assertEquals(ScanMethod.RANGE, request.getMethod()); + Assert.assertEquals(ScanOrderType.ORDER_BY_KEY, + request.getOrderType()); + Assert.assertEquals(64, request.getPageSize()); + Assert.assertEquals(limit, request.getLimit()); + } + } + + private static void addPagedNode(HgStoreNodeManager manager, String graph, + long nodeId, PagedStream stream) { + PagedRecordingSession handler = new PagedRecordingSession( + graph, nodeId, stream); + HgStoreNodeSession session = handler.proxy(); + manager.addNode(graph, new RecordingStoreNode(nodeId, session)); + } + private static void restoreNodePartitioner(HgStoreNodeManager manager, HgStoreNodePartitioner old) throws Exception { @@ -179,6 +342,22 @@ private static int key(HgKvEntry entry) { return entry.key()[0] & 0xff; } + private static List keys(HgKvIterator iterator) { + List keys = new ArrayList<>(); + while (iterator.hasNext()) { + keys.add(key(iterator.next())); + } + return keys; + } + + private static List range(int start, int end, int step) { + List keys = new ArrayList<>(); + for (int key = start; key < end; key += step) { + keys.add(key); + } + return keys; + } + private static byte[] keyBytes(int key) { return new byte[]{(byte) key}; } @@ -395,4 +574,125 @@ public Object invoke(Object proxy, Method method, Object[] args) { throw new UnsupportedOperationException(method.toString()); } } + + private static final class PagedRecordingSession + implements InvocationHandler { + + private final String graph; + private final long nodeId; + private final PagedStream stream; + + private PagedRecordingSession(String graph, long nodeId, + PagedStream stream) { + this.graph = graph; + this.nodeId = nodeId; + this.stream = stream; + } + + private HgStoreNodeSession proxy() { + return (HgStoreNodeSession) Proxy.newProxyInstance( + HgStoreNodeSession.class.getClassLoader(), + new Class[]{HgStoreNodeSession.class}, this); + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + if ("scanIterator".equals(method.getName()) && + args != null && args.length == 1 && + args[0] instanceof Builder) { + return KvPageScannerTestSupport.iterator( + (HgStoreNodeSession) proxy, + ((Builder) args[0]).clone(), this.stream::open); + } + if ("getGraphName".equals(method.getName())) { + return this.graph; + } + if ("getStoreNode".equals(method.getName())) { + return new RecordingStoreNode( + this.nodeId, (HgStoreSession) proxy); + } + if ("isTx".equals(method.getName())) { + return false; + } + if ("toString".equals(method.getName())) { + return "PagedRecordingSession"; + } + throw new UnsupportedOperationException(method.toString()); + } + } + + private static final class PagedStream { + + private static final int ORDERED_SCAN_VERSION = + ScanStreamVersion.SCAN_STREAM_VERSION_ORDERED_BY_KEY_VALUE; + + private final List> pages; + private final int failPage; + private final int responseVersion; + private final List requests; + private int page; + private int dataRequests; + private int closeRequests; + private boolean clientCompleted; + + private PagedStream(List> pages, int failPage) { + this(pages, failPage, ORDERED_SCAN_VERSION); + } + + private PagedStream(List> pages, int failPage, + int responseVersion) { + this.pages = pages; + this.failPage = failPage; + this.responseVersion = responseVersion; + this.requests = new ArrayList<>(); + this.page = 0; + this.dataRequests = 0; + this.closeRequests = 0; + this.clientCompleted = false; + } + + private StreamObserver open( + StreamObserver response) { + return new StreamObserver() { + @Override + public void onNext(ScanStreamReq request) { + if (request.getCloseFlag() != 0) { + closeRequests++; + response.onCompleted(); + return; + } + requests.add(request); + dataRequests++; + if (page == failPage) { + response.onError(new IllegalStateException( + "injected page failure")); + return; + } + List keys = pages.get(page++); + KvPageRes.Builder result = KvPageRes.newBuilder() + .setVersion(responseVersion) + .setOver(page == + pages.size() && + failPage < 0); + for (int key : keys) { + ByteString bytes = ByteString.copyFrom(keyBytes(key)); + result.addData(Kv.newBuilder() + .setKey(bytes) + .setValue(bytes)); + } + response.onNext(result.build()); + } + + @Override + public void onError(Throwable throwable) { + clientCompleted = true; + } + + @Override + public void onCompleted() { + clientCompleted = true; + } + }; + } + } } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java index 104ee4530b..ce34696aea 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java @@ -19,12 +19,16 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.NoSuchElementException; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; @@ -134,6 +138,92 @@ public void testMergeClosesAllSourcesWhenConcurrentInitializeFails() { } } + @Test + public void testConcurrentInitializeFailsWithoutWaitingForSlowSource() + throws Exception { + ExecutorService initializer = Executors.newFixedThreadPool(2); + ExecutorService caller = Executors.newSingleThreadExecutor(); + CountDownLatch slowStarted = new CountDownLatch(1); + CountDownLatch releaseSlow = new CountDownLatch(1); + TestIterator slow = new TestIterator(1); + TestIterator failed = new TestIterator(2); + slow.blockFirstHasNext(slowStarted, releaseSlow); + failed.failOnHasNextAfter(0); + + Future result = caller.submit(() -> { + OrderedKvIterator iterator = new OrderedKvIterator( + Arrays.asList(slow, failed), 0L, initializer); + return iterator.hasNext(); + }); + try { + Assert.assertTrue(slowStarted.await(3L, TimeUnit.SECONDS)); + try { + result.get(3L, TimeUnit.SECONDS); + Assert.fail("Expected initialization failure"); + } catch (ExecutionException e) { + Assert.assertTrue(e.getCause() instanceof + IllegalStateException); + } + Assert.assertTrue(slow.closed); + Assert.assertTrue(failed.closed); + } finally { + releaseSlow.countDown(); + result.cancel(true); + caller.shutdownNow(); + initializer.shutdownNow(); + } + } + + @Test + public void testSlowQueriesDoNotBlockIndependentQuery() throws Exception { + int concurrentQueries = 8; + ExecutorService callers = Executors.newFixedThreadPool( + concurrentQueries + 1); + CountDownLatch slowQueriesStarted = new CountDownLatch( + concurrentQueries); + CountDownLatch releaseSlowQueries = new CountDownLatch(1); + List> slowResults = new ArrayList<>(); + try { + for (int i = 0; i < concurrentQueries; i++) { + TestIterator slow = new TestIterator(i); + slow.blockFirstHasNext(slowQueriesStarted, + releaseSlowQueries); + slowResults.add(callers.submit(() -> { + OrderedKvIterator iterator = new OrderedKvIterator( + Collections.singletonList(slow), 0L); + try { + return iterator.hasNext(); + } finally { + iterator.close(); + } + })); + } + Assert.assertTrue(slowQueriesStarted.await(3L, + TimeUnit.SECONDS)); + + Future fastResult = callers.submit(() -> { + OrderedKvIterator iterator = new OrderedKvIterator( + Collections.singletonList(new TestIterator(100)), + 0L); + try { + return iterator.hasNext(); + } finally { + iterator.close(); + } + }); + Assert.assertTrue(fastResult.get(3L, TimeUnit.SECONDS)); + } catch (TimeoutException e) { + Assert.fail("An independent ordered scan was starved by slow " + + "queries"); + } finally { + releaseSlowQueries.countDown(); + for (Future result : slowResults) { + result.cancel(true); + } + callers.shutdownNow(); + } + } + private static List keys(HgKvIterator iterator) { List keys = new ArrayList<>(); while (iterator.hasNext()) { @@ -155,6 +245,8 @@ private static final class TestIterator implements HgKvIterator { private boolean closed; private int failOnHasNextAfter; private CountDownLatch firstHasNextBarrier; + private CountDownLatch firstHasNextStarted; + private CountDownLatch firstHasNextRelease; private boolean firstHasNextBlocked; private TestIterator(Integer... keys) { @@ -168,6 +260,8 @@ private TestIterator(Integer... keys) { this.closed = false; this.failOnHasNextAfter = -1; this.firstHasNextBarrier = null; + this.firstHasNextStarted = null; + this.firstHasNextRelease = null; this.firstHasNextBlocked = false; } @@ -179,6 +273,12 @@ private void blockFirstHasNext(CountDownLatch barrier) { this.firstHasNextBarrier = barrier; } + private void blockFirstHasNext(CountDownLatch started, + CountDownLatch release) { + this.firstHasNextStarted = started; + this.firstHasNextRelease = release; + } + @Override public boolean hasNext() { if (this.nextCalls == this.failOnHasNextAfter) { @@ -199,6 +299,21 @@ public boolean hasNext() { throw new IllegalStateException(e); } } + if (this.firstHasNextRelease != null && + !this.firstHasNextBlocked) { + this.firstHasNextBlocked = true; + this.firstHasNextStarted.countDown(); + try { + if (!this.firstHasNextRelease.await(5L, + TimeUnit.SECONDS)) { + throw new IllegalStateException( + "Timed out waiting for source release"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } return this.offset < this.entries.size(); } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java index 89fc8f0662..c2c70e2d50 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTest.java @@ -25,6 +25,7 @@ import org.apache.hugegraph.store.client.HgStoreNodeSession; import org.apache.hugegraph.store.grpc.common.Kv; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.KvPageRes; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.junit.Assert; @@ -66,6 +67,41 @@ public void testKeepsExplicitPageSizeAndRequestsNextPageOnDemand() { Assert.assertEquals(2, requests.size()); } + @Test + public void testRejectsOrderedScanWithoutCapabilityAck() { + ScanStreamReq.Builder builder = ScanStreamReq.newBuilder() + .setMethod(ScanMethod.RANGE) + .setTable("table") + .setPageSize(2) + .setOrderType( + ScanOrderType.ORDER_BY_KEY); + KvPageScanner scanner = new KvPageScanner( + session(), builder, response -> onePageObserver(response, 0)); + + try { + scanner.hasNext(); + Assert.fail("Expected ordered scan capability failure"); + } catch (RuntimeException e) { + Assert.assertTrue(e.getMessage().contains("ordered scan")); + } + } + + @Test + public void testAcceptsOrderedScanWithCapabilityAck() { + ScanStreamReq.Builder builder = ScanStreamReq.newBuilder() + .setMethod(ScanMethod.RANGE) + .setTable("table") + .setPageSize(2) + .setOrderType( + ScanOrderType.ORDER_BY_KEY); + KvPageScanner scanner = new KvPageScanner( + session(), builder, response -> onePageObserver(response, 1)); + + Assert.assertTrue(scanner.hasNext()); + Assert.assertEquals(1, scanner.next().getKey().byteAt(0) & 0xff); + Assert.assertFalse(scanner.hasNext()); + } + private static HgStoreNodeSession session() { return (HgStoreNodeSession) Proxy.newProxyInstance( HgStoreNodeSession.class.getClassLoader(), @@ -114,6 +150,27 @@ public void onCompleted() { }; } + private static StreamObserver onePageObserver( + StreamObserver response, int version) { + return new StreamObserver() { + @Override + public void onNext(ScanStreamReq request) { + if (request.getCloseFlag() == 0) { + response.onNext(versionedPage(true, version, 1)); + } + } + + @Override + public void onError(Throwable throwable) { + // The client cancels the incompatible ordered stream + } + + @Override + public void onCompleted() { + } + }; + } + private static KvPageRes page(boolean over, int... keys) { KvPageRes.Builder page = KvPageRes.newBuilder().setOver(over); for (int key : keys) { @@ -122,4 +179,9 @@ private static KvPageRes page(boolean over, int... keys) { } return page.build(); } + + private static KvPageRes versionedPage(boolean over, int version, + int... keys) { + return page(over, keys).toBuilder().setVersion(version).build(); + } } diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTestSupport.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTestSupport.java new file mode 100644 index 0000000000..8255a6b533 --- /dev/null +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/grpc/KvPageScannerTestSupport.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.store.client.grpc; + +import java.util.function.Function; + +import org.apache.hugegraph.store.HgKvEntry; +import org.apache.hugegraph.store.HgKvIterator; +import org.apache.hugegraph.store.client.HgStoreNodeSession; +import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; + +import io.grpc.stub.StreamObserver; + +public final class KvPageScannerTestSupport { + + private KvPageScannerTestSupport() { + } + + public static HgKvIterator iterator( + HgStoreNodeSession session, ScanStreamReq.Builder builder, + Function, + StreamObserver> streamFactory) { + KvPageScanner scanner = new KvPageScanner(session, builder, + streamFactory); + return GrpcKvIteratorImpl.of(session, scanner); + } +} From b426dc701a276cb0d7991db330411691df74f6fd Mon Sep 17 00:00:00 2001 From: contrueCT Date: Fri, 7 Aug 2026 16:50:55 +0800 Subject: [PATCH 5/6] fix(hstore): harden ordered range scans --- .../apache/hugegraph/core/VertexCoreTest.java | 81 +++++++++++++++++++ .../apache/hugegraph/store/HgScanQuery.java | 2 + .../store/client/OrderedKvIterator.java | 38 ++------- .../store/node/grpc/ScanQueryProducer.java | 3 + .../store/node/grpc/ScanUtilTest.java | 17 ++++ .../store/client/NodeTxSessionProxyTest.java | 19 +++++ .../store/client/OrderedKvIteratorTest.java | 17 +++- 7 files changed, 144 insertions(+), 33 deletions(-) diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/VertexCoreTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/VertexCoreTest.java index 6118deee2b..9aa144542e 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/VertexCoreTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/VertexCoreTest.java @@ -8477,6 +8477,54 @@ public void testQueryByPropertyInPageWithLimitLtePageSize() { }); } + @Test + public void testQueryByRangeIndexKeepsOrderAcrossStorePages() { + Assume.assumeTrue("Not support paging", + storeFeatures().supportsQueryByPage()); + + initRangeIndexOrderTestData(); + + GraphTraversalSource g = graph().traversal(); + List vertices = g.V().hasLabel("ranked") + .has("rank", P.between(0, 130)) + .limit(70) + .toList(); + assertRanks(vertices, 0, 70); + + GraphTraversal firstPage = + g.V().hasLabel("ranked") + .has("rank", P.between(0, 130)) + .has("~page", "") + .limit(70); + vertices = firstPage.toList(); + assertRanks(vertices, 0, 70); + + String page = TraversalUtil.page(firstPage); + Assert.assertNotNull(page); + Assert.assertFalse(page.isEmpty()); + + vertices = g.V().hasLabel("ranked") + .has("rank", P.between(0, 130)) + .has("~page", page) + .limit(70) + .toList(); + assertRanks(vertices, 70, 60); + } + + @Test + public void testQueryByRangeIndexKeepsOffsetOrderInHstore() { + Assume.assumeTrue("Only run for hstore", + Objects.equals("hstore", graph().backend())); + + initRangeIndexOrderTestData(); + + List vertices = graph().traversal().V().hasLabel("ranked") + .has("rank", P.between(0, 130)) + .range(65, 75) + .toList(); + assertRanks(vertices, 65, 10); + } + @Test public void testQueryByPropertyInPageWithLimitGtPageSize() { // FIXME: The legacy HStore guard and related coverage debt are tracked in @@ -9504,6 +9552,26 @@ private void initPageTestData() { this.commitTx(); } + private void initRangeIndexOrderTestData() { + SchemaManager schema = graph().schema(); + schema.propertyKey("rank").asInt().create(); + schema.vertexLabel("ranked") + .properties("rank") + .useCustomizeStringId() + .create(); + schema.indexLabel("rankedByRank") + .onV("ranked") + .by("rank") + .range() + .create(); + + for (int rank = 129; rank >= 0; rank--) { + graph().addVertex(T.label, "ranked", T.id, "ranked-" + rank, + "rank", rank); + } + this.commitTx(); + } + private Vertex vertex(String label, String pkName, Object pkValue) { List vertices = graph().traversal().V() .hasLabel(label).has(pkName, pkValue) @@ -9512,6 +9580,19 @@ private Vertex vertex(String label, String pkName, Object pkValue) { return vertices.size() == 1 ? vertices.get(0) : null; } + private static void assertRanks(List vertices, int firstRank, + int expectedSize) { + List actualRanks = new ArrayList<>(vertices.size()); + for (Vertex vertex : vertices) { + actualRanks.add(vertex.value("rank")); + } + Assert.assertEquals(expectedSize, vertices.size()); + for (int i = 0; i < expectedSize; i++) { + Assert.assertEquals("Unexpected ranks: " + actualRanks, + firstRank + i, (int) actualRanks.get(i)); + } + } + private static void assertContains(List vertices, Object... keyValues) { Assert.assertTrue(Utils.contains(vertices, new FakeObjects.FakeVertex(keyValues))); diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgScanQuery.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgScanQuery.java index cc64ba945b..ed0f527368 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgScanQuery.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/HgScanQuery.java @@ -185,6 +185,8 @@ public ScanBuilder setScanType(int scanType) { } public ScanBuilder setOrderType(ScanOrderType orderType) { + HgAssert.isFalse(orderType == ScanOrderType.ORDER_BY_KEY, + "ORDER_BY_KEY is not supported by batch scan"); this.orderType = orderType; return this; } diff --git a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java index b234706caa..9008e0ea25 100644 --- a/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java +++ b/hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/OrderedKvIterator.java @@ -28,7 +28,6 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorCompletionService; -import java.util.concurrent.Executors; import java.util.concurrent.Future; import org.apache.hugegraph.store.HgKvEntry; @@ -39,13 +38,17 @@ final class OrderedKvIterator implements HgKvIterator { private static final int INITIALIZE_THREADS = 8; + private static final long INITIALIZE_KEEP_ALIVE_SECONDS = 60L; + private static final ExecutorService INITIALIZER = + ExecutorPool.createExecutor("ordered-scan-init", + INITIALIZE_KEEP_ALIVE_SECONDS, + 0, INITIALIZE_THREADS); private final List> iterators; private final PriorityQueue queue; private final boolean[] sourceClosed; private final long limit; private final ExecutorService initializer; - private final boolean ownsInitializer; private boolean initialized; private boolean closed; @@ -55,19 +58,12 @@ final class OrderedKvIterator implements HgKvIterator { OrderedKvIterator(List> iterators, long limit) { - this(iterators, limit, createInitializer(iterators), true); + this(iterators, limit, INITIALIZER); } OrderedKvIterator(List> iterators, long limit, ExecutorService initializer) { - this(iterators, limit, initializer, false); - } - - private OrderedKvIterator( - List> iterators, - long limit, ExecutorService initializer, - boolean ownsInitializer) { - this.iterators = iterators; + this.iterators = Objects.requireNonNull(iterators); this.queue = new PriorityQueue<>((left, right) -> { int result = Arrays.compareUnsigned(left.entry.key(), right.entry.key()); @@ -80,7 +76,6 @@ private OrderedKvIterator( this.limit = limit <= HgStoreClientConst.NO_LIMIT ? Long.MAX_VALUE : limit; this.initializer = Objects.requireNonNull(initializer); - this.ownsInitializer = ownsInitializer; this.initialized = false; this.closed = false; this.count = 0L; @@ -159,7 +154,6 @@ public void close() { } } this.queue.clear(); - this.shutdownInitializer(); if (failure instanceof RuntimeException) { throw (RuntimeException) failure; } @@ -208,8 +202,6 @@ private void initialize() { this.cancel(futures); this.closeAfterFailure(e); throw e; - } finally { - this.shutdownInitializer(); } } @@ -222,16 +214,6 @@ private SourceEntry firstEntry(int source) { return new SourceEntry(source, iterator.next()); } - private static ExecutorService createInitializer( - List> iterators) { - Objects.requireNonNull(iterators); - int threads = Math.max(1, Math.min(INITIALIZE_THREADS, - iterators.size())); - return Executors.newFixedThreadPool( - threads, - ExecutorPool.newThreadFactory("ordered-scan-init")); - } - private void cancel(List> futures) { for (Future future : futures) { if (!future.isDone()) { @@ -240,12 +222,6 @@ private void cancel(List> futures) { } } - private void shutdownInitializer() { - if (this.ownsInitializer) { - this.initializer.shutdownNow(); - } - } - private void addNext(int source) { HgKvIterator iterator = this.iterators.get(source); diff --git a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanQueryProducer.java b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanQueryProducer.java index 204c32d9c0..d25843539e 100644 --- a/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanQueryProducer.java +++ b/hugegraph-store/hg-store-node/src/main/java/org/apache/hugegraph/store/node/grpc/ScanQueryProducer.java @@ -25,6 +25,7 @@ import javax.annotation.concurrent.NotThreadSafe; import org.apache.hugegraph.store.grpc.common.ScanMethod; +import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.ScanCondition; import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; import org.apache.hugegraph.store.node.util.HgAssert; @@ -58,6 +59,8 @@ public static ScanQueryProducer requestOf(String graph, String[] tables, HgAssert.isArgumentValid(graph, "graph"); HgAssert.isArgumentNotNull(tables, "tables"); HgAssert.isArgumentNotNull(request, "ScanQueryRequest"); + HgAssert.isFalse(request.getOrderType() == ScanOrderType.ORDER_BY_KEY, + "ORDER_BY_KEY is not supported by batch scan"); ScanQueryProducer res = new ScanQueryProducer(); res.graph = graph; diff --git a/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java b/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java index 9764d7264d..6f2df672c4 100644 --- a/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java +++ b/hugegraph-store/hg-store-node/src/test/java/org/apache/hugegraph/store/node/grpc/ScanUtilTest.java @@ -28,6 +28,7 @@ import org.apache.hugegraph.store.grpc.common.ScanMethod; import org.apache.hugegraph.store.grpc.common.ScanOrderType; import org.apache.hugegraph.store.grpc.stream.KvPageRes; +import org.apache.hugegraph.store.grpc.stream.ScanQueryRequest; import org.apache.hugegraph.store.grpc.stream.ScanStreamReq; import org.apache.hugegraph.store.node.AppConfig; import org.junit.Assert; @@ -37,6 +38,22 @@ public class ScanUtilTest { + @Test + public void testBatchScanRejectsOrderByKey() { + ScanQueryRequest request = ScanQueryRequest.newBuilder() + .setMethod(ScanMethod.RANGE) + .setTable("table") + .setOrderType(ScanOrderType.ORDER_BY_KEY) + .build(); + + IllegalArgumentException error = Assert.assertThrows( + IllegalArgumentException.class, + () -> ScanQueryProducer.requestOf("graph", + new String[]{"table"}, + request)); + Assert.assertTrue(error.getMessage().contains("batch scan")); + } + @Test public void testOrderedAllPartitionRangeUsesOrderedScan() { RecordingWrapper wrapper = new RecordingWrapper(); diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java index a35838f328..3fb1a907b4 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java @@ -34,6 +34,7 @@ import org.apache.hugegraph.store.HgKvEntry; import org.apache.hugegraph.store.HgKvIterator; import org.apache.hugegraph.store.HgOwnerKey; +import org.apache.hugegraph.store.HgScanQuery; import org.apache.hugegraph.store.HgStoreSession; import org.apache.hugegraph.store.client.grpc.KvPageScannerTestSupport; import org.apache.hugegraph.store.grpc.common.Kv; @@ -52,6 +53,24 @@ public class NodeTxSessionProxyTest { + @Test + public void testBatchQueryRejectsOrderByKey() { + HgOwnerKey start = HgOwnerKey.of(keyBytes(9), keyBytes(1)); + HgOwnerKey end = HgOwnerKey.of(keyBytes(9), keyBytes(5)); + HgScanQuery.ScanBuilder builder = HgScanQuery.ScanBuilder.rangeOf( + "table", Collections.singletonList(start), + Collections.singletonList(end)); + + IllegalArgumentException error = Assert.assertThrows( + IllegalArgumentException.class, + () -> builder.setOrderType(ScanOrderType.ORDER_BY_KEY)); + Assert.assertTrue(error.getMessage().contains("batch scan")); + + HgScanQuery query = builder.setOrderType(ScanOrderType.ORDER_STRICT) + .build(); + Assert.assertEquals(ScanOrderType.ORDER_STRICT, query.getOrderType()); + } + @Test public void testNodeTkvDoesNotMutateSharedOwnerKeys() { HgOwnerKey start = HgOwnerKey.of(keyBytes(9), keyBytes(1)); diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java index ce34696aea..34ec7bbaa3 100644 --- a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java +++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/client/OrderedKvIteratorTest.java @@ -175,8 +175,9 @@ public void testConcurrentInitializeFailsWithoutWaitingForSlowSource() } @Test - public void testSlowQueriesDoNotBlockIndependentQuery() throws Exception { - int concurrentQueries = 8; + public void testSlowQueriesUseBoundedWorkersAndDoNotBlockIndependentQuery() + throws Exception { + int concurrentQueries = 16; ExecutorService callers = Executors.newFixedThreadPool( concurrentQueries + 1); CountDownLatch slowQueriesStarted = new CountDownLatch( @@ -200,6 +201,10 @@ public void testSlowQueriesDoNotBlockIndependentQuery() throws Exception { } Assert.assertTrue(slowQueriesStarted.await(3L, TimeUnit.SECONDS)); + long workers = initializerThreads(); + Assert.assertTrue("Expected at most 8 ordered scan initializer " + + "threads, but found " + workers, + workers <= 8L); Future fastResult = callers.submit(() -> { OrderedKvIterator iterator = new OrderedKvIterator( @@ -224,6 +229,14 @@ public void testSlowQueriesDoNotBlockIndependentQuery() throws Exception { } } + private static long initializerThreads() { + return Thread.getAllStackTraces().keySet().stream() + .filter(Thread::isAlive) + .filter(thread -> thread.getName().startsWith( + "ordered-scan-init-")) + .count(); + } + private static List keys(HgKvIterator iterator) { List keys = new ArrayList<>(); while (iterator.hasNext()) { From ea9cce81584b61070974ab4a1eddf54050da782d Mon Sep 17 00:00:00 2001 From: contrueCT Date: Sat, 8 Aug 2026 01:42:33 +0800 Subject: [PATCH 6/6] fix(core): count offset across index batches --- .../apache/hugegraph/backend/query/Query.java | 10 ++++------ .../apache/hugegraph/unit/core/QueryTest.java | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/Query.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/Query.java index 1816bff237..b8fc8e2964 100644 --- a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/Query.java +++ b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/query/Query.java @@ -251,13 +251,11 @@ public Set skipOffsetIfNeeded(Set elems) { if (fromIndex < 0L) { // Skipping offset is overhead, no need to skip fromIndex = 0L; - } else if (fromIndex > 0L) { - this.goOffset(fromIndex); } - if (fromIndex > Integer.MAX_VALUE) { - E.checkArgument(false, - "Offset must be <= 0x7fffffff, but got '%s'", - fromIndex); + // An index holder yields ids in batches, only count this batch's ids + fromIndex = Math.min(fromIndex, elems.size()); + if (fromIndex > 0L) { + this.goOffset(fromIndex); } if (fromIndex >= elems.size()) { diff --git a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryTest.java b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryTest.java index 7d48084dbf..5778ceba7c 100644 --- a/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryTest.java +++ b/hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/unit/core/QueryTest.java @@ -95,6 +95,22 @@ public void testConditionWithMultipleMatchedInValues() { e.getMessage())); } + @Test + public void testSkipOffsetAcrossBatches() { + Query query = new Query(HugeType.VERTEX); + query.offset(6L); + query.limit(2L); + + Assert.assertTrue(query.skipOffsetIfNeeded( + ImmutableSet.of(0, 1, 2, 3)).isEmpty()); + Assert.assertEquals(4L, query.actualOffset()); + + Assert.assertEquals(ImmutableSet.of(6, 7), + query.skipOffsetIfNeeded( + ImmutableSet.of(4, 5, 6, 7))); + Assert.assertEquals(6L, query.actualOffset()); + } + @Test public void testToString() { Query query = new Query(HugeType.VERTEX);