Examples with KillableCompletionStage io.crate.concurrent.KillableCompletionStage used on opensource projects

Example 1 with KillableCompletionStage

use of io.crate.concurrent.KillableCompletionStage in project crate by crate.

the class BatchPagingIteratorTest method testBatchPagingIteratorWithPagedSource.

@Test
public void testBatchPagingIteratorWithPagedSource() throws Exception {
    List<Object[]> expectedResult = StreamSupport.stream(RowGenerator.range(0, 10).spliterator(), false).map(Row::materialize).collect(Collectors.toList());
    Supplier<BatchIterator<Row>> batchIteratorSupplier = () -> {
        BatchSimulatingIterator<Row> source = new BatchSimulatingIterator<>(TestingBatchIterators.range(0, 10), 2, 5, executor);
        Function<Integer, KillableCompletionStage<? extends Iterable<? extends KeyIterable<Integer, Row>>>> fetchMore = exhausted -> {
            List<Row> rows = new ArrayList<>();
            while (source.moveNext()) {
                rows.add(new RowN(source.currentElement().materialize()));
            }
            if (source.allLoaded()) {
                return KillableCompletionStage.whenKilled(CompletableFuture.completedFuture(singletonList(new KeyIterable<>(1, rows))), t -> {
                });
            }
            // but to simulate multiple pages and fetchMore calls
            try {
                return KillableCompletionStage.whenKilled(source.loadNextBatch().toCompletableFuture().thenApply(ignored -> {
                    while (source.moveNext()) {
                        rows.add(new RowN(source.currentElement().materialize()));
                    }
                    return singleton(new KeyIterable<>(1, rows));
                }), t -> {
                });
            } catch (Exception e) {
                return KillableCompletionStage.failed(e);
            }
        };
        return new BatchPagingIterator<>(PassThroughPagingIterator.repeatable(), fetchMore, source::allLoaded, throwable -> {
            if (throwable == null) {
                source.close();
            } else {
                source.kill(throwable);
            }
        });
    };
    BatchIteratorTester tester = new BatchIteratorTester(batchIteratorSupplier);
    tester.verifyResultAndEdgeCaseBehaviour(expectedResult);
}