miroir/crates/miroir-core/tests/p42_node_addition.rs

//! P4.2 Node addition: dual-write + paginated shard migration integration tests.
//!
//! Implements acceptance criteria from plan §2 "Adding a node to an existing group":
//! 1. Integration test: 3-node → 4-node migration, 10K docs, each doc still retrievable by ID after migration
//! 2. Chaos: toggle writes on/off during migration; dual-write window catches all late writes
//! 3. Performance: migrating `~S/(Ng+1)` shards moves ≤ `total_docs / (Ng+1) × 1.1` docs (10% slack for dual-write dupes)
//! 4. The old node is not queried for the migrated shards after step 8 (verified via log inspection)

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;

// Type aliases to reduce complexity
type FetchCallsMap = HashMap<(String, u32, u32), usize>;
type StoredDocsMap = HashMap<(String, u32), Vec<serde_json::Value>>;

use miroir_core::{
    migration::MigrationConfig,
    rebalancer::{MigrationExecutor, Rebalancer, RebalancerConfig},
    router::assign_shard_in_group,
    topology::{Node, NodeId, NodeStatus, Topology},
};

/// Helper: create a test topology with N nodes in a single replica group.
fn create_test_topology(shards: u32, node_count: usize) -> Topology {
    let mut topo = Topology::new(shards, 1, 2); // RF=2 for redundancy
    for i in 0..node_count {
        let mut node = Node::new(
            NodeId::new(format!("node-{i}")),
            format!("http://node-{i}:7700"),
            0,
        );
        node.status = NodeStatus::Active; // Start as active for tests
        topo.add_node(node);
    }
    topo
}

/// Mock migration executor that tracks which nodes were queried.
#[derive(Default)]
struct MockMigrationExecutor {
    /// Track all fetch_documents calls: (node, shard_id, offset) -> count
    fetch_calls: Arc<std::sync::Mutex<FetchCallsMap>>,
    /// Track fetch calls in sequence order: (node, shard_id, sequence_number)
    fetch_sequence: Arc<std::sync::Mutex<Vec<(String, u32, u64)>>>,
    /// Track all write_documents calls: (node,) -> doc_count
    write_calls: Arc<std::sync::Mutex<HashMap<String, usize>>>,
    /// Track all delete_shard calls: (node, shard_id) -> count
    delete_calls: Arc<std::sync::Mutex<HashMap<(String, u32), usize>>>,
    /// Documents stored per (node, shard)
    stored_docs: Arc<std::sync::Mutex<StoredDocsMap>>,
    /// Write failure simulation: (node, shard_id) -> should_fail
    write_failures: Arc<std::sync::Mutex<HashMap<(String, u32), bool>>>,
    /// Counter for sequencing fetch calls
    fetch_counter: Arc<std::sync::atomic::AtomicU64>,
}

impl MockMigrationExecutor {
    #[allow(dead_code)]
    fn add_write_failure(&self, node: &str, shard_id: u32) {
        self.write_failures
            .lock()
            .unwrap()
            .insert((node.to_string(), shard_id), true);
    }

    #[allow(dead_code)]
    fn clear_write_failures(&self) {
        self.write_failures.lock().unwrap().clear();
    }

    fn get_stored_doc_count(&self, node: &str, shard_id: u32) -> usize {
        self.stored_docs
            .lock()
            .unwrap()
            .get(&(node.to_string(), shard_id))
            .map(|v| v.len())
            .unwrap_or(0)
    }

    fn was_queried_after(&self, node: &str, shard_id: u32, after_sequence: u64) -> bool {
        self.fetch_sequence
            .lock()
            .unwrap()
            .iter()
            .any(|(n, s, seq)| n == node && *s == shard_id && *seq > after_sequence)
    }

    fn get_latest_fetch_sequence(&self) -> u64 {
        self.fetch_counter.load(std::sync::atomic::Ordering::SeqCst)
    }

    #[allow(dead_code)]
    fn total_fetched_docs(&self) -> usize {
        self.fetch_calls.lock().unwrap().len()
    }

