[bug] Sender-side snapshot shard transfer keeps the shards holder read lock across the whole transfer

//! PoC: the sender of a snapshot shard transfer keeps the collection-wide `shards_holder` read
//! guard for the whole transfer, so a consensus-applied writer queues behind it and the
//! write-preferring lock then fences every collection reader until the transfer ends.
//!
//! Integration location: lib/collection/tests/snapshot_transfer_sender_holder_guard.rs
//! Run: cargo +nightly-2026-03-21 test -p collection --test snapshot_transfer_sender_holder_guard -- --nocapture

use std::collections::HashSet;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, Instant};

use api::grpc::qdrant::collections_internal_server::{
    CollectionsInternal, CollectionsInternalServer,
};
use api::grpc::qdrant::qdrant_server::{Qdrant, QdrantServer};
use api::grpc::qdrant::shard_snapshots_server::{ShardSnapshots, ShardSnapshotsServer};
use api::grpc::qdrant::{
    CollectionOperationResponse, CreateShardSnapshotRequest, CreateSnapshotResponse,
    DeleteShardSnapshotRequest, DeleteSnapshotResponse, GetCollectionInfoRequestInternal,
    GetCollectionInfoResponse, GetShardMemoryReportRequest, GetShardMemoryReportResponse,
    GetShardOptimizationsRequest, GetShardOptimizationsResponse, GetShardRecoveryPointRequest,
    GetShardRecoveryPointResponse, HealthCheckReply, HealthCheckRequest,
    InitiateShardTransferRequest, ListShardSnapshotsRequest, ListSnapshotsResponse,
    RecoverShardSnapshotRequest, RecoverSnapshotResponse, UpdateShardCutoffPointRequest,
    WaitForShardStateRequest,
};
use collection::collection::{Collection, RequestShardTransfer};
use collection::config::{CollectionConfigInternal, CollectionParams, WalConfig};
use collection::operations::types::CollectionResult;
use collection::operations::vector_params_builder::VectorParamsBuilder;
use collection::optimizers_builder::OptimizersConfig;
use collection::shards::CollectionId;
use collection::shards::channel_service::ChannelService;
use collection::shards::collection_shard_distribution::CollectionShardDistribution;
use collection::shards::replica_set::replica_set_state::ReplicaState;
use collection::shards::replica_set::{AbortShardTransfer, ChangePeerFromState};
use collection::shards::resharding::ReshardKey;
use collection::shards::shard::{PeerId, ShardId};
use collection::shards::transfer::{
    ShardTransfer, ShardTransferConsensus, ShardTransferKey, ShardTransferMethod,
};
use common::budget::ResourceBudget;
use segment::types::{Distance, StrictModeConfig};
use tempfile::Builder;
use tokio::sync::watch;
use tonic::transport::{Server, Uri};
use tonic::{Request, Response, Status};

const THIS_PEER: PeerId = 0;
const REMOTE_PEER: PeerId = 1;
const SHARD: ShardId = 0;

/// Mock remote peer: `initiate` succeeds, `recover` blocks until released.
#[derive(Clone)]
struct MockPeer {
    released: watch::Receiver<bool>,
}

#[async_trait::async_trait]
impl Qdrant for MockPeer {
    async fn health_check(
        &self,
        _request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckReply>, Status> {
        Ok(Response::new(HealthCheckReply {
            title: "qdrant - vector search engine".to_string(),
            version: "1.19.2".to_string(),
            commit: None,
        }))
    }
}

#[async_trait::async_trait]
impl CollectionsInternal for MockPeer {
    async fn get(
        &self,
        _request: Request<GetCollectionInfoRequestInternal>,
    ) -> Result<Response<GetCollectionInfoResponse>, Status> {
        Err(Status::unimplemented("get"))
    }

    async fn initiate(
        &self,
        _request: Request<InitiateShardTransferRequest>,
    ) -> Result<Response<CollectionOperationResponse>, Status> {
        Ok(Response::new(CollectionOperationResponse {
            result: true,
            time: 0.0,
        }))
    }

    async fn wait_for_shard_state(
        &self,
        _request: Request<WaitForShardStateRequest>,
    ) -> Result<Response<CollectionOperationResponse>, Status> {
        Ok(Response::new(CollectionOperationResponse {
            result: true,
            time: 0.0,
        }))
    }

    async fn get_shard_recovery_point(
        &self,
        _request: Request<GetShardRecoveryPointRequest>,
    ) -> Result<Response<GetShardRecoveryPointResponse>, Status> {
        Err(Status::unimplemented("get_shard_recovery_point"))
    }

    async fn update_shard_cutoff_point(
        &self,
        _request: Request<UpdateShardCutoffPointRequest>,
    ) -> Result<Response<CollectionOperationResponse>, Status> {
        Err(Status::unimplemented("update_shard_cutoff_point"))
    }

