From 5fccee3a130c2d124a2add4aec0fd4f3350b954c Mon Sep 17 00:00:00 2001
From: Tomasz Grabiec <tgrabiec@scylladb.com>
Date: Sun, 21 Jan 2024 00:20:01 +0100
Subject: [PATCH] storage_service, topology: Track excluded state in
 locator::topology

Will be used by tablet load balancer to avoid excluded nodes in
scheduling.
---
 locator/topology.cc        |  6 ++++++
 locator/topology.hh        | 12 ++++++++++++
 service/storage_service.cc |  8 ++++++++
 3 files changed, 26 insertions(+)

diff --git a/locator/topology.cc b/locator/topology.cc
index 4f8c75cb97..4ffbe6792a 100644
--- a/locator/topology.cc
+++ b/locator/topology.cc
@@ -423,6 +423,12 @@ const node* topology::find_node(host_id id) const noexcept {
     return nullptr;
 }
 
+// Finds a node by its host_id
+// Returns nullptr if not found
+node* topology::find_node(host_id id) noexcept {
+    return make_mutable(const_cast<const topology*>(this)->find_node(id));
+}
+
 // Finds a node by its endpoint
 // Returns nullptr if not found
 const node* topology::find_node(const inet_address& ep) const noexcept {
diff --git a/locator/topology.hh b/locator/topology.hh
index 77bb367332..3a4cdf2ced 100644
--- a/locator/topology.hh
+++ b/locator/topology.hh
@@ -63,6 +63,7 @@ private:
     endpoint_dc_rack _dc_rack;
     state _state;
     shard_id _shard_count = 0;
+    bool _excluded = false;
 
     // Is this node the `localhost` instance
     this_node _is_this_node;
@@ -119,6 +120,16 @@ public:
         return _state == state::normal;
     }
 
+    // Excluded nodes are still part of topology, possibly present in replica sets, but
+    // are not going to be up again and can be ignored in barriers.
+    bool is_excluded() const {
+        return _excluded;
+    }
+
+    void set_excluded(bool excluded) {
+        _excluded = excluded;
+    }
+
     bool is_leaving() const noexcept {
         switch (_state) {
         case state::being_decommissioned:
@@ -208,6 +219,7 @@ public:
     // Looks up a node by its host_id.
     // Returns a pointer to the node if found, or nullptr otherwise.
     const node* find_node(host_id id) const noexcept;
+    node* find_node(host_id id) noexcept;
 
     const node& get_node(host_id id) const {
         auto n = find_node(id);
diff --git a/service/storage_service.cc b/service/storage_service.cc
index cc3b925564..0d46fb096f 100644
--- a/service/storage_service.cc
+++ b/service/storage_service.cc
@@ -534,6 +534,12 @@ future<> storage_service::sync_raft_topology_nodes(mutable_token_metadata_ptr tm
     for (const auto& [id, rs]: t.transition_nodes) {
         co_await process_transition_node(id, rs);
     }
+    for (auto id : t.get_excluded_nodes()) {
+        locator::node* n = tmptr->get_topology().find_node(locator::host_id(id.uuid()));
+        if (n) {
+            n->set_excluded(true);
+        }
+    }
 }
 
 future<> storage_service::topology_state_load() {
@@ -634,6 +640,8 @@ future<> storage_service::topology_state_load() {
         rtlogger.debug("topology_state_load: current CDC generation ID: {}", *gen_id);
         co_await _cdc_gens.local().handle_cdc_generation(*gen_id);
     }
+
+    slogger.debug("topology_state_load: excluded nodes: {}", _topology_state_machine._topology.get_excluded_nodes());
 }
 
 future<> storage_service::topology_transition() {