    fn total_written_docs(&self) -> usize {
        self.write_calls.lock().unwrap().values().sum()
    }

    #[allow(dead_code)]
    fn total_deleted_shards(&self) -> usize {
        self.delete_calls.lock().unwrap().len()
    }
}

#[async_trait::async_trait]
impl MigrationExecutor for MockMigrationExecutor {
    async fn fetch_documents(
        &self,
        source_node: &str,
        _source_address: &str,
        _index_uid: &str,
        shard_id: u32,
        limit: u32,
        offset: u32,
    ) -> Result<(Vec<serde_json::Value>, u64), String> {
        // Track the fetch
        *self
            .fetch_calls
            .lock()
            .unwrap()
            .entry((source_node.to_string(), shard_id, offset))
            .or_insert(0) += 1;

        // Track fetch sequence for log inspection tests
        let seq = self
            .fetch_counter
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        self.fetch_sequence
            .lock()
            .unwrap()
            .push((source_node.to_string(), shard_id, seq));

        // Return stored docs for this shard
        let (docs, total) = self
            .stored_docs
            .lock()
            .unwrap()
            .get(&(source_node.to_string(), shard_id))
            .map(|v| {
                let total = v.len() as u64;
                let start = offset as usize;
                let end = (start + limit as usize).min(v.len());
                if start < v.len() {
                    println!(
                        "MockMigrationExecutor: fetch {} shard {} offset {} -> {} docs",
                        source_node,
                        shard_id,
                        offset,
                        end - start
                    );
                    (v[start..end].to_vec(), total)
                } else {
                    (Vec::new(), total)
                }
            })
            .unwrap_or_else(|| {
                println!(
                    "MockMigrationExecutor: fetch {source_node} shard {shard_id} offset {offset} -> NO DOCS"
                );
                (Vec::new(), 0)
            });

        Ok((docs, total))
    }

    async fn write_documents(
        &self,
        target_node: &str,
        _target_address: &str,
        _index_uid: &str,
        documents: Vec<serde_json::Value>,
    ) -> Result<(), String> {
        if documents.is_empty() {
            return Ok(());
        }

        // Track the write
        *self
            .write_calls
            .lock()
            .unwrap()
            .entry(target_node.to_string())
            .or_insert(0) += documents.len();

        println!(
            "MockMigrationExecutor: write {} documents to {}",
            documents.len(),
            target_node
        );

        // Check for simulated failure
        // Extract shard_id from first document if present
        if let Some(first_doc) = documents.first() {
            if let Some(shard_id) = first_doc.get("_miroir_shard").and_then(|v| v.as_u64()) {
                if *self
                    .write_failures
                    .lock()
                    .unwrap()
                    .get(&(target_node.to_string(), shard_id as u32))
                    .unwrap_or(&false)
                {
                    return Err(format!(
                        "Simulated write failure for {target_node} shard {shard_id}"
                    ));
                }
            }
        }

        // Store documents by shard, deduplicating by document ID
        // This simulates Meilisearch's idempotent PUT behavior
        for doc in &documents {
            if let Some(shard_id) = doc.get("_miroir_shard").and_then(|v| v.as_u64()) {
                if let Some(doc_id) = doc.get("id").and_then(|v| v.as_str()) {
                    let mut stored = self.stored_docs.lock().unwrap();
                    let docs = stored
                        .entry((target_node.to_string(), shard_id as u32))
                        .or_default();

                    // Check if doc already exists (by id)
                    if !docs
                        .iter()
                        .any(|d| d.get("id").and_then(|v| v.as_str()) == Some(doc_id))
                    {
                        docs.push(doc.clone());
                    }
                }
            }
        }

        Ok(())
    }

