struct LoopBody {

struct Callback {

Result<ControlFlow<>> operator()(const T& next) {

if (IsIterationEnd(next)) {

return Break();

} else {

auto visited = visitor(next);

if (visited.ok()) {

return Continue();

} else {

return visited;

}

}

}

Visitor visitor;

};

Future<ControlFlow<>> operator()() {

Callback callback{visitor};

auto next = generator();

return next.Then(std::move(callback));

}

AsyncGenerator<T> generator;

Visitor visitor;

};

return Loop(LoopBody{std::move(generator), std::move(visitor)});

auto vec = std::make_shared<std::vector<T>>();

auto loop_body = [generator = std::move(generator),

vec = std::move(vec)]() -> Future<ControlFlow<std::vector<T>>> {

auto next = generator();

return next.Then([vec](const T& result) -> Result<ControlFlow<std::vector<T>>> {

if (IsIterationEnd(result)) {

return Break(*vec);

} else {

vec->push_back(result);

return Continue();

}

});

};

return Loop(std::move(loop_body));

public:

MappingGenerator(AsyncGenerator<T> source, std::function<Future<V>(const T&)> map)

: state_(std::make_shared<State>(std::move(source), std::move(map))) {}

Future<V> operator()() {

auto future = Future<V>::Make();

bool should_trigger;

{

auto guard = state_->mutex.Lock();

if (state_->finished) {

return AsyncGeneratorEnd<V>();

}

should_trigger = state_->waiting_jobs.empty();

state_->waiting_jobs.push_back(future);

}

if (should_trigger) {

state_->source().AddCallback(Callback{state_});

}

return future;

}

private:

struct State {

State(AsyncGenerator<T> source, std::function<Future<V>(const T&)> map)

: source(std::move(source)),

map(std::move(map)),

waiting_jobs(),

mutex(),

finished(false) {}

void Purge() {

// This might be called by an original callback (if the source iterator fails or

// ends) or by a mapped callback (if the map function fails or ends prematurely).

// Either way it should only be called once and after finished is set so there is no

// need to guard access to `waiting_jobs`.

while (!waiting_jobs.empty()) {

waiting_jobs.front().MarkFinished(IterationTraits<V>::End());

waiting_jobs.pop_front();

}

}

AsyncGenerator<T> source;

std::function<Future<V>(const T&)> map;

std::deque<Future<V>> waiting_jobs;

util::Mutex mutex;

bool finished;

};

struct Callback;

struct MappedCallback {

void operator()(const Result<V>& maybe_next) {

bool end = !maybe_next.ok() || IsIterationEnd(*maybe_next);

bool should_purge = false;

if (end) {

{

auto guard = state->mutex.Lock();

should_purge = !state->finished;

state->finished = true;

}

}

sink.MarkFinished(maybe_next);

if (should_purge) {

state->Purge();

}

}

std::shared_ptr<State> state;

Future<V> sink;

};

struct Callback {

void operator()(const Result<T>& maybe_next) {

Future<V> sink;

bool end = !maybe_next.ok() || IsIterationEnd(*maybe_next);

bool should_purge = false;

bool should_trigger;

{

auto guard = state->mutex.Lock();

// A MappedCallback may have purged or be purging the queue;

// we shouldn't do anything here.

if (state->finished) return;

if (end) {

should_purge = !state->finished;

state->finished = true;

}

sink = state->waiting_jobs.front();

state->waiting_jobs.pop_front();

should_trigger = !end && !state->waiting_jobs.empty();

}

if (should_purge) {

state->Purge();

}

if (should_trigger) {

state->source().AddCallback(Callback{state});

}

if (maybe_next.ok()) {

const T& val = maybe_next.ValueUnsafe();

if (IsIterationEnd(val)) {

sink.MarkFinished(IterationTraits<V>::End());

} else {

Future<V> mapped_fut = state->map(val);

mapped_fut.AddCallback(MappedCallback{std::move(state), std::move(sink)});

}

} else {

sink.MarkFinished(maybe_next.status());

}

}

std::shared_ptr<State> state;

};

std::shared_ptr<State> state_;

auto map_callback = [map = std::move(map)](const T& val) mutable -> Future<V> {

return ToFuture(map(val));

};

