/usr/local/lib64/python3.6/site-packages/pyarrow/include/arrow/util
Edit: /usr/local/lib64/python3.6/site-packages/pyarrow/include/arrow/util/async_generator.h (64088B)
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include
#include
#include
#include
#include
#include
#include "arrow/util/async_util.h"
#include "arrow/util/functional.h"
#include "arrow/util/future.h"
#include "arrow/util/io_util.h"
#include "arrow/util/iterator.h"
#include "arrow/util/mutex.h"
#include "arrow/util/optional.h"
#include "arrow/util/queue.h"
#include "arrow/util/thread_pool.h"
namespace arrow {
// The methods in this file create, modify, and utilize AsyncGenerator which is an
// iterator of futures. This allows an asynchronous source (like file input) to be run
// through a pipeline in the same way that iterators can be used to create pipelined
// workflows.
//
// In order to support pipeline parallelism we introduce the concept of asynchronous
// reentrancy. This is different than synchronous reentrancy. With synchronous code a
// function is reentrant if the function can be called again while a previous call to that
// function is still running. Unless otherwise specified none of these generators are
// synchronously reentrant. Care should be taken to avoid calling them in such a way (and
// the utilities Visit/Collect/Await take care to do this).
//
// Asynchronous reentrancy on the other hand means the function is called again before the
// future returned by the function is marked finished (but after the call to get the
// future returns). Some of these generators are async-reentrant while others (e.g.
// those that depend on ordered processing like decompression) are not. Read the MakeXYZ
// function comments to determine which generators support async reentrancy.
//
// Note: Generators that are not asynchronously reentrant can still support readahead
// (\see MakeSerialReadaheadGenerator).
//
// Readahead operators, and some other operators, may introduce queueing. Any operators
// that introduce buffering should detail the amount of buffering they introduce in their
// MakeXYZ function comments.
template
using AsyncGenerator = std::function()>;
template
struct IterationTraits> {
/// \brief by default when iterating through a sequence of AsyncGenerator,
/// an empty function indicates the end of iteration.
static AsyncGenerator End() { return AsyncGenerator(); }
static bool IsEnd(const AsyncGenerator& val) { return !val; }
};
template
Future AsyncGeneratorEnd() {
return Future::MakeFinished(IterationTraits::End());
}
/// returning a future that completes when all have been visited
template
Future<> VisitAsyncGenerator(AsyncGenerator generator, Visitor visitor) {
struct LoopBody {
struct Callback {
Result> 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> operator()() {
Callback callback{visitor};
auto next = generator();
return next.Then(std::move(callback));
}
AsyncGenerator generator;
Visitor visitor;
};
return Loop(LoopBody{std::move(generator), std::move(visitor)});
}
/// \brief Wait for an async generator to complete, discarding results.
template
Future<> DiscardAllFromAsyncGenerator(AsyncGenerator generator) {
std::function visitor = [](const T&) { return Status::OK(); };
return VisitAsyncGenerator(generator, visitor);
}
/// \brief Collect the results of an async generator into a vector
template
Future> CollectAsyncGenerator(AsyncGenerator generator) {
auto vec = std::make_shared>();
struct LoopBody {
Future>> operator()() {
auto next = generator_();
auto vec = vec_;
return next.Then([vec](const T& result) -> Result>> {
if (IsIterationEnd(result)) {
return Break(*vec);
} else {
vec->push_back(result);
return Continue();
}
});
}
AsyncGenerator generator_;
std::shared_ptr> vec_;
};
return Loop(LoopBody{std::move(generator), std::move(vec)});
}
/// \see MakeMappedGenerator
template
class MappingGenerator {
public:
MappingGenerator(AsyncGenerator source, std::function(const T&)> map)
: state_(std::make_shared(std::move(source), std::move(map))) {}
Future operator()() {
auto future = Future::Make();
bool should_trigger;
{
auto guard = state_->mutex.Lock();
if (state_->finished) {
return AsyncGeneratorEnd();
}
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 source, std::function(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::End());
waiting_jobs.pop_front();
}
}
AsyncGenerator source;
std::function(const T&)> map;
std::deque> waiting_jobs;
util::Mutex mutex;
bool finished;
};
struct Callback;
struct MappedCallback {
void operator()(const Result& 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;
Future sink;
};
struct Callback {
void operator()(const Result& maybe_next) {
Future 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::End());
} else {
Future 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;
};
std::shared_ptr state_;
};
/// \brief Create a generator that will apply the map function to each element of
/// source. The map function is not called on the end token.
///
/// Note: This function makes a copy of `map` for each item
/// Note: Errors returned from the `map` function will be propagated
///
/// If the source generator is async-reentrant then this generator will be also
template ,
typename V = typename EnsureFuture::type::ValueType>
AsyncGenerator MakeMappedGenerator(AsyncGenerator source_generator, MapFn map) {
struct MapCallback {
MapFn map_;
Future operator()(const T& val) { return ToFuture(map_(val)); }
};
return MappingGenerator(std::move(source_generator), MapCallback{std::move(map)});
}
/// \brief Create a generator that will apply the map function to
/// each element of source. The map function is not called on the end
/// token. The result of the map function should be another
/// generator; all these generators will then be flattened to produce
/// a single stream of items.
///
/// Note: This function makes a copy of `map` for each item
/// Note: Errors returned from the `map` function will be propagated
///
/// If the source generator is async-reentrant then this generator will be also
template ,
typename V = typename EnsureFuture::type::ValueType>
AsyncGenerator MakeFlatMappedGenerator(AsyncGenerator source_generator, MapFn map) {
return MakeConcatenatedGenerator(
MakeMappedGenerator(std::move(source_generator), std::move(map)));
}
/// \see MakeSequencingGenerator
template
class SequencingGenerator {
public:
SequencingGenerator(AsyncGenerator source, ComesAfter compare, IsNext is_next,
T initial_value)
: state_(std::make_shared(std::move(source), std::move(compare),
std::move(is_next), std::move(initial_value))) {}
Future 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::MakeFinished(result);
}
if (state_->finished) {
return AsyncGeneratorEnd();
}
// The next item is not in the queue so we will need to wait
auto new_waiting_fut = Future::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& left, const Result& right) {
if (!left.ok() || !right.ok()) {
// Should never happen
return false;
}
return compare(*left, *right);
}
ComesAfter compare;
};
struct State {
State(AsyncGenerator 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 source;
IsNext is_next;
T previous_value;
Future waiting_future;
std::priority_queue, std::vector>, WrappedComesAfter> queue;
bool finished;
util::Mutex mutex;
};
class Callback {
public:
explicit Callback(std::shared_ptr state) : state_(std::move(state)) {}
void operator()(const Result result) {
Future 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(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_;
};
const std::shared_ptr state_;
};
/// \brief Buffer an AsyncGenerator to return values in sequence order ComesAfter
/// and IsNext determine the sequence order.
///
/// ComesAfter should be a BinaryPredicate that only returns true if a comes after b
///
/// IsNext should be a BinaryPredicate that returns true, given `a` and `b`, only if
/// `b` follows immediately after `a`. It should return true given `initial_value` and
/// `b` if `b` is the first item in the sequence.
///
/// This operator will queue unboundedly while waiting for the next item. It is intended
/// for jittery sources that might scatter an ordered sequence. It is NOT intended to
/// sort. Using it to try and sort could result in excessive RAM usage. This generator
/// will queue up to N blocks where N is the max "out of order"ness of the source.
///
/// For example, if the source is 1,6,2,5,4,3 it will queue 3 blocks because 3 is 3
/// blocks beyond where it belongs.
///
/// This generator is not async-reentrant but it consists only of a simple log(n)
/// insertion into a priority queue.
template
AsyncGenerator MakeSequencingGenerator(AsyncGenerator source_generator,
ComesAfter compare, IsNext is_next,
T initial_value) {
return SequencingGenerator(
std::move(source_generator), std::move(compare), std::move(is_next),
std::move(initial_value));
}
/// \see MakeTransformedGenerator
template
class TransformingGenerator {
// 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(AsyncGenerator generator, Transformer transformer)
: generator_(std::move(generator)),
transformer_(std::move(transformer)),
last_value_(),
finished_() {}
Future operator()() {
while (true) {
auto maybe_next_result = Pump();
if (!maybe_next_result.ok()) {
return Future::MakeFinished(maybe_next_result.status());
}
auto maybe_next = std::move(maybe_next_result).ValueUnsafe();
if (maybe_next.has_value()) {
return Future::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::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> Pump() {
if (!finished_ && last_value_.has_value()) {
ARROW_ASSIGN_OR_RAISE(TransformFlow 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::End();
}
return util::nullopt;
}
AsyncGenerator generator_;
Transformer transformer_;
util::optional last_value_;
bool finished_;
};
public:
explicit TransformingGenerator(AsyncGenerator generator,
Transformer transformer)
: state_(std::make_shared(std::move(generator),
std::move(transformer))) {}
Future operator()() { return (*state_)(); }
protected:
std::shared_ptr state_;
};
/// \brief Transform an async generator using a transformer function returning a new
/// AsyncGenerator
///
/// The transform function here behaves exactly the same as the transform function in
/// MakeTransformedIterator and you can safely use the same transform function to
/// transform both synchronous and asynchronous streams.
///
/// This generator is not async-reentrant
///
/// This generator may queue up to 1 instance of T but will not delay
template
AsyncGenerator MakeTransformedGenerator(AsyncGenerator generator,
Transformer transformer) {
return TransformingGenerator(generator, transformer);
}
/// \see MakeSerialReadaheadGenerator
template
class SerialReadaheadGenerator {
public:
SerialReadaheadGenerator(AsyncGenerator source_generator, int max_readahead)
: state_(std::make_shared(std::move(source_generator), max_readahead)) {}
Future 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();
}
std::shared_ptr> 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 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& 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>();
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 source_;
std::atomic 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 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>> readahead_queue_;
};
struct Callback {
Result 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_;
};
struct ErrCallback {
Result operator()(const Status& st) {
state_->finished_.store(true);
return st;
}
std::shared_ptr state_;
};
std::shared_ptr state_;
};
/// \see MakeFromFuture
template
class FutureFirstGenerator {
public:
explicit FutureFirstGenerator(Future> future)
: state_(std::make_shared(std::move(future))) {}
Future operator()() {
if (state_->source_) {
return state_->source_();
} else {
auto state = state_;
return state_->future_.Then([state](const AsyncGenerator& source) {
state->source_ = source;
return state->source_();
});
}
}
private:
struct State {
explicit State(Future> future) : future_(future), source_() {}
Future> future_;
AsyncGenerator source_;
};
std::shared_ptr state_;
};
/// \brief Transform a Future> into an AsyncGenerator
/// that waits for the future to complete as part of the first item.
///
/// This generator is not async-reentrant (even if the generator yielded by future is)
///
/// This generator does not queue
template
AsyncGenerator MakeFromFuture(Future> future) {
return FutureFirstGenerator(std::move(future));
}
/// \brief Create a generator that will pull from the source into a queue. Unlike
/// MakeReadaheadGenerator this will not pull reentrantly from the source.
///
/// The source generator does not need to be async-reentrant
///
/// This generator is not async-reentrant (even if the source is)
///
/// This generator may queue up to max_readahead additional instances of T
template
AsyncGenerator MakeSerialReadaheadGenerator(AsyncGenerator source_generator,
int max_readahead) {
return SerialReadaheadGenerator(std::move(source_generator), max_readahead);
}
/// \brief Create a generator that immediately pulls from the source
///
/// Typical generators do not pull from their source until they themselves
/// are pulled. This generator does not follow that convention and will call
/// generator() once before it returns. The returned generator will otherwise
/// mirror the source.
///
/// This generator forwards aysnc-reentrant pressure to the source
/// This generator buffers one item (the first result) until it is delivered.
template
AsyncGenerator MakeAutoStartingGenerator(AsyncGenerator generator) {
struct AutostartGenerator {
Future operator()() {
if (first_future->is_valid()) {
Future result = *first_future;
*first_future = Future();
return result;
}
return source();
}
std::shared_ptr> first_future;
AsyncGenerator source;
};
std::shared_ptr> first_future = std::make_shared>(generator());
return AutostartGenerator{std::move(first_future), std::move(generator)};
}
/// \see MakeReadaheadGenerator
template
class ReadaheadGenerator {
public:
ReadaheadGenerator(AsyncGenerator source_generator, int max_readahead)
: state_(std::make_shared(std::move(source_generator), max_readahead)) {}
Future AddMarkFinishedContinuation(Future fut) {
auto state = state_;
return fut.Then(
[state](const T& result) -> Result {
state->MarkFinishedIfDone(result);
return result;
},
[state](const Status& err) -> Result {
state->finished.store(true);
return err;
});
}
Future operator()() {
if (state_->readahead_queue.empty()) {
// This is the first request, let's pump the underlying queue
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());
} else {
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 source_generator, int max_readahead)
: source_generator(std::move(source_generator)), max_readahead(max_readahead) {
finished.store(false);
}
void MarkFinishedIfDone(const T& next_result) {
if (IsIterationEnd(next_result)) {
finished.store(true);
}
}
AsyncGenerator source_generator;
int max_readahead;
std::atomic finished;
std::queue> readahead_queue;
};
std::shared_ptr state_;
};
/// \brief A generator where the producer pushes items on a queue.
///
/// No back-pressure is applied, so this generator is mostly useful when
/// producing the values is neither CPU- nor memory-expensive (e.g. fetching
/// filesystem metadata).
///
/// This generator is not async-reentrant.
template
class PushGenerator {
struct State {
explicit State(util::BackpressureOptions backpressure)
: backpressure(std::move(backpressure)) {}
void OpenBackpressureIfFreeUnlocked(util::Mutex::Guard&& guard) {
if (backpressure.toggle && result_q.size() < backpressure.resume_if_below) {
// Open might trigger callbacks so release the lock first
guard.Unlock();
backpressure.toggle->Open();
}
}
void CloseBackpressureIfFullUnlocked() {
if (backpressure.toggle && result_q.size() > backpressure.pause_if_above) {
backpressure.toggle->Close();
}
}
util::BackpressureOptions backpressure;
util::Mutex mutex;
std::deque> result_q;
util::optional> consumer_fut;
bool finished = false;
};
public:
/// Producer API for PushGenerator
class Producer {
public:
explicit Producer(const std::shared_ptr& 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 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));
state->CloseBackpressureIfFullUnlocked();
}
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::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 weak_state_;
};
explicit PushGenerator(util::BackpressureOptions backpressure = {})
: state_(std::make_shared(std::move(backpressure))) {}
/// Read an item from the queue
Future operator()() const {
auto lock = state_->mutex.Lock();
assert(!state_->consumer_fut.has_value()); // Non-reentrant
if (!state_->result_q.empty()) {
auto fut = Future::MakeFinished(std::move(state_->result_q.front()));
state_->result_q.pop_front();
state_->OpenBackpressureIfFreeUnlocked(std::move(lock));
return fut;
}
if (state_->finished) {
return AsyncGeneratorEnd();
}
auto fut = Future::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_;
};
/// \brief Create a generator that pulls reentrantly from a source
/// This generator will pull reentrantly from a source, ensuring that max_readahead
/// requests are active at any given time.
///
/// The source generator must be async-reentrant
///
/// This generator itself is async-reentrant.
///
/// This generator may queue up to max_readahead instances of T
template
AsyncGenerator MakeReadaheadGenerator(AsyncGenerator source_generator,
int max_readahead) {
return ReadaheadGenerator(std::move(source_generator), max_readahead);
}
/// \brief Creates a generator that will yield finished futures from a vector
///
/// This generator is async-reentrant
template
AsyncGenerator MakeVectorGenerator(std::vector vec) {
struct State {
explicit State(std::vector vec_) : vec(std::move(vec_)), vec_idx(0) {}
std::vector vec;
std::atomic vec_idx;
};
auto state = std::make_shared(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();
}
return Future::MakeFinished(state->vec[idx]);
};
}
/// \see MakeMergedGenerator
template
class MergedGenerator {
public:
explicit MergedGenerator(AsyncGenerator> source,
int max_subscriptions)
: state_(std::make_shared(std::move(source), max_subscriptions)) {}
Future operator()() {
Future waiting_future;
std::shared_ptr delivered_job;
{
auto guard = state_->mutex.Lock();
if (!state_->delivered_jobs.empty()) {
delivered_job = std::move(state_->delivered_jobs.front());
state_->delivered_jobs.pop_front();
} else if (state_->finished) {
return IterationTraits::End();
} else {
waiting_future = Future::Make();
state_->waiting_jobs.push_back(std::make_shared>(waiting_future));
}
}
if (delivered_job) {
// deliverer will be invalid if outer callback encounters an error and delivers a
// failed result
if (delivered_job->deliverer) {
delivered_job->deliverer().AddCallback(
InnerCallback{state_, delivered_job->index});
}
return std::move(delivered_job->value);
}
if (state_->first) {
state_->first = false;
for (std::size_t i = 0; i < state_->active_subscriptions.size(); i++) {
state_->PullSource().AddCallback(OuterCallback{state_, i});
}
}
return waiting_future;
}
private:
struct DeliveredJob {
explicit DeliveredJob(AsyncGenerator deliverer_, Result value_,
std::size_t index_)
: deliverer(deliverer_), value(std::move(value_)), index(index_) {}
AsyncGenerator deliverer;
Result value;
std::size_t index;
};
struct State {
State(AsyncGenerator> source, int max_subscriptions)
: source(std::move(source)),
active_subscriptions(max_subscriptions),
delivered_jobs(),
waiting_jobs(),
mutex(),
first(true),
source_exhausted(false),
finished(false),
num_active_subscriptions(max_subscriptions) {}
Future> PullSource() {
// Need to guard access to source() so we don't pull sync-reentrantly which
// is never valid.
auto lock = mutex.Lock();
return source();
}
AsyncGenerator> source;
// active_subscriptions and delivered_jobs will be bounded by max_subscriptions
std::vector> active_subscriptions;
std::deque> delivered_jobs;
// waiting_jobs is unbounded, reentrant pulls (e.g. AddReadahead) will provide the
// backpressure
std::deque>> waiting_jobs;
util::Mutex mutex;
bool first;
bool source_exhausted;
bool finished;
int num_active_subscriptions;
};
struct InnerCallback {
void operator()(const Result& maybe_next_ref) {
Future next_fut;
const Result* maybe_next = &maybe_next_ref;
while (true) {
Future sink;
bool sub_finished = maybe_next->ok() && IsIterationEnd(**maybe_next);
{
auto guard = state->mutex.Lock();
if (state->finished) {
// We've errored out so just ignore this result and don't keep pumping
return;
}
if (!sub_finished) {
if (state->waiting_jobs.empty()) {
state->delivered_jobs.push_back(std::make_shared(
state->active_subscriptions[index], *maybe_next, index));
} else {
sink = std::move(*state->waiting_jobs.front());
state->waiting_jobs.pop_front();
}
}
}
if (sub_finished) {
state->PullSource().AddCallback(OuterCallback{state, index});
} else if (sink.is_valid()) {
sink.MarkFinished(*maybe_next);
if (!maybe_next->ok()) return;
next_fut = state->active_subscriptions[index]();
if (next_fut.TryAddCallback([this]() { return *this; })) {
return;
}
// Already completed. Avoid very deep recursion by looping
// here instead of relying on the callback.
maybe_next = &next_fut.result();
continue;
}
return;
}
}
std::shared_ptr state;
std::size_t index;
};
struct OuterCallback {
void operator()(const Result>& maybe_next) {
bool should_purge = false;
bool should_continue = false;
Future error_sink;
{
auto guard = state->mutex.Lock();
if (!maybe_next.ok() || IsIterationEnd(*maybe_next)) {
state->source_exhausted = true;
if (!maybe_next.ok() || --state->num_active_subscriptions == 0) {
state->finished = true;
should_purge = true;
}
if (!maybe_next.ok()) {
if (state->waiting_jobs.empty()) {
state->delivered_jobs.push_back(std::make_shared(
AsyncGenerator(), maybe_next.status(), index));
} else {
error_sink = std::move(*state->waiting_jobs.front());
state->waiting_jobs.pop_front();
}
}
} else {
state->active_subscriptions[index] = *maybe_next;
should_continue = true;
}
}
if (error_sink.is_valid()) {
error_sink.MarkFinished(maybe_next.status());
}
if (should_continue) {
(*maybe_next)().AddCallback(InnerCallback{state, index});
} else if (should_purge) {
// At this point state->finished has been marked true so no one else
// will be interacting with waiting_jobs and we can iterate outside lock
while (!state->waiting_jobs.empty()) {
state->waiting_jobs.front()->MarkFinished(IterationTraits::End());
state->waiting_jobs.pop_front();
}
}
}
std::shared_ptr state;
std::size_t index;
};
std::shared_ptr state_;
};
/// \brief Create a generator that takes in a stream of generators and pulls from up to
/// max_subscriptions at a time
///
/// Note: This may deliver items out of sequence. For example, items from the third
/// AsyncGenerator generated by the source may be emitted before some items from the first
/// AsyncGenerator generated by the source.
///
/// This generator will pull from source async-reentrantly unless max_subscriptions is 1
/// This generator will not pull from the individual subscriptions reentrantly. Add
/// readahead to the individual subscriptions if that is desired.
/// This generator is async-reentrant
///
/// This generator may queue up to max_subscriptions instances of T
template
AsyncGenerator MakeMergedGenerator(AsyncGenerator> source,
int max_subscriptions) {
return MergedGenerator(std::move(source), max_subscriptions);
}
template
Result> MakeSequencedMergedGenerator(
AsyncGenerator> source, int max_subscriptions) {
if (max_subscriptions < 0) {
return Status::Invalid("max_subscriptions must be a positive integer");
}
if (max_subscriptions == 1) {
return Status::Invalid("Use MakeConcatenatedGenerator if max_subscriptions is 1");
}
AsyncGenerator> autostarting_source = MakeMappedGenerator(
std::move(source),
[](const AsyncGenerator& sub) { return MakeAutoStartingGenerator(sub); });
AsyncGenerator> sub_readahead =
MakeSerialReadaheadGenerator(std::move(autostarting_source), max_subscriptions - 1);
return MakeConcatenatedGenerator(std::move(sub_readahead));
}
/// \brief Create a generator that takes in a stream of generators and pulls from each
/// one in sequence.
///
/// This generator is async-reentrant but will never pull from source reentrantly and
/// will never pull from any subscription reentrantly.
///
/// This generator may queue 1 instance of T
///
/// TODO: Could potentially make a bespoke implementation instead of MergedGenerator that
/// forwards async-reentrant requests instead of buffering them (which is what
/// MergedGenerator does)
template
AsyncGenerator MakeConcatenatedGenerator(AsyncGenerator> source) {
return MergedGenerator(std::move(source), 1);
}
template
struct Enumerated {
T value;
int index;
bool last;
};
template
struct IterationTraits> {
static Enumerated End() { return Enumerated{IterationEnd