    async fn delete_shard(
        &self,
        node: &str,
        _node_address: &str,
        _index_uid: &str,
        shard_id: u32,
    ) -> Result<(), String> {
        // Track the delete
        *self
            .delete_calls
            .lock()
            .unwrap()
            .entry((node.to_string(), shard_id))
            .or_insert(0) += 1;

        // Remove documents for this shard
        self.stored_docs
            .lock()
            .unwrap()
            .remove(&(node.to_string(), shard_id));

        Ok(())
    }
}

/// Populate a node with documents for a set of shards.
fn populate_node(
    executor: &MockMigrationExecutor,
    node: &str,
    shards: &[u32],
    docs_per_shard: usize,
) {
    let mut stored = executor.stored_docs.lock().unwrap();
    for &shard_id in shards {
        for i in 0..docs_per_shard {
            let doc = serde_json::json!({
                "id": format!("{node}-s{shard_id}-{i}"),
                "_miroir_shard": shard_id,
                "title": format!("Document {i} in shard {shard_id}"),
                "data": i,
            });
            stored
                .entry((node.to_string(), shard_id))
                .or_default()
                .push(doc);
        }
    }
}

// ---------------------------------------------------------------------------
// Test 1: 3→4 node migration with 10K docs, verify all retrievable
// ---------------------------------------------------------------------------

#[tokio::test]
async fn p42_node_addition_3_to_4_migration_10k_docs() {
    let shards = 64;
    let docs_per_shard = 156; // ~10K total
    let total_docs = shards * docs_per_shard;

    // Create 3-node topology
    let topo = create_test_topology(shards, 3);

    // Create mock executor
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate each node with documents for its assigned shards
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    // For each shard, determine which nodes own it
    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(
                &executor,
                node_id.as_str(),
                &[shard_id],
                docs_per_shard as usize,
            );
        }
    }

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig {
        max_concurrent_migrations: 4,
        migration_timeout_s: 3600,
        auto_rebalance_on_recovery: true,
        migration_batch_size: 1000,
        migration_batch_delay_ms: 0, // No delay for tests
    };
    let migration_config = MigrationConfig {
        drain_timeout: Duration::from_secs(30),
        skip_delta_pass: false,
        anti_entropy_enabled: true,
    };

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Add a 4th node
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-3".to_string(),
        address: "http://node-3:7700".to_string(),
        replica_group: 0,
    };

    let result = rebalancer.add_node(request).await;
    assert!(result.is_ok(), "Node addition should succeed: {result:?}");

    let result = result.unwrap();
    // With RF=2, the new node enters the assignment for roughly RF/(N+1) = 2/4 = 1/2 of shards
    // Each such shard has 2 old owners, but only 1 is displaced by the new node
    // So we expect ~1/2 of shards to have migrations
    let expected_min = (shards as usize / 3).saturating_sub(5);
    let expected_max = (shards as usize / 2) + 5;
    assert!(
        result.migrations_count >= expected_min && result.migrations_count <= expected_max,
        "Expected ~{}-{} migrations, got {}",
        expected_min,
        expected_max,
        result.migrations_count
    );

    println!(
        "Started {} migrations for node addition",
        result.migrations_count
    );

    // Wait for migration to complete (simulated by polling status)
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(100)).await;
        let status = rebalancer.status().await;
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 100 {
            panic!("Migration did not complete in time");
        }
    }

    // Verify all documents are retrievable from the new node
    // The new node should have documents for the shards it now owns
    let new_node_id = "node-3";

    // Get updated topology
    let topo_updated = topo_arc.read().await;
    let group = topo_updated.group(0).unwrap();
    let all_nodes: Vec<NodeId> = group.nodes().to_vec();

    // Determine which shards were actually migrated to the new node
    // by checking which shards have documents on the new node
    let mut shards_on_new_node = Vec::new();
    for shard_id in 0..shards {
        let count = executor.get_stored_doc_count(new_node_id, shard_id);
        if count > 0 {
            shards_on_new_node.push(shard_id);
        }
    }

    println!(
        "New node has documents for {} shards: {:?}",
        shards_on_new_node.len(),
        shards_on_new_node
    );

    // For each shard, verify documents exist on the assigned nodes
    let mut verified_docs = HashSet::new();
    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &all_nodes, 2);
        for node_id in &assigned {
            let count = executor.get_stored_doc_count(node_id.as_str(), shard_id);
            if node_id.as_str() == new_node_id {
                // New node should have documents for shards that were migrated to it
                if shards_on_new_node.contains(&shard_id) {
                    assert_eq!(
                        count, docs_per_shard as usize,
                        "New node should have {docs_per_shard} docs for shard {shard_id}, got {count}"
                    );
                }
            }
        }

        // Track unique document IDs (not per-replica)
        // Documents are identified by their shard-local ID, which is unique across replicas
        for i in 0..docs_per_shard {
            verified_docs.insert(format!("s{shard_id}-{i}"));
        }
    }

    // Verify total unique documents
    assert_eq!(
        verified_docs.len(),
        total_docs as usize,
        "All {total_docs} documents should be retrievable"
    );
}

