Skip to content

Commit 9eb75e8

Browse files
contrueCTcontrueCT
committed
fix(hstore): stream ordered range scans lazily
1 parent 3646a14 commit 9eb75e8

2 files changed

Lines changed: 201 additions & 24 deletions

File tree

‎hugegraph-store/hg-store-client/src/main/java/org/apache/hugegraph/store/client/NodeTxSessionProxy.java‎

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -54,12 +54,17 @@
5454
import org.apache.hugegraph.store.client.util.HgAssert;
5555
import org.apache.hugegraph.store.client.util.HgStoreClientConst;
5656
import org.apache.hugegraph.store.client.util.HgStoreClientUtil;
57+
import org.apache.hugegraph.store.grpc.common.Header;
58+
import org.apache.hugegraph.store.grpc.common.ScanMethod;
59+
import org.apache.hugegraph.store.grpc.stream.ScanStreamReq;
5760
import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder;
5861
import org.apache.hugegraph.store.query.StoreQueryParam;
5962
import org.apache.hugegraph.store.term.HgPair;
6063
import org.apache.hugegraph.store.term.HgTriple;
6164
import org.apache.hugegraph.structure.BaseElement;
6265

66+
import com.google.protobuf.ByteString;
67+
6368
import lombok.extern.slf4j.Slf4j;
6469