    async fn get_shard_optimizations(
        &self,
        _request: Request<GetShardOptimizationsRequest>,
    ) -> Result<Response<GetShardOptimizationsResponse>, Status> {
        Err(Status::unimplemented("get_shard_optimizations"))
    }

    async fn get_shard_memory_report(
        &self,
        _request: Request<GetShardMemoryReportRequest>,
    ) -> Result<Response<GetShardMemoryReportResponse>, Status> {
        Err(Status::unimplemented("get_shard_memory_report"))
    }
}

#[async_trait::async_trait]
impl ShardSnapshots for MockPeer {
    async fn create(
        &self,
        _request: Request<CreateShardSnapshotRequest>,
    ) -> Result<Response<CreateSnapshotResponse>, Status> {
        Err(Status::unimplemented("create"))
    }

    async fn list(
        &self,
        _request: Request<ListShardSnapshotsRequest>,
    ) -> Result<Response<ListSnapshotsResponse>, Status> {
        Err(Status::unimplemented("list"))
    }

    async fn delete(
        &self,
        _request: Request<DeleteShardSnapshotRequest>,
    ) -> Result<Response<DeleteSnapshotResponse>, Status> {
        Err(Status::unimplemented("delete"))
    }

    async fn recover(
        &self,
        _request: Request<RecoverShardSnapshotRequest>,
    ) -> Result<Response<RecoverSnapshotResponse>, Status> {
        let mut released = self.released.clone();
        while !*released.borrow() {
            if released.changed().await.is_err() {
                break;
            }
        }
        Err(Status::cancelled("released by PoC"))
    }
}

/// Consensus stand-in: a leader is available and proposals are accepted.
struct TestConsensus;

#[async_trait::async_trait]
impl ShardTransferConsensus for TestConsensus {
    fn this_peer_id(&self) -> PeerId {
        THIS_PEER
    }

    fn peers(&self) -> Vec<PeerId> {
        vec![THIS_PEER, REMOTE_PEER]
    }

    fn consensus_commit_term(&self) -> (u64, u64) {
        (1, 1)
    }

    fn is_leader_established(&self) -> bool {
        true
    }