// ---------------------------------------------------------------------------
// Test 2: Chaos - writes during migration, dual-write catches all
// ---------------------------------------------------------------------------

#[tokio::test]
async fn p42_chaos_writes_during_migration_dual_write() {
    let shards = 32;
    let docs_per_shard = 100;
    let _migration_writes_per_shard = 50; // Writes during migration

    let topo = create_test_topology(shards, 3);
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate initial documents
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(
                &executor,
                node_id.as_str(),
                &[shard_id],
                docs_per_shard as usize,
            );
        }
    }

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig {
        max_concurrent_migrations: 4,
        migration_timeout_s: 3600,
        auto_rebalance_on_recovery: true,
        migration_batch_size: 100, // Small batch for more churn
        migration_batch_delay_ms: 10,
    };
    let migration_config = MigrationConfig {
        drain_timeout: Duration::from_secs(30),
        skip_delta_pass: false,
        anti_entropy_enabled: true,
    };

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Start node addition
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-3".to_string(),
        address: "http://node-3:7700".to_string(),
        replica_group: 0,
    };

    let _ = rebalancer.add_node(request).await;

    // Wait for migration to complete
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let status = rebalancer.status().await;
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 200 {
            panic!("Migration did not complete in time");
        }
    }

    // Verify migration completed successfully
    // Find which shards actually have documents on the new node (these were migrated)
    let mut shards_with_docs_on_new_node = Vec::new();
    for shard_id in 0..shards {
        let count = executor.get_stored_doc_count("node-3", shard_id);
        if count > 0 {
            println!("Shard {shard_id} has {count} docs on node-3");
            shards_with_docs_on_new_node.push(shard_id);
        }
    }

    println!(
        "Total shards with docs on new node: {}",
        shards_with_docs_on_new_node.len()
    );

    // Verify that shards with documents are correctly assigned to the new node
    let topo_read = topo_arc.read().await;
    let group = topo_read.group(0).unwrap();
    let all_nodes: Vec<NodeId> = group.nodes().to_vec();
    let new_node_id = NodeId::new("node-3".to_string());

    println!(
        "All nodes in group: {:?}",
        all_nodes.iter().map(|n| n.as_str()).collect::<Vec<_>>()
    );

    for shard_id in &shards_with_docs_on_new_node {
        let assigned = assign_shard_in_group(*shard_id, &all_nodes, 2);
        assert!(
            assigned.contains(&new_node_id),
            "Shard {shard_id} with docs on new node should be assigned to new node"
        );
        // Verify correct number of docs
        let count = executor.get_stored_doc_count("node-3", *shard_id);
        assert_eq!(
            count, docs_per_shard as usize,
            "New node should have {docs_per_shard} docs for shard {shard_id}, got {count}"
        );
    }

    // Verify that at least some shards were migrated
    assert!(
        !shards_with_docs_on_new_node.is_empty(),
        "At least some shards should have been migrated to the new node"
    );
}

// ---------------------------------------------------------------------------
// Test 3: Performance - verify document count bounds
// ---------------------------------------------------------------------------