6570
/**
@@ -497,11 +502,11 @@ public HgKvIterator<HgKvEntry> scanIteratorOrdered(String table, HgOwnerKey star
497502
.parallelStream()
498503
.map(e -> this.getStoreNode(e.getNodeId())
499504
.openSession(this.graphName)
500-
.scanIterator(e.getTable(), e.getKey(),
501-
e.getEndKey(), limit,
502-
scanType, query))
505+
.scanIterator(this.orderedRangeScanBuilder(e,
506+
limit,
507+
scanType,
508+
query)))
503509
.collect(Collectors.toList());
504-
prefetchOrderedRangeScanIterators(iterators);
505510
return this.toOrderedHgKvIteratorProxy(iterators, limit);
506511
}
507512

@@ -677,16 +682,6 @@ static HgKvIterator<HgKvEntry> mergeOrderedRangeScanIterators(
677682
return new OrderedKvIterator(iterators, limit);
678683
}
679684

680-
static void prefetchOrderedRangeScanIterators(
681-
List<? extends HgKvIterator> iteratorList) {
682-
try {
683-
iteratorList.parallelStream().forEach(iterator -> iterator.hasNext());
684-
} catch (RuntimeException | Error e) {
685-
iteratorList.forEach(HgKvIterator::close);
686-
throw e;
687-
}
688-
}
689-
690685
private HgKvIterator toHgKvIteratorProxy(List<HgKvIterator> iteratorList, long limit) {
691686
return mergeRangeScanIterators(iteratorList, limit);
692687
}
@@ -806,6 +801,29 @@ private List<NodeTkv> toOrderedRangeNodeTkvList(String table,
806801
return this.toNodeTkvList(table, startKey, endKey);
807802
}
808803

804+
private Builder orderedRangeScanBuilder(NodeTkv nodeTkv, long limit,
805+
int scanType, byte[] query) {
806+
long scanLimit = limit <= HgStoreClientConst.NO_LIMIT ?
807+
Integer.MAX_VALUE : limit;
808+
return ScanStreamReq.newBuilder()
809+
.setHeader(Header.newBuilder()
810+
.setGraph(this.graphName)
811+
.build())
812+
.setMethod(ScanMethod.RANGE)
813+
.setTable(nodeTkv.getTable())
814+
.setStart(toByteString(nodeTkv.getKey().getKey()))
815+
.setEnd(toByteString(nodeTkv.getEndKey().getKey()))
816+
.setLimit(scanLimit)
817+
.setCode(nodeTkv.getKey().getKeyCode())
818+
.setScanType(scanType)
819+
.setQuery(toByteString(query));
820+
}
821+
822+
private static ByteString toByteString(byte[] bytes) {
823+
return ByteString.copyFrom(bytes != null ? bytes :
824+
HgStoreClientConst.EMPTY_BYTES);
825+
}
826+
809827
private List<NodeTkv> toNodeTkvList(String table, int startCode, int endCode) {
810828
Collection<HgNodePartition> partitions = this.doPartition(table, startCode, endCode);
811829
ArrayList<NodeTkv> nodeTkvs = new ArrayList<>(partitions.size());

‎hugegraph-store/hg-store-client/src/test/java/org/apache/hugegraph/store/client/NodeTxSessionProxyTest.java‎

Lines changed: 169 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,12 @@
1818
package org.apache.hugegraph.store.client;
1919

2020
import java.lang.reflect.Field;
21+
import java.lang.reflect.InvocationHandler;
2122
import java.lang.reflect.Method;
23+
import java.lang.reflect.Proxy;
24+
import java.util.ArrayList;
2225
import java.util.Arrays;
26+
import java.util.Collections;
2327
import java.util.LinkedHashSet;
2428
import java.util.List;
2529
import java.util.NoSuchElementException;
@@ -28,6 +32,9 @@
2832
import org.apache.hugegraph.store.HgKvEntry;
2933
import org.apache.hugegraph.store.HgKvIterator;
3034
import org.apache.hugegraph.store.HgOwnerKey;
35+
import org.apache.hugegraph.store.HgStoreSession;
36+
import org.apache.hugegraph.store.grpc.common.ScanMethod;
37+
import org.apache.hugegraph.store.grpc.stream.ScanStreamReq.Builder;
3138
import org.junit.Assert;
3239
import org.junit.Test;
3340

@@ -78,17 +85,92 @@ public void testMergeOrderedRangeScanIteratorsSortsByKey() {
7885
}
7986

8087
@Test
81-
public void testPrefetchOrderedRangeScanIteratorsDoesNotConsumeEntries() {
88+
public void testMergeOrderedRangeScanIteratorsIsLazy() {
8289
TestIterator first = new TestIterator(3, 4);
8390
TestIterator second = new TestIterator(1, 2);
84-
NodeTxSessionProxy.prefetchOrderedRangeScanIterators(
85-
Arrays.asList(first, second));
8691

87-
Assert.assertTrue(first.hasNextCalls > 0);
88-
Assert.assertTrue(second.hasNextCalls > 0);
92+
HgKvIterator<HgKvEntry> iterator =
93+
NodeTxSessionProxy.mergeOrderedRangeScanIterators(
94+
Arrays.asList(first, second), 0L);
95+
96+
Assert.assertEquals(0, first.hasNextCalls);
97+
Assert.assertEquals(0, second.hasNextCalls);
8998
Assert.assertEquals(0, first.nextCalls);
9099
Assert.assertEquals(0, second.nextCalls);
91100

101+
Assert.assertEquals(1, key(iterator.next()));
102+
Assert.assertEquals(2, key(iterator.next()));
103+
Assert.assertEquals(3, key(iterator.next()));
104+
Assert.assertEquals(4, key(iterator.next()));
105+
Assert.assertFalse(iterator.hasNext());
106+
}
107+
108+
@Test
109+
public void testScanIteratorOrderedUsesStreamingBuildersLazily()
110+
throws Exception {
111+
HgStoreNodeManager manager = HgStoreNodeManager.getInstance();
112+
HgStoreNodePartitioner oldPartitioner = manager.getNodePartitioner();
113+
long firstNodeId = System.nanoTime();
114+
long secondNodeId = firstNodeId + 1L;
115+
RecordingPartitioner partitioner =
116+
new RecordingPartitioner(firstNodeId, secondNodeId);
117+
TestIterator firstIterator = new TestIterator(3, 4);
118+
TestIterator secondIterator = new TestIterator(1, 2);
119+
RecordingSession firstSession =
120+
new RecordingSession(firstIterator);
121+
RecordingSession secondSession =
122+
new RecordingSession(secondIterator);
123+
String graph = "graph-" + firstNodeId;
124+
manager.addNode(graph, new RecordingStoreNode(firstNodeId,
125+
firstSession.proxy()));
126+
manager.addNode(graph, new RecordingStoreNode(secondNodeId,
127+
secondSession.proxy()));
128+
manager.setNodePartitioner(partitioner);
129+
try {
130+
NodeTxSessionProxy proxy = new NodeTxSessionProxy(graph, manager);
131+
HgKvIterator<HgKvEntry> iterator = proxy.scanIteratorOrdered(
132+
"table", HgOwnerKey.of(keyBytes(9), keyBytes(1)),
133+
HgOwnerKey.of(keyBytes(9), keyBytes(5)), 5L, 123,
134+
keyBytes(7));
135+
136+
Assert.assertEquals(1, firstSession.builders.size());
137+
Assert.assertEquals(1, secondSession.builders.size());
138+
Assert.assertEquals(0, firstSession.rangeScanCalls);
139+
Assert.assertEquals(0, secondSession.rangeScanCalls);
140+
Assert.assertEquals(0, firstIterator.hasNextCalls);
141+
Assert.assertEquals(0, secondIterator.hasNextCalls);
142+
assertOrderedRangeBuilder(firstSession.builders.get(0), 5L, 123,
143+
keyBytes(7));
144+
assertOrderedRangeBuilder(secondSession.builders.get(0), 5L, 123,
145+
keyBytes(7));
146+
147+
Assert.assertEquals(1, key(iterator.next()));
148+
Assert.assertEquals(2, key(iterator.next()));
149+
Assert.assertEquals(3, key(iterator.next()));
150+
Assert.assertEquals(4, key(iterator.next()));
151+
Assert.assertFalse(iterator.hasNext());
152+
} finally {
153+
restoreNodePartitioner(manager, oldPartitioner);
154+
}
155+
}
156+
157+
private static void assertOrderedRangeBuilder(Builder builder, long limit,
158+
int scanType,
159+
byte[] query) {
160+
Assert.assertEquals(ScanMethod.RANGE, builder.getMethod());
161+
Assert.assertEquals("table", builder.getTable());
162+
Assert.assertEquals(limit, builder.getLimit());
163+
Assert.assertEquals(scanType, builder.getScanType());
164+
Assert.assertArrayEquals(query, builder.getQuery().toByteArray());
165+
}
166+
167+
@Test
168+
public void testMergeOrderedRangeScanIteratorsSortsPrefetchedSources() {
169+
TestIterator first = new TestIterator(3, 4);
170+
TestIterator second = new TestIterator(1, 2);
171+
first.hasNext();
172+
second.hasNext();
173+
92174
HgKvIterator<HgKvEntry> iterator =
93175
NodeTxSessionProxy.mergeOrderedRangeScanIterators(
94176
Arrays.asList(first, second), 0L);
@@ -225,19 +307,30 @@ public byte[] value() {
225307
private static final class RecordingPartitioner
226308
implements HgStoreNodePartitioner {
227309

310+
private final long firstNodeId;
311+
private final long secondNodeId;
228312
private int byteRangeCalls;
229313
private int codeRangeCalls;
230314
private byte[] lastStartKey;
231315
private byte[] lastEndKey;
232316

317+
private RecordingPartitioner() {
318+
this(1L, 2L);
319+
}
320+
321+
private RecordingPartitioner(long firstNodeId, long secondNodeId) {
322+
this.firstNodeId = firstNodeId;
323+
this.secondNodeId = secondNodeId;
324+
}
325+
233326
@Override
234327
public int partition(HgNodePartitionerBuilder builder,
235328
String graphName, byte[] startKey,
236329
byte[] endKey) {
237330
this.byteRangeCalls++;
238331
this.lastStartKey = startKey;
239332
this.lastEndKey = endKey;
240-
builder.setPartitions(partitions());
333+
builder.setPartitions(this.partitions());
241334
return 0;
242335
}
243336

@@ -246,15 +339,81 @@ public int partition(HgNodePartitionerBuilder builder,
246339
String graphName, int startCode,
247340
int endCode) {
248341
this.codeRangeCalls++;
249-
builder.setPartitions(partitions());
342+
builder.setPartitions(this.partitions());
250343
return 0;
251344
}
252345

253-
private static Set<HgNodePartition> partitions() {
346+
private Set<HgNodePartition> partitions() {
254347
Set<HgNodePartition> partitions = new LinkedHashSet<>();
255-
partitions.add(HgNodePartition.of(1L, 10, 10, 20));
256-
partitions.add(HgNodePartition.of(2L, 30, 30, 40));
348+
partitions.add(HgNodePartition.of(this.firstNodeId, 10, 10, 20));
349+
partitions.add(HgNodePartition.of(this.secondNodeId, 30, 30, 40));
257350
return partitions;
258351
}
259352
}
353+
354+
private static final class RecordingStoreNode implements HgStoreNode {
355+
356+
private final Long nodeId;
357+
private final HgStoreSession session;
358+
359+
private RecordingStoreNode(Long nodeId, HgStoreSession session) {
360+
this.nodeId = nodeId;
361+
this.session = session;
362+
}
363+
364+
@Override
365+
public Long getNodeId() {
366+
return this.nodeId;
367+
}
368+
369+
@Override
370+
public String getAddress() {
371+
return "127.0.0.1:" + this.nodeId;
372+
}
373+
374+
@Override
375+
public HgStoreSession openSession(String graphName) {
376+
return this.session;
377+
}
378+
}
379+
380+
private static final class RecordingSession
381+
implements InvocationHandler {
382+
383+
private final List<Builder> builders;
384+
private final TestIterator iterator;
385+
private int rangeScanCalls;
386+
387+
private RecordingSession(TestIterator iterator) {
388+
this.builders = Collections.synchronizedList(new ArrayList<>());
389+
this.iterator = iterator;
390+
this.rangeScanCalls = 0;
391+
}
392+
393+
private HgStoreSession proxy() {
394+
return (HgStoreSession) Proxy.newProxyInstance(
395+
HgStoreSession.class.getClassLoader(),
396+
new Class[]{HgStoreSession.class}, this);
397+
}
398+
399+
@Override
400+
public Object invoke(Object proxy, Method method, Object[] args) {
401+
if ("scanIterator".equals(method.getName())) {
402+
if (args != null && args.length == 1 &&
403+
args[0] instanceof Builder) {
404+
this.builders.add(((Builder) args[0]).clone());
405+
return this.iterator;
406+
}
407+
this.rangeScanCalls++;
408+
return this.iterator;
409+
}
410+
if ("isTx".equals(method.getName())) {
411+
return false;
412+
}
413+
if ("toString".equals(method.getName())) {
414+
return "RecordingSession";
415+
}
416+
throw new UnsupportedOperationException(method.toString());
417+
}
418+
}
260419
}

0 commit comments

Comments
 (0)