GH-48137: [C++] Restore ThreadPool state when a worker fails to start

cpp/src/arrow/util/thread_pool.cc

@@ -23,6 +23,7 @@
 #include <list>
 #include <mutex>
 #include <string>
+#include <system_error>
 #include <thread>
 #include <vector>
 
@@ -580,7 +581,7 @@ Status ThreadPool::SetCapacity(int threads) {
  threads - static_cast<int>(state_->workers_.size()));
  if (required > 0) {
  // Some tasks are pending, spawn the number of needed threads immediately
- LaunchWorkersUnlocked(required);
+RETURN_NOT_OK(LaunchWorkersUnlocked(required));
  } else if (required < 0) {
  // Excess threads are running, wake them so that they stop
  state_->cv_.notify_all();
@@ -692,17 +693,23 @@ static void SetCurrentThreadPool(ThreadPool* pool) { current_thread_pool_ = pool
 
 bool ThreadPool::OwnsThisThread() { return GetCurrentThreadPool() == this; }
 
-void ThreadPool::LaunchWorkersUnlocked(int threads) {
+Status ThreadPool::LaunchWorkersUnlocked(int threads) {
  std::shared_ptr<State> state = sp_state_;
 
  for (int i = 0; i < threads; i++) {
  state_->workers_.emplace_back();
  auto it = --(state_->workers_.end());
- *it = std::thread([this, state, it] {
- SetCurrentThreadPool(this);
- WorkerLoop(state, it);
- });
+ try {
+ *it = std::thread([this, state, it] {
+ SetCurrentThreadPool(this);
+ WorkerLoop(state, it);
+ });
+ } catch (const std::exception& e) {
+ state_->workers_.erase(it);
+ return Status::UnknownError("Failed to launch worker thread: ", e.what());
+ }
  }
+ return Status::OK();
 }
 
 Status ThreadPool::SpawnReal(TaskHints hints, FnOnce<void()> task, StopToken stop_token,
@@ -729,12 +736,12 @@ Status ThreadPool::SpawnReal(TaskHints hints, FnOnce<void()> task, StopToken sto
  return Status::Invalid("operation forbidden during or after shutdown");
  }
  CollectFinishedWorkersUnlocked();
- state_->tasks_queued_or_running_++;
- if (static_cast<int>(state_->workers_.size()) < state_->tasks_queued_or_running_ &&
+ if (static_cast<int>(state_->workers_.size()) <= state_->tasks_queued_or_running_ &&
  state_->desired_capacity_ > static_cast<int>(state_->workers_.size())) {
  // We can still spin up more workers so spin up a new worker
- LaunchWorkersUnlocked(/*threads=*/1);
+RETURN_NOT_OK(LaunchWorkersUnlocked(/*threads=*/1));
  }
+ state_->tasks_queued_or_running_++;
  state_->pending_tasks_.push(
  QueuedTask{{std::move(task), std::move(stop_token), std::move(stop_callback)},
  hints.priority,

cpp/src/arrow/util/thread_pool.h

@@ -496,6 +496,7 @@ class ARROW_EXPORT ThreadPool : public Executor {
 
  protected:
  FRIEND_TEST(TestThreadPool, SetCapacity);
+ FRIEND_TEST(TestThreadPool, FailedWorkerLaunch);
  FRIEND_TEST(TestGlobalThreadPool, Capacity);
  ARROW_FRIEND_EXPORT friend ThreadPool* GetCpuThreadPool();
 
@@ -507,7 +508,7 @@ class ARROW_EXPORT ThreadPool : public Executor {
  // Collect finished worker threads, making sure the OS threads have exited
  void CollectFinishedWorkersUnlocked();
  // Launch a given number of additional workers
-void LaunchWorkersUnlocked(int threads);
+Status LaunchWorkersUnlocked(int threads);
  // Get the current actual capacity
  int GetActualCapacity();
 

cpp/src/arrow/util/thread_pool_test.cc

@@ -16,6 +16,7 @@
 // under the License.
 
 #ifndef _WIN32
+# include <sys/resource.h>
 # include <sys/types.h>
 # include <unistd.h>
 #endif
@@ -832,6 +833,36 @@ TEST_F(TestThreadPool, SetCapacity) {
  ASSERT_EQ(pool->GetCapacity(), 7);
 }
 #endif
+
+#if defined(ARROW_ENABLE_THREADING) && !defined(_WIN32)
+TEST_F(TestThreadPool, FailedWorkerLaunch) {
+# ifdef __APPLE__
+ GTEST_SKIP() << "RLIMIT_NPROC does not limit thread creation on macOS";
+# else
+ auto pool = this->MakeThreadPool(4);
+
+ struct rlimit limit;
+ ASSERT_EQ(getrlimit(RLIMIT_NPROC, &limit), 0);
+ const rlim_t soft_limit = limit.rlim_cur;
+ limit.rlim_cur = 1;
+ if (setrlimit(RLIMIT_NPROC, &limit) != 0) {
+ GTEST_SKIP() << "Could not lower RLIMIT_NPROC";
+ }
+ const Status st = pool->Spawn([] {});
+ limit.rlim_cur = soft_limit;
+ ASSERT_EQ(setrlimit(RLIMIT_NPROC, &limit), 0);
+
+ if (st.ok()) {
+ GTEST_SKIP() << "Lowering RLIMIT_NPROC did not prevent thread creation";
+ }
+ ASSERT_RAISES(UnknownError, st);
+ ASSERT_EQ(pool->GetActualCapacity(), 0);
+ ASSERT_EQ(pool->GetNumTasks(), 0);
+ ASSERT_OK(pool->Shutdown());
+# endif
+}
+#endif
+
 // Test Submit() functionality
 
 TEST_F(TestThreadPool, Submit) {