service: tablet_allocator: Introduce tablet load balancer

Will be invoked by the topology coordinator later to decide
which tablets to migrate.

locator/load_sketch.hh

@@ -12,6 +12,7 @@
 #include "locator/token_metadata.hh"
 #include "locator/tablets.hh"
 #include "utils/stall_free.hh"
+#include "utils/div_ceil.hh"
 
 #include <seastar/core/smp.hh>
 #include <seastar/coroutine/maybe_yield.hh>
@@ -45,6 +46,14 @@ class load_sketch {
                 s.load = 0;
             }
         }
+
+        uint64_t load() const {
+            uint64_t result = 0;
+            for (auto&& s : _shards) {
+                result += s.load;
+            }
+            return result;
+        }
     };
     std::unordered_map<host_id, node_load> _nodes;
     token_metadata_ptr _tm;
@@ -95,6 +104,21 @@ public:
         std::push_heap(n._shards.begin(), n._shards.end(), shard_load_cmp());
         return shard;
     }
+
+    uint64_t get_load(host_id node) const {
+        if (!_nodes.contains(node)) {
+            return 0;
+        }
+        return _nodes.at(node).load();
+    }
+
+    uint64_t get_avg_shard_load(host_id node) const {
+        if (!_nodes.contains(node)) {
+            return 0;
+        }
+        auto& n = _nodes.at(node);
+        return div_ceil(n.load(), n._shards.size());
+    }
 };
 
 } // namespace locator

locator/tablets.hh

@@ -80,6 +80,22 @@ std::ostream& operator<<(std::ostream&, const tablet_replica&);
 
 using tablet_replica_set = utils::small_vector<tablet_replica, 3>;
 
+/// Creates a new replica set with old_replica replaced by new_replica.
+/// If there is no old_replica, the set is returned unchanged.
+inline
+tablet_replica_set replace_replica(const tablet_replica_set& rs, tablet_replica old_replica, tablet_replica new_replica) {
+    tablet_replica_set result;
+    result.reserve(rs.size());
+    for (auto&& r : rs) {
+        if (r == old_replica) {
+            result.push_back(new_replica);
+        } else {
+            result.push_back(r);
+        }
+    }
+    return result;
+}
+
 /// Stores information about a single tablet.
 struct tablet_info {
     tablet_replica_set replicas;