locator/topology: do_sort_by_proximity: shuffle equal-distance replicas

To improve balancing when reading in 1 < CL < ALL

This implementation has a moderate impact on
the function performance in contrast to full
std::shuffle of the vector before sorting it
(especially with large number of nodes to sort).

Before:
test                                               iterations      median         mad         min         max      allocs       tasks        inst      cycles
sort_by_proximity_topology.perf_sort_by_proximity    21133744    49.504ns     0.207ns    49.266ns    49.738ns       0.000       0.000       661.3       148.1

After:
sort_by_proximity_topology.perf_sort_by_proximity    17888163    55.934ns     0.265ns    55.583ns    56.300ns       0.000       0.000       715.8       167.4

Signed-off-by: Benny Halevy <[email protected]>

locator/topology.cc

@@ -6,11 +6,14 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
+#include <bit>
+#include <ranges>
+#include <utility>
+
 #include <seastar/core/coroutine.hh>
 #include <seastar/coroutine/maybe_yield.hh>
 #include <seastar/core/on_internal_error.hh>
 #include <seastar/util/lazy.hh>
-#include <utility>
 
 #include <seastar/core/shard_id.hh>
 #include "utils/log.hh"
@@ -108,6 +111,7 @@ topology::topology(config cfg)
         : _shard(this_shard_id())
         , _cfg(cfg)
         , _sort_by_proximity(!cfg.disable_proximity_sorting)
+        , _random_engine(std::random_device{}())
 {
     tlogger.trace("topology[{}]: constructing using config: endpoint={} id={} dc={} rack={}", fmt::ptr(this),
             cfg.this_endpoint, cfg.this_host_id, cfg.local_dc_rack.dc, cfg.local_dc_rack.rack);
@@ -127,6 +131,7 @@ topology::topology(topology&& o) noexcept
     , _dc_racks(std::move(o._dc_racks))
     , _sort_by_proximity(o._sort_by_proximity)
     , _datacenters(std::move(o._datacenters))
+    , _random_engine(std::move(o._random_engine))
 {
     SCYLLA_ASSERT(_shard == this_shard_id());
     tlogger.trace("topology[{}]: move from [{}]", fmt::ptr(this), fmt::ptr(&o));
@@ -178,6 +183,7 @@ future<topology> topology::clone_gently() const {
         co_await coroutine::maybe_yield();
     }
     ret._sort_by_proximity = _sort_by_proximity;
+    ret._random_engine = _random_engine;
     co_return ret;
 }
 
@@ -585,6 +591,19 @@ void topology::do_sort_by_proximity(locator::host_id address, host_id_vector_rep
         return info{ id, distance(address, loc, id, loc1) };
     }), std::back_inserter(host_infos));
     std::ranges::sort(host_infos, std::ranges::less{}, std::mem_fn(&info::distance));
+    auto it = host_infos.begin();
+    auto prev = it;
+    auto shuffler = _random_engine();
+    for (++it; it < host_infos.end(); ++it) {
+        // Shuffle equal-distance replicas to improve load-balancing
+        if (prev->distance == it->distance) {
+            if (shuffler & 1) {
+                std::swap(prev->id, it->id);
+            }
+            shuffler = std::rotr(shuffler, 1);
+        }
+        prev = it;
+    }
     std::ranges::copy(host_infos | std::ranges::views::transform(std::mem_fn(&info::id)), addresses.begin());
 }
 
@@ -637,6 +656,10 @@ topology::get_datacenter_host_ids() const {
     return ret;
 }
 
+void topology::seed_random_engine(random_engine_type::result_type value) {
+    _random_engine.seed(value);
+}
+
 } // namespace locator
 
 namespace std {

locator/topology.hh

@@ -17,6 +17,7 @@
 #include <unordered_map>
 #include <compare>
 #include <iostream>
+#include <random>
 
 #include <seastar/core/future.hh>
 #include <seastar/core/sstring.hh>
@@ -27,6 +28,8 @@
 
 using namespace seastar;
 
+struct sort_by_proximity_topology;
+
 namespace locator {
 class topology;
 }
@@ -409,6 +412,8 @@ public:
     }
 
 private:
+    using random_engine_type = std::mt19937_64;
+
     bool is_configured_this_node(const node&) const;
     const node& add_node(node_holder node);
     void remove_node(const node& node);
@@ -423,6 +428,8 @@ private:
         return const_cast<node*>(nptr);
     }
 
+    void seed_random_engine(random_engine_type::result_type);
+
     unsigned _shard;
     config _cfg;
     const node* _this_node = nullptr;
@@ -455,7 +462,10 @@ private:
         return _nodes_by_endpoint;
     };
 
+    mutable random_engine_type _random_engine;
+
     friend class token_metadata_impl;
+    friend struct ::sort_by_proximity_topology;
 public:
     void test_compare_endpoints(const locator::host_id& address, const locator::host_id& a1, const locator::host_id& a2) const;
     void test_sort_by_proximity(const locator::host_id& address, const host_id_vector_replica_set& nodes) const;

test/perf/perf_sort_by_proximity.cc

@@ -19,6 +19,8 @@
 #include <seastar/testing/test_runner.hh>
 
 #include "locator/token_metadata.hh"
+#include "test/lib/log.hh"
+#include "test/lib/random_utils.hh"
 
 struct sort_by_proximity_topology {
     static constexpr size_t DCS = 1;
@@ -59,6 +61,8 @@ struct sort_by_proximity_topology {
                     }
                 }
             }
+            auto seed = tests::random::get_int<locator::topology::random_engine_type::result_type>();
+            topology.seed_random_engine(seed);
         });
     }
 };