    fn recovered_switch_to_partial(
        &self,
        _transfer_config: &ShardTransfer,
        _collection_id: CollectionId,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn start_shard_transfer(
        &self,
        _transfer_config: ShardTransfer,
        _collection_id: CollectionId,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn restart_shard_transfer(
        &self,
        _transfer_config: ShardTransfer,
        _collection_id: CollectionId,
        _default_method: ShardTransferMethod,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn abort_shard_transfer(
        &self,
        _transfer: ShardTransferKey,
        _collection_id: CollectionId,
        _reason: &str,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn set_shard_replica_set_state(
        &self,
        _peer_id: Option<PeerId>,
        _collection_id: CollectionId,
        _shard_id: ShardId,
        _state: ReplicaState,
        _from_state: Option<ReplicaState>,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn commit_read_hashring(
        &self,
        _collection_id: CollectionId,
        _reshard_key: ReshardKey,
    ) -> CollectionResult<()> {
        Ok(())
    }

    async fn commit_write_hashring(
        &self,
        _collection_id: CollectionId,
        _reshard_key: ReshardKey,
    ) -> CollectionResult<()> {
        Ok(())
    }
}

fn dummy_on_replica_failure() -> ChangePeerFromState {
    Arc::new(|_, _, _| {})
}

fn dummy_request_shard_transfer() -> RequestShardTransfer {
    Arc::new(|_| {})
}

fn dummy_abort_shard_transfer() -> AbortShardTransfer {
    Arc::new(|_, _| {})
}

fn test_collection_config() -> CollectionConfigInternal {
    CollectionConfigInternal {
        params: CollectionParams {
            vectors: VectorParamsBuilder::new(4, Distance::Dot).build().into(),
            shard_number: NonZeroU32::new(1).unwrap(),
            ..CollectionParams::empty()
        },
        optimizer_config: OptimizersConfig {
            deleted_threshold: 0.9,
            vacuum_min_vector_number: 1000,
            default_segment_number: 2,
            max_segment_size: None,
            #[expect(deprecated)]
            memmap_threshold: None,
            indexing_threshold: Some(50_000),
            flush_interval_sec: 30,
            max_optimization_threads: Some(2),
            prevent_unoptimized: None,
        },
        wal_config: WalConfig {
            wal_capacity_mb: 1,
            wal_segments_ahead: 0,
            wal_retain_closed: 1,
        },
        hnsw_config: Default::default(),
        quantization_config: Default::default(),
        strict_mode_config: Default::default(),
        uuid: None,
        metadata: None,
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn snapshot_transfer_sender_holds_shards_holder_read_guard() {
    let dir = Builder::new().prefix("collection").tempdir().unwrap();
    let snapshots_path = dir.path().join("snapshots");
    std::fs::create_dir_all(&snapshots_path).unwrap();

    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let remote_addr = listener.local_addr().unwrap();
    drop(listener);
    let (released_tx, released_rx) = watch::channel(false);
    let mock = MockPeer {
        released: released_rx,
    };
    tokio::spawn(async move {
        let _ = Server::builder()
            .add_service(QdrantServer::new(mock.clone()))
            .add_service(CollectionsInternalServer::new(mock.clone()))
            .add_service(ShardSnapshotsServer::new(mock))
            .serve(remote_addr)
            .await;
    });
    for _ in 0..100 {
        if tokio::net::TcpStream::connect(remote_addr).await.is_ok() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }

    let channel_service = ChannelService::new(6333, false, None, None);
    channel_service
        .id_to_address
        .write()
        .insert(THIS_PEER, "http://127.0.0.1:6333".parse::<Uri>().unwrap());
    channel_service.id_to_address.write().insert(
        REMOTE_PEER,
        format!("http://127.0.0.1:{}", remote_addr.port())
            .parse::<Uri>()
            .unwrap(),
    );

    let distribution = CollectionShardDistribution {
        shards: [(SHARD, HashSet::from([THIS_PEER, REMOTE_PEER]))]
            .into_iter()
            .collect(),
    };

    let collection = Collection::new(
        "test".to_string(),
        THIS_PEER,
        dir.path(),
        &snapshots_path,
        &test_collection_config(),
        Default::default(),
        distribution,
        None,
        channel_service,
        dummy_on_replica_failure(),
        dummy_request_shard_transfer(),
        dummy_abort_shard_transfer(),
        None,
        None,
        ResourceBudget::default(),
        None,
    )
    .await
    .expect("collection must be created");
    collection
        .set_shard_replica_state(SHARD, THIS_PEER, ReplicaState::Active, None)
        .await
        .unwrap();
    let collection = Arc::new(collection);

    let transfer = ShardTransfer {
        shard_id: SHARD,
        to_shard_id: None,
        from: THIS_PEER,
        to: REMOTE_PEER,
        sync: false,
        method: Some(ShardTransferMethod::Snapshot),
        filter: None,
    };
    let started = Instant::now();
    collection
        .start_shard_transfer(
            transfer,
            Box::new(TestConsensus),
            dir.path().to_path_buf(),
            std::future::ready(()),
            std::future::ready(()),
        )
        .await
        .expect("transfer must be started");

    let diff = StrictModeConfig {
        enabled: Some(true),
        ..Default::default()
    };
    let mut guard_seen = None;
    let mut consecutive_blocks = 0usize;
    let deadline = Instant::now() + Duration::from_secs(60);
    while Instant::now() < deadline {
        let blocked = tokio::time::timeout(
            Duration::from_millis(50),
            collection.update_strict_mode_config(diff.clone()),
        )
        .await
        .is_err();
        consecutive_blocks = if blocked { consecutive_blocks + 1 } else { 0 };
        if consecutive_blocks >= 3 {
            guard_seen = Some(started.elapsed());
            break;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    let guard_seen = guard_seen.expect("sender transfer must hold the shards holder read guard");

    let writer = tokio::spawn({
        let collection = Arc::clone(&collection);
        async move { collection.update_strict_mode_config(diff).await }
    });

    let mut reader_blocked = false;
    for _ in 0..50 {
        assert!(
            !writer.is_finished(),
            "writer must stay queued while the transfer holds the guard",
        );
        if tokio::time::timeout(Duration::from_millis(500), collection.state())
            .await
            .is_err()
        {
            reader_blocked = true;
            break;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    assert!(
        reader_blocked,
        "collection readers must be fenced behind the queued writer",
    );

    let hold_started = Instant::now();
    while hold_started.elapsed() < Duration::from_secs(3) {
        assert!(!writer.is_finished(), "writer must stay queued");
    }
    let blocked_for = started.elapsed();

    let _ = released_tx.send(true);
    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 guard is released");

    assert!(
        blocked_for >= Duration::from_secs(3),
        "writer and readers must stay blocked for the whole observed hold",
    );

    println!("transfer reached the remote snapshot recovery after: {guard_seen:?}");
    println!("writer and readers stayed blocked for: {blocked_for:?}");
    println!("snapshot transfer sender holder-guard test done");
}

$ cargo +nightly-2026-03-21 test -p collection --test snapshot_transfer_sender_holder_guard -- --nocapture
    Finished `test` profile [unoptimized + debuginfo] target(s) in 0.46s
     Running tests/snapshot_transfer_sender_holder_guard.rs (target/debug/deps/snapshot_transfer_sender_holder_guard-c426223c3840892e)

running 1 test
transfer reached the remote snapshot recovery after: 234.221137ms
writer and readers stayed blocked for: 3.786511458s
snapshot transfer sender holder-guard test done
test snapshot_transfer_sender_holds_shards_holder_read_guard ... ok

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