Merge 'test: tablets_test: Create proper schema in load balancer tests' from Tomasz Grabiec

This PR converts boost load balancer tests in preparation for load balancer changes
which add per-table tablet hints. After those changes, load balancer consults with the replication
strategy in the database, so we need to create proper schema in the
database. To do that, we need proper topology for replication
strategies which use RF > 1, otherwise keyspace creation will fail.

Topology is created in tests via group0 commands, which is abstracted by
the new `topology_builder` class.

Tests cannot modify token_metadata only in memory now as it needs to be
consistent with the schema and on-disk metadata. That's why modifications to
tablet metadata are now made under group0 guard and save back metadata to disk.

Closes scylladb/scylladb#22648

* github.com:scylladb/scylladb:
  test: tablets: Drop keyspace after do_test_load_balancing_merge_colocation() scenario
  tests: tablets: Set initial tablets to 1 to exit growing mode
  test: tablets_test: Create proper schema in load balancer tests
  test: lib: Introduce topology_builder
  test: cql_test_env: Expose topology_state_machine
  topology_state_machine: Introduce lock transition

docs/dev/topology-over-raft.md

@@ -124,6 +124,9 @@ Additionally to specific node states, there entire topology can also be in a tra
     it from group 0. We also use this state to rollback a failed bootstrap or decommission.
 - `rollback_to_normal` - the decommission or removenode operation failed. Rollback the operation by
     moving the node we tried to decommission/remove back to the normal state.
+- `lock` - the topology stays in this state until externally changed (to null state), preventing topology
+    requests from starting. Intended to be used in tests which want to prevent internally-triggered topology
+    operations during the test.
 
 When a node bootstraps, we create new tokens for it and a new CDC generation
 and enter the `commit_cdc_generation` state. Once the generation is committed,

service/storage_service.cc

