This is an automated email from the ASF dual-hosted git repository.
zanmato1984 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 065db5f6e6 GH-51495: [C++] Fix race in MergedGenerator that could drop
an error and end the stream early (#51498)
065db5f6e6 is described below
commit 065db5f6e669befd11632c06bfd697b9b9f94fe2
Author: Kita Renji <[email protected]>
AuthorDate: Tue Sep 29 00:18:51 2026 +0700
GH-51495: [C++] Fix race in MergedGenerator that could drop an error and
end the stream early (#51498)
### Rationale for this change
`MergedGenerator` could turn an error from a subscription into a normal
end-of-stream. When an inner or outer subscription failed while no caller was
waiting, the callback set `broken = true` under the mutex but only stored the
error in `final_error` after releasing it. A call to `operator()` in that
window saw `broken` with an OK `final_error` and returned `IterationEnd`, and
the error was never delivered.
In the dataset scanner this shows up as `ToTable()` returning a truncated
table without raising when a fragment cannot be opened (#51495).
### What changes are included in this PR?
- The first error is recorded in the same locked section that sets
`broken`, via `SetFinalErrorUnlocked` (replacing `MarkFinalError`), together
with the caller that will receive it, if one is already waiting. Otherwise the
error goes to the next caller that asks, so a concurrent pull that sees
`broken` also sees the error.
- `all_finished` is removed, and the remaining futures are no longer
completed through callbacks on it. `Future` callbacks have no guaranteed order
(`future.h`): one added while a future is being marked finished may run
immediately. With every terminal item chained on `all_finished`, a pull made
during that dispatch could complete before an earlier future: the one receiving
the error (pointed out in review), or an earlier `IterationEnd`.
- Terminal pulls made once the generator is broken or exhausted are
queued in `waiting_jobs`.
- `MarkFinishedAndPurge` completes the remaining futures itself, in the
order they were handed out: the future receiving the error first, then the
waiting callers and terminal items, which get `IterationEnd`. Pulls that arrive
meanwhile are queued behind them, and the queue is drained until it stays empty.
- A caller that asks from the callbacks of the last pending future still
gets an already-finished `IterationEnd`, as before.
- A private, test-only static hook,
`MergedGenerator<T>::error_signaled_hook_for_testing`, that runs right after
the generator enters its error state and the mutex is released. It is empty by
default and only reachable from the `MergedGeneratorErrorHookTest` fixture,
which is a friend. Without it I couldn't find a way to test this
deterministically, since no user code runs inside the window.
- Tests, for both the inner and outer error paths:
-
`MergedGeneratorErrorHookTest.{Inner,Outer}ErrorNotLostToConcurrentPull`:
nobody waiting; a pull from inside the hook must raise.
-
`MergedGeneratorErrorHookTest.{Inner,Outer}ErrorToWaiterNotOvertakenByLaterPull`:
a caller already waiting must get the error before a later pull gets
end-of-stream.
-
`MergedGeneratorErrorHookTest.{Inner,Outer}ErrorToWaiterNotOvertakenDuringCompletion`
and `ClaimedErrorNotOvertakenDuringCompletion`: the later pull is made while
the generator is completing its futures, with the earlier future's completion
held up (via `TryAddCallback`, as in the review).
- `MergedGeneratorTest.PullFromLastFutureCallbackCompletesAtOnce`.
- `MergedGeneratorStressTest.TerminalNeverOvertakesEarlierFutures`: inner
and outer failures on the thread pool; checks that no terminal item completes
while an earlier future is pending and that exactly one error is delivered.
### Are these changes tested?
Yes. `arrow-async-utility-test`.
Latest commit (macOS arm64): all 138 tests pass. The `MergedGenerator`
tests pass 300/300 with `--gtest_repeat=300` (Debug) and 30/30 under TSAN with
no reports. The three `*DuringCompletion` tests fail every time on the previous
commit, and the stress test found the `IterationEnd` ordering issue before the
fix.
Original commit (Linux):
- With this PR: all 131 tests pass. The two new tests also passed 1000/1000
with `--gtest_repeat=1000`.
- With only the `final_error` change reverted (hook kept): both new tests
fail every time (2000/2000 over `--gtest_repeat=1000`) with `Expected
'_fut.status()' to fail with Invalid, but got OK`. The other 129 tests still
pass.
I also checked the fix end to end outside the test suite. I built
`libarrow_dataset` from the 25.0.1 tag with the same header change and swapped
it into the 25.0.1 wheel. The Python reproducer from the issue (going through a
`PyFileSystem` wrapper, which is what made it reproducible for me) went from
57/4500 truncated tables to 0 in about 11.6k runs.
### Are there any user-facing changes?
A scan that hits a fragment error now always raises. Before, it could
sometimes return a partial result.
**This PR contains a "Critical Fix".** It fixes a bug where a dataset scan
could return incorrect (truncated) results with no error.
### Was AI used for this PR?
In accordance to the [AI generation
guidelines](https://arrow.apache.org/docs/dev/developers/overview.html#ai-generated-code),
please disclose below whether and how AI was used in this PR.
**PR code and description written by:**
- [ ] Human
- [x] AI
**Reviewed before submission by:**
- [x] Human
- [ ] AI
- [ ] Not reviewed
🤖 Generated with [Claude Code](https://claude.com/claude-code)
* GitHub Issue: #51495
Lead-authored-by: Panchawat <[email protected]>
Co-authored-by: Panchawat <[email protected]>
Signed-off-by: Rossi Sun <[email protected]>
---
cpp/src/arrow/util/async_generator.h | 195 ++++++++++++----
cpp/src/arrow/util/async_generator_test.cc | 349 +++++++++++++++++++++++++++++
2 files changed, 499 insertions(+), 45 deletions(-)
diff --git a/cpp/src/arrow/util/async_generator.h
b/cpp/src/arrow/util/async_generator.h
index 8b2c4dc54d..1e8c72d2a5 100644
--- a/cpp/src/arrow/util/async_generator.h
+++ b/cpp/src/arrow/util/async_generator.h
@@ -1080,13 +1080,28 @@ class MergedGenerator {
} else if (state_->broken ||
(!state_->first && state_->num_running_subscriptions == 0)) {
// If we are broken or exhausted then prepare a terminal item but
- // we won't complete it until we've finished.
- Result<T> end_res = IterationEnd<T>();
- if (!state_->final_error.ok()) {
- end_res = state_->final_error;
- state_->final_error = Status::OK();
+ // we won't complete it until we've finished, and not before any
future we
+ // returned earlier.
+ if (!state_->final_error.ok() && !state_->error_sink.is_valid()) {
+ // Nobody has claimed the error yet, so this caller receives it.
Nothing
+ // returned earlier is still waiting: that would have received the
error.
+ if (state_->finishing) {
+ Status err = std::move(state_->final_error);
+ state_->final_error = Status::OK();
+ return Future<T>::MakeFinished(std::move(err));
+ }
+ state_->error_sink = Future<T>::Make();
+ return state_->error_sink;
+ }
+ if (state_->AllCompletedUnlocked(guard)) {
+ // Every future returned earlier has completed
+ return Future<T>::MakeFinished(IterationEnd<T>());
}
- return state_->all_finished.Then([end_res]() -> Result<T> { return
end_res; });
+ // Queue it behind the waiting callers, MarkFinishedAndPurge completes
them in
+ // order
+ auto terminal = Future<T>::Make();
+ state_->waiting_jobs.push_back(std::make_shared<Future<T>>(terminal));
+ return terminal;
} else {
// Otherwise we just queue the request and it will be completed when
one of the
// ongoing inner subscriptions delivers a result
@@ -1108,7 +1123,7 @@ class MergedGenerator {
// subscription that delivered it (deliverer).
if (delivered_job) {
if (mark_generator_complete) {
- state_->all_finished.MarkFinished();
+ state_->MarkFinishedAndPurge();
} else {
delivered_job->deliverer().AddCallback(
InnerCallback(state_, delivered_job->index));
@@ -1149,6 +1164,14 @@ class MergedGenerator {
}
private:
+ friend class MergedGeneratorErrorHookTest;
+
+ // Test-only hook. If set, it is called right after an error from an inner
or outer
+ // subscription has put the generator in its errored state and the mutex has
been
+ // released, so tests can pull from the generator at that exact point. Only
set it
+ // while no MergedGenerator<T> is in use.
+ static inline std::function<void()> error_signaled_hook_for_testing;
+
struct DeliveredJob {
explicit DeliveredJob(AsyncGenerator<T> deliverer_, Result<T> value_,
std::size_t index_)
@@ -1193,30 +1216,97 @@ class MergedGenerator {
}
}
- // This function is called outside the mutex but it will only ever be
- // called once
+ // This function is called outside the mutex, once all outstanding work is
done, and
+ // it will only ever be called once.
+ //
+ // It completes the remaining futures in the order they were handed out:
the one
+ // that receives the error, then the waiting callers (including terminal
items
+ // requested since), which get the end of the stream. Callers that ask
while this
+ // runs are queued behind them, so the queue is drained until it stays
empty. This
+ // ordering cannot be left to callbacks: Future callbacks do not run in a
guaranteed
+ // order, and one added while a future is being marked finished may run at
once.
void MarkFinishedAndPurge() {
- all_finished.MarkFinished();
- while (!waiting_jobs.empty()) {
- waiting_jobs.front()->MarkFinished(IterationEnd<T>());
- waiting_jobs.pop_front();
+ std::deque<std::shared_ptr<Future<T>>> to_complete;
+ Status err;
+ {
+ auto guard = mutex.Lock();
+ finishing = true;
+ if (error_sink.is_valid()) {
+
to_complete.push_back(std::make_shared<Future<T>>(std::move(error_sink)));
+ error_sink = Future<T>();
+ err = std::move(final_error);
+ final_error = Status::OK();
+ }
+ for (auto& waiting_job : waiting_jobs) {
+ to_complete.push_back(std::move(waiting_job));
+ }
+ waiting_jobs.clear();
+ if (to_complete.empty()) {
+ // Nothing is pending
+ purged = true;
+ return;
+ }
+ }
+ auto complete = [&err](const std::shared_ptr<Future<T>>& fut) {
+ if (err.ok()) {
+ fut->MarkFinished(IterationEnd<T>());
+ } else {
+ fut->MarkFinished(std::move(err));
+ err = Status::OK();
+ }
+ };
+ while (true) {
+ std::shared_ptr<Future<T>> last;
+ if (!to_complete.empty()) {
+ last = std::move(to_complete.back());
+ to_complete.pop_back();
+ }
+ for (const auto& fut : to_complete) {
+ complete(fut);
+ }
+ to_complete.clear();
+ if (last) {
+ {
+ // If nobody asked meanwhile, `last` is the only future still
pending, so a
+ // caller that asks once it has completed, e.g. from its
callbacks, can be
+ // given its terminal item at once, as it would be after we are
done
+ auto guard = mutex.Lock();
+ if (waiting_jobs.empty()) {
+ completing_last = last;
+ }
+ }
+ complete(last);
+ }
+ auto guard = mutex.Lock();
+ completing_last.reset();
+ if (waiting_jobs.empty()) {
+ purged = true;
+ return;
+ }
+ to_complete.swap(waiting_jobs);
}
}
- // This is called outside the mutex but it is only ever called
- // once and Future<>::AddCallback is thread-safe
- void MarkFinalError(const Status& err, Future<T> maybe_sink) {
- if (maybe_sink.is_valid()) {
- // Someone is waiting for this error so lets mark it complete when
- // all the work is done
- all_finished.AddCallback([maybe_sink, err](const Status& status)
mutable {
- maybe_sink.MarkFinished(err);
- });
- } else {
- // No one is waiting for this error right now so it will be delivered
- // next.
- final_error = err;
- }
+ // True if every future handed out so far has completed. Must be called
with the
+ // mutex held, once we are broken or exhausted.
+ bool AllCompletedUnlocked(const util::Mutex::Guard& guard) {
+ return purged ||
+ (waiting_jobs.empty() && completing_last &&
completing_last->is_finished());
+ }
+
+ // Must be called with the mutex held, when the first error arrives, in
the same
+ // locked section that sets `broken`: a concurrent caller that sees
`broken` must
+ // also see the error, or it would get a plain end of stream and the error
would be
+ // silently dropped (GH-51495).
+ //
+ // `sink` is the caller waiting for the item that failed, if any. If
there is none,
+ // the error goes to the next caller. Either way it is delivered once all
+ // outstanding work is done, by MarkFinishedAndPurge (or directly, to a
caller that
+ // asks after that has started).
+ void SetFinalErrorUnlocked(const util::Mutex::Guard& guard, const Status&
err,
+ Future<T> sink) {
+ final_error = err;
+ error_sink = std::move(sink);
}
bool IsCompleteUnlocked(const util::Mutex::Guard& guard) {
@@ -1239,12 +1329,9 @@ class MergedGenerator {
// caller
std::deque<std::shared_ptr<DeliveredJob>> delivered_jobs;
// waiting_jobs is unbounded, reentrant pulls (e.g. AddReadahead) will
provide the
- // backpressure
+ // backpressure. Once we are broken or exhausted it also holds the
terminal items
+ // requested until MarkFinishedAndPurge completes them.
std::deque<std::shared_ptr<Future<T>>> waiting_jobs;
- // A future that will be marked complete when the terminal item has
arrived and all
- // outstanding futures have completed. It is used to hold off emission of
an error
- // until all outstanding work is done.
- Future<> all_finished = Future<>::Make();
util::Mutex mutex;
// A flag cleared when the caller firsts asks for a future. Used to start
polling.
bool first;
@@ -1255,7 +1342,7 @@ class MergedGenerator {
// are finishing up.
bool source_exhausted;
// The number of futures that we have requested from either the outer or
inner
- // subscriptions that have not yet completed. We cannot mark all_finished
until this
+ // subscriptions that have not yet completed. We cannot finish until this
// reaches 0. This will never be greater than max_subscriptions
int outstanding_requests;
// The number of running subscriptions. We ramp this up to
`max_subscriptions` as
@@ -1264,12 +1351,28 @@ class MergedGenerator {
// subscription is exhausted at which point this descends to 0 (and
source_exhausted)
// is then set to true.
int num_running_subscriptions;
- // If an error arrives, and the caller hasn't asked for that item, we
store the error
- // here. It is analagous to delivered_jobs but for errors instead of
finished
- // results.
+ // The first error, until it is delivered. It is analagous to
delivered_jobs but for
+ // errors instead of finished results. Guarded by `mutex`.
Status final_error;
+ // The caller that will receive `final_error`, once one has asked for it.
Guarded
+ // by `mutex`.
+ Future<T> error_sink;
+ // Set once MarkFinishedAndPurge has started. Guarded by `mutex`.
+ bool finishing = false;
+ // Set once MarkFinishedAndPurge has completed every future handed out so
far.
+ // Guarded by `mutex`.
+ bool purged = false;
+ // While MarkFinishedAndPurge completes what it knows to be the last
pending future,
+ // that future. Guarded by `mutex`.
+ std::shared_ptr<Future<T>> completing_last;
};
+ static void RunErrorSignaledHookForTesting() {
+ if (error_signaled_hook_for_testing) {
+ error_signaled_hook_for_testing();
+ }
+ }
+
struct InnerCallback {
InnerCallback(std::shared_ptr<State> state, std::size_t index, bool
recursive = false)
: state(std::move(state)), index(index), recursive(recursive) {}
@@ -1289,7 +1392,7 @@ class MergedGenerator {
bool pull_next_sub = false;
bool was_broken = false;
bool should_mark_gen_complete = false;
- bool should_mark_final_error = false;
+ bool signaled_error = false;
{
auto guard = state->mutex.Lock();
if (state->broken) {
@@ -1311,8 +1414,9 @@ class MergedGenerator {
// If this is the first error then we transition the state to a
broken state
if (!maybe_next->ok()) {
- should_mark_final_error = true;
+ signaled_error = true;
state->SignalErrorUnlocked(guard);
+ state->SetFinalErrorUnlocked(guard, maybe_next->status(),
std::move(sink));
}
}
@@ -1338,10 +1442,9 @@ class MergedGenerator {
// Now we have given up the lock and we can take all the actions we
decided we
// need to take.
- if (should_mark_final_error) {
- state->MarkFinalError(maybe_next->status(), std::move(sink));
+ if (signaled_error) {
+ RunErrorSignaledHookForTesting();
}
-
if (should_mark_gen_complete) {
state->MarkFinishedAndPurge();
}
@@ -1391,7 +1494,7 @@ class MergedGenerator {
// We have been given a new inner subscription
bool should_continue = false;
bool should_mark_gen_complete = false;
- bool should_deliver_error = false;
+ bool signaled_error = false;
bool source_exhausted = maybe_next.ok() && IsIterationEnd(*maybe_next);
Future<T> error_sink;
{
@@ -1399,13 +1502,15 @@ class MergedGenerator {
if (!maybe_next.ok() || source_exhausted || state->broken) {
// If here then we will not pull any more from the outer source
if (!state->broken && !maybe_next.ok()) {
+ signaled_error = true;
state->SignalErrorUnlocked(guard);
// If here then we are the first error so we need to deliver it
- should_deliver_error = true;
if (!state->waiting_jobs.empty()) {
error_sink = std::move(*state->waiting_jobs.front());
state->waiting_jobs.pop_front();
}
+ state->SetFinalErrorUnlocked(guard, maybe_next.status(),
+ std::move(error_sink));
}
if (source_exhausted) {
state->source_exhausted = true;
@@ -1419,8 +1524,8 @@ class MergedGenerator {
should_continue = true;
}
}
- if (should_deliver_error) {
- state->MarkFinalError(maybe_next.status(), std::move(error_sink));
+ if (signaled_error) {
+ RunErrorSignaledHookForTesting();
}
if (should_mark_gen_complete) {
state->MarkFinishedAndPurge();
diff --git a/cpp/src/arrow/util/async_generator_test.cc
b/cpp/src/arrow/util/async_generator_test.cc
index afb03b6720..9d965dbc4e 100644
--- a/cpp/src/arrow/util/async_generator_test.cc
+++ b/cpp/src/arrow/util/async_generator_test.cc
@@ -18,12 +18,15 @@
#include <atomic>
#include <chrono>
#include <condition_variable>
+#include <functional>
#include <mutex>
#include <optional>
#include <random>
+#include <string>
#include <thread>
#include <unordered_set>
#include <utility>
+#include <vector>
#include "arrow/io/slow.h"
#include "arrow/testing/async_test_util.h"
@@ -791,6 +794,352 @@ TEST_P(MergedGeneratorTestFixture,
DeepOuterGeneratorStackOverflow) {
INSTANTIATE_TEST_SUITE_P(MergedGeneratorTests, MergedGeneratorTestFixture,
::testing::Values(false, true));
+// GH-51495: when an inner or outer subscription fails while no caller is
waiting, the
+// error must be visible to the next pull as soon as the generator is in its
errored
+// state. Previously the error was stored only after the internal mutex had
been
+// released, so a pull landing in between got a plain end-of-stream and the
error was
+// lost. The test hook runs exactly in that window, which makes these tests
+// deterministic.
+class MergedGeneratorErrorHookTest : public ::testing::Test {
+ protected:
+ void TearDown() override {
+ MergedGenerator<TestInt>::error_signaled_hook_for_testing = nullptr;
+ }
+
+ // Pull from `merged` once, from inside the hook
+ void PullFromErrorHook(AsyncGenerator<TestInt>* merged, Future<TestInt>*
pulled) {
+ MergedGenerator<TestInt>::error_signaled_hook_for_testing = [merged,
pulled]() {
+ if (!pulled->is_valid()) {
+ *pulled = (*merged)();
+ }
+ };
+ }
+
+ // Fails a pending item by calling `fail` on another thread, and pulls again
while
+ // `gen` is completing `*earlier`: `earlier`'s internal mutex is held
meanwhile, so
+ // completing it blocks, and the later pull is made once `gen` has started
completing
+ // its remaining futures. If `pull_earlier_in_hook`, `*earlier` is pulled
from the
+ // error hook, i.e. it is the caller that claims an error nobody was waiting
for.
+ // Returns true if the later pull completed while `*earlier` was still
pending.
+ bool LaterPullCompletesFirst(const MergedGenerator<TestInt>& gen,
+ AsyncGenerator<TestInt>* merged,
Future<TestInt>* earlier,
+ std::function<void()> fail, bool
pull_earlier_in_hook,
+ Future<TestInt>* later) {
+ std::atomic<bool> signaled{false};
+ std::atomic<bool> holding{false};
+ MergedGenerator<TestInt>::error_signaled_hook_for_testing = [&]() {
+ if (pull_earlier_in_hook) {
+ *earlier = (*merged)();
+ }
+ signaled = true;
+ while (!holding) {
+ std::this_thread::yield();
+ }
+ };
+ std::thread failer(std::move(fail));
+ while (!signaled) {
+ std::this_thread::yield();
+ }
+ bool later_completed_first = false;
+ bool held = earlier->TryAddCallback([&]() {
+ holding = true;
+ // Relies on TryAddCallback running this factory under `earlier`'s
internal mutex
+ const auto deadline = std::chrono::steady_clock::now() +
std::chrono::seconds(10);
+ while (!IsFinishing(gen)) {
+ if (std::chrono::steady_clock::now() > deadline) {
+ ADD_FAILURE() << "generator never started completing its futures";
+ break;
+ }
+ std::this_thread::yield();
+ }
+ *later = (*merged)();
+ later_completed_first = later->is_finished();
+ return [](const Result<TestInt>&) {};
+ });
+ if (!held) {
+ holding = true;
+ }
+ failer.join();
+ EXPECT_TRUE(held) << "earlier future completed before the later pull";
+ return later_completed_first;
+ }
+
+ static bool IsFinishing(const MergedGenerator<TestInt>& gen) {
+ auto guard = gen.state_->mutex.Lock();
+ return gen.state_->finishing;
+ }
+
+ // Like PullFromErrorHook, but also records "terminal" in `order` when the
pulled
+ // future completes
+ void PullFromErrorHookAndRecord(AsyncGenerator<TestInt>* merged,
+ Future<TestInt>* pulled,
+ std::vector<std::string>* order) {
+ MergedGenerator<TestInt>::error_signaled_hook_for_testing = [=]() {
+ if (!pulled->is_valid()) {
+ *pulled = (*merged)();
+ pulled->AddCallback(
+ [order](const Result<TestInt>&) { order->push_back("terminal"); });
+ }
+ };
+ }
+};
+
+TEST_F(MergedGeneratorErrorHookTest, InnerErrorNotLostToConcurrentPull) {
+ auto failing = Future<TestInt>::Make();
+ AsyncGenerator<TestInt> failing_sub = [failing]() { return failing; };
+ std::vector<AsyncGenerator<TestInt>> subs =
{MakeVectorGenerator<TestInt>({TestInt(1)}),
+ failing_sub};
+ auto merged = MakeMergedGenerator(MakeVectorGenerator(std::move(subs)), 1);
+ // Delivers 1 and then subscribes to failing_sub, whose first item is
pending. Now
+ // there is one outstanding request and nobody waiting.
+ ASSERT_FINISHES_OK_AND_EQ(TestInt(1), merged());
+
+ Future<TestInt> pulled;
+ PullFromErrorHook(&merged, &pulled);
+ failing.MarkFinished(Status::Invalid("XYZ"));
+ ASSERT_TRUE(pulled.is_valid());
+ ASSERT_FINISHES_AND_RAISES(Invalid, pulled);
+ AssertGeneratorExhausted(merged);
+}
+
+TEST_F(MergedGeneratorErrorHookTest, OuterErrorNotLostToConcurrentPull) {
+ auto failing = Future<AsyncGenerator<TestInt>>::Make();
+ int num_pulls = 0;
+ AsyncGenerator<AsyncGenerator<TestInt>> source =
+ [&]() -> Future<AsyncGenerator<TestInt>> {
+ if (num_pulls++ == 0) {
+ return Future<AsyncGenerator<TestInt>>::MakeFinished(
+ MakeVectorGenerator<TestInt>({TestInt(1)}));
+ }
+ return failing;
+ };
+ auto merged = MakeMergedGenerator(std::move(source), 1);
+ // Delivers 1 and then pulls the next subscription from the source, which is
pending.
+ // Now there is one outstanding request and nobody waiting.
+ ASSERT_FINISHES_OK_AND_EQ(TestInt(1), merged());
+
+ Future<TestInt> pulled;
+ PullFromErrorHook(&merged, &pulled);
+ failing.MarkFinished(Status::Invalid("XYZ"));
+ ASSERT_TRUE(pulled.is_valid());
+ ASSERT_FINISHES_AND_RAISES(Invalid, pulled);
+ AssertGeneratorExhausted(merged);
+}
+
+// When a caller is already waiting for the item that fails, it must receive
the error
+// before any later pull receives its end-of-stream: the error callback has to
be
+// registered in the same locked transition that marks the generator errored,
or a
+// pull landing in between registers its terminal continuation first and
overtakes it.
+TEST_F(MergedGeneratorErrorHookTest,
InnerErrorToWaiterNotOvertakenByLaterPull) {
+ auto failing = Future<TestInt>::Make();
+ AsyncGenerator<TestInt> failing_sub = [failing]() { return failing; };
+ std::vector<AsyncGenerator<TestInt>> subs = {failing_sub};
+ auto merged = MakeMergedGenerator(MakeVectorGenerator(std::move(subs)), 1);
+ std::vector<std::string> order;
+ // Subscribes to failing_sub, whose first item is pending, so this caller
waits
+ Future<TestInt> waiting = merged();
+ waiting.AddCallback([&order](const Result<TestInt>&) {
order.push_back("error"); });
+
+ Future<TestInt> pulled;
+ PullFromErrorHookAndRecord(&merged, &pulled, &order);
+ failing.MarkFinished(Status::Invalid("XYZ"));
+ ASSERT_TRUE(pulled.is_valid());
+ ASSERT_FINISHES_AND_RAISES(Invalid, waiting);
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto terminal, pulled);
+ ASSERT_TRUE(IsIterationEnd(terminal));
+ ASSERT_EQ(order, (std::vector<std::string>{"error", "terminal"}));
+}
+
+TEST_F(MergedGeneratorErrorHookTest,
OuterErrorToWaiterNotOvertakenByLaterPull) {
+ auto failing = Future<AsyncGenerator<TestInt>>::Make();
+ AsyncGenerator<AsyncGenerator<TestInt>> source = [failing]() { return
failing; };
+ auto merged = MakeMergedGenerator(std::move(source), 1);
+ std::vector<std::string> order;
+ // Pulls the first subscription from the source, which is pending, so this
caller
+ // waits
+ Future<TestInt> waiting = merged();
+ waiting.AddCallback([&order](const Result<TestInt>&) {
order.push_back("error"); });
+
+ Future<TestInt> pulled;
+ PullFromErrorHookAndRecord(&merged, &pulled, &order);
+ failing.MarkFinished(Status::Invalid("XYZ"));
+ ASSERT_TRUE(pulled.is_valid());
+ ASSERT_FINISHES_AND_RAISES(Invalid, waiting);
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto terminal, pulled);
+ ASSERT_TRUE(IsIterationEnd(terminal));
+ ASSERT_EQ(order, (std::vector<std::string>{"error", "terminal"}));
+}
+
+// Callbacks on a Future do not run in a guaranteed order: one added while the
future is
+// being marked finished may run immediately. So a terminal item must not
merely be
+// registered after the error, it must not be able to complete before the
future that
+// receives the error has completed, even when it is requested while the
generator is
+// already completing.
+TEST_F(MergedGeneratorErrorHookTest,
InnerErrorToWaiterNotOvertakenDuringCompletion) {
+ auto failing = Future<TestInt>::Make();
+ AsyncGenerator<TestInt> failing_sub = [failing]() { return failing; };
+ std::vector<AsyncGenerator<TestInt>> subs = {failing_sub};
+ MergedGenerator<TestInt> gen(MakeVectorGenerator(std::move(subs)), 1);
+ AsyncGenerator<TestInt> merged = gen;
+ Future<TestInt> waiting = merged();
+
+ Future<TestInt> later;
+ ASSERT_FALSE(LaterPullCompletesFirst(
+ gen, &merged, &waiting,
+ [failing]() mutable { failing.MarkFinished(Status::Invalid("XYZ")); },
+ /*pull_earlier_in_hook=*/false, &later));
+ ASSERT_FINISHES_AND_RAISES(Invalid, waiting);
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto terminal, later);
+ ASSERT_TRUE(IsIterationEnd(terminal));
+}
+
+TEST_F(MergedGeneratorErrorHookTest,
OuterErrorToWaiterNotOvertakenDuringCompletion) {
+ auto failing = Future<AsyncGenerator<TestInt>>::Make();
+ AsyncGenerator<AsyncGenerator<TestInt>> source = [failing]() { return
failing; };
+ MergedGenerator<TestInt> gen(std::move(source), 1);
+ AsyncGenerator<TestInt> merged = gen;
+ Future<TestInt> waiting = merged();
+
+ Future<TestInt> later;
+ ASSERT_FALSE(LaterPullCompletesFirst(
+ gen, &merged, &waiting,
+ [failing]() mutable { failing.MarkFinished(Status::Invalid("XYZ")); },
+ /*pull_earlier_in_hook=*/false, &later));
+ ASSERT_FINISHES_AND_RAISES(Invalid, waiting);
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto terminal, later);
+ ASSERT_TRUE(IsIterationEnd(terminal));
+}
+
+TEST_F(MergedGeneratorErrorHookTest, ClaimedErrorNotOvertakenDuringCompletion)
{
+ auto failing = Future<TestInt>::Make();
+ AsyncGenerator<TestInt> failing_sub = [failing]() { return failing; };
+ std::vector<AsyncGenerator<TestInt>> subs =
{MakeVectorGenerator<TestInt>({TestInt(1)}),
+ failing_sub};
+ MergedGenerator<TestInt> gen(MakeVectorGenerator(std::move(subs)), 1);
+ AsyncGenerator<TestInt> merged = gen;
+ // Delivers 1 and then subscribes to failing_sub, whose first item is
pending. Now
+ // there is one outstanding request and nobody waiting, so the error goes to
the
+ // next pull, made from the hook.
+ ASSERT_FINISHES_OK_AND_EQ(TestInt(1), merged());
+
+ Future<TestInt> claimed;
+ Future<TestInt> later;
+ ASSERT_FALSE(LaterPullCompletesFirst(
+ gen, &merged, &claimed,
+ [failing]() mutable { failing.MarkFinished(Status::Invalid("XYZ")); },
+ /*pull_earlier_in_hook=*/true, &later));
+ ASSERT_FINISHES_AND_RAISES(Invalid, claimed);
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto terminal, later);
+ ASSERT_TRUE(IsIterationEnd(terminal));
+}
+
+// A caller that asks from the callbacks of the last pending future gets its
terminal item
+// at once, so it may block on it.
+TEST(MergedGeneratorTest, PullFromLastFutureCallbackCompletesAtOnce) {
+ for (bool fail : {false, true}) {
+ ARROW_SCOPED_TRACE("fail = ", fail);
+ auto pending = Future<TestInt>::Make();
+ AsyncGenerator<TestInt> sub = [pending]() { return pending; };
+ std::vector<AsyncGenerator<TestInt>> subs = {sub};
+ auto merged = MakeMergedGenerator(MakeVectorGenerator(std::move(subs)), 1);
+ Future<TestInt> waiting = merged();
+ Future<TestInt> pulled;
+ waiting.AddCallback([&](const Result<TestInt>&) { pulled = merged(); });
+ if (fail) {
+ pending.MarkFinished(Status::Invalid("XYZ"));
+ ASSERT_FINISHES_AND_RAISES(Invalid, waiting);
+ } else {
+ pending.MarkFinished(IterationEnd<TestInt>());
+ ASSERT_FINISHES_OK_AND_ASSIGN(auto end, waiting);
+ ASSERT_TRUE(IsIterationEnd(end));
+ }
+ ASSERT_TRUE(pulled.is_valid());
+ ASSERT_TRUE(pulled.is_finished());
+ ASSERT_OK_AND_ASSIGN(auto terminal, pulled.result());
+ ASSERT_TRUE(IsIterationEnd(terminal));
+ }
+}
+
+// Stress test of the AsyncGenerator contract with inner items and errors
arriving on
+// other threads: a terminal item (an error or the end of the stream) must not
complete
+// while any future returned before it is still pending, and an error must
never be
+// lost.
+TEST(MergedGeneratorStressTest, TerminalNeverOvertakesEarlierFutures) {
+ constexpr int kIterations = 300;
+ constexpr int kNumSubs = 6;
+ constexpr int kItemsPerSub = 4;
+ constexpr int kNumPulls = kNumSubs * kItemsPerSub + 8;
+ auto pool = internal::GetCpuThreadPool();
+ for (int iter = 0; iter < kIterations; iter++) {
+ const bool outer_fails = iter % 3 == 0;
+ std::vector<AsyncGenerator<TestInt>> subs;
+ for (int i = 0; i < kNumSubs; i++) {
+ auto sub = NewBackgroundAsyncVectorIt({1, 2, 3, 4}, /*sleep=*/false);
+ if (!outer_fails && i == iter % kNumSubs) {
+ sub = MakeTransferredGenerator(util::FailAt(std::move(sub), iter %
kItemsPerSub),
+ pool);
+ }
+ subs.push_back(std::move(sub));
+ }
+ AsyncGenerator<AsyncGenerator<TestInt>> source =
MakeVectorGenerator(std::move(subs));
+ if (outer_fails) {
+ source =
+ MakeTransferredGenerator(util::FailAt(std::move(source), 1 + iter %
4), pool);
+ }
+ auto merged = MakeMergedGenerator(std::move(source), 3);
+
+ // Shared with the callbacks, which may still be running after the futures
are
+ // marked finished
+ struct Shared {
+ std::mutex mutex;
+ std::vector<Future<TestInt>> futures;
+ int overtaken = 0;
+ int callbacks_run = 0;
+ };
+ auto shared = std::make_shared<Shared>();
+ for (int i = 0; i < kNumPulls; i++) {
+ Future<TestInt> fut = merged();
+ {
+ std::lock_guard<std::mutex> lock(shared->mutex);
+ shared->futures.push_back(fut);
+ }
+ fut.AddCallback([shared, i](const Result<TestInt>& result) {
+ std::lock_guard<std::mutex> lock(shared->mutex);
+ shared->callbacks_run++;
+ if (result.ok() && !IsIterationEnd(*result)) {
+ return;
+ }
+ for (int j = 0; j < i; j++) {
+ if (!shared->futures[j].is_finished()) {
+ shared->overtaken++;
+ }
+ }
+ });
+ }
+ std::vector<Future<TestInt>> futures;
+ {
+ std::lock_guard<std::mutex> lock(shared->mutex);
+ futures = shared->futures;
+ }
+ int num_errors = 0;
+ for (const auto& fut : futures) {
+ ASSERT_TRUE(fut.Wait(kDefaultAssertFinishesWaitSeconds));
+ if (!fut.status().ok()) {
+ ASSERT_TRUE(fut.status().IsInvalid()) << fut.status();
+ num_errors++;
+ }
+ }
+ BusyWait(kDefaultAssertFinishesWaitSeconds, [&]() {
+ std::lock_guard<std::mutex> lock(shared->mutex);
+ return shared->callbacks_run == kNumPulls;
+ });
+ std::lock_guard<std::mutex> lock(shared->mutex);
+ ASSERT_EQ(shared->callbacks_run, kNumPulls) << "iteration " << iter;
+ ASSERT_EQ(shared->overtaken, 0) << "iteration " << iter;
+ ASSERT_EQ(num_errors, 1) << "iteration " << iter;
+ }
+}
+
class AutoStartingGeneratorTestFixture : public GeneratorTestFixture {};
TEST_P(AutoStartingGeneratorTestFixture, Basic) {
