[bug] Recovering a shard snapshot holds the collection shards holder read lock across unpack and shard replacement

//! PoC: restoring a shard snapshot holds the collection-wide `shards_holder` read guard for the
//! whole archive unpack and shard replacement; a consensus-applied config update queues on the
//! write side, and the write-preferring lock then fences every collection reader until the
//! restore finishes. This is the receiver side of a snapshot-based shard transfer.
//!
//! Integration location: lib/storage/tests/restore_shard_snapshot_holder_guard.rs
//! Run: cargo +nightly-2026-03-21 test -p storage --test restore_shard_snapshot_holder_guard -- --nocapture
//! Target: dev (1.19.2-dev).

use std::io::Read as _;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use cancel::CancellationToken;
use collection::shards::channel_service::ChannelService;
use common::budget::ResourceBudget;
use segment::types::StrictModeConfig;
use shard::snapshots::snapshot_data::SnapshotData;
use shard::snapshots::snapshot_manifest::RecoveryType;
use storage::content_manager::CollectionContainer;
use storage::content_manager::collection_meta_ops::{
    CollectionMetaOperations, CreateCollection, CreateCollectionOperation,
};
use storage::content_manager::consensus::operation_sender::OperationSender;
use storage::content_manager::toc::TableOfContent;
use storage::rbac::{Access, AccessRequirements};
use storage::types::StorageConfig;
use tempfile::Builder;
use tokio::runtime::Runtime;

/// Padding appended to the snapshot archive so the unpack phase lasts long enough to observe the
/// lock hold. A real large shard produces a snapshot of this order (see issue #8671: 178 GB).
const SNAPSHOT_PADDING_BYTES: u64 = 1024 * 1024 * 1024;

/// Built from JSON so field additions do not break the PoC across versions.
fn test_storage_config(storage_dir: &Path) -> StorageConfig {
    let json = serde_json::json!({
        "storage_path": storage_dir.to_string_lossy().into_owned(),
        "snapshots_path": storage_dir.join("snapshots").to_string_lossy().into_owned(),
        "temp_path": null,
        "on_disk_payload": false,
        "optimizers": {
            "deleted_threshold": 0.5,
            "vacuum_min_vector_number": 100,
            "default_segment_number": 2,
            "max_segment_size": null,
            "memmap_threshold": 100,
            "indexing_threshold": 100,
            "flush_interval_sec": 2,
            "max_optimization_threads": 2,
        },
        "wal": { "wal_capacity_mb": 32, "wal_segments_ahead": 0 },
        "performance": {
            "max_search_threads": 1,
            "max_optimization_runtime_threads": 2,
            "optimizer_cpu_budget": 0,
            "optimizer_io_budget": 0,
            "incoming_shard_transfers_limit": 1,
            "outgoing_shard_transfers_limit": 1,
        },
        "hnsw_index": { "m": 16, "ef_construct": 100, "full_scan_threshold": 10000 },
        "update_concurrency": 2,
    });
    serde_json::from_value(json).expect("valid storage config")
}

/// Rebuild a snapshot archive with one extra large root-level member.
///
/// The original archive is unpacked and re-archived so all entries are plain regular files and
/// directories, then a padding member is appended. The restore ignores the padding (it only
/// installs the known shard paths), but unwinding it takes a disk-bound amount of time
/// proportional to its size, which is the point: a real large shard produces a snapshot of this
/// order.
fn pad_snapshot(src: &Path, dst: &Path, temp_dir: &Path) {
    let unpacked = Builder::new()
        .prefix("snapshot-pad")
        .tempdir_in(temp_dir)
        .unwrap();
    common::tar_unpack::tar_unpack_file(src, unpacked.path()).unwrap();

    let mut builder = tar::Builder::new(std::fs::File::create(dst).unwrap());
    for entry in std::fs::read_dir(unpacked.path()).unwrap() {
        let entry = entry.unwrap();
        let path = entry.path();
        let name = entry.file_name();
        if path.is_dir() {
            builder.append_dir_all(&name, &path).unwrap();
        } else {
            builder.append_path_with_name(&path, &name).unwrap();
        }
    }

    let mut header = tar::Header::new_gnu();
    header.set_path("padding.bin").unwrap();
    header.set_mode(0o644);
    header.set_size(SNAPSHOT_PADDING_BYTES);
    header.set_cksum();
    builder
        .append(&header, std::io::repeat(0).take(SNAPSHOT_PADDING_BYTES))
        .unwrap();

    builder.finish().unwrap();
    builder.into_inner().unwrap();
}