return MappingGenerator<T, V>(std::move(source_generator), std::move(map_callback));

public:

SequencingGenerator(AsyncGenerator<T> source, ComesAfter compare, IsNext is_next,

T initial_value)

: state_(std::make_shared<State>(std::move(source), std::move(compare),

std::move(is_next), std::move(initial_value))) {}

Future<T> operator()() {

{

auto guard = state_->mutex.Lock();

// We can send a result immediately if the top of the queue is either an

// error or the next item

if (!state_->queue.empty() &&

(!state_->queue.top().ok() ||

state_->is_next(state_->previous_value, *state_->queue.top()))) {

auto result = std::move(state_->queue.top());

if (result.ok()) {

state_->previous_value = *result;

}

state_->queue.pop();

return Future<T>::MakeFinished(result);

}

if (state_->finished) {

return AsyncGeneratorEnd<T>();

}

// The next item is not in the queue so we will need to wait

auto new_waiting_fut = Future<T>::Make();

state_->waiting_future = new_waiting_fut;

guard.Unlock();

state_->source().AddCallback(Callback{state_});

return new_waiting_fut;

}

}

private:

struct WrappedComesAfter {

bool operator()(const Result<T>& left, const Result<T>& right) {

if (!left.ok() || !right.ok()) {

// Should never happen

return false;

}

return compare(*left, *right);

}

ComesAfter compare;

};

struct State {

State(AsyncGenerator<T> source, ComesAfter compare, IsNext is_next, T initial_value)

: source(std::move(source)),

is_next(std::move(is_next)),

previous_value(std::move(initial_value)),

waiting_future(),

queue(WrappedComesAfter{compare}),

finished(false),

mutex() {}

AsyncGenerator<T> source;

IsNext is_next;

T previous_value;

Future<T> waiting_future;

std::priority_queue<Result<T>, std::vector<Result<T>>, WrappedComesAfter> queue;

bool finished;

util::Mutex mutex;

};

class Callback {

public:

explicit Callback(std::shared_ptr<State> state) : state_(std::move(state)) {}

void operator()(const Result<T> result) {

Future<T> to_deliver;

bool finished;

{

auto guard = state_->mutex.Lock();

bool ready_to_deliver = false;

if (!result.ok()) {

// Clear any cached results

while (!state_->queue.empty()) {

state_->queue.pop();

}

ready_to_deliver = true;

state_->finished = true;

} else if (IsIterationEnd<T>(result.ValueUnsafe())) {

ready_to_deliver = state_->queue.empty();

state_->finished = true;

} else {

ready_to_deliver = state_->is_next(state_->previous_value, *result);

}

if (ready_to_deliver && state_->waiting_future.is_valid()) {

to_deliver = state_->waiting_future;

if (result.ok()) {

state_->previous_value = *result;

}

} else {

state_->queue.push(result);

}

// Capture state_->finished so we can access it outside the mutex

finished = state_->finished;

}

// Must deliver result outside of the mutex

if (to_deliver.is_valid()) {

to_deliver.MarkFinished(result);

} else {

// Otherwise, if we didn't get the next item (or a terminal item), we

// need to keep looking

if (!finished) {

state_->source().AddCallback(Callback{state_});

}

}

}

private:

const std::shared_ptr<State> state_;

};

const std::shared_ptr<State> state_;

ComesAfter compare, IsNext is_next,

