diff --git a/google/cloud/odbc/bq_driver/internal/utils.h b/google/cloud/odbc/bq_driver/internal/utils.h index a557ce6725..fdc77c527f 100644 --- a/google/cloud/odbc/bq_driver/internal/utils.h +++ b/google/cloud/odbc/bq_driver/internal/utils.h @@ -33,6 +33,7 @@ extern HINSTANCE g_hDllInstance; #include "google/cloud/status_or.h" #include "re2/re2.h" #include +#include #include #include #include @@ -42,9 +43,12 @@ extern HINSTANCE g_hDllInstance; #include #include #include +#include +#include #include #include #include +#include #include namespace google::cloud::odbc_bq_driver_internal { @@ -102,75 +106,91 @@ size_t BufferSizeForType(SQLSMALLINT type, size_t requested); // TaskInput: The type of a single item in the input vector. // TaskResult: The type of data returned by the function on success. // -// Returns: A vector containing the results of all successful tasks, or the -// StatusRecord of the first error encountered. +// Returns: A vector containing the results of all successful tasks in input +// order, or the StatusRecord of the first (in input order) error encountered. +// After a task fails, tasks not yet started are skipped, but tasks already +// in flight run to completion before this function returns. +// +// Implemented as a fixed pool of max_threads workers pulling the next input +// from a shared atomic index. Each worker moves on to the next input the +// moment its current task finishes, so one slow task (a straggler REST call) +// never stalls the dispatch of the remaining work -- a previous +// dispatcher-based implementation blocked on the oldest in-flight task to +// free a slot, which serialized the whole batch behind stragglers and made +// catalog enumeration latency vary heavily from run to run. template odbc_internal::StatusRecordOr> ExecuteParallelTasks( std::uint32_t max_threads, std::vector const& inputs, std::function(TaskInput const&)> task_func) { - using FutureType = std::future>; - - std::vector active_futures; - std::vector aggregated_results; - odbc_internal::StatusRecord error_status = odbc_internal::StatusRecord::Ok(); - bool error_occurred = false; - - // Helper to collect result from a finished future - auto process_future = [&](FutureType& f) { - auto result = f.get(); - if (!result) { - // Record the first error that occurs, but keep draining threads - if (!error_occurred) { - error_status = result.GetStatusRecord(); - error_occurred = true; + if (inputs.empty()) { + return std::vector{}; + } + // A misconfigured MaxThreads of 0 previously spun the dispatch loop forever; + // treat it as serial execution. Never spawn more workers than inputs. + std::size_t const num_workers = (std::min)( + static_cast((std::max)(max_threads, std::uint32_t{1})), + inputs.size()); + + // One pre-sized slot per input: workers write disjoint indices, so no + // synchronization is needed, and results come back in input order. + std::vector>> slots( + inputs.size()); + std::atomic next_index{0}; + std::atomic error_occurred{false}; + std::exception_ptr first_exception; + std::mutex exception_mutex; + + auto worker = [&] { + while (!error_occurred.load(std::memory_order_relaxed)) { + std::size_t const index = + next_index.fetch_add(1, std::memory_order_relaxed); + if (index >= inputs.size()) { + return; } - } else if (!error_occurred) { - // Only store results if we are still in a success state - aggregated_results.push_back(std::move(*result)); - } - }; - - for (auto const& input : inputs) { - // If we have hit the thread limit, wait for at least one thread to finish - while (active_futures.size() >= max_threads) { - bool slot_freed = false; - for (auto it = active_futures.begin(); it != active_futures.end();) { - // Check if ready without blocking - if (it->wait_for(std::chrono::milliseconds(1)) == - std::future_status::ready) { - process_future(*it); - it = active_futures.erase(it); - slot_freed = true; - break; // We freed a slot, proceed to launch next task + try { + auto result = task_func(inputs[index]); + if (!result) { + error_occurred.store(true, std::memory_order_relaxed); } - ++it; - } - - // If no threads finished yet, block on the oldest one to prevent spinning - if (!slot_freed && !active_futures.empty()) { - process_future(active_futures.front()); - active_futures.erase(active_futures.begin()); + slots[index] = std::move(result); + } catch (...) { + // Propagate the exception on the calling thread (std::async did this + // via future::get); letting it escape a worker would call terminate. + std::lock_guard lock(exception_mutex); + if (!first_exception) { + first_exception = std::current_exception(); + } + error_occurred.store(true, std::memory_order_relaxed); + return; } } + }; - // If an error occurred previously, we stop spawning new tasks, - // but the loop continues to ensure we drain existing futures safely. - if (!error_occurred) { - active_futures.push_back( - std::async(std::launch::async, task_func, input)); - } + std::vector workers; + workers.reserve(num_workers); + for (std::size_t i = 0; i < num_workers; ++i) { + workers.emplace_back(worker); } - - // Wait for and collect all remaining threads - for (auto& f : active_futures) { - process_future(f); + for (auto& thread : workers) { + thread.join(); } - if (error_occurred) { - return error_status; + if (first_exception) { + std::rethrow_exception(first_exception); } + std::vector aggregated_results; + aggregated_results.reserve(inputs.size()); + for (auto& slot : slots) { + if (!slot.has_value()) { + continue; // Skipped after an error; an errored slot below reports it. + } + if (!*slot) { + return slot->GetStatusRecord(); + } + aggregated_results.push_back(std::move(**slot)); + } return aggregated_results; } diff --git a/google/cloud/odbc/bq_driver/internal/utils_test.cc b/google/cloud/odbc/bq_driver/internal/utils_test.cc index f134d0a0ca..41fe818c6b 100644 --- a/google/cloud/odbc/bq_driver/internal/utils_test.cc +++ b/google/cloud/odbc/bq_driver/internal/utils_test.cc @@ -27,6 +27,7 @@ using ::google::cloud::odbc_internal::SQLStates; using ::google::cloud::odbc_internal::StatusRecord; using ::google::cloud::odbc_internal::StatusRecordOr; using google::cloud::odbc_testing_utils::StatusRecordIs; +using ::testing::ElementsAre; using ::testing::HasSubstr; using ::testing::IsEmpty; using ::testing::UnorderedElementsAre; @@ -1016,4 +1017,69 @@ TEST(ExecuteParallelTasksTest, RespectsSlidingWindow) { EXPECT_GE(duration, (task_count / max_threads) * min_sleep_ms); } +TEST(ExecuteParallelTasksTest, ResultsPreserveInputOrder) { + std::vector inputs = {5, 1, 4, 2, 3}; + + // Sleep inversely to the input so tasks complete in the reverse of their + // submission order; the results must still come back in input order. + auto task = [](int input) -> StatusRecordOr { + std::this_thread::sleep_for(std::chrono::milliseconds(input * 10)); + return input; + }; + + auto result = ExecuteParallelTasks(5, inputs, task); + + ASSERT_STATUS_RECORD_OK(result); + EXPECT_THAT(*result, ElementsAre(5, 1, 4, 2, 3)); +} + +TEST(ExecuteParallelTasksTest, ZeroMaxThreadsRunsSerially) { + // A misconfigured MaxThreads of 0 must not hang; it degrades to serial + // execution. + std::vector inputs = {1, 2, 3}; + auto task = [](int input) -> StatusRecordOr { return input * 2; }; + + auto result = ExecuteParallelTasks(0, inputs, task); + + ASSERT_STATUS_RECORD_OK(result); + EXPECT_THAT(*result, ElementsAre(2, 4, 6)); +} + +TEST(ExecuteParallelTasksTest, StragglerDoesNotStallDispatch) { + // One slow task at the head of the queue must not prevent the other worker + // from chewing through the remaining short tasks. The previous + // dispatcher-based implementation blocked on the oldest in-flight future, + // so this workload took ~straggler + sum(short tasks) instead of + // ~max(straggler, sum(short tasks)). + int const straggler_ms = 300; + int const short_task_ms = 20; + int const short_task_count = 8; + + std::vector inputs; + inputs.push_back(straggler_ms); + for (int i = 0; i < short_task_count; ++i) { + inputs.push_back(short_task_ms); + } + + auto sleeping_task = [](int sleep_ms) -> StatusRecordOr { + std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms)); + return sleep_ms; + }; + + auto start_time = std::chrono::steady_clock::now(); + auto result = + ExecuteParallelTasks(2, inputs, sleeping_task); + auto duration = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time) + .count(); + + ASSERT_STATUS_RECORD_OK(result); + EXPECT_EQ(result->size(), inputs.size()); + // Ideal: worker 1 takes the straggler (300ms) while worker 2 runs the eight + // 20ms tasks (160ms) => ~300ms total. The old scheduler needed ~460ms + // (straggler + all short tasks serialized behind it). Use a generous bound + // to stay robust on loaded CI machines while still distinguishing the two. + EXPECT_LT(duration, straggler_ms + short_task_count * short_task_ms - 60); +} + } // namespace google::cloud::odbc_bq_driver_internal