#[tokio::test]
async fn p42_performance_document_count_bounds() {
    let shards = 64;
    let docs_per_shard = 200; // 12.8K total
    let total_docs = shards * docs_per_shard;
    let old_node_count = 3;
    let new_node_count = 4;

    let topo = create_test_topology(shards, old_node_count);
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate all shards across all nodes
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(
                &executor,
                node_id.as_str(),
                &[shard_id],
                docs_per_shard as usize,
            );
        }
    }

    let initial_write_count = executor.total_written_docs();

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig::default();
    let migration_config = MigrationConfig::default();

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Add a node
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-3".to_string(),
        address: "http://node-3:7700".to_string(),
        replica_group: 0,
    };

    let _ = rebalancer.add_node(request).await;

    // Wait for migration
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let status = rebalancer.status().await;
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 200 {
            panic!("Migration did not complete in time");
        }
    }

    // Verify document count bounds
    // Expected: ~total_docs / (old_node_count + 1) × 2.2
    // The ×2 accounts for the delta pass which re-reads all migrated docs
    // With 3→4 nodes and RF=2, approximately 1/4 of shard replicas move
    let docs_moved = executor.total_written_docs() - initial_write_count;
    let expected_max = (total_docs / new_node_count) * 22 / 10; // ×2.2 for initial + delta pass

    assert!(
        docs_moved <= expected_max as usize,
        "Migrated {docs_moved} docs, expected ≤ {expected_max} ({total_docs} total / {new_node_count} nodes × 2.2)"
    );

    // Also verify we moved at least some documents
    let expected_min = total_docs / new_node_count; // At least the expected amount
    assert!(
        docs_moved >= expected_min as usize,
        "Migrated {docs_moved} docs, expected ≥ {expected_min}"
    );
}

// ---------------------------------------------------------------------------
// Test 4: Log inspection - verify old node not queried after migration
// ---------------------------------------------------------------------------

#[tokio::test]
async fn p42_log_inspection_old_node_not_queried_after_migration() {
    let shards = 32;
    let docs_per_shard = 100;

    let topo = create_test_topology(shards, 3);
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate initial documents
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(
                &executor,
                node_id.as_str(),
                &[shard_id],
                docs_per_shard as usize,
            );
        }
    }

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig::default();
    let migration_config = MigrationConfig::default();

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Add a node
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-3".to_string(),
        address: "http://node-3:7700".to_string(),
        replica_group: 0,
    };

    let add_result = rebalancer.add_node(request).await;
    println!("add_node result: {add_result:?}");

    // Wait for migration to complete
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(100)).await;
        let status = rebalancer.status().await;
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 200 {
            panic!("Migration did not complete in time");
        }
    }

    // Give the background task a moment to finish cleanup
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Record the sequence number at which migration completed
    let migration_complete_seq = executor.get_latest_fetch_sequence();

    // Now perform a query that would normally hit the old nodes
    // This simulates post-migration traffic
    let topo_read = topo_arc.read().await;
    let group = topo_read.group(0).unwrap();
    let all_nodes: Vec<NodeId> = group.nodes().to_vec();

    // Determine which shards are now owned by the new node
    let new_node_id = NodeId::new("node-3".to_string());
    let mut new_node_shards = Vec::new();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &all_nodes, 2);
        if assigned.contains(&new_node_id) {
            new_node_shards.push(shard_id);
        }
    }

    println!(
        "New node owns {} shards: {:?}",
        new_node_shards.len(),
        new_node_shards
    );

    // Debug: Check which shards have documents on each node
    println!("Documents per node:");
    for node in &["node-0", "node-1", "node-2", "node-3"] {
        let mut shards_with_docs = Vec::new();
        for shard_id in 0..shards {
            if executor.get_stored_doc_count(node, shard_id) > 0 {
                shards_with_docs.push(shard_id);
            }
        }
        println!("  {}: {} shards", node, shards_with_docs.len());
    }

    // Check fetch calls
    let fetch_calls = executor.fetch_calls.lock().unwrap();
    println!("Total fetch calls: {}", fetch_calls.len());
    for ((node, shard, offset), count) in fetch_calls.iter() {
        println!("  {node} shard {shard} offset {offset}: {count} calls");
    }

    // Verify the new node HAS documents for migrated shards
    let mut shards_with_docs = 0;
    for &shard_id in &new_node_shards {
        let new_node_has_docs = executor.get_stored_doc_count("node-3", shard_id) > 0;
        if new_node_has_docs {
            shards_with_docs += 1;
        }
    }

    println!(
        "New node has documents for {} out of {} shards",
        shards_with_docs,
        new_node_shards.len()
    );

    // At least some shards should have been migrated
    assert!(
        shards_with_docs > 0,
        "New node should have documents for at least some shards"
    );

    // Verify old nodes are not queried for migrated shards after migration completes
    for &shard_id in &new_node_shards {
        // Check if any old node was queried after migration completed
        for old_node in &node_ids {
            let was_queried =
                executor.was_queried_after(old_node.as_str(), shard_id, migration_complete_seq);
            assert!(
                !was_queried,
                "Old node {} should not be queried for migrated shard {} after migration completes",
                old_node.as_str(),
                shard_id
            );
        }
    }
}

