fix(nemesis): handle new keyspace being written to during AddRemoveDCNemesis

sdcm/nemesis/monkey/add_remove_dc.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterator
+from collections.abc import Generator
 from contextlib import ExitStack, contextmanager
 
 from sdcm.cluster import BaseNode
@@ -51,7 +51,7 @@ def num_nodes_in_new_dc(self) -> int:
         if self.runner.cluster.is_features_enabled_on_node(
             node=self.runner.target_node, feature_list=["KEYSPACE_MULTI_RF_CHANGE"]
         ):
-            return 2
+            return 3
         return 1
 
     @property
@@ -65,7 +65,7 @@ def system_keyspaces(self) -> list[str]:
         return system_keyspaces
 
     @contextmanager
-    def temporary_system_keyspaces_network_topology_strategy(self) -> Iterator[None]:
+    def temporary_system_keyspaces_network_topology_strategy(self) -> Generator[None]:
         """Temporarily switch system keyspaces to NetworkTopologyStrategy."""
         with (
             temporary_replication_strategy_setter(self.runner.target_node) as ntrs_setter,
@@ -78,8 +78,45 @@ def temporary_system_keyspaces_network_topology_strategy(self) -> Iterator[None]
                     ntrs_setter(**{keyspace: new_rs})
             yield
 
+    def _preserve_keyspaces_with_new_dc_replication(self, new_dc_setter: temporary_replication_strategy_setter) -> None:
+        """Add rollback records for keyspaces created while the temporary DC exists."""
+        new_dc_name = self.new_dc_name
+        if not new_dc_name:
+            return
+
+        try:
+            with self.runner.cluster.cql_connection_patient(self.runner.target_node) as session:
+                result = session.execute("SELECT keyspace_name, replication FROM system_schema.keyspaces")
+                keyspace_rows = result.current_rows
+        except Exception as exc:  # noqa: BLE001 - best-effort cleanup must not mask nemesis failure
+            self.runner.log.warning(f"Failed to discover keyspaces with replication to {new_dc_name}: {exc}")
+            return
+
+        for row in keyspace_rows:
+            keyspace = row.keyspace_name
+            replication = row.replication
+            self.runner.log.info(f"Found {keyspace=} with {replication=}")
+            if keyspace in new_dc_setter.preserved:
+                continue
+            if "NetworkTopologyStrategy" not in replication.get("class", ""):
+                continue
+            if new_dc_name not in replication:
+                continue
+
+            try:
+                strategy = ReplicationStrategy.get(self.runner.target_node, keyspace)
+                rollback_strategy = NetworkTopologyReplicationStrategy(
+                    **{**strategy.replication_factors_per_dc, new_dc_name: 0}
+                )
+            except Exception as exc:  # noqa: BLE001 - keep rollback of other keyspaces best-effort
+                self.runner.log.warning(f"Failed to preserve rollback strategy for keyspace {keyspace}: {exc}")
+                continue
+
+            self.runner.log.info(f"Preserve rollback strategy for keyspace {keyspace} with RF=0 in {new_dc_name}")
+            new_dc_setter.preserved[keyspace] = rollback_strategy
+
     @contextmanager
-    def temporary_new_dc_replication_factors(self) -> Iterator[None]:
+    def temporary_new_dc_replication_factors(self) -> Generator[None]:
         """Temporarily update replication factors for the new datacenter."""
         with (
             temporary_replication_strategy_setter(self.runner.target_node) as new_dc_setter,
@@ -99,7 +136,10 @@ def temporary_new_dc_replication_factors(self) -> Iterator[None]:
                     )
                     new_dc_setter.preserved[keyspace] = rollback_strategy
 
-            yield
+            try:
+                yield
+            finally:
+                self._preserve_keyspaces_with_new_dc_replication(new_dc_setter)
 
     @property
     def datacenters(self) -> list[str]:

sdcm/utils/replication_strategy_utils.py

@@ -5,6 +5,8 @@
 from contextlib import ContextDecorator
 from typing import Callable, Dict, TYPE_CHECKING
 
+from cassandra import InvalidRequest
+
 from sdcm.exceptions import DatacenterNotResolvedError
 from sdcm.utils.cql_utils import cql_quote_if_needed
 from sdcm.utils.database_query_utils import is_system_keyspace, LOGGER
@@ -126,8 +128,13 @@ def _preserve_replication_strategy(self, keyspace: str) -> None:
 
     def __call__(self, **keyspaces: ReplicationStrategy) -> None:
         for keyspace, strategy in keyspaces.items():
-            self._preserve_replication_strategy(keyspace)
-            strategy.apply(self.node, keyspace)
+            try:
+                self._preserve_replication_strategy(keyspace)
+                strategy.apply(self.node, keyspace)
+            except InvalidRequest as exc:
+                # A concurrent nemesis may have dropped this keyspace in the meantime:
+                # skip it instead of aborting rollback of the remaining keyspaces.
+                LOGGER.warning("Skipping replication strategy update for keyspace %s: %s", keyspace, exc)
 
 
 class DataCenterTopologyRfControl:

unit_tests/unit/nemesis/monkey/test_add_remove_dc.py

@@ -1,6 +1,7 @@
 """Tests for sdcm.nemesis.monkey.add_remove_dc module."""
 
 from threading import Event
+from types import SimpleNamespace
 from unittest.mock import MagicMock, call, patch
 
 import pytest
@@ -18,6 +19,11 @@
 pytestmark = pytest.mark.usefixtures("events")
 
 
+def _keyspace_row(name, replication):
+    """Build a system_schema.keyspaces row for replication discovery tests."""
+    return SimpleNamespace(keyspace_name=name, replication=replication)
+
+
 def _clone_strategy(strategy):
     """Creates a deep copy of a replication strategy for comparison.
 
@@ -155,8 +161,31 @@ def test_disrupt_updates_replication_and_cleans_new_dc(runner):
     monkey = AddRemoveDcNemesis(runner)
     new_nodes = [MagicMock(name="node-new-1"), MagicMock(name="node-new-2")]
     events = []
+    created_keyspace = "created_during_nemesis"
     first_setter = FakeReplicationStrategySetter("system-keyspaces", events)
     second_setter = FakeReplicationStrategySetter("new-dc-rf", events)
+    session = runner.cluster.cql_connection_patient.return_value.__enter__.return_value
+    session.execute.side_effect = None
+    session.execute.return_value = MagicMock(
+        current_rows=[
+            _keyspace_row(
+                created_keyspace,
+                {
+                    "class": "org.apache.cassandra.locator.NetworkTopologyStrategy",
+                    monkey.initial_dc_name: str(monkey.new_ks_rf),
+                    "dc1_nemesis_dc": str(monkey.num_nodes_in_new_dc),
+                },
+            ),
+            _keyspace_row("local_strategy_keyspace", {"class": "org.apache.cassandra.locator.LocalStrategy"}),
+            _keyspace_row(
+                "other_dc_keyspace",
+                {
+                    "class": "org.apache.cassandra.locator.NetworkTopologyStrategy",
+                    monkey.initial_dc_name: str(monkey.new_ks_rf),
+                },
+            ),
+        ]
+    )
 
     def add_nodes():
         monkey.new_nodes = new_nodes
@@ -180,6 +209,10 @@ def fake_get(node, keyspace):
             return NetworkTopologyReplicationStrategy(**{monkey.initial_dc_name: monkey.new_ks_rf})
         if keyspace in ("system_distributed", "system_traces"):
             return SimpleReplicationStrategy(monkey.new_ks_rf)
+        if keyspace == created_keyspace:
+            return NetworkTopologyReplicationStrategy(
+                **{monkey.initial_dc_name: monkey.new_ks_rf, "dc1_nemesis_dc": monkey.num_nodes_in_new_dc}
+            )
         return NetworkTopologyReplicationStrategy(**{monkey.initial_dc_name: monkey.new_ks_rf})
 
     with (
@@ -234,12 +267,56 @@ def fake_get(node, keyspace):
         }
 
     assert len(second_setter.rollback_calls) == 1
-    assert list(second_setter.rollback_calls[0]) == updated_keyspaces
+    assert list(second_setter.rollback_calls[0]) == updated_keyspaces + [created_keyspace]
     assert events.index("new-dc-rf:rollback") < events.index("decommission")
 
     assert monkey.new_dc_name is None
 
 
+def test_temporary_new_dc_replication_factors_preserves_created_keyspaces_on_error(runner):
+    """Keyspaces created during the temporary DC window should be rolled back on errors too."""
+    monkey = AddRemoveDcNemesis(runner)
+    runner.status_by_dc["dc1_nemesis_dc"] = object()
+    created_keyspace = "created_during_error"
+    setter = FakeReplicationStrategySetter("new-dc-rf", [])
+    session = runner.cluster.cql_connection_patient.return_value.__enter__.return_value
+    session.execute.side_effect = None
+    session.execute.return_value = MagicMock(
+        current_rows=[
+            _keyspace_row(
+                created_keyspace,
+                {
+                    "class": "org.apache.cassandra.locator.NetworkTopologyStrategy",
+                    monkey.initial_dc_name: str(monkey.new_ks_rf),
+                    "dc1_nemesis_dc": str(monkey.num_nodes_in_new_dc),
+                },
+            )
+        ]
+    )
+
+    def fake_get(node, keyspace):
+        if keyspace == created_keyspace:
+            return NetworkTopologyReplicationStrategy(
+                **{monkey.initial_dc_name: monkey.new_ks_rf, "dc1_nemesis_dc": monkey.num_nodes_in_new_dc}
+            )
+        return NetworkTopologyReplicationStrategy(**{monkey.initial_dc_name: monkey.new_ks_rf})
+
+    with pytest.raises(RuntimeError, match="boom"):
+        with (
+            patch(f"{_MODULE}.temporary_replication_strategy_setter", return_value=setter),
+            patch(f"{_MODULE}.ReplicationStrategy.get", side_effect=fake_get),
+        ):
+            with monkey.temporary_new_dc_replication_factors():
+                raise RuntimeError("boom")
+
+    expected_keyspaces = ["system_distributed", "system_traces", monkey.new_ks_name, created_keyspace]
+    assert list(setter.rollback_calls[0]) == expected_keyspaces
+    assert setter.preserved[created_keyspace].replication_factors_per_dc == {
+        monkey.initial_dc_name: monkey.new_ks_rf,
+        "dc1_nemesis_dc": 0,
+    }
+
+
 def test_assert_new_dc_registered_retries_until_status_contains_new_dc(runner):
     """assert_new_dc_registered() should retry until nodetool status shows the temporary DC."""
     monkey = AddRemoveDcNemesis(runner)
@@ -343,7 +420,7 @@ def test_finalizer_skips_decommission_when_no_new_nodes(runner):
 @pytest.mark.parametrize(
     "feature_enabled,expected_count",
     [
-        pytest.param(True, 2, id="multi-rf-change-enabled"),
+        pytest.param(True, 3, id="multi-rf-change-enabled"),
         pytest.param(False, 1, id="multi-rf-change-disabled"),
     ],
 )