#[test]
fn restore_shard_snapshot_holds_shards_holder_read_guard() {
    let storage_dir = Builder::new().prefix("storage").tempdir().unwrap();
    let config = test_storage_config(storage_dir.path());

    let (propose_sender, _propose_receiver) = std::sync::mpsc::channel();
    let toc = Arc::new(
        TableOfContent::new(
            &config,
            ResourceBudget::default(),
            ChannelService::default(),
            0,
            Some(OperationSender::new(propose_sender)),
        )
        .unwrap(),
    );

    // Seed one collection with a single local shard.
    let create_collection: CreateCollection = serde_json::from_value(serde_json::json!({
        "vectors": { "size": 4, "distance": "Dot" },
        "shard_number": 1,
        "hnsw_config": null,
        "wal_config": null,
        "optimizers_config": null,
        "sparse_vectors": null,
        "strict_mode_config": null,
    }))
    .expect("valid create collection request");
    CollectionContainer::perform_collection_meta_op(
        toc.as_ref(),
        CollectionMetaOperations::CreateCollection(
            CreateCollectionOperation::new("seed".to_string(), create_collection).unwrap(),
        ),
    )
    .unwrap();

    let rt = Runtime::new().unwrap();
    let access = Access::full("restore shard snapshot holder guard PoC");
    let pass = access
        .check_collection_access("seed", AccessRequirements::new())
        .unwrap();
    let collection = rt.block_on(toc.get_collection(&pass)).unwrap();

    let temp_dir = storage_dir.path().join("temp");
    std::fs::create_dir_all(&temp_dir).unwrap();

    // Snapshot the local shard and pad the archive. Shard snapshots are stored under the
    // collection's snapshots path in a `shards/<shard_id>/` subdirectory.
    let shard_snapshots_dir = collection.snapshots_path().join("shards/0");
    let snapshot = rt.block_on(collection.create_shard_snapshot(0, &temp_dir)).unwrap();
    let snapshot_path = std::fs::read_dir(&shard_snapshots_dir)
        .unwrap()
        .filter_map(|entry| entry.ok())
        .map(|entry| entry.path())
        .find(|path| {
            path.extension()
                .is_some_and(|ext| ext.to_string_lossy() == "snapshot")
        })
        .expect("shard snapshot file must exist");
    let padded_path = temp_dir.join("padded.snapshot");
    println!("snapshot: {}", snapshot.name);
    pad_snapshot(&snapshot_path, &padded_path, &temp_dir);

    // Start the restore. The guard is taken before the restore future is returned, before the
    // unpack even begins.
    let restore = rt
        .block_on(collection.restore_shard_snapshot(
            0,
            SnapshotData::new_packed_persistent(&padded_path),
            RecoveryType::Full,
            0,
            false,
            &temp_dir,
            None,
            CancellationToken::new(),
        ))
        .unwrap();
    let restore_started = Instant::now();
    let restore = rt.spawn(restore);

    // Wait until the restore holds the read guard: a write-taking call must block. This is the
    // same call the consensus apply of `UpdateCollection` makes for a strict mode change.
    let diff = StrictModeConfig {
        enabled: Some(true),
        ..Default::default()
    };
    let guard_held = rt.block_on(async {
        for _ in 0..300 {
            let blocked = tokio::time::timeout(
                Duration::from_millis(50),
                collection.update_strict_mode_config(diff.clone()),
            )
            .await
            .is_err();
            if blocked {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
        false
    });
    let guard_seen_after = restore_started.elapsed();
    assert!(
        guard_held,
        "snapshot restore must hold the shards holder read guard while it unpacks",
    );

    // Queue that write-taking call and probe a reader. The write-preferring lock keeps every
    // reader fenced behind the queued writer.
    let writer = rt.spawn({
        let collection = collection.clone();
        async move { collection.update_strict_mode_config(diff).await }
    });
    let mut reader_blocked = false;
    for _ in 0..20 {
        let done = rt.block_on(async {
            tokio::time::timeout(Duration::from_millis(200), collection.state())
                .await
                .is_ok()
        });
        if !done {
            assert!(!writer.is_finished(), "writer must stay queued");
            reader_blocked = true;
            break;
        }
        if writer.is_finished() || restore.is_finished() {
            break;
        }
        rt.block_on(async { tokio::time::sleep(Duration::from_millis(20)).await });
    }
    assert!(
        reader_blocked,
        "collection reader must be fenced behind the queued writer while the restore runs",
    );

    // The restore finishes on its own schedule; only then are the writer and readers released.
    let restore_result = rt.block_on(async {
        tokio::time::timeout(Duration::from_secs(600), restore)
            .await
            .expect("restore must finish")
            .expect("restore task must not panic")
    });
    let held_for = restore_started.elapsed();
    assert!(
        restore_result.is_ok(),
        "snapshot restore must succeed, got {restore_result:?}",
    );
    rt.block_on(async {
        tokio::time::timeout(Duration::from_secs(60), writer)
            .await
            .expect("writer must complete after the guard is released")
            .unwrap()
            .unwrap();
        tokio::time::timeout(Duration::from_secs(60), collection.state())
            .await
            .expect("reader must complete after the writer");
    });

    println!("guard observed held after: {guard_seen_after:?}");
    println!("restore held the shards holder read guard for: {held_for:?}");
    println!("reader fenced behind the queued writer: {reader_blocked}");
    println!("restore result: {restore_result:?}");
    println!("restore shard snapshot holder guard test done");
}

$ cargo +nightly-2026-03-21 test -p storage --test restore_shard_snapshot_holder_guard -- --nocapture
   Compiling storage v0.2.0
    Finished `test` profile [unoptimized + debuginfo] target(s) in 4.54s
     Running tests/restore_shard_snapshot_holder_guard.rs (target/debug/deps/restore_shard_snapshot_holder_guard-779f76d63c6bb72c)

running 1 test
snapshot: seed-shard-0-2026-09-29-15-44-43.snapshot
guard observed held after: 51.540561ms
restore held the shards holder read guard for: 2.861365081s
reader fenced behind the queued writer: true
restore result: Ok(())
restore shard snapshot holder guard test done
test restore_shard_snapshot_holds_shards_holder_read_guard ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 8.02s

// src/common/snapshots.rs:177-181
    // - `recover_shard_snapshot_impl` is *not* cancel safe
    //   - but the task is *spawned* on the runtime and won't be cancelled, if request is cancelled

    cancel::future::spawn_cancel_on_drop(async move |cancel| {

// src/common/snapshots.rs:193-204 (abridged)
        // Waits for a recovery of this shard already in progress, then holds it excluded
        // for the whole recovery below and reports progress meanwhile. ...
        let recovery_guard = cancel::future::cancel_on_token(
            cancel.clone(),
            collection.start_shard_recovery(shard_id),
        )
        .await??;

// src/common/snapshots.rs:290-306 (abridged)
        let snapshot_data =
            cancel::future::cancel_on_token(cancel.clone(), download_task).await??;

        // `recover_shard_snapshot_impl` is *not* cancel safe
        let result = recover_shard_snapshot_impl(
            &toc,
            &collection,
            shard_id,
            snapshot_data,
            snapshot_priority,
            RecoveryType::Full,
            Some(recovery_guard.progress_handle()),
            cancel,
        )
        .await;

// src/common/snapshots.rs:341-360 (abridged)
    // TODO: `Collection::restore_shard_snapshot` *is* cancel-safe, but `recover_shard_snapshot_impl` is *not* cancel-safe (yet)
    collection
        .restore_shard_snapshot(
            shard,
            snapshot_data,
            recovery_type,
            toc.this_peer_id,
            toc.is_distributed(),
            &toc.optional_temp_or_storage_temp_path()?,
            recovery_progress,
            cancel,
        )
        .await?
        .await?;

// lib/collection/src/collection/snapshots.rs:342-370 (abridged)
        // `ShardHolder::validate_shard_snapshot` is cancel safe, so we explicitly cancel it
        // when token is triggered
        let shard_holder = self.shards_holder.clone().read_owned().await;

        let collection_path = self.path.clone();
        let collection_name = self.name().to_string();

        let temp_dir = temp_dir.to_path_buf();

        // `ShardHolder::restore_shard_snapshot` is *not* cancel safe, so we spawn it onto runtime,
        // so that it won't be cancelled if current future is dropped
        let restore = self.update_runtime.spawn(async move {
            shard_holder
                .restore_shard_snapshot(
                    snapshot_data,
                    recovery_type,
                    &collection_path,
                    &collection_name,
                    shard_id,
                    this_peer_id,
                    is_distributed,
                    &temp_dir,
                    recovery_progress,
                    cancel,
                )
                .await?;

            CollectionResult::Ok(())
        });

// lib/collection/src/shards/shard_holder/mod.rs:1410-1461 (abridged)
        let extract = {
            let snapshot_temp_dir = snapshot_temp_dir.path().to_path_buf();

            cancel::blocking::spawn_cancel_on_token(
                cancel.child_token(),
                move |cancel| -> CollectionResult<_> {
                    match snapshot_data {
                        SnapshotData::Packed(snapshot_path) => {
                            ...
                            tar_unpack_file(&snapshot_path, &snapshot_temp_dir)?;
                            snapshot_path.close()?;
                        }
                        SnapshotData::Unpacked(snapshot_dir) => {