cloud_storage_clients: rewrite pick_two_random_shards for O(1) sampling

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<ss::shard_id, ss::shard_id> pick_two_random_shards() {
+    using dist_t = std::uniform_int_distribution<ss::shard_id>;
+
+    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<unsigned int, unsigned int> pick_two_random_shards() {
-    static thread_local std::vector<unsigned> shards = [] {
-        std::vector<unsigned> 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