@Override
    public <U> StateFuture<Collection<U>> onNext(Function<T, StateFuture<? extends U>> iterating) {
        // Public interface implementation, this is on task thread.
        // We perform the user code on cache, and create a new request and chain with it.
        if (isEmpty()) {
            return StateFutureUtils.completedFuture(Collections.emptyList());
        }
        Collection<StateFuture<? extends U>> resultFutures = new ArrayList<>();

        for (T item : cache) {
            resultFutures.add(iterating.apply(item));
        }
        if (hasNext()) {
            return StateFutureUtils.combineAll(resultFutures)
                    .thenCombine(
                            asyncNextLoad().thenCompose(itr -> itr.onNext(iterating)),
                            (a, b) -> {
                                // TODO optimization: Avoid results copy.
                                Collection<U> result = new ArrayList<>(a.size() + b.size());
                                result.addAll(a);
                                result.addAll(b);
                                return result;
                            });
        } else {
            return StateFutureUtils.combineAll(resultFutures);
        }
    }
@Test
    @SuppressWarnings({"unchecked", "rawtypes"})
    public void testPartialLoading() {
        TestIteratorStateExecutor stateExecutor = new TestIteratorStateExecutor(100, 3);
        AsyncExecutionController aec =
                new AsyncExecutionController(
                        new SyncMailboxExecutor(), (a, b) -> {}, stateExecutor, 1, 100, 1000, 1);
        stateExecutor.bindAec(aec);
        RecordContext<String> recordContext = aec.buildContext("1", "key1");
        aec.setCurrentContext(recordContext);

        AtomicInteger processed = new AtomicInteger();

        aec.handleRequest(null, StateRequestType.MAP_ITER, null)
                .thenAccept(
                        (iter) -> {
                            assertThat(iter).isInstanceOf(StateIterator.class);
                            ((StateIterator<Integer>) iter)
                                    .onNext(
                                            (item) -> {
                                                assertThat(item)
                                                        .isEqualTo(processed.getAndIncrement());
                                            })
                                    .thenAccept(
                                            (v) -> {
                                                assertThat(processed.get()).isEqualTo(100);
                                            });
                        });
        aec.drainInflightRecords(0);
    }