tidaldb/tidal-net/tests/ship_timeout_breaker.rs

// Integration-test exemptions (same posture as the other tidal-net integration tests).
#![allow(clippy::unwrap_used, clippy::significant_drop_tightening)]
//! The write-burst false-partition fix, proven END-TO-END through REAL gRPC.
//!
//! Under a sustained 1536-D ack=quorum burst a follower applies every write on its
//! single segment-receiver thread (CPU-heavy HNSW insert), which can starve its
//! transport runtime so a leader→follower `ship_segment` RPC misses the client's
//! `request_timeout` — even though the follower is up and still acking other ships.
//! Before the fix, `send_to` called `record_failure` on that deadline exactly as on
//! a genuine severance, so 5 such opened the breaker, both followers tripped under
//! the uniform burst, the commit index stalled, and ack=quorum 503'd with no
//! self-heal. The fix routes a ship DEADLINE (tonic `DeadlineExceeded`/`Cancelled`)
//! to `record_timeout`, which opens the breaker ONLY when the peer shows no recent
//! proof of life.
//!
//! These tests use a REAL tonic `WalShipping` server over a REAL established HTTP/2
//! connection (not a synthetic `tonic::Status`): the handler succeeds once (proving
//! liveness) then hangs past the client deadline, so the client observes a genuine
//! request timeout on a live connection — exactly the production shape.

use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use tidal_net::{
    circuit_breaker::BreakerState,
    client::PeerPool,
    config::GrpcTransportConfig,
    error::GrpcTransportError,
    proto::{
        AppliedReport, AppliedReportAck, HeartbeatRequest, HeartbeatResponse, JoinRequest,
        JoinResponse, ShipSegmentRequest, ShipSegmentResponse, SnapshotChunk, SnapshotRequest,
        StreamRequest, TimeoutNowRequest, TimeoutNowResponse, VoteRequest, VoteResponse,
        wal_shipping_server::{WalShipping, WalShippingServer},
    },
};
use tidaldb::replication::{
    WalSegmentId,
    shard::{RegionId, ShardId},
    transport::WalSegmentPayload,
};
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status, transport::Server};

/// A real WAL-shipping server whose `ship_segment` handler succeeds for the first
/// `hang_after` calls (fast `accepted=true`, refreshing the leader's liveness stamp)
/// and then HANGS `hang_for` on every subsequent call — long enough to exceed the
/// client's `request_timeout`, so the leader observes a genuine deadline on a LIVE,
/// already-established connection. Every other RPC is unimplemented (the client only
/// ships in these tests).
struct FlakyShipServer {
    calls: Arc<AtomicUsize>,
    hang_after: usize,
    hang_for: Duration,
}

#[tonic::async_trait]
impl WalShipping for FlakyShipServer {
    async fn ship_segment(
        &self,
        _request: Request<ShipSegmentRequest>,
    ) -> Result<Response<ShipSegmentResponse>, Status> {
        let n = self.calls.fetch_add(1, Ordering::SeqCst);
        if n >= self.hang_after {
            // Slow-but-ALIVE: the server is up and processing, just late — the
            // client's request_timeout fires first and surfaces a real deadline.
            tokio::time::sleep(self.hang_for).await;
        }
        Ok(Response::new(ShipSegmentResponse {
            accepted: true,
            ..Default::default()
        }))
    }

    type StreamSegmentsStream = ReceiverStream<Result<ShipSegmentRequest, Status>>;
    async fn stream_segments(
        &self,
        _request: Request<StreamRequest>,
    ) -> Result<Response<Self::StreamSegmentsStream>, Status> {
        Err(Status::unimplemented("mock"))
    }

    type FetchSnapshotStream = ReceiverStream<Result<SnapshotChunk, Status>>;
    async fn fetch_snapshot(
        &self,
        _request: Request<SnapshotRequest>,
    ) -> Result<Response<Self::FetchSnapshotStream>, Status> {
        Err(Status::unimplemented("mock"))
    }

    async fn heartbeat(
        &self,
        _request: Request<HeartbeatRequest>,
    ) -> Result<Response<HeartbeatResponse>, Status> {
        Err(Status::unimplemented("mock"))
    }

    async fn report_applied(
        &self,
        _request: Request<AppliedReport>,
    ) -> Result<Response<AppliedReportAck>, Status> {
        Err(Status::unimplemented("mock"))
    }

    async fn request_vote(
        &self,
        _request: Request<VoteRequest>,
    ) -> Result<Response<VoteResponse>, Status> {
        Err(Status::unimplemented("mock"))
    }

    async fn timeout_now(
        &self,
        _request: Request<TimeoutNowRequest>,
    ) -> Result<Response<TimeoutNowResponse>, Status> {
        Err(Status::unimplemented("mock"))
    }

    async fn join_cluster(
        &self,
        _request: Request<JoinRequest>,
    ) -> Result<Response<JoinResponse>, Status> {
        Err(Status::unimplemented("mock"))
    }
}

