improve(bq_driver): SQLtables improvement

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 <algorithm>
+#include <atomic>
 #include <chrono>
 #include <codecvt>
 #include <fstream>
@@ -42,9 +43,12 @@ extern HINSTANCE g_hDllInstance;
 #include <locale>
 #include <map>
 #include <memory>
+#include <mutex>
+#include <optional>
 #include <stdexcept>
 #include <string>
 #include <string_view>
+#include <thread>
 #include <vector>
 
 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 <typename TaskInput, typename TaskResult>
 odbc_internal::StatusRecordOr<std::vector<TaskResult>> ExecuteParallelTasks(
     std::uint32_t max_threads, std::vector<TaskInput> const& inputs,
     std::function<odbc_internal::StatusRecordOr<TaskResult>(TaskInput const&)>
         task_func) {
-  using FutureType = std::future<odbc_internal::StatusRecordOr<TaskResult>>;
-
-  std::vector<FutureType> active_futures;
-  std::vector<TaskResult> 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<TaskResult>{};
+  }
+  // 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::size_t>((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<std::optional<odbc_internal::StatusRecordOr<TaskResult>>> slots(
+      inputs.size());
+  std::atomic<std::size_t> next_index{0};
+  std::atomic<bool> 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<std::mutex> 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<std::thread> 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<TaskResult> 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;
 }
 

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<int> 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<int> {
+    std::this_thread::sleep_for(std::chrono::milliseconds(input * 10));
+    return input;
+  };
+
+  auto result = ExecuteParallelTasks<int, int>(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<int> inputs = {1, 2, 3};
+  auto task = [](int input) -> StatusRecordOr<int> { return input * 2; };
+
+  auto result = ExecuteParallelTasks<int, int>(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<int> 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<int> {
+    std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms));
+    return sleep_ms;
+  };
+
+  auto start_time = std::chrono::steady_clock::now();
+  auto result =
+      ExecuteParallelTasks<int, int>(2, inputs, sleeping_task);
+  auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
+                      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