diff --git a/src/v/cloud_storage_clients/client_pool.cc b/src/v/cloud_storage_clients/client_pool.cc index 251821739e346..4dea16a174a13 100644 --- a/src/v/cloud_storage_clients/client_pool.cc +++ b/src/v/cloud_storage_clients/client_pool.cc @@ -34,6 +34,47 @@ namespace { constexpr auto self_configure_attempts = 3; constexpr auto self_configure_backoff = 1s; constexpr auto pool_ready_timeout = 15s; + +/// Picks two distinct random shards, excluding self. Used for "power of two +/// choices" load balancing. +/// +/// Uses rejection-free sampling: instead of picking from [0,n) and retrying on +/// excluded values, we pick from a smaller range [0,n-k) and map to valid +/// values by skipping over excluded ones. This guarantees O(1) with exactly one +/// random draw per selection. +[[nodiscard]] std::pair pick_two_random_shards() { + using dist_t = std::uniform_int_distribution; + + static thread_local std::mt19937 gen{std::random_device{}()}; + + const ss::shard_id n = ss::smp::count; + const ss::shard_id self = ss::this_shard_id(); + + vassert(n > 1, "At least two shards are required"); + + if (n == 2) { + ss::shard_id other = (self == 0) ? 1 : 0; + return {other, other}; + } + + // Pick cpu_a from [0, n), excluding self + dist_t dist(0, n - 2); + auto r = dist(gen); + const ss::shard_id cpu_a = r < self ? r : r + 1; + + // Pick cpu_b from [0, n), excluding self and cpu_a + dist.param(dist_t::param_type{0, n - 3}); + auto cpu_b = dist(gen); + const auto [lo, hi] = std::minmax(self, cpu_a); + if (cpu_b >= lo) { + ++cpu_b; + } + if (cpu_b >= hi) { + ++cpu_b; + } + + return {cpu_a, cpu_b}; +} } // namespace namespace cloud_storage_clients { @@ -247,26 +288,6 @@ void client_pool::shutdown_connections() { bool client_pool::shutdown_initiated() { return _as.abort_requested(); } -std::tuple pick_two_random_shards() { - static thread_local std::vector shards = [] { - std::vector res; - for (auto i = 0UL; i < ss::smp::count; i++) { - if (i != ss::this_shard_id()) { - res.push_back(i); - } - } - return res; - }(); - vassert(ss::smp::count > 1, "At least two shards are required"); - if (shards.size() == 1) { - return std::tie(shards.at(0), shards.at(0)); - } - std::random_device rd; - std::mt19937 gen(rd()); - std::shuffle(shards.begin(), shards.end(), gen); - return std::tie(shards.at(0), shards.at(1)); -} - /// \brief Acquire http client from the pool. /// /// as: An abort source which must outlive the lease, that will