fn free_addr() -> SocketAddr {
    std::net::TcpListener::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
}

fn make_payload(seqno: u64) -> WalSegmentPayload {
    WalSegmentPayload {
        id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seqno),
        bytes: vec![0xCD; 50],
        event_count: 1,
        leader_last_seq: seqno,
        stream_baseline: 0,
        term: 0,
        leader_region: 0,
    }
}

fn config_for(peer: SocketAddr, threshold: u32) -> GrpcTransportConfig {
    let mut peers = HashMap::new();
    peers.insert(ShardId(0), peer.to_string());
    GrpcTransportConfig {
        local_shard: ShardId(0),
        listen_addr: free_addr(),
        peers,
        insecure: true,
        circuit_breaker_threshold: threshold,
        circuit_breaker_reset: Duration::from_secs(30),
        // Short request deadline so a hung handler trips it fast; generous connect
        // timeout so the first (fast) ship establishes the connection cleanly.
        request_timeout: Duration::from_millis(500),
        connect_timeout: Duration::from_secs(5),
        ..Default::default()
    }
}

/// Spawn the flaky server and wait until it is actually accepting TCP.
async fn spawn_server(addr: SocketAddr, server: FlakyShipServer) {
    tokio::spawn(async move {
        let _ = Server::builder()
            .add_service(WalShippingServer::new(server))
            .serve(addr)
            .await;
    });
    for _ in 0..50 {
        if tokio::net::TcpStream::connect(addr).await.is_ok() {
            return;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    panic!("mock WAL server never started listening on {addr}");
}

fn assert_ship_timeout(result: &Result<(u64, u64), GrpcTransportError>) {
    match result {
        Err(GrpcTransportError::Grpc(status)) => assert!(
            matches!(
                status.code(),
                tonic::Code::DeadlineExceeded | tonic::Code::Cancelled
            ),
            "an established-connection ship deadline must surface as DeadlineExceeded|Cancelled \
             (the codes routed to record_timeout), got {:?}: {status}",
            status.code()
        ),
        other => panic!("expected a gRPC deadline status from a hung ship, got {other:?}"),
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn recent_liveness_keeps_breaker_closed_under_real_ship_timeouts() {
    // THE fix, end-to-end: a follower that proved liveness (ship #1 succeeded) then
    // times out a burst of ships must NOT trip the breaker — a slow-but-alive peer
    // is not a partition.
    let addr = free_addr();
    let calls = Arc::new(AtomicUsize::new(0));
    spawn_server(
        addr,
        FlakyShipServer {
            calls: calls.clone(),
            hang_after: 1, // first ship fast; the rest hang
            hang_for: Duration::from_secs(3),
        },
    )
    .await;

    let pool = PeerPool::new(&config_for(addr, 3)).unwrap();

    // Ship #1: real success over a freshly-established H2 connection → record_success.
    let first = pool.send_to(ShardId(0), make_payload(1)).await;
    assert!(first.is_ok(), "the first ship must succeed: {first:?}");
    assert!(
        pool.peer_grpc_fresh(ShardId(0), Duration::from_secs(5)),
        "an accepted ship must refresh the peer's liveness stamp"
    );

    // Ships #2..#6: each hangs past the 500ms deadline → REAL ship timeouts (far more
    // than the threshold of 3) on the SAME live connection.
    for seq in 2..=6 {
        let r = pool.send_to(ShardId(0), make_payload(seq)).await;
        assert_ship_timeout(&r);
    }

    assert_eq!(
        pool.breaker_state(ShardId(0)),
        BreakerState::Closed,
        "5 real ship timeouts from a peer with recent liveness must NOT open the breaker \
         (the write-burst false-partition fix)"
    );
    // And replication is not wedged: the breaker still admits the next ship.
    assert!(
        pool.peer_grpc_fresh(ShardId(0), Duration::from_secs(60)),
        "the slow-but-alive peer's liveness stamp is still fresh from ship #1"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn ship_timeouts_with_no_proof_of_life_still_open_breaker() {
    // Genuine-blackhole detection preserved END-TO-END: a peer that only ever times
    // out (never round-trips a reply) has no liveness stamp, so real ship timeouts
    // still open the breaker at the threshold.
    let addr = free_addr();
    let calls = Arc::new(AtomicUsize::new(0));
    spawn_server(
        addr,
        FlakyShipServer {
            calls: calls.clone(),
            hang_after: 0, // every ship hangs — no proof of life ever
            hang_for: Duration::from_secs(3),
        },
    )
    .await;

    let pool = PeerPool::new(&config_for(addr, 3)).unwrap();

    for seq in 1..=3 {
        let r = pool.send_to(ShardId(0), make_payload(seq)).await;
        assert_ship_timeout(&r);
    }

    assert_eq!(
        pool.breaker_state(ShardId(0)),
        BreakerState::Open,
        "ship timeouts with no proof of life must still open the breaker (blackhole detection)"
    );
}