tidaldb/tidal-net/tests/transport_contract.rs

// Integration-test exemptions (same posture as the tidaldb integration tests):
// unwrap/unwrap_err on known-good fixtures and short-lived read guards are
// idiomatic here.
#![allow(clippy::unwrap_used, clippy::significant_drop_tightening)]
//! Transport trait contract tests for `GrpcTransport`.
//!
//! Mirrors the test patterns from `InProcessTransport` but over gRPC on localhost.

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

use tidal_net::{GrpcTransport, config::GrpcTransportConfig};
use tidaldb::replication::{
    WalSegmentId,
    shard::{RegionId, ShardId},
    transport::{Transport, TransportError, WalSegmentPayload},
};

/// Get a unique listen address using port 0 (OS-assigned).
/// Since tonic doesn't support port 0, we bind a `TcpListener` to find a free port.
fn free_addr() -> SocketAddr {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    listener.local_addr().unwrap()
}

fn make_config(
    shard: ShardId,
    listen: SocketAddr,
    peers: HashMap<ShardId, SocketAddr>,
) -> GrpcTransportConfig {
    GrpcTransportConfig {
        local_shard: shard,
        listen_addr: listen,
        peers,
        insecure: true,
        ..Default::default()
    }
}

fn make_payload(shard: ShardId, seqno: u64) -> WalSegmentPayload {
    WalSegmentPayload {
        id: WalSegmentId::new(RegionId::SINGLE, shard, seqno),
        bytes: vec![0xAB; 100],
        event_count: 5,
        leader_last_seq: seqno,
    }
}

/// Build two `GrpcTransports` on localhost that can talk to each other.
fn build_pair() -> (GrpcTransport, GrpcTransport) {
    let addr0 = free_addr();
    let addr1 = free_addr();

    let config0 = make_config(ShardId(0), addr0, HashMap::from([(ShardId(1), addr1)]));
    let config1 = make_config(ShardId(1), addr1, HashMap::from([(ShardId(0), addr0)]));

    let t0 = GrpcTransport::new(config0).expect("transport 0");
    let t1 = GrpcTransport::new(config1).expect("transport 1");

    // Give servers a moment to start.
    thread::sleep(Duration::from_millis(100));

    (t0, t1)
}

#[test]
fn send_and_receive_between_shards() {
    let (t0, t1) = build_pair();

    // Shard 0 sends to Shard 1.
    let payload = make_payload(ShardId(0), 42);
    t0.send_segment(ShardId(1), payload).unwrap();

    // Shard 1 receives it.
    let received = t1.recv_segment().unwrap();
    assert_eq!(received.id.seqno, 42);
    assert_eq!(received.event_count, 5);
    assert_eq!(received.bytes.len(), 100);
}

#[test]
fn send_to_unknown_peer_fails() {
    let addr0 = free_addr();
    let config0 = make_config(ShardId(0), addr0, HashMap::new());
    let t0 = GrpcTransport::new(config0).expect("transport 0");
    thread::sleep(Duration::from_millis(50));

    let result = t0.send_segment(ShardId(99), make_payload(ShardId(0), 1));
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        TransportError::UnknownPeer(_)
    ));
}

#[test]
fn payload_too_large_rejected() {
    let addr0 = free_addr();
    let addr1 = free_addr();
    let config0 = make_config(ShardId(0), addr0, HashMap::from([(ShardId(1), addr1)]));
    let t0 = GrpcTransport::new(config0).expect("transport 0");
    thread::sleep(Duration::from_millis(50));

    let payload = WalSegmentPayload {
        id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
        bytes: vec![0u8; 64 * 1024 * 1024 + 1],
        event_count: 0,
        leader_last_seq: 1,
    };
    let result = t0.send_segment(ShardId(1), payload);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        TransportError::PayloadTooLarge { .. }
    ));
}

#[test]
fn local_shard_returns_correct_id() {
    let addr = free_addr();
    let config = make_config(ShardId(7), addr, HashMap::new());
    let t = GrpcTransport::new(config).expect("transport");
    assert_eq!(t.local_shard(), ShardId(7));
}

#[test]
fn shutdown_unparks_recv_segment_and_thread_exits() {
    // Models the real receiver wiring: a worker thread holds its own
    // Arc<GrpcTransport> clone and parks in recv_segment. Without a cancellation
    // path that thread never returns (recv_segment only yields None when every
    // server-side inbound_tx sender drops), so it would leak until process exit
    // AND its Arc would keep GrpcTransport::drop from ever running. shutdown_receivers
    // breaks the cycle: a parked recv_segment returns None deterministically.
    let addr = free_addr();
    let config = make_config(ShardId(0), addr, HashMap::new());
    let transport = Arc::new(GrpcTransport::new(config).expect("transport"));
    thread::sleep(Duration::from_millis(50));

    let returned_none = Arc::new(AtomicBool::new(false));
    let worker_transport = Arc::clone(&transport);
    let worker_flag = Arc::clone(&returned_none);
    let receiver = thread::spawn(move || {
        // No segment will ever arrive (no peers); this parks until shutdown.
        let segment = worker_transport.recv_segment();
        worker_flag.store(segment.is_none(), Ordering::SeqCst);
        // Drop the worker's Arc clone on return so only the test's Arc remains.
    });

    // Let the worker reach the parked recv_segment.
    thread::sleep(Duration::from_millis(100));
    assert!(
        !receiver.is_finished(),
        "receiver should still be parked in recv_segment before shutdown"
    );

    // Trip the cancellation path.
    transport.shutdown_receivers();

    // The parked recv_segment must return None promptly and the thread must exit.
    let deadline = Instant::now() + Duration::from_secs(5);
    while !receiver.is_finished() && Instant::now() < deadline {
        thread::sleep(Duration::from_millis(10));
    }
    assert!(
        receiver.is_finished(),
        "recv_segment must return after shutdown_receivers so the thread can exit"
    );
    receiver.join().expect("receiver thread joins cleanly");
    assert!(
        returned_none.load(Ordering::SeqCst),
        "recv_segment must return None on shutdown (the receiver's exit condition)"
    );

    // Only the test's Arc remains; dropping it runs GrpcTransport::drop without leak.
    assert_eq!(Arc::strong_count(&transport), 1);
    drop(transport);
}

#[test]
fn multiple_segments_fifo_order() {
    let (t0, t1) = build_pair();

    for seq in 1..=5u64 {
        t0.send_segment(ShardId(1), make_payload(ShardId(0), seq))
            .unwrap();
    }

    // Brief pause to let segments flow through gRPC.
    thread::sleep(Duration::from_millis(50));

    for expected_seq in 1..=5u64 {
        let received = t1.recv_segment().unwrap();
        assert_eq!(received.id.seqno, expected_seq);
    }
}