T initial_value) {

return SequencingGenerator<T, ComesAfter, IsNext>(

std::move(source_generator), std::move(compare), std::move(is_next),

std::move(initial_value));

// The transforming generator state will be referenced as an async generator but will

// also be referenced via callback to various futures. If the async generator owner

// moves it around we need the state to be consistent for future callbacks.

struct TransformingGeneratorState

: std::enable_shared_from_this<TransformingGeneratorState> {

TransformingGeneratorState(AsyncGenerator<T> generator, Transformer<T, V> transformer)

: generator_(std::move(generator)),

transformer_(std::move(transformer)),

last_value_(),

finished_() {}

Future<V> operator()() {

while (true) {

auto maybe_next_result = Pump();

if (!maybe_next_result.ok()) {

return Future<V>::MakeFinished(maybe_next_result.status());

}

auto maybe_next = std::move(maybe_next_result).ValueUnsafe();

if (maybe_next.has_value()) {

return Future<V>::MakeFinished(*std::move(maybe_next));

}

auto next_fut = generator_();

// If finished already, process results immediately inside the loop to avoid

// stack overflow

if (next_fut.is_finished()) {

auto next_result = next_fut.result();

if (next_result.ok()) {

last_value_ = *next_result;

} else {

return Future<V>::MakeFinished(next_result.status());

}

// Otherwise, if not finished immediately, add callback to process results

} else {

auto self = this->shared_from_this();

return next_fut.Then([self](const T& next_result) {

self->last_value_ = next_result;

return (*self)();

});

}

}

}

// See comment on TransformingIterator::Pump

Result<std::optional<V>> Pump() {

if (!finished_ && last_value_.has_value()) {

ARROW_ASSIGN_OR_RAISE(TransformFlow<V> next, transformer_(*last_value_));

if (next.ReadyForNext()) {

if (IsIterationEnd(*last_value_)) {

finished_ = true;

}

last_value_.reset();

}

if (next.Finished()) {

finished_ = true;

}

if (next.HasValue()) {

return next.Value();

}

}

if (finished_) {

return IterationTraits<V>::End();

}

return std::nullopt;

}

AsyncGenerator<T> generator_;

Transformer<T, V> transformer_;

std::optional<T> last_value_;

bool finished_;

};

public:

explicit TransformingGenerator(AsyncGenerator<T> generator,

Transformer<T, V> transformer)

: state_(std::make_shared<TransformingGeneratorState>(std::move(generator),

std::move(transformer))) {}

Future<V> operator()() { return (*state_)(); }

protected:

std::shared_ptr<TransformingGeneratorState> state_;

public:

SerialReadaheadGenerator(AsyncGenerator<T> source_generator, int max_readahead)

: state_(std::make_shared<State>(std::move(source_generator), max_readahead)) {}

Future<T> operator()() {

if (state_->first_) {

// Lazy generator, need to wait for the first ask to prime the pump

state_->first_ = false;

auto next = state_->source_();

return next.Then(Callback{state_}, ErrCallback{state_});

}

// This generator is not async-reentrant. We won't be called until the last

// future finished so we know there is something in the queue

auto finished = state_->finished_.load();

if (finished && state_->readahead_queue_.IsEmpty()) {

return AsyncGeneratorEnd<T>();

}

std::shared_ptr<Future<T>> next;

if (!state_->readahead_queue_.Read(next)) {

return Status::UnknownError("Could not read from readahead_queue");

}

auto last_available = state_->spaces_available_.fetch_add(1);

if (last_available == 0 && !finished) {

// Reader idled out, we need to restart it

ARROW_RETURN_NOT_OK(state_->Pump(state_));

}

return *next;

}

private:

struct State {

State(AsyncGenerator<T> source, int max_readahead)

: first_(true),

source_(std::move(source)),

finished_(false),

// There is one extra "space" for the in-flight request

spaces_available_(max_readahead + 1),

// The SPSC queue has size-1 "usable" slots so we need to overallocate 1

readahead_queue_(max_readahead + 1) {}

Status Pump(const std::shared_ptr<State>& self) {

// Can't do readahead_queue.write(source().Then(...)) because then the

// callback might run immediately and add itself to the queue before this gets added

// to the queue messing up the order.

auto next_slot = std::make_shared<Future<T>>();

auto written = readahead_queue_.Write(next_slot);

if (!written) {

return Status::UnknownError("Could not write to readahead_queue");

}

// If this Pump is being called from a callback it is possible for the source to

// poll and read from the queue between the Write and this spot where we fill the

// value in. However, it is not possible for the future to read this value we are

// writing. That is because this callback (the callback for future X) must be

// finished before future X is marked complete and this source is not pulled

// reentrantly so it will not poll for future X+1 until this callback has completed.

*next_slot = source_().Then(Callback{self}, ErrCallback{self});

return Status::OK();

}

// Only accessed by the consumer end

bool first_;

// Accessed by both threads

AsyncGenerator<T> source_;

std::atomic<bool> finished_;

// The queue has a size but it is not atomic. We keep track of how many spaces are

// left in the queue here so we know if we've just written the last value and we need

// to stop reading ahead or if we've just read from a full queue and we need to

// restart reading ahead

std::atomic<uint32_t> spaces_available_;

// Needs to be a queue of shared_ptr and not Future because we set the value of the

// future after we add it to the queue

util::SpscQueue<std::shared_ptr<Future<T>>> readahead_queue_;

};

struct Callback {

Result<T> operator()(const T& next) {

if (IsIterationEnd(next)) {

state_->finished_.store(true);

return next;

}

auto last_available = state_->spaces_available_.fetch_sub(1);

if (last_available > 1) {

ARROW_RETURN_NOT_OK(state_->Pump(state_));

}

return next;

}

std::shared_ptr<State> state_;

};

struct ErrCallback {

Result<T> operator()(const Status& st) {

state_->finished_.store(true);

return st;

}

std::shared_ptr<State> state_;

};

std::shared_ptr<State> state_;

public:

explicit FutureFirstGenerator(Future<AsyncGenerator<T>> future)

: state_(std::make_shared<State>(std::move(future))) {}

Future<T> operator()() {

if (state_->source_) {

return state_->source_();

} else {

auto state = state_;

return state_->future_.Then([state](const AsyncGenerator<T>& source) {

state->source_ = source;

return state->source_();

});

}

}

private:

struct State {

explicit State(Future<AsyncGenerator<T>> future) : future_(future), source_() {}

Future<AsyncGenerator<T>> future_;

AsyncGenerator<T> source_;

};

std::shared_ptr<State> state_;

struct AutostartGenerator {

Future<T> operator()() {

if (first_future->is_valid()) {

Future<T> result = *first_future;

*first_future = Future<T>();

return result;

}

return source();

}

std::shared_ptr<Future<T>> first_future;

AsyncGenerator<T> source;

};

std::shared_ptr<Future<T>> first_future = std::make_shared<Future<T>>(generator());

return AutostartGenerator{std::move(first_future), std::move(generator)};

public:

ReadaheadGenerator(AsyncGenerator<T> source_generator, int max_readahead)

: state_(std::make_shared<State>(std::move(source_generator), max_readahead)) {}

Future<T> AddMarkFinishedContinuation(Future<T> fut) {

auto state = state_;

return fut.Then(

[state](const T& result) -> Future<T> {

state->MarkFinishedIfDone(result);

if (state->finished.load()) {

if (state->num_running.fetch_sub(1) == 1) {

state->final_future.MarkFinished();

}

} else {

state->num_running.fetch_sub(1);

}

return result;

},

[state](const Status& err) -> Future<T> {

// If there is an error we need to make sure all running

// tasks finish before we return the error.

state->finished.store(true);

if (state->num_running.fetch_sub(1) == 1) {

state->final_future.MarkFinished();

}

return state->final_future.Then([err]() -> Result<T> { return err; });

});

}

Future<T> operator()() {

if (state_->readahead_queue.empty()) {

// This is the first request, let's pump the underlying queue

state_->num_running.store(state_->max_readahead);

for (int i = 0; i < state_->max_readahead; i++) {

auto next = state_->source_generator();

auto next_after_check = AddMarkFinishedContinuation(std::move(next));

state_->readahead_queue.push(std::move(next_after_check));

}

}

// Pop one and add one

auto result = state_->readahead_queue.front();

state_->readahead_queue.pop();

if (state_->finished.load()) {

state_->readahead_queue.push(AsyncGeneratorEnd<T>());

} else {

state_->num_running.fetch_add(1);

auto back_of_queue = state_->source_generator();

auto back_of_queue_after_check =

AddMarkFinishedContinuation(std::move(back_of_queue));

state_->readahead_queue.push(std::move(back_of_queue_after_check));

}

return result;

}

private:

struct State {

State(AsyncGenerator<T> source_generator, int max_readahead)

: source_generator(std::move(source_generator)), max_readahead(max_readahead) {}

void MarkFinishedIfDone(const T& next_result) {

if (IsIterationEnd(next_result)) {

finished.store(true);

}

}

AsyncGenerator<T> source_generator;

int max_readahead;

Future<> final_future = Future<>::Make();

std::atomic<int> num_running{0};

std::atomic<bool> finished{false};

std::queue<Future<T>> readahead_queue;

};

std::shared_ptr<State> state_;

struct State {

State() {}

util::Mutex mutex;

std::deque<Result<T>> result_q;

std::optional<Future<T>> consumer_fut;

bool finished = false;

};

public:

/// Producer API for PushGenerator

class Producer {

public:

explicit Producer(const std::shared_ptr<State>& state) : weak_state_(state) {}

/// \brief Push a value on the queue

///

/// True is returned if the value was pushed, false if the generator is

/// already closed or destroyed. If the latter, it is recommended to stop

/// producing any further values.

bool Push(Result<T> result) {

auto state = weak_state_.lock();

if (!state) {

// Generator was destroyed

return false;

}

auto lock = state->mutex.Lock();

if (state->finished) {

// Closed early

return false;

}

if (state->consumer_fut.has_value()) {

auto fut = std::move(state->consumer_fut.value());

state->consumer_fut.reset();

lock.Unlock(); // unlock before potentially invoking a callback

fut.MarkFinished(std::move(result));

} else {

state->result_q.push_back(std::move(result));

}

return true;

}

/// \brief Tell the consumer we have finished producing

///

/// It is allowed to call this and later call Push() again ("early close").

/// In this case, calls to Push() after the queue is closed are silently

/// ignored. This can help implementing non-trivial cancellation cases.

///

/// True is returned on success, false if the generator is already closed

/// or destroyed.

bool Close() {

auto state = weak_state_.lock();

if (!state) {

// Generator was destroyed

return false;

}

auto lock = state->mutex.Lock();

if (state->finished) {

// Already closed

return false;

}

state->finished = true;

if (state->consumer_fut.has_value()) {

auto fut = std::move(state->consumer_fut.value());

state->consumer_fut.reset();

lock.Unlock(); // unlock before potentially invoking a callback

fut.MarkFinished(IterationTraits<T>::End());

}

return true;

}

/// Return whether the generator was closed or destroyed.

bool is_closed() const {

auto state = weak_state_.lock();

if (!state) {

// Generator was destroyed

return true;

}

auto lock = state->mutex.Lock();

return state->finished;

}

private:

const std::weak_ptr<State> weak_state_;

};

PushGenerator() : state_(std::make_shared<State>()) {}

/// Read an item from the queue

Future<T> operator()() const {

auto lock = state_->mutex.Lock();

assert(!state_->consumer_fut.has_value()); // Non-reentrant

if (!state_->result_q.empty()) {

auto fut = Future<T>::MakeFinished(std::move(state_->result_q.front()));

state_->result_q.pop_front();

return fut;

}

if (state_->finished) {

return AsyncGeneratorEnd<T>();

}

auto fut = Future<T>::Make();

state_->consumer_fut = fut;

return fut;

}

/// \brief Return producer-side interface

///

/// The returned object must be used by the producer to push values on the queue.

/// Only a single Producer object should be instantiated.

Producer producer() { return Producer{state_}; }

private:

const std::shared_ptr<State> state_;

struct State {

explicit State(std::vector<T> vec_) : vec(std::move(vec_)), vec_idx(0) {}

std::vector<T> vec;

std::atomic<std::size_t> vec_idx;

};

auto state = std::make_shared<State>(std::move(vec));

return [state]() {

auto idx = state->vec_idx.fetch_add(1);

if (idx >= state->vec.size()) {

// Eagerly return memory

state->vec.clear();

return AsyncGeneratorEnd<T>();

}

return Future<T>::MakeFinished(state->vec[idx]);

};

// Note, the implementation of this class is quite complex at the moment (PRs to

// simplify are always welcome)

//

// Terminology is borrowed from rxjs. This is a pull based implementation of the

// mergeAll operator. The "outer subscription" refers to the async

// generator that the caller provided when creating this. The outer subscription

// yields generators.

//

// Each of these generators is then subscribed to (up to max_subscriptions) and these

// are referred to as "inner subscriptions".

//

// As soon as we start we try and establish `max_subscriptions` inner subscriptions. For

// each inner subscription we will cache up to 1 value. This means we may have more

// values than we have been asked for. In our example, if a caller asks for one record

// batch we will start scanning `max_subscriptions` different files. For each file we

// will only queue up to 1 batch (so a separate readahead is needed on the file if batch