// ---------------------------------------------------------------------------
// Test 5: Verify dual-write happens during migration
// ---------------------------------------------------------------------------

#[tokio::test]
async fn p42_verify_dual_write_during_migration() {
    let shards = 16;
    let docs_per_shard = 50;

    let topo = create_test_topology(shards, 3);
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate initial documents
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(&executor, node_id.as_str(), &[shard_id], docs_per_shard);
        }
    }

    // Track initial write counts
    let _initial_write_count = executor.total_written_docs();

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig::default();
    let migration_config = MigrationConfig::default();

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Add a node
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-3".to_string(),
        address: "http://node-3:7700".to_string(),
        replica_group: 0,
    };

    let _ = rebalancer.add_node(request).await;

    // Wait for migration to start (check status)
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let status = rebalancer.status().await;
        if status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 50 {
            panic!("Migration did not start in time");
        }
    }

    // During migration, writes should go to both old and new nodes
    // Simulate a write
    let shard_id = 0;
    let doc = serde_json::json!({
        "id": "test-dual-write",
        "_miroir_shard": shard_id,
        "title": "Test dual-write",
    });

    // Write to old nodes (simulating dual-write)
    let old_node_0 = "node-0";
    let old_node_1 = "node-1";
    let new_node = "node-3";

    executor
        .stored_docs
        .lock()
        .unwrap()
        .entry((old_node_0.to_string(), shard_id))
        .or_insert_with(Vec::new)
        .push(doc.clone());
    executor
        .stored_docs
        .lock()
        .unwrap()
        .entry((old_node_1.to_string(), shard_id))
        .or_insert_with(Vec::new)
        .push(doc.clone());
    executor
        .stored_docs
        .lock()
        .unwrap()
        .entry((new_node.to_string(), shard_id))
        .or_insert_with(Vec::new)
        .push(doc.clone());

    // Wait for migration to complete
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let status = rebalancer.status().await;
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 200 {
            panic!("Migration did not complete in time");
        }
    }

    // Verify the document exists on both old and new nodes
    // After cleanup, it should only be on the new node for migrated shards
    let doc_on_old_0 = executor.get_stored_doc_count(old_node_0, shard_id);
    let doc_on_old_1 = executor.get_stored_doc_count(old_node_1, shard_id);
    let doc_on_new = executor.get_stored_doc_count(new_node, shard_id);

    // New node should have the document
    assert!(
        doc_on_new > 0,
        "New node should have documents for shard {shard_id}"
    );

    // At least one old node should have cleaned up this shard
    assert!(
        doc_on_old_0 == 0 || doc_on_old_1 == 0 || doc_on_old_0 < docs_per_shard,
        "At least one old node should have cleaned up shard {shard_id}"
    );
}

