StagedStream.java

package com.macvalves.util.concurrent;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;

import java.util.Collection;
import java.util.List;
import java.util.Spliterators;
import java.util.concurrent.CyclicBarrier;
import java.util.function.Function;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import com.google.common.collect.AbstractIterator;
import com.macvalves.util.Exceptions;

public class StagedStream<S, T> {

  // when all parties initially arrive that increments the index and populates the Ts
  private int idx = -1;
  private List<T> currentTs = null;

  private final CyclicBarrier barrier;

  private final List<S> seeds;
  private final Function<S, List<T>> expansionFunction;
  private int takenStreams = 0;

  public StagedStream(int streams, List<S> seeds, Function<S, List<T>> expansion){
    this.seeds = requireNonNull(seeds);
    this.expansionFunction = requireNonNull(expansion);

    checkArgument(streams > 0);
    this.barrier = new CyclicBarrier(streams, () -> {
      if(idx < seeds.size())
        currentTs = expansionFunction.apply(seeds.get(++idx));
      else
        currentTs = null;
    });
  }

  public Stream<T> takeStream(){
    checkState(takenStreams < barrier.getParties(), "Too many streams taken");
    takenStreams++;

    return StreamSupport.stream(Spliterators.spliterator(
      new ItrWithAcquire(), seeds.size(), 0), false)
      .flatMap(Collection::stream);
  }

  private class ItrWithAcquire extends AbstractIterator<List<T>>{
    @Override
    protected List<T> computeNext(){
      Exceptions.runR(barrier::await);
      if(null == currentTs){
        assert idx >= seeds.size();
        return this.endOfData();
      }
      assert idx < seeds.size();
      return currentTs;
    }
  }

}