use of com.facebook.presto.thrift.api.connector.PrestoThriftNullableToken in project presto by prestodb.
the class TestThriftIndexPageSource method testGetNextPageTwoConcurrentRequests.
@Test
public void testGetNextPageTwoConcurrentRequests() throws Exception {
final int splits = 3;
final int lookupRequestsConcurrency = 2;
final int rowsPerSplit = 1;
List<SettableFuture<PrestoThriftPageResult>> futures = IntStream.range(0, splits).mapToObj(i -> SettableFuture.<PrestoThriftPageResult>create()).collect(toImmutableList());
List<CountDownLatch> signals = IntStream.range(0, splits).mapToObj(i -> new CountDownLatch(1)).collect(toImmutableList());
TestingThriftService client = new TestingThriftService(rowsPerSplit, false, false) {
@Override
public ListenableFuture<PrestoThriftPageResult> getRows(PrestoThriftId splitId, List<String> columns, long maxBytes, PrestoThriftNullableToken nextToken) {
int key = Ints.fromByteArray(splitId.getId());
signals.get(key).countDown();
return futures.get(key);
}
};
ThriftConnectorStats stats = new ThriftConnectorStats();
long pageSizeReceived = 0;
ThriftIndexPageSource pageSource = new ThriftIndexPageSource((context, headers) -> client, ImmutableMap.of(), stats, new ThriftIndexHandle(new SchemaTableName("default", "table1"), TupleDomain.all()), ImmutableList.of(column("a", INTEGER)), ImmutableList.of(column("b", INTEGER)), new InMemoryRecordSet(ImmutableList.of(INTEGER), generateKeys(0, splits)), MAX_BYTES_PER_RESPONSE, lookupRequestsConcurrency);
assertNull(pageSource.getNextPage());
assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
signals.get(0).await(1, SECONDS);
signals.get(1).await(1, SECONDS);
signals.get(2).await(1, SECONDS);
assertEquals(signals.get(0).getCount(), 0, "first request wasn't sent");
assertEquals(signals.get(1).getCount(), 0, "second request wasn't sent");
assertEquals(signals.get(2).getCount(), 1, "third request shouldn't be sent");
// at this point first two requests were sent
assertFalse(pageSource.isFinished());
assertNull(pageSource.getNextPage());
assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
// completing the second request
futures.get(1).set(pageResult(20, null));
Page page = pageSource.getNextPage();
pageSizeReceived += page.getSizeInBytes();
assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
assertNotNull(page);
assertEquals(page.getPositionCount(), 1);
assertEquals(page.getBlock(0).getInt(0), 20);
// not complete yet
assertFalse(pageSource.isFinished());
// once one of the requests completes the next one should be sent
signals.get(2).await(1, SECONDS);
assertEquals(signals.get(2).getCount(), 0, "third request wasn't sent");
// completing the first request
futures.get(0).set(pageResult(10, null));
page = pageSource.getNextPage();
assertNotNull(page);
pageSizeReceived += page.getSizeInBytes();
assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
assertEquals(page.getPositionCount(), 1);
assertEquals(page.getBlock(0).getInt(0), 10);
// still not complete
assertFalse(pageSource.isFinished());
// completing the third request
futures.get(2).set(pageResult(30, null));
page = pageSource.getNextPage();
assertNotNull(page);
pageSizeReceived += page.getSizeInBytes();
assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
assertEquals(page.getPositionCount(), 1);
assertEquals(page.getBlock(0).getInt(0), 30);
// finished now
assertTrue(pageSource.isFinished());
// after completion
assertNull(pageSource.getNextPage());
pageSource.close();
}
use of com.facebook.presto.thrift.api.connector.PrestoThriftNullableToken in project presto by prestodb.
the class ThriftIndexPageSource method sendSplitRequest.
private ListenableFuture<PrestoThriftSplitBatch> sendSplitRequest(@Nullable PrestoThriftId nextToken) {
long start = System.nanoTime();
ListenableFuture<PrestoThriftSplitBatch> future = client.get(thriftHeaders).getIndexSplits(schemaTableName, lookupColumnNames, outputColumnNames, keys, outputConstraint, MAX_SPLIT_COUNT, new PrestoThriftNullableToken(nextToken));
future = catchingThriftException(future);
future.addListener(() -> readTimeNanos.addAndGet(System.nanoTime() - start), directExecutor());
return future;
}
use of com.facebook.presto.thrift.api.connector.PrestoThriftNullableToken in project presto by prestodb.
the class ThriftIndexPageSource method sendDataRequest.
private void sendDataRequest(RunningSplitContext context, @Nullable PrestoThriftId nextToken) {
long start = System.nanoTime();
ListenableFuture<PrestoThriftPageResult> future = context.getClient().getRows(context.getSplit().getSplitId(), outputColumnNames, maxBytesPerResponse, new PrestoThriftNullableToken(nextToken));
future = catchingThriftException(future);
future.addListener(() -> readTimeNanos.addAndGet(System.nanoTime() - start), directExecutor());
dataRequests.add(future);
contexts.put(future, context);
}
use of com.facebook.presto.thrift.api.connector.PrestoThriftNullableToken in project presto by prestodb.
the class ThriftPageSource method sendDataRequestInternal.
private CompletableFuture<PrestoThriftPageResult> sendDataRequestInternal() {
long start = System.nanoTime();
ListenableFuture<PrestoThriftPageResult> rowsBatchFuture = client.getRows(splitId, columnNames, maxBytesPerResponse, new PrestoThriftNullableToken(nextToken));
rowsBatchFuture = catchingThriftException(rowsBatchFuture);
rowsBatchFuture.addListener(() -> readTimeNanos.addAndGet(System.nanoTime() - start), directExecutor());
return toCompletableFuture(nonCancellationPropagating(rowsBatchFuture));
}
Aggregations