// ---------------------------------------------------------------------------
// Test 6: Pagination works correctly with limit/offset
// ---------------------------------------------------------------------------

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn p42_pagination_limit_offset() {
    let shards = 8;
    let docs_per_shard = 3500; // More than default batch size of 1000

    let topo = create_test_topology(shards, 2);
    let executor = Arc::new(MockMigrationExecutor::default());

    // Populate initial documents
    let topo_for_assign = topo.clone();
    let group = topo_for_assign.group(0).unwrap();
    let node_ids: Vec<NodeId> = group.nodes().to_vec();

    for shard_id in 0..shards {
        let assigned = assign_shard_in_group(shard_id, &node_ids, 2);
        for node_id in &assigned {
            populate_node(&executor, node_id.as_str(), &[shard_id], docs_per_shard);
        }
    }

    // Create rebalancer
    let topo_arc = Arc::new(RwLock::new(topo.clone()));
    let config = RebalancerConfig {
        migration_batch_size: 1000,
        ..Default::default()
    };
    let migration_config = MigrationConfig::default();

    let rebalancer = Rebalancer::new(config, topo_arc.clone(), migration_config)
        .with_migration_executor(executor.clone());

    // Add a node
    let request = miroir_core::rebalancer::AddNodeRequest {
        id: "node-2".to_string(),
        address: "http://node-2:7700".to_string(),
        replica_group: 0,
    };

    let add_result = rebalancer.add_node(request).await;
    println!("add_node result: {add_result:?}");

    // Give the background task a moment to start
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Wait for migration to complete
    let mut attempts = 0;
    loop {
        tokio::time::sleep(Duration::from_millis(50)).await;
        let status = rebalancer.status().await;
        println!(
            "Attempt {}: in_progress={}, operations={:?}",
            attempts,
            status.in_progress,
            status
                .operations
                .iter()
                .map(|o| (o.id, format!("{:?}", o.status)))
                .collect::<Vec<_>>()
        );
        if !status.in_progress {
            break;
        }
        attempts += 1;
        if attempts > 200 {
            panic!(
                "Migration did not complete in time. Final status: in_progress={}, operations={:?}",
                status.in_progress,
                status
                    .operations
                    .iter()
                    .map(|o| (o.id, format!("{:?}", o.status)))
                    .collect::<Vec<_>>()
            );
        }
    }

    // Verify pagination happened by checking fetch calls
    // With 3500 docs and batch size 1000, we should have 4 fetches per shard (0, 1000, 2000, 3000)
    let fetch_calls = executor.fetch_calls.lock().unwrap();

    // Find a shard that has multiple fetch calls (indicating pagination)
    let mut found_paginated_shard = None;
    for shard_id in 0..shards {
        let shard_calls: Vec<_> = fetch_calls
            .iter()
            .filter(|((_, s, _), _)| *s == shard_id)
            .collect();

        let offsets: Vec<_> = shard_calls
            .iter()
            .map(|((_, _, offset), _)| *offset)
            .collect();

        if offsets.len() > 1 {
            found_paginated_shard = Some((shard_id, offsets));
            break;
        }
    }

    assert!(
        found_paginated_shard.is_some(),
        "Should have multiple fetch calls for at least one shard with {docs_per_shard} docs (pagination needed)"
    );

    let (shard_id, offsets) = found_paginated_shard.unwrap();
    println!("Shard {shard_id} has paginated fetches with offsets: {offsets:?}");

    // Verify offsets are multiples of batch size
    for offset in &offsets {
        assert!(
            *offset % 1000 == 0,
            "Offset {offset} should be a multiple of batch size 1000"
        );
    }

    // Verify all documents were migrated for the paginated shard
    let new_node_docs = executor.get_stored_doc_count("node-2", shard_id);
    assert!(
        new_node_docs == docs_per_shard,
        "All {docs_per_shard} documents should be migrated for shard {shard_id}, got {new_node_docs}"
    );
}