From 97e8e1787b7337e92650fdedf9fce04792b985a3 Mon Sep 17 00:00:00 2001 From: Panchawat Date: Sat, 26 Sep 2026 17:58:35 +0700 Subject: [PATCH 1/3] GH-51495: [C++] Fix race in MergedGenerator that could drop an error and end the stream early When an inner or outer subscription failed while no caller was waiting, MergedGenerator set `broken` under the mutex but stored the error in `final_error` only after releasing it. A pull in that window saw `broken` with an OK `final_error` and returned end-of-stream, so the error was lost. In the dataset scanner this made ToTable() return a truncated table instead of raising when a fragment could not be opened. Store `final_error` while the mutex is held. Add a test-only hook that runs right after the error state is entered, and use it for deterministic regression tests of the inner and outer error paths. Co-Authored-By: Claude Opus 5.5 --- cpp/src/arrow/util/async_generator.h | 64 +++++++++++++++++----- cpp/src/arrow/util/async_generator_test.cc | 64 ++++++++++++++++++++++ 2 files changed, 113 insertions(+), 15 deletions(-) diff --git a/cpp/src/arrow/util/async_generator.h b/cpp/src/arrow/util/async_generator.h index 8b2c4dc54d9c..d50ab908e39b 100644 --- a/cpp/src/arrow/util/async_generator.h +++ b/cpp/src/arrow/util/async_generator.h @@ -1056,6 +1056,15 @@ class MergedGenerator { int max_subscriptions) : state_(std::make_shared(std::move(source), max_subscriptions)) {} + /// \brief Test-only hook, not part of the public API + /// + /// If set, it is called without holding the internal mutex, right after an error from + /// an inner or outer subscription has put the generator in its errored state, and + /// before that error is handed to a caller that was already waiting for it. This lets + /// tests pull from the generator at that exact point. Only set it while no + /// MergedGenerator is in use. + static inline std::function error_signaled_hook_for_testing; + Future operator()() { // A caller has requested a future Future waiting_future; @@ -1203,22 +1212,29 @@ class MergedGenerator { } } - // 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 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); - }); + // Must be called with the mutex held, when the first error arrives. + void SetFinalErrorUnlocked(const util::Mutex::Guard& guard, const Status& err, + Future* sink, bool* should_deliver_to_sink) { + if (sink->is_valid()) { + // Someone is waiting for this error, it will be delivered to them (outside + // the lock) once all outstanding work is done + *should_deliver_to_sink = true; } else { // No one is waiting for this error right now so it will be delivered - // next. + // next. This must happen under the lock: 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). final_error = err; } } + // This is called outside the mutex but it is only ever called + // once and Future<>::AddCallback is thread-safe + void DeliverFinalError(const Status& err, Future sink) { + all_finished.AddCallback( + [sink, err](const Status& status) mutable { sink.MarkFinished(err); }); + } + bool IsCompleteUnlocked(const util::Mutex::Guard& guard) { return outstanding_requests == 0 && (broken || (source_exhausted && num_running_subscriptions == 0 && @@ -1266,10 +1282,16 @@ class MergedGenerator { 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. + // results. Guarded by `mutex`. Status final_error; }; + static void RunErrorSignaledHookForTesting() { + if (error_signaled_hook_for_testing) { + error_signaled_hook_for_testing(); + } + } + struct InnerCallback { InnerCallback(std::shared_ptr state, std::size_t index, bool recursive = false) : state(std::move(state)), index(index), recursive(recursive) {} @@ -1290,6 +1312,7 @@ class MergedGenerator { 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 +1334,10 @@ 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(), &sink, + &should_mark_final_error); } } @@ -1338,8 +1363,11 @@ class MergedGenerator { // Now we have given up the lock and we can take all the actions we decided we // need to take. + if (signaled_error) { + RunErrorSignaledHookForTesting(); + } if (should_mark_final_error) { - state->MarkFinalError(maybe_next->status(), std::move(sink)); + state->DeliverFinalError(maybe_next->status(), std::move(sink)); } if (should_mark_gen_complete) { @@ -1392,6 +1420,7 @@ class MergedGenerator { 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 error_sink; { @@ -1399,13 +1428,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(), &error_sink, + &should_deliver_error); } if (source_exhausted) { state->source_exhausted = true; @@ -1419,8 +1450,11 @@ class MergedGenerator { should_continue = true; } } + if (signaled_error) { + RunErrorSignaledHookForTesting(); + } if (should_deliver_error) { - state->MarkFinalError(maybe_next.status(), std::move(error_sink)); + state->DeliverFinalError(maybe_next.status(), std::move(error_sink)); } 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 afb03b67209a..3b4637271a71 100644 --- a/cpp/src/arrow/util/async_generator_test.cc +++ b/cpp/src/arrow/util/async_generator_test.cc @@ -791,6 +791,70 @@ 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::error_signaled_hook_for_testing = nullptr; + } + + // Pull from `merged` once, from inside the hook + void PullFromErrorHook(AsyncGenerator* merged, Future* pulled) { + MergedGenerator::error_signaled_hook_for_testing = [merged, pulled]() { + if (!pulled->is_valid()) { + *pulled = (*merged)(); + } + }; + } +}; + +TEST_F(MergedGeneratorErrorHookTest, InnerErrorNotLostToConcurrentPull) { + auto failing = Future::Make(); + AsyncGenerator failing_sub = [failing]() { return failing; }; + std::vector> subs = {MakeVectorGenerator({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 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>::Make(); + int num_pulls = 0; + AsyncGenerator> source = + [&]() -> Future> { + if (num_pulls++ == 0) { + return Future>::MakeFinished( + MakeVectorGenerator({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 pulled; + PullFromErrorHook(&merged, &pulled); + failing.MarkFinished(Status::Invalid("XYZ")); + ASSERT_TRUE(pulled.is_valid()); + ASSERT_FINISHES_AND_RAISES(Invalid, pulled); + AssertGeneratorExhausted(merged); +} + class AutoStartingGeneratorTestFixture : public GeneratorTestFixture {}; TEST_P(AutoStartingGeneratorTestFixture, Basic) { From 9a48233e66aecbcb6d7f5cec16b2c65a21861414 Mon Sep 17 00:00:00 2001 From: Panchawat Date: Mon, 28 Sep 2026 11:11:39 +0700 Subject: [PATCH 2/3] GH-51495: [C++] Register waiting error before later terminal pulls in MergedGenerator When a caller was already waiting for the item that failed, the error callback was attached to `all_finished` only after the mutex was released. A pull in that window saw `broken`, attached its end-of-stream continuation to `all_finished` first, and since callbacks run in registration order the later terminal item completed before the earlier waiting future got its error. Attach the waiting caller's error callback in the same locked section that sets `broken`, for both the inner and outer error paths. This removes `DeliverFinalError`. Add regression tests for the waiter-present path. Also make `error_signaled_hook_for_testing` private and friend the test fixture instead of exposing a mutable static in the public header. Co-Authored-By: Claude Opus 5.5 (1M context) --- cpp/src/arrow/util/async_generator.h | 59 +++++++++------------ cpp/src/arrow/util/async_generator_test.cc | 60 ++++++++++++++++++++++ 2 files changed, 84 insertions(+), 35 deletions(-) diff --git a/cpp/src/arrow/util/async_generator.h b/cpp/src/arrow/util/async_generator.h index d50ab908e39b..74441692df68 100644 --- a/cpp/src/arrow/util/async_generator.h +++ b/cpp/src/arrow/util/async_generator.h @@ -1056,15 +1056,6 @@ class MergedGenerator { int max_subscriptions) : state_(std::make_shared(std::move(source), max_subscriptions)) {} - /// \brief Test-only hook, not part of the public API - /// - /// If set, it is called without holding the internal mutex, right after an error from - /// an inner or outer subscription has put the generator in its errored state, and - /// before that error is handed to a caller that was already waiting for it. This lets - /// tests pull from the generator at that exact point. Only set it while no - /// MergedGenerator is in use. - static inline std::function error_signaled_hook_for_testing; - Future operator()() { // A caller has requested a future Future waiting_future; @@ -1158,6 +1149,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 is in use. + static inline std::function error_signaled_hook_for_testing; + struct DeliveredJob { explicit DeliveredJob(AsyncGenerator deliverer_, Result value_, std::size_t index_) @@ -1212,13 +1211,20 @@ class MergedGenerator { } } - // Must be called with the mutex held, when the first error arrives. + // Must be called with the mutex held, when the first error arrives, in the same + // locked section that sets `broken`. void SetFinalErrorUnlocked(const util::Mutex::Guard& guard, const Status& err, - Future* sink, bool* should_deliver_to_sink) { - if (sink->is_valid()) { - // Someone is waiting for this error, it will be delivered to them (outside - // the lock) once all outstanding work is done - *should_deliver_to_sink = true; + Future sink) { + if (sink.is_valid()) { + // Someone is waiting for this error, it will be delivered to them once all + // outstanding work is done. The callback must be registered under the lock: + // a concurrent caller that sees `broken` registers its terminal continuation + // on `all_finished` too, and callbacks run in registration order, so this one + // must come first or the terminal item would overtake the error. + // The request that failed is still outstanding, so `all_finished` cannot be + // finished yet and the callback does not run here, under the lock. + all_finished.AddCallback( + [sink, err](const Status& status) mutable { sink.MarkFinished(err); }); } else { // No one is waiting for this error right now so it will be delivered // next. This must happen under the lock: a concurrent caller that sees @@ -1228,13 +1234,6 @@ class MergedGenerator { } } - // This is called outside the mutex but it is only ever called - // once and Future<>::AddCallback is thread-safe - void DeliverFinalError(const Status& err, Future sink) { - all_finished.AddCallback( - [sink, err](const Status& status) mutable { sink.MarkFinished(err); }); - } - bool IsCompleteUnlocked(const util::Mutex::Guard& guard) { return outstanding_requests == 0 && (broken || (source_exhausted && num_running_subscriptions == 0 && @@ -1311,7 +1310,6 @@ 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(); @@ -1336,8 +1334,7 @@ class MergedGenerator { if (!maybe_next->ok()) { signaled_error = true; state->SignalErrorUnlocked(guard); - state->SetFinalErrorUnlocked(guard, maybe_next->status(), &sink, - &should_mark_final_error); + state->SetFinalErrorUnlocked(guard, maybe_next->status(), std::move(sink)); } } @@ -1366,10 +1363,6 @@ class MergedGenerator { if (signaled_error) { RunErrorSignaledHookForTesting(); } - if (should_mark_final_error) { - state->DeliverFinalError(maybe_next->status(), std::move(sink)); - } - if (should_mark_gen_complete) { state->MarkFinishedAndPurge(); } @@ -1419,7 +1412,6 @@ 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 error_sink; @@ -1435,8 +1427,8 @@ class MergedGenerator { error_sink = std::move(*state->waiting_jobs.front()); state->waiting_jobs.pop_front(); } - state->SetFinalErrorUnlocked(guard, maybe_next.status(), &error_sink, - &should_deliver_error); + state->SetFinalErrorUnlocked(guard, maybe_next.status(), + std::move(error_sink)); } if (source_exhausted) { state->source_exhausted = true; @@ -1453,9 +1445,6 @@ class MergedGenerator { if (signaled_error) { RunErrorSignaledHookForTesting(); } - if (should_deliver_error) { - state->DeliverFinalError(maybe_next.status(), std::move(error_sink)); - } 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 3b4637271a71..7f0ff55bc4be 100644 --- a/cpp/src/arrow/util/async_generator_test.cc +++ b/cpp/src/arrow/util/async_generator_test.cc @@ -21,9 +21,11 @@ #include #include #include +#include #include #include #include +#include #include "arrow/io/slow.h" #include "arrow/testing/async_test_util.h" @@ -811,6 +813,20 @@ class MergedGeneratorErrorHookTest : public ::testing::Test { } }; } + + // Like PullFromErrorHook, but also records "terminal" in `order` when the pulled + // future completes + void PullFromErrorHookAndRecord(AsyncGenerator* merged, + Future* pulled, + std::vector* order) { + MergedGenerator::error_signaled_hook_for_testing = [=]() { + if (!pulled->is_valid()) { + *pulled = (*merged)(); + pulled->AddCallback( + [order](const Result&) { order->push_back("terminal"); }); + } + }; + } }; TEST_F(MergedGeneratorErrorHookTest, InnerErrorNotLostToConcurrentPull) { @@ -855,6 +871,50 @@ TEST_F(MergedGeneratorErrorHookTest, OuterErrorNotLostToConcurrentPull) { 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::Make(); + AsyncGenerator failing_sub = [failing]() { return failing; }; + std::vector> subs = {failing_sub}; + auto merged = MakeMergedGenerator(MakeVectorGenerator(std::move(subs)), 1); + std::vector order; + // Subscribes to failing_sub, whose first item is pending, so this caller waits + Future waiting = merged(); + waiting.AddCallback([&order](const Result&) { order.push_back("error"); }); + + Future 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{"error", "terminal"})); +} + +TEST_F(MergedGeneratorErrorHookTest, OuterErrorToWaiterNotOvertakenByLaterPull) { + auto failing = Future>::Make(); + AsyncGenerator> source = [failing]() { return failing; }; + auto merged = MakeMergedGenerator(std::move(source), 1); + std::vector order; + // Pulls the first subscription from the source, which is pending, so this caller + // waits + Future waiting = merged(); + waiting.AddCallback([&order](const Result&) { order.push_back("error"); }); + + Future 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{"error", "terminal"})); +} + class AutoStartingGeneratorTestFixture : public GeneratorTestFixture {}; TEST_P(AutoStartingGeneratorTestFixture, Basic) { From 364639c414074b15569db3a340466af5e6783184 Mon Sep 17 00:00:00 2001 From: Panchawat Date: Mon, 28 Sep 2026 12:37:03 +0700 Subject: [PATCH 3/3] GH-51495: [C++] Complete MergedGenerator futures in order instead of relying on callback order Future callbacks have no guaranteed execution order: one added while the future is being marked finished may run immediately. MergedGenerator chained every terminal item, and the delivery of an error to a waiting caller, on `all_finished`, so a pull made while `all_finished` was dispatching could get its terminal item before an earlier future (the one receiving the error, or an earlier IterationEnd) had completed. Remove `all_finished`. Terminal pulls made once the generator is broken or exhausted are queued in `waiting_jobs`, and 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. 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 a finished IterationEnd. Add deterministic tests for pulls made while the generator is completing (inner error, outer error, error nobody was waiting for), a test for pulling from the last future's callback, and a stress test that checks the AsyncGenerator ordering contract directly. Co-Authored-By: Claude Opus 5.5 (1M context) --- cpp/src/arrow/util/async_generator.h | 162 +++++++++++---- cpp/src/arrow/util/async_generator_test.cc | 225 +++++++++++++++++++++ 2 files changed, 347 insertions(+), 40 deletions(-) diff --git a/cpp/src/arrow/util/async_generator.h b/cpp/src/arrow/util/async_generator.h index 74441692df68..1e8c72d2a5eb 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 end_res = IterationEnd(); - 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::MakeFinished(std::move(err)); + } + state_->error_sink = Future::Make(); + return state_->error_sink; + } + if (state_->AllCompletedUnlocked(guard)) { + // Every future returned earlier has completed + return Future::MakeFinished(IterationEnd()); } - return state_->all_finished.Then([end_res]() -> Result { return end_res; }); + // Queue it behind the waiting callers, MarkFinishedAndPurge completes them in + // order + auto terminal = Future::Make(); + state_->waiting_jobs.push_back(std::make_shared>(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)); @@ -1201,37 +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()); - waiting_jobs.pop_front(); + std::deque>> to_complete; + Status err; + { + auto guard = mutex.Lock(); + finishing = true; + if (error_sink.is_valid()) { + to_complete.push_back(std::make_shared>(std::move(error_sink))); + error_sink = Future(); + 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>& fut) { + if (err.ok()) { + fut->MarkFinished(IterationEnd()); + } else { + fut->MarkFinished(std::move(err)); + err = Status::OK(); + } + }; + while (true) { + std::shared_ptr> 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); } } + // 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`. + // 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 sink) { - if (sink.is_valid()) { - // Someone is waiting for this error, it will be delivered to them once all - // outstanding work is done. The callback must be registered under the lock: - // a concurrent caller that sees `broken` registers its terminal continuation - // on `all_finished` too, and callbacks run in registration order, so this one - // must come first or the terminal item would overtake the error. - // The request that failed is still outstanding, so `all_finished` cannot be - // finished yet and the callback does not run here, under the lock. - all_finished.AddCallback( - [sink, err](const Status& status) mutable { sink.MarkFinished(err); }); - } else { - // No one is waiting for this error right now so it will be delivered - // next. This must happen under the lock: 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). - final_error = err; - } + final_error = err; + error_sink = std::move(sink); } bool IsCompleteUnlocked(const util::Mutex::Guard& guard) { @@ -1254,12 +1329,9 @@ class MergedGenerator { // caller std::deque> 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>> 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; @@ -1270,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 @@ -1279,10 +1351,20 @@ 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. Guarded by `mutex`. + // 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 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> completing_last; }; static void RunErrorSignaledHookForTesting() { diff --git a/cpp/src/arrow/util/async_generator_test.cc b/cpp/src/arrow/util/async_generator_test.cc index 7f0ff55bc4be..9d965dbc4e5f 100644 --- a/cpp/src/arrow/util/async_generator_test.cc +++ b/cpp/src/arrow/util/async_generator_test.cc @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -814,6 +815,60 @@ class MergedGeneratorErrorHookTest : public ::testing::Test { }; } + // 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& gen, + AsyncGenerator* merged, Future* earlier, + std::function fail, bool pull_earlier_in_hook, + Future* later) { + std::atomic signaled{false}; + std::atomic holding{false}; + MergedGenerator::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&) {}; + }); + 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& 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* merged, @@ -915,6 +970,176 @@ TEST_F(MergedGeneratorErrorHookTest, OuterErrorToWaiterNotOvertakenByLaterPull) ASSERT_EQ(order, (std::vector{"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::Make(); + AsyncGenerator failing_sub = [failing]() { return failing; }; + std::vector> subs = {failing_sub}; + MergedGenerator gen(MakeVectorGenerator(std::move(subs)), 1); + AsyncGenerator merged = gen; + Future waiting = merged(); + + Future 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>::Make(); + AsyncGenerator> source = [failing]() { return failing; }; + MergedGenerator gen(std::move(source), 1); + AsyncGenerator merged = gen; + Future waiting = merged(); + + Future 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::Make(); + AsyncGenerator failing_sub = [failing]() { return failing; }; + std::vector> subs = {MakeVectorGenerator({TestInt(1)}), + failing_sub}; + MergedGenerator gen(MakeVectorGenerator(std::move(subs)), 1); + AsyncGenerator 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 claimed; + Future 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::Make(); + AsyncGenerator sub = [pending]() { return pending; }; + std::vector> subs = {sub}; + auto merged = MakeMergedGenerator(MakeVectorGenerator(std::move(subs)), 1); + Future waiting = merged(); + Future pulled; + waiting.AddCallback([&](const Result&) { pulled = merged(); }); + if (fail) { + pending.MarkFinished(Status::Invalid("XYZ")); + ASSERT_FINISHES_AND_RAISES(Invalid, waiting); + } else { + pending.MarkFinished(IterationEnd()); + 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> 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> 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> futures; + int overtaken = 0; + int callbacks_run = 0; + }; + auto shared = std::make_shared(); + for (int i = 0; i < kNumPulls; i++) { + Future fut = merged(); + { + std::lock_guard lock(shared->mutex); + shared->futures.push_back(fut); + } + fut.AddCallback([shared, i](const Result& result) { + std::lock_guard 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> futures; + { + std::lock_guard 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 lock(shared->mutex); + return shared->callbacks_run == kNumPulls; + }); + std::lock_guard 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) {