test: tablets: Add test for failed streaming being fenced away

tgrabiec · tgrabiec · commit 7d0f4c10a264 · 2023-12-06T18:37:01.000+01:00
diff --git a/streaming/stream_session.cc b/streaming/stream_session.cc
@@ -38,6 +38,7 @@
 #include "consumer.hh"
 #include "readers/generating_v2.hh"
 #include "service/topology_guard.hh"
+#include "utils/error_injection.hh"
 
 namespace streaming {
 
@@ -127,18 +128,27 @@ void stream_manager::init_messaging_service_handler(abort_source& as) {
             return make_exception_future<rpc::sink<int>>(std::runtime_error(format("Node {} is not fully initialized for streaming, try again later",
                     utils::fb_utilities::get_broadcast_address())));
         }
-        return _mm.local().get_schema_for_write(schema_id, from, _ms.local(), as).then([this, from, estimated_partitions, plan_id, cf_id, source, reason, topo_guard] (schema_ptr s) mutable {
-          return _db.local().obtain_reader_permit(s, "stream-session", db::no_timeout, {}).then([this, from, estimated_partitions, plan_id, cf_id, source, reason, topo_guard, s] (reader_permit permit) mutable {
+        return _mm.local().get_schema_for_write(schema_id, from, _ms.local(), as).then([this, from, estimated_partitions, plan_id, cf_id, source, reason, topo_guard, &as] (schema_ptr s) mutable {
+          return _db.local().obtain_reader_permit(s, "stream-session", db::no_timeout, {}).then([this, from, estimated_partitions, plan_id, cf_id, source, reason, topo_guard, s, &as] (reader_permit permit) mutable {
             struct stream_mutation_fragments_cmd_status {
                 bool got_cmd = false;
                 bool got_end_of_stream = false;
             };
             auto cmd_status = make_lw_shared<stream_mutation_fragments_cmd_status>();
             auto offstrategy_update = make_lw_shared<offstrategy_trigger>(_db, cf_id, plan_id);
             auto guard = service::topology_guard(s->table(), topo_guard);
-            auto get_next_mutation_fragment = [guard = std::move(guard), &sm = container(), source, plan_id, from, s, cmd_status, offstrategy_update, permit] () mutable {
+
+            // Will log a message when streaming is done. Used to synchronize tests.
+            lw_shared_ptr<std::any> log_done;
+            if (utils::get_local_injector().is_enabled("stream_mutation_fragments")) {
+                log_done = make_lw_shared<std::any>(seastar::make_shared(seastar::defer([] {
+                    sslog.info("stream_mutation_fragments: done");
+                })));
+            }
+
+            auto get_next_mutation_fragment = [guard = std::move(guard), &as, &sm = container(), source, plan_id, from, s, cmd_status, offstrategy_update, permit] () mutable {
                 guard.check();
-                return source().then([&sm, plan_id, from, s, cmd_status, offstrategy_update, permit] (std::optional<std::tuple<frozen_mutation_fragment, rpc::optional<stream_mutation_fragments_cmd>>> opt) mutable {
+                return source().then([&sm, &guard, &as, plan_id, from, s, cmd_status, offstrategy_update, permit] (std::optional<std::tuple<frozen_mutation_fragment, rpc::optional<stream_mutation_fragments_cmd>>> opt) mutable {
                     if (opt) {
                         auto cmd = std::get<1>(*opt);
                         if (cmd) {
@@ -160,7 +170,19 @@ void stream_manager::init_messaging_service_handler(abort_source& as) {
                         auto mf = fmf.unfreeze(*s, permit);
                         sm.local().update_progress(plan_id, from.addr, progress_info::direction::IN, sz);
                         offstrategy_update->update();
-                        return make_ready_future<mutation_fragment_opt>(std::move(mf));
+
+                        return utils::get_local_injector().inject_with_handler("stream_mutation_fragments", [&guard, &as] (auto& handler) -> future<> {
+                            auto& guard_ = guard;
+                            auto& as_ = as;
+                            sslog.info("stream_mutation_fragments: waiting");
+                            while (!handler.poll_for_message()) {
+                                guard_.check();
+                                co_await sleep_abortable(std::chrono::milliseconds(5), as_);
+                            }
+                            sslog.info("stream_mutation_fragments: released");
+                        }).then([mf = std::move(mf)] () mutable {
+                            return mutation_fragment_opt(std::move(mf));
+                        });
                     } else {
                         // If the sender has sent stream_mutation_fragments_cmd it means it is
                         // a node that understands the new protocol. It must send end_of_stream
@@ -185,7 +207,7 @@ void stream_manager::init_messaging_service_handler(abort_source& as) {
                 make_generating_reader_v1(s, permit, std::move(get_next_mutation_fragment)),
                 make_streaming_consumer("streaming", _db, _sys_dist_ks, _view_update_generator, estimated_partitions, reason, is_offstrategy_supported(reason), topo_guard),
                 std::move(op)
-            ).then_wrapped([s, plan_id, from, sink, estimated_partitions, sh_ptr = std::move(sharder_ptr)] (future<uint64_t> f) mutable {
+            ).then_wrapped([s, plan_id, from, sink, estimated_partitions, log_done, sh_ptr = std::move(sharder_ptr)] (future<uint64_t> f) mutable {
                 int32_t status = 0;
                 uint64_t received_partitions = 0;
                 if (f.failed()) {
diff --git a/test/pylib/manager_client.py b/test/pylib/manager_client.py
@@ -337,6 +337,10 @@ async def get_host_id(self, server_id: ServerNum) -> HostID:
             raise Exception(f"Failed to get local host id address for server {server_id}") from exc
         return HostID(host_id)
 
+    async def get_table_id(self, keyspace: str, table: str):
+        rows = await self.cql.run_async(f"select id from system_schema.tables where keyspace_name = '{keyspace}' and table_name = '{table}'")
+        return rows[0].id
+
     async def server_sees_others(self, server_id: ServerNum, count: int, interval: float = 45.):
         """Wait till a server sees a minimum given count of other servers"""
         if count < 1:
diff --git a/test/topology_experimental_raft/test_tablets.py b/test/topology_experimental_raft/test_tablets.py
@@ -3,11 +3,15 @@
 #
 # SPDX-License-Identifier: AGPL-3.0-or-later
 #
+from uuid import UUID
+
+from cassandra.query import SimpleStatement, ConsistencyLevel
 
 from test.pylib.manager_client import ManagerClient
-from test.pylib.rest_client import inject_error_one_shot
+from test.pylib.rest_client import inject_error_one_shot, HTTPError
 from test.pylib.rest_client import inject_error
 from test.pylib.util import wait_for_cql_and_get_hosts
+from test.topology.conftest import skip_mode
 from test.topology.util import reconnect_driver
 
 import pytest
@@ -29,6 +33,21 @@ async def inject_error_on(manager, error_name, servers):
     await asyncio.gather(*errs)
 
 
+async def get_tablet_replicas(manager, keyspace_name, table_name, token):
+    table_id = await manager.get_table_id(keyspace_name, table_name)
+    rows = await manager.cql.run_async(f"SELECT last_token, replicas FROM system.tablets where "
+                                       f"keyspace_name = '{keyspace_name}' and "
+                                       f"table_id = {table_id}")
+    for row in rows:
+        if row.last_token >= token:
+            return row.replicas
+
+
+async def get_tablet_replica(manager, keyspace_name, table_name, token):
+    replicas = await get_tablet_replicas(manager, keyspace_name, table_name, token)
+    return replicas[0]
+
+
 @pytest.mark.asyncio
 async def test_tablet_metadata_propagates_with_schema_changes_in_snapshot_mode(manager: ManagerClient):
     """Test that you can create a table and insert and query data"""
@@ -199,3 +218,77 @@ async def check():
     await check()
 
     await cql.run_async("DROP KEYSPACE test;")
+
+
+@pytest.mark.asyncio
+@skip_mode('release', 'error injections are not supported in release mode')
+async def test_streaming_is_guarded_by_topology_guard(manager: ManagerClient):
+    logger.info("Bootstrapping cluster")
+    cmdline = [
+        '--logger-log-level', 'storage_service=trace',
+    ]
+    servers = [await manager.server_add(cmdline=cmdline)]
+
+    await manager.api.disable_tablet_balancing(servers[0].ip_addr)
+
+    cql = manager.get_cql()
+    await cql.run_async("CREATE KEYSPACE test WITH replication = {'class': 'NetworkTopologyStrategy', "
+                        "'replication_factor': 1, 'initial_tablets': 1};")
+    await cql.run_async("CREATE TABLE test.test (pk int PRIMARY KEY, c int);")
+
+    servers.append(await manager.server_add(cmdline=cmdline))
+
+    key = 7 # Whatever
+    tablet_token = 0 # Doesn't matter since there is one tablet
+    await cql.run_async(f"INSERT INTO test.test (pk, c) VALUES ({key}, 0)")
+    rows = await cql.run_async("SELECT pk from test.test")
+    assert len(list(rows)) == 1
+
+    replica = await get_tablet_replica(manager, 'test', 'test', tablet_token)
+
+    s0_host_id = await manager.get_host_id(servers[0].server_id)
+    s1_host_id = await manager.get_host_id(servers[1].server_id)
+    dst_shard = 0
+
+    await manager.api.enable_injection(servers[1].ip_addr, "stream_mutation_fragments", one_shot=True)
+    s1_log = await manager.server_open_log(servers[1].server_id)
+    s1_mark = await s1_log.mark()
+
+    migration_task = asyncio.create_task(
+        manager.api.move_tablet(servers[0].ip_addr, "test", "test", replica[0], replica[1], s1_host_id, dst_shard, tablet_token))
+
+    # Wait for the replica-side writer of streaming to reach a place where it already
+    # received writes from the leaving replica but haven't applied them yet.
+    # Once the writer reaches this place, it will wait for the message_injection() call below before proceeding.
+    # The place we block the writer in should not hold to erm or topology_guard because that will block the migration
+    # below and prevent test from proceeding.
+    await s1_log.wait_for('stream_mutation_fragments: waiting', from_mark=s1_mark)
+    s1_mark = await s1_log.mark()
+
+    # Should cause streaming to fail and be retried while leaving behind the replica-side writer.
+    await manager.api.inject_disconnect(servers[0].ip_addr, servers[1].ip_addr)
+
+    logger.info("Waiting for migration to finish")
+    await migration_task
+    logger.info("Migration done")
+
+    # Sanity test
+    rows = await cql.run_async("SELECT pk from test.test")
+    assert len(list(rows)) == 1
+
+    await cql.run_async("TRUNCATE test.test")
+    rows = await cql.run_async("SELECT pk from test.test")
+    assert len(list(rows)) == 0
+
+    # Release abandoned streaming
+    await manager.api.message_injection(servers[1].ip_addr, "stream_mutation_fragments")
+    await s1_log.wait_for('stream_mutation_fragments: done', from_mark=s1_mark)
+
+    # Verify that there is no data resurrection
+    rows = await cql.run_async("SELECT pk from test.test")
+    assert len(list(rows)) == 0
+
+    # Verify that moving the tablet back works
+    await manager.api.move_tablet(servers[0].ip_addr, "test", "test", s1_host_id, dst_shard, replica[0], replica[1], tablet_token)
+    rows = await cql.run_async("SELECT pk from test.test")
+    assert len(list(rows)) == 0
