diff --git a/cpp/src/arrow/util/thread_pool.cc b/cpp/src/arrow/util/thread_pool.cc index 4fbce97c2ff..00e09b6573f 100644 --- a/cpp/src/arrow/util/thread_pool.cc +++ b/cpp/src/arrow/util/thread_pool.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -580,7 +581,7 @@ Status ThreadPool::SetCapacity(int threads) { threads - static_cast(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 = 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 task, StopToken stop_token, @@ -729,12 +736,12 @@ Status ThreadPool::SpawnReal(TaskHints hints, FnOnce task, StopToken sto return Status::Invalid("operation forbidden during or after shutdown"); } CollectFinishedWorkersUnlocked(); - state_->tasks_queued_or_running_++; - if (static_cast(state_->workers_.size()) < state_->tasks_queued_or_running_ && + if (static_cast(state_->workers_.size()) <= state_->tasks_queued_or_running_ && state_->desired_capacity_ > static_cast(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, diff --git a/cpp/src/arrow/util/thread_pool.h b/cpp/src/arrow/util/thread_pool.h index 201b8cef790..b106b577abb 100644 --- a/cpp/src/arrow/util/thread_pool.h +++ b/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(); diff --git a/cpp/src/arrow/util/thread_pool_test.cc b/cpp/src/arrow/util/thread_pool_test.cc index c1391c8be88..7c7498838a6 100644 --- a/cpp/src/arrow/util/thread_pool_test.cc +++ b/cpp/src/arrow/util/thread_pool_test.cc @@ -16,6 +16,7 @@ // under the License. #ifndef _WIN32 +# include # include # include #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) {