[bug] Receiver-side shard transfer initiation keeps the shards holder read lock across consensus dependent waits

//! PoC: receiver-side shard transfer initiation holds the collection-wide `shards_holder` read
//! guard while it waits for the transfer record and replica state driven through consensus.
//! A consensus-applied config update queues on the write side and fences every collection reader
//! behind it, and the receiver can never see the notification it waits for.
//!
//! Integration location: lib/storage/tests/initiate_shard_transfer_holder_guard.rs
//! Run: cargo +nightly-2026-03-21 test -p storage --test initiate_shard_transfer_holder_guard -- --nocapture
//! Target: dev (1.19.2-dev).

use std::time::{Duration, Instant};

use collection::shards::channel_service::ChannelService;
use common::budget::ResourceBudget;
use segment::types::StrictModeConfig;
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;

/// Built from JSON so field additions do not break the PoC across versions.
fn test_storage_config(storage_dir: &std::path::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")
}

#[test]
fn initiate_shard_transfer_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 = std::sync::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("initiate shard transfer holder lock PoC");
    let pass = access
        .check_collection_access("seed", AccessRequirements::new())
        .unwrap();
    let collection = rt.block_on(toc.get_collection(&pass)).unwrap();

    // Receiver-side transfer initiation. Without a consensus notification it parks for 60 seconds,
    // holding the collection-wide shards holder read guard the whole time.
    let started = Instant::now();
    let initiate = rt.spawn({
        let collection = collection.clone();
        async move { collection.initiate_shard_transfer(0, None).await }
    });

    // Wait until the receiver holds the read guard: a write-taking call must block.
    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 = started.elapsed();
    assert!(
        guard_held,
        "initiation must hold the shards holder read guard while waiting for consensus",
    );

    // A config update, in the shape applied through consensus, queues on the write side. Retry the
    // reader probe until the writer task has been polled and parked on the write guard; after that
    // the write-preferring lock keeps every later reader queued.
    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..10 {
        let done = rt.block_on(async {
            tokio::time::timeout(Duration::from_millis(500), collection.state())
                .await
                .is_ok()
        });
        if !done {
            reader_blocked = true;
            break;
        }
        assert!(!writer.is_finished(), "writer must stay queued");
        rt.block_on(async { tokio::time::sleep(Duration::from_millis(50)).await });
    }
    assert!(
        reader_blocked,
        "collection reader must be fenced behind the queued writer",
    );

    // Both stay blocked while the receiver waits for the notification it can never receive.
    let hold_started = Instant::now();
    while hold_started.elapsed() < Duration::from_secs(5) {
        assert!(!writer.is_finished(), "writer must stay queued");
        assert!(!initiate.is_finished(), "initiation must stay parked");
    }

    // The receiver gives up after its 60 second wait; the guard is released and the writer plus
    // readers proceed.
    let initiate_result = rt.block_on(async {
        tokio::time::timeout(Duration::from_secs(90), initiate)
            .await
            .expect("initiation must finish")
            .expect("initiation task must not panic")
    });
    let blocked_for = started.elapsed();
    assert!(
        initiate_result.is_err(),
        "without a consensus notification the initiation must time out",
    );
    rt.block_on(async {
        tokio::time::timeout(Duration::from_secs(30), writer)
            .await
            .expect("writer must complete after the guard is released")
            .unwrap()
            .unwrap();
        collection.state().await;
    });

    println!("guard observed held after: {guard_seen_after:?}");
    println!("writer and readers stayed blocked for: {blocked_for:?}");
    println!("initiate result: {initiate_result:?}");
    println!("initiate shard transfer test done");
}

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

running 1 test
test initiate_shard_transfer_holds_shards_holder_read_guard has been running for over 60 seconds
guard observed held after: 67.473964ms
writer and readers stayed blocked for: 60.001507251s
initiate result: Err(Timeout { description: "Failed to initiate shard transfer: Didn't receive shard transfer notification from consensus in 60 seconds" })
initiate shard transfer test done
test initiate_shard_transfer_holds_shards_holder_read_guard ... ok

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

