diff --git a/include/cucascade/exec/thread_pool.hpp b/include/cucascade/exec/thread_pool.hpp index 64b8b5eb..b84fad7e 100644 --- a/include/cucascade/exec/thread_pool.hpp +++ b/include/cucascade/exec/thread_pool.hpp @@ -19,9 +19,9 @@ #pragma once #include +#include #include -#include #include #include #include @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -38,42 +39,51 @@ namespace cucascade::exec { class static_thread_pool { public: + /// Start all workers before returning. Names and best-effort affinity are + /// applied before per_thread_init. An empty CPU list uses the process main + /// thread's mask, rather than inheriting a potentially pinned creator. + /// Throws for a non-positive width or failed initialization, after joining + /// any started workers. The initializer must be safe to call concurrently. explicit static_thread_pool(int num_threads, const std::string& name = "thread_pool", std::vector cpu_ids = {}, std::function per_thread_init = nullptr) { + if (num_threads <= 0) { throw std::invalid_argument("thread pool width must be positive"); } threads_.reserve(static_cast(num_threads)); - - cpu_set_t cpuset; - CPU_ZERO(&cpuset); - std::for_each(cpu_ids.begin(), cpu_ids.end(), [&cpuset](int cpu_id) { - CPU_SET(static_cast(cpu_id), &cpuset); - }); - - std::unique_ptr init_latch; - if (per_thread_init) { init_latch = std::make_unique(num_threads); } - - auto* init_fn_ptr = per_thread_init ? &per_thread_init : nullptr; - auto* latch_ptr = init_latch.get(); - - for (int i = 0; i < num_threads; ++i) { - auto& t = threads_.emplace_back([this, init_fn_ptr, latch_ptr]() { - if (init_fn_ptr) { - (*init_fn_ptr)(); - latch_ptr->count_down(); - } - work_loop(); - }); - if (!name.empty()) { - pthread_setname_np(t.native_handle(), (name + "_" + std::to_string(i)).c_str()); + if (cpu_ids.empty()) { cpu_ids = thread_util::get_process_affinity(); } + + std::latch initialized{num_threads}; + std::exception_ptr init_error; + try { + for (int i = 0; i < num_threads; ++i) { + auto worker_name = name.empty() ? std::string{} : name + "_" + std::to_string(i); + threads_.emplace_back([&, worker_name = std::move(worker_name)] { + bool ready = false; + try { + if (!worker_name.empty()) { (void)thread_util::set_current_thread_name(worker_name); } + (void)thread_util::set_current_thread_affinity(cpu_ids); + if (per_thread_init) { per_thread_init(); } + ready = true; + } catch (...) { + std::lock_guard lock(mu_); + if (!init_error) { init_error = std::current_exception(); } + } + initialized.count_down(); + if (ready) { work_loop(); } + }); } - if (!cpu_ids.empty()) { - pthread_setaffinity_np(t.native_handle(), sizeof(cpu_set_t), &cpuset); + initialized.wait(); + if (init_error) { std::rethrow_exception(init_error); } + } catch (...) { + // The destructor does not run if construction fails. Keep the startup + // state above alive until every started worker has stopped using it. + stop(); + for (auto& t : threads_) { + if (t.joinable()) { t.join(); } } + throw; } - - if (init_latch) { init_latch->wait(); } } static_thread_pool(const static_thread_pool&) = delete; diff --git a/include/cucascade/exec/thread_util.hpp b/include/cucascade/exec/thread_util.hpp index a1c3b35c..f1382395 100644 --- a/include/cucascade/exec/thread_util.hpp +++ b/include/cucascade/exec/thread_util.hpp @@ -93,6 +93,12 @@ inline constexpr std::size_t max_thread_name_len = 15; [[nodiscard]] std::vector get_thread_affinity(pthread_t handle) noexcept; [[nodiscard]] std::vector get_current_thread_affinity() noexcept; +/// Read the main thread's CPU mask, even when called from a pinned worker. +/// Pools and reactors use this as their default instead of inheriting the +/// creator's affinity. Returns an empty vector on syscall failure; allocation +/// failures propagate so callers can recover before starting any threads. +[[nodiscard]] std::vector get_process_affinity(); + // --------------------------------------------------------------------------- // scheduling policy and priority // --------------------------------------------------------------------------- diff --git a/src/exec/thread_util.cpp b/src/exec/thread_util.cpp index ee34266f..c41ded1f 100644 --- a/src/exec/thread_util.cpp +++ b/src/exec/thread_util.cpp @@ -179,6 +179,19 @@ std::vector get_current_thread_affinity() noexcept return get_thread_affinity(::pthread_self()); } +std::vector get_process_affinity() +{ + // On Linux the process ID identifies the main thread, not the caller. + cpu_set_t mask{}; + if (::sched_getaffinity(::getpid(), sizeof(mask), &mask) != 0) { return {}; } + std::vector ids; + ids.reserve(static_cast(CPU_COUNT(&mask))); + for (int i = 0; i < CPU_SETSIZE; ++i) { + if (CPU_ISSET(static_cast(i), &mask)) { ids.push_back(i); } + } + return ids; +} + // --------------------------------------------------------------------------- // scheduling policy and priority // --------------------------------------------------------------------------- diff --git a/src/io/rest/rest_reactor.cpp b/src/io/rest/rest_reactor.cpp index b9db8e33..e42cd889 100644 --- a/src/io/rest/rest_reactor.cpp +++ b/src/io/rest/rest_reactor.cpp @@ -582,17 +582,22 @@ void rest_reactor::start() _running = true; _accepting = true; try { - _worker = std::jthread([this](std::stop_token const& st) { worker_loop(st); }, - _stop_source.get_token()); + auto cpus = cucascade::exec::thread_util::get_process_affinity(); + auto name = _tname.empty() ? std::string{} : _tname + "_worker"; + _worker = std::jthread( + [this, cpus = std::move(cpus), name = std::move(name)](std::stop_token const& st) { + std::ignore = cucascade::exec::thread_util::set_current_thread_affinity(cpus); + if (!name.empty()) { + std::ignore = cucascade::exec::thread_util::set_current_thread_name(name); + } + worker_loop(st); + }, + _stop_source.get_token()); } catch (...) { _running = false; _accepting = false; throw; } - if (!_tname.empty()) { - auto const full_name = _tname + "_worker"; - std::ignore = cucascade::exec::thread_util::set_thread_name(_worker, full_name); - } } rest_reactor::~rest_reactor() { shutdown(); } diff --git a/src/io/uring/uring_reactor.cpp b/src/io/uring/uring_reactor.cpp index f4e1fc57..458c1c93 100644 --- a/src/io/uring/uring_reactor.cpp +++ b/src/io/uring/uring_reactor.cpp @@ -493,18 +493,24 @@ void uring_reactor::start() std::lock_guard lock(_enqueue_mutex); try { + auto cpus = cucascade::exec::thread_util::get_process_affinity(); + auto name = _tname.empty() ? std::string{} : _tname + "_worker"; _accepting.store(true, std::memory_order_release); - _worker = std::jthread([this](std::stop_token stop_token) { worker_loop(stop_token); }, - _stop_source.get_token()); + _worker = std::jthread( + [this, cpus = std::move(cpus), name = std::move(name)](std::stop_token stop_token) { + // io_uring's own workers inherit this mask on the first submission. + std::ignore = cucascade::exec::thread_util::set_current_thread_affinity(cpus); + if (!name.empty()) { + std::ignore = cucascade::exec::thread_util::set_current_thread_name(name); + } + worker_loop(stop_token); + }, + _stop_source.get_token()); } catch (...) { _accepting.store(false, std::memory_order_release); _bounce_storage.reset(); throw; } - if (!_tname.empty()) { - auto const name = _tname + "_worker"; - std::ignore = cucascade::exec::thread_util::set_thread_name(_worker, name); - } } void uring_reactor::interrupt() {} diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index c8ec8e4b..884809d6 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -64,6 +64,7 @@ if(NOT CUCASCADE_TOPOLOGY_ONLY AND CUCASCADE_BUILD_IO) # IO test executable - links the cudf-free cucascade-io datasource layer. add_executable( cucascade_io_tests + exec/test_thread_pool.cpp io/test_datasource_registry.cpp io/test_dispatch_failure_hook.cpp io/test_uri_parser.cpp diff --git a/test/exec/test_thread_pool.cpp b/test/exec/test_thread_pool.cpp new file mode 100644 index 00000000..63edbfeb --- /dev/null +++ b/test/exec/test_thread_pool.cpp @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using cucascade::exec::static_thread_pool; +namespace thread_util = cucascade::exec::thread_util; + +TEST_CASE("thread pool names fit the Linux thread name limit", "[thread_pool]") +{ + for (auto const& name : {std::string{"io"}, std::string{"footer-one-worker"}}) { + CAPTURE(name); + std::promise observed; + auto result = observed.get_future(); + static_thread_pool pool{1, name}; + pool.schedule([&observed] { observed.set_value(thread_util::get_current_thread_name()); }); + REQUIRE(result.wait_for(std::chrono::seconds{5}) == std::future_status::ready); + CHECK(result.get() == (name + "_0").substr(0, 15)); + } +} + +TEST_CASE("thread pool affinity uses valid CPU IDs and defaults to process CPUs", "[thread_pool]") +{ + auto const inherited = thread_util::get_current_thread_affinity(); + REQUIRE_FALSE(inherited.empty()); + auto const cpu = inherited.back(); + + for (auto const& requested : + {std::vector{}, std::vector{cpu}, std::vector{-1, cpu, CPU_SETSIZE}}) { + CAPTURE(requested); + std::promise> observed; + auto result = observed.get_future(); + static_thread_pool pool{1, "affinity", requested}; + pool.schedule([&observed] { observed.set_value(thread_util::get_current_thread_affinity()); }); + REQUIRE(result.wait_for(std::chrono::seconds{5}) == std::future_status::ready); + CHECK(result.get() == (requested.empty() ? inherited : std::vector{cpu})); + } +} + +TEST_CASE("thread pool applies naming and affinity before initialization", "[thread_pool]") +{ + auto const cpus = thread_util::get_current_thread_affinity(); + REQUIRE_FALSE(cpus.empty()); + std::string observed_name; + std::vector observed_cpus; + static_thread_pool pool{1, "worker-bootstrap-long", {cpus.back()}, [&] { + observed_name = thread_util::get_current_thread_name(); + observed_cpus = thread_util::get_current_thread_affinity(); + }}; + CHECK(observed_name == "worker-bootstra"); + CHECK(observed_cpus == std::vector{cpus.back()}); +} + +TEST_CASE("thread pool does not inherit a pinned creator's CPU restriction", "[thread_pool]") +{ + // Read the test main thread's actual mask, independently of get_process_affinity. + auto const process_cpus = thread_util::get_current_thread_affinity(); + REQUIRE_FALSE(process_cpus.empty()); + if (process_cpus.size() < 2) { SKIP("requires at least two allowed CPUs"); } + + auto result = std::async(std::launch::async, [&] { + std::array pinned{process_cpus.front()}; + if (!thread_util::set_current_thread_affinity(pinned)) { + throw std::runtime_error("cannot pin the creator thread"); + } + auto creator_cpus = thread_util::get_current_thread_affinity(); + auto queried_cpus = thread_util::get_process_affinity(); + std::vector initialized_cpus; + std::promise> observed; + auto worker = observed.get_future(); + static_thread_pool pool{1, "process-cpus", {}, [&] { + initialized_cpus = thread_util::get_current_thread_affinity(); + }}; + pool.schedule([&] { observed.set_value(thread_util::get_current_thread_affinity()); }); + return std::make_tuple(creator_cpus, queried_cpus, initialized_cpus, worker.get()); + }); + auto const [creator, queried, initialized, worker] = result.get(); + CHECK(creator == std::vector{process_cpus.front()}); + CHECK(queried == process_cpus); + CHECK(initialized == process_cpus); + CHECK(worker == process_cpus); +} + +TEST_CASE("thread pool propagates initialization failures after joining every worker", + "[thread_pool]") +{ + std::atomic initialized{0}; + std::atomic exited{0}; + struct exit_counter { + std::atomic& count; + ~exit_counter() { ++count; } + }; + auto initialize = [&] { + thread_local exit_counter on_exit{exited}; + if (initialized.fetch_add(1) == 1) { throw std::runtime_error("worker initialization failed"); } + }; + + CHECK_THROWS_WITH((static_thread_pool{3, "startup", {}, initialize}), + "worker initialization failed"); + CHECK(initialized == 3); + CHECK(exited == 3); +} + +TEST_CASE("thread pool rejects widths that cannot execute work", "[thread_pool]") +{ + CHECK_THROWS_AS(static_thread_pool{0}, std::invalid_argument); + CHECK_THROWS_AS(static_thread_pool{-1}, std::invalid_argument); +} diff --git a/test/io/rest/test_rest_footer_resolve.cpp b/test/io/rest/test_rest_footer_resolve.cpp index 64a6ca67..c28f4221 100644 --- a/test/io/rest/test_rest_footer_resolve.cpp +++ b/test/io/rest/test_rest_footer_resolve.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -1116,3 +1117,56 @@ TEST_CASE("batched footer resolve reuses at most the configured in-flight connec CHECK(server.get_count() == paths.size()); CHECK(server.accepted_connection_count() <= inflight); } + +TEST_CASE("REST worker startup escapes a pinned creator before authorizing requests", + "[rest][affinity]") +{ + namespace thread_util = cucascade::exec::thread_util; + namespace rest = cucascade::io::rest; + auto const process_cpus = thread_util::get_current_thread_affinity(); + REQUIRE_FALSE(process_cpus.empty()); + if (process_cpus.size() < 2) { SKIP("requires at least two allowed CPUs"); } + + // Warmup authorizes on the real reactor worker. Observe its OS state there, + // then decline the request so this test needs no network or credentials. + struct observing_authorizer : rest::request_authorizer { + std::promise>> observed; + + rest::authorized_request authorize(rest::object_ref const&, + rest::request_method, + std::chrono::seconds) override + { + throw std::runtime_error("unexpected object request"); + } + + rest::authorized_request authorize_list(std::string_view, + std::string_view, + std::chrono::seconds) override + { + observed.set_value( + {thread_util::get_current_thread_name(), thread_util::get_current_thread_affinity()}); + throw std::runtime_error("warmup observation complete"); + } + }; + + auto result = std::async(std::launch::async, [&] { + std::array pinned{process_cpus.front()}; + if (!thread_util::set_current_thread_affinity(pinned)) { + throw std::runtime_error("cannot pin the creator thread"); + } + auto authorizer = std::make_shared(); + auto observed = authorizer->observed.get_future(); + auto context = + std::make_shared(test_config(), authorizer, nullptr); + rest_reactor reactor{std::move(context), "rest"}; + reactor.start(); + reactor.warmup("bucket"); + if (observed.wait_for(5s) != std::future_status::ready) { + throw std::runtime_error("reactor did not authorize warmup"); + } + return observed.get(); + }); + auto const [name, cpus] = result.get(); + CHECK(name == "rest_worker"); + CHECK(cpus == process_cpus); +}