test: add test for snapshoting during migration (#4108)

* test: add test for snapshotting during migration

* test: add test to run replication after migration

Author: BorysTheDev

tests/dragonfly/cluster_test.py

@@ -1567,6 +1567,150 @@ async def test_cluster_replication_migration(
     assert await seeder.compare(r1_capture, r2_node.instance.port)
 
 
+@dfly_args({"proactor_threads": 4, "cluster_mode": "yes"})
+async def test_start_replication_during_migration(
+    df_factory: DflyInstanceFactory, df_seeder_factory: DflySeederFactory
+):
+    """
+    Test replication with migration. Create the following setup:
+
+    master_1 do migration to master_2 and we start replication for master_1 during this migration
+
+    in the end master_1 and replica_1 should have the same data
+    """
+    instances = [
+        df_factory.create(port=BASE_PORT + i, admin_port=BASE_PORT + 1000 + i) for i in range(3)
+    ]
+    df_factory.start_all(instances)
+
+    nodes = [await create_node_info(n) for n in instances]
+    m1_node, r1_node, m2_node = nodes
+    master_nodes = [m1_node, m2_node]
+
+    m1_node.slots = [(0, 16383)]
+    m1_node.replicas = [r1_node]
+    m2_node.slots = []
+
+    logging.debug("Push initial config")
+    await push_config(
+        json.dumps(generate_config(master_nodes)), [node.admin_client for node in nodes]
+    )
+
+    logging.debug("create data")
+    seeder = df_seeder_factory.create(keys=10000, port=nodes[0].instance.port, cluster_mode=True)
+    await seeder.run(target_deviation=0.1)
+
+    logging.debug("start migration")
+    m1_node.migrations = [
+        MigrationInfo("127.0.0.1", m2_node.instance.admin_port, [(2001, 16383)], m2_node.id)
+    ]
+    await push_config(
+        json.dumps(generate_config(master_nodes)), [node.admin_client for node in nodes]
+    )
+
+    logging.debug("start replication")
+    await r1_node.admin_client.execute_command(f"replicaof localhost {m1_node.instance.port}")
+
+    await wait_available_async(r1_node.admin_client)
+
+    await wait_for_status(m1_node.admin_client, m2_node.id, "FINISHED")
+
+    logging.debug("finish migration")
+    m1_node.migrations = []
+    m1_node.slots = [(0, 2000)]
+    m2_node.migrations = []
+    m2_node.slots = [(2001, 16383)]
+
+    await push_config(
+        json.dumps(generate_config(master_nodes)), [node.admin_client for node in nodes]
+    )
+
+    await check_all_replicas_finished([r1_node.client], m1_node.client)
+
+    m1_capture = await seeder.capture(m1_node.instance.port)
+
+    assert await seeder.compare(m1_capture, r1_node.instance.port)
+
+
+@pytest.mark.parametrize("migration_first", [False, True])
+@dfly_args({"proactor_threads": 4, "cluster_mode": "yes", "dbfilename": "snap_during_migration"})
+async def test_snapshoting_during_migration(
+    df_factory: DflyInstanceFactory, df_seeder_factory: DflySeederFactory, migration_first: bool
+):
+    """
+    Test saving snapshot during migration. Create the following setups:
+
+    1) Start saving and then run migration simultaneously
+    2) Run migration and start saving simultaneously
+
+    The result should be the same: snapshot contains all the data that existed before migration
+    """
+    instances = [
+        df_factory.create(port=BASE_PORT + i, admin_port=BASE_PORT + 1000 + i) for i in range(2)
+    ]
+    df_factory.start_all(instances)
+
+    nodes = [await create_node_info(n) for n in instances]
+
+    nodes[0].slots = [(0, 16383)]
+    nodes[1].slots = []
+
+    logging.debug("Push initial config")
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    logging.debug("create data")
+    seeder = df_seeder_factory.create(keys=10000, port=nodes[0].instance.port, cluster_mode=True)
+    await seeder.run(target_deviation=0.1)
+
+    capture_before_migration = await seeder.capture(nodes[0].instance.port)
+
+    nodes[0].migrations = [
+        MigrationInfo("127.0.0.1", nodes[1].instance.admin_port, [(0, 16383)], nodes[1].id)
+    ]
+
+    async def start_migration():
+        logging.debug("start migration")
+        await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    async def start_save():
+        logging.debug("BGSAVE")
+        await nodes[0].client.execute_command(f"BGSAVE")
+
+    if migration_first:
+        await start_migration()
+        await asyncio.sleep(random.randint(0, 10) / 100)
+        await start_save()
+    else:
+        await start_save()
+        await asyncio.sleep(random.randint(0, 10) / 100)
+        await start_migration()
+
+    logging.debug("wait for snapshot")
+    while await is_saving(nodes[0].client):
+        await asyncio.sleep(0.1)
+
+    logging.debug("wait migration finish")
+    await wait_for_status(nodes[0].admin_client, nodes[1].id, "FINISHED")
+
+    logging.debug("finish migration")
+    nodes[0].migrations = []
+    nodes[0].slots = []
+    nodes[0].migrations = []
+    nodes[0].slots = [(0, 16383)]
+
+    await push_config(json.dumps(generate_config(nodes)), [node.admin_client for node in nodes])
+
+    assert await seeder.compare(capture_before_migration, nodes[1].instance.port)
+
+    await nodes[1].client.execute_command(
+        "DFLY",
+        "LOAD",
+        "snap_during_migration-summary.dfs",
+    )
+
+    assert await seeder.compare(capture_before_migration, nodes[1].instance.port)
+
+
 @dfly_args({"proactor_threads": 4, "cluster_mode": "yes"})
 @pytest.mark.asyncio
 async def test_cluster_migration_cancel(df_factory: DflyInstanceFactory):