// lib/collection/src/collection/shard_transfer.rs:679-689
pub fn initiate_shard_transfer(
    &self,
    shard_id: ShardId,
    from_peer_id: Option<PeerId>,
) -> impl Future<Output = CollectionResult<()>> + 'static {
    let shards_holder = self.shards_holder.clone();

    let collection_path = self.path.clone();

    async move {
        let shards_holder_guard = shards_holder.clone().read_owned().await;

// lib/collection/src/collection/shard_transfer.rs:696-701
// Wait for the replica set to have the local shard initialized
// This can take some time as this is arranged through consensus
replica_set
    .wait_for_local(defaults::CONSENSUS_META_OP_WAIT)
    .await?;

// lib/collection/src/collection/shard_transfer.rs:706-756 (abridged)
let shard_transfer_requested = tokio::task::spawn_blocking(move || {
    // Wait for transfer targeting this replica
    let shard_transfer_registered = shards_holder_guard.shard_transfers.wait_for(
        |shard_transfers| {
            shard_transfers
                .iter()
                .any(|shard_transfer| shard_transfer.is_target(this_peer_id, shard_id))
        },
        Duration::from_secs(60),
    );

    // Reject request, if no transfer was registered within 60 seconds
    if !shard_transfer_registered {
        return Ok(false);
    }
    ...
    // Transfer is registered before the replica state is set.
    // Wait until the local replica switches to a shard transfer state.
    Ok(replica_set.wait_for_state_condition_sync(
        |state| {
            state
                .get_peer_state(this_peer_id)
                .is_some_and(|peer_state| peer_state.is_partial_or_recovery())
        },
        defaults::CONSENSUS_META_OP_WAIT,
    ))
});

// lib/collection/src/collection/shard_transfer.rs:193
shards_holder.register_start_shard_transfer(shard_transfer.clone())?;

// lib/collection/src/collection/shard_transfer.rs:219-226
let old_shard = to_replica_set.set_local(shard, Some(initial_state)).await?;
...
to_replica_set
    .ensure_replica_with_state(shard_transfer.to, initial_state)
    .await?;

// lib/collection/src/collection/collection_ops.rs:199-204
let shard_holder = self.shards_holder.write().await;
let updates = shard_holder.all_shards().map(|replica_set| {
    replica_set.on_strict_mode_config_update(&new_strict_mode_config)
});
let results = future::join_all(updates).await;

// lib/storage/src/content_manager/consensus_manager.rs:593-632 (abridged)
pub fn apply_normal_entry(&self, entry: &RaftEntry) -> Result<bool, StorageError> {
    ...
    ConsensusOperations::CollectionMeta(operation) => {
        self.toc.perform_collection_meta_op(*operation)?
    }
    ...
}

receiver handler:        holds shards_holder.read()
                          └─ waits for the transfer record (60s) and replica state (10s)

consensus thread:        applies entry M (UpdateCollection / resharding / snapshot apply)
                          └─ blocks on shards_holder.write()          <- queued behind the receiver

consensus entry N:       StartTransfer (registers the record, sets the replica state)
                          └─ cannot be applied, the thread is still in M

external readers:        queries, updates, internal APIs
                          └─ block on shards_holder.read()            <- fenced behind the writer

// lib/collection/src/collection/shard_transfer.rs:751-764 (abridged)
match AbortOnDropHandle::new(shard_transfer_requested).await {
    Ok(Ok(true)) => Ok(()),

    Ok(Ok(false)) => {
        let description = "\
            Failed to initiate shard transfer: \
            Didn't receive shard transfer notification from consensus in 60 seconds";

        Err(CollectionError::Timeout {
            description: description.into(),
        })
    }
    ...
}