@@ -744,6 +744,8 @@ future<> storage_service::topology_state_load(state_change_hint hint) {
                 return read_new_t::no;
             }
             switch (*state) {
+                case topology::transition_state::lock:
+                    [[fallthrough]];
                 case topology::transition_state::join_group0:
                     [[fallthrough]];
                 case topology::transition_state::tablet_migration:

service/topology_coordinator.cc

@@ -2341,6 +2341,10 @@ class topology_coordinator : public endpoint_lifecycle_subscriber {
             case topology::transition_state::tablet_resize_finalization:
                 co_await handle_tablet_resize_finalization(std::move(guard));
                 break;
+            case topology::transition_state::lock:
+                release_guard(std::move(guard));
+                co_await await_event();
+                break;
             case topology::transition_state::left_token_ring: {
                 auto node = get_node_to_work_on(std::move(guard));
 

service/topology_state_machine.cc

@@ -153,6 +153,7 @@ static std::unordered_map<topology::transition_state, sstring> transition_state_
     {topology::transition_state::left_token_ring, "left token ring"},
     {topology::transition_state::rollback_to_normal, "rollback to normal"},
     {topology::transition_state::truncate_table, "truncate table"},
+    {topology::transition_state::lock, "lock"},
 };
 
 // Allows old deprecated names to be recognized and point to the correct transition.

service/topology_state_machine.hh

@@ -120,6 +120,7 @@ struct topology {
         left_token_ring,
         rollback_to_normal,
         truncate_table,
+        lock,
     };
 
     std::optional<transition_state> tstate;

test/lib/cql_test_env.cc

@@ -411,6 +411,10 @@ public:
         return _token_metadata;
     }
 
+    virtual sharded<service::topology_state_machine>& get_topology_state_machine() override {
+        return _topology_state_machine;
+    }
+
     virtual future<> refresh_client_state() override {
         return _core_local.invoke_on_all([] (core_local_state& state) {
             return state.client_state.maybe_update_per_service_level_params();

test/lib/cql_test_env.hh

@@ -189,6 +189,8 @@ public:
 
     virtual sharded<locator::shared_token_metadata>& get_shared_token_metadata() = 0;
 
+    virtual sharded<service::topology_state_machine>& get_topology_state_machine() = 0;
+
     data_dictionary::database data_dictionary();
 
     virtual sharded<qos::service_level_controller>& service_level_controller_service() = 0;

test/lib/topology_builder.hh

@@ -0,0 +1,210 @@
+/*
+ * Copyright (C) 2025-present ScyllaDB
+ *
+ */
+
+/*
+ * SPDX-License-Identifier: (LicenseRef-ScyllaDB-Source-Available-1.0 and Apache-2.0)
+ */
+
+#pragma once
+
+#include "cql_test_env.hh"
+#include "locator/topology.hh"
+#include "gms/inet_address.hh"
+#include "service/topology_mutation.hh"
+#include "service/topology_state_machine.hh"
+#include "service/raft/raft_group0_client.hh"
+#include "locator/host_id.hh"
+#include "test/lib/log.hh"
+#include "version.hh"
+
+#include <atomic>
+
+
+/// Modifies topology inside a given cql_test_env.
+/// Local node's membership is not affected, but it belongs to a different DC than those produced by this builder.
+///
+/// Creating the builder locks the topology state machine so there are no concurrent topology operations
+/// and load balancing.
+/// The built topology is not removed when the builder is destroyed and the state machine is left locked.
+///
+/// All methods expect to be run in a seastar thread.
+///
+/// Examples usage:
+///
+///     topology_builder topo(e);
+///     auto host1 = topo.add_node(); // dc1 rack1
+///     auto host2 = topo.add_node(); // dc1 rack1
+///     topo.start_new_dc();
+///     auto host3 = topo.add_node(); // dc2 rack1
+///     auto host4 = topo.add_node(); // dc2 rack1
+///     topo.start_new_rack();
+///     auto host5 = topo.add_node(); // dc2 rack2
+///     auto host6 = topo.add_node(); // dc2 rack2
+///
+class topology_builder {
+public:
+    using inet_address = locator::inet_address;
+    using endpoint_dc_rack = locator::endpoint_dc_rack;
+private:
+    cql_test_env& _env;
+    int _nr_nodes = 0;
+    int _rack_id;
+    sstring _dc;
+    sstring _rack;
+private:
+    inet_address make_node_address(int n) {
+        assert(n > 0);
+        int a = n % 256;
+        n /= 256;
+        int b = n % 256;
+        n /= 256;
+        assert(n < 256);
+        return inet_address(fmt::format("10.{}.{}.{}", n, b, a));
+    }
+
+    // Locks the topology to prevent concurrent topology operations and load balancing.
+    // Setting transition_state to "lock" blocks background load-balancing which could interfere with the test
+    // and prevents errors from load_topology_state() complaining about nodes in transition with no transition state.
+    void lock_topology() {
+        abort_source as;
+        auto& client = _env.get_raft_group0_client();
+        service::topology& topo = _env.get_topology_state_machine().local()._topology;
+
+        while (true) {
+            if (topo.tstate && *topo.tstate == service::topology::transition_state::lock) {
+                testlog.info("Topology is locked");
+                return;
+            }
+
+            auto guard = client.start_operation(as).get();
+
+            if (topo.tstate) {
+                testlog.info("Waiting for topology state machine to be idle");
+                release_guard(std::move(guard));
+                _env.get_topology_state_machine().local().await_not_busy().get();
+                testlog.info("Woken up");
+                continue;
+            }
+
+            service::topology_change change({service::topology_mutation_builder(guard.write_timestamp())
+                                                 .set_transition_state(service::topology::transition_state::lock)
+                                                 .build()});
+            service::group0_command g0_cmd = client.prepare_command(std::move(change), guard, "locking topology");
+            try {
+                client.add_entry(std::move(g0_cmd), std::move(guard), as).get();
+            } catch (service::group0_concurrent_modification&) {
+                testlog.info("Concurrent modification detected while locking topology, retrying");
+            }
+        }
+    }
+public:
+    topology_builder(cql_test_env& e)
+        : _env(e)
+    {
+        start_new_dc();
+        lock_topology();
+    }
+
+    // Returns a new token from some sequence of unique tokens.
+    // Uniqueness is in the scope of the process, not just this object.
+    dht::token new_token() {
+        static std::atomic<int64_t> next_token = 1;
+        return dht::token(next_token.fetch_add(1));
+    }
+
+    // Returns the name of the currently built DC.
+    const sstring& dc() const {
+        return _dc;
+    }
+
+    // Returns location of the currently built rack.
+    endpoint_dc_rack rack() const {
+        return {_dc, _rack};
+    }
+
+    // Starts building a new rack in the current DC.
+    // Returns location of the new rack.
+    endpoint_dc_rack start_new_rack() {
+        _rack_id++;
+        _rack = fmt::format("rack{}", _rack_id);
+        return rack();
+    }
+
+    // Starts building a new DC.
+    // DC is named uniquely in the scope of the process, not just this object.
+    endpoint_dc_rack start_new_dc() {
+        static std::atomic<int> next_id = 1;
+        _dc = fmt::format("dc{}", next_id.fetch_add(1));
+        _rack_id = 0;
+        return start_new_rack();
+    }
+
+    locator::host_id add_node(service::node_state state = service::node_state::normal,
+                              unsigned shard_count = 1,
+                              std::optional<endpoint_dc_rack> rack_override = {})
+    {
+        ++_nr_nodes;
+
+        auto ip = make_node_address(_nr_nodes);
+        auto id = locator::host_id(utils::UUID_gen::get_time_UUID());
+        auto dc_rack = rack_override.value_or(rack());
+        dht::token token = new_token();
+        std::unordered_set<dht::token> tokens({token});
+
+        abort_source as;
+        auto& client = _env.get_raft_group0_client();
+
+        while (true) {
+            auto guard = client.start_operation(as).get();
+
+            service::topology_mutation_builder builder(guard.write_timestamp());
+            builder.with_node(raft::server_id(id.uuid()))
+                    .set("datacenter", dc_rack.dc)
+                    .set("rack", dc_rack.rack)
+                    .set("node_state", state)
+                    .set("shard_count", (uint32_t) shard_count)
+                    .set("cleanup_status", service::cleanup_status::clean)
+                    .set("release_version", version::release())
+                    .set("num_tokens", (uint32_t) 1)
+                    .set("tokens_string", fmt::format("{}", token))
+                    .set("tokens", tokens)
+                    .set("supported_features", std::set<sstring>())
+                    .set("request_id", utils::UUID())
+                    .set("ignore_msb", (uint32_t) 0);
+            service::topology_change change({builder.build()});
+            service::group0_command g0_cmd = client.prepare_command(std::move(change), guard,
+                                                                    format("adding node {} to topology", id));
+            testlog.info("Adding node {}/{} dc={} rack={} to topology", id, ip, dc_rack.dc, dc_rack.rack);
+            try {
+                client.add_entry(std::move(g0_cmd), std::move(guard), as).get();
+                break;
+            } catch (service::group0_concurrent_modification&) {
+                testlog.warn("Concurrent modification detected, retrying");
+            }
+        }
+        return id;
+    }
+
+    void set_node_state(locator::host_id id, service::node_state state) {
+        abort_source as;
+        auto& client = _env.get_raft_group0_client();
+        while (true) {
+            auto guard = client.start_operation(as).get();
+            service::topology_mutation_builder builder(guard.write_timestamp());
+            builder.with_node(raft::server_id(id.uuid()))
+                    .set("node_state", state);
+            service::topology_change change({builder.build()});
+            service::group0_command g0_cmd = client.prepare_command(std::move(change), guard,
+                                                                    format("node {} state={}", id, state));
+            testlog.info("Changing node {} state={}", id, state);
+            try {
+                client.add_entry(std::move(g0_cmd), std::move(guard), as).get();
+                break;
+            } catch (service::group0_concurrent_modification&) {
+                testlog.warn("Concurrent modification detected, retrying");
+            }
+        }
+    }
+};