tidaldb/tidal-net/tests/election_rpc.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)]
//! m11p4 election RPC + term-fencing contract tests over real sockets.
//!
//! Proves the wire half of the election design:
//! 1. `RequestVote` round-trips through the late-bound [`ElectionHooks`]
//!    (grant and refusal, with the voter's term in the reply).
//! 2. A node WITHOUT hooks answers `Unimplemented` — what a pre-m11p4 binary
//!    looks like to a candidate (counts as not-granted, never an error).
//! 3. A stale-term ship is fenced with `FAILED_PRECONDITION` BEFORE the
//!    inbound queue; a current-term ship passes and acks the responder term.
//! 4. The heartbeat exchange carries the responder's term + acceptance.
//! 5. `ElectionNet` fan-outs deliver events (reply or unreachable) for every
//!    peer — the driver's inbox never starves on a dead peer.

use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::{Arc, Mutex, mpsc},
    thread,
    time::Duration,
};

use tidal_net::{
    ClaimRejection, ElectionHooks, ElectionNetEvent, GrpcTransport, HeartbeatExchange,
    config::GrpcTransportConfig, proto, sources::ServingSources,
};
use tidaldb::replication::{
    VoteReply, VoteRpc, WalSegmentId,
    shard::{RegionId, ShardId},
    transport::{Transport, TransportError, WalSegmentPayload},
};

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()
    }
}

/// A scriptable hooks impl: a fixed current term + a recorded log of inbound
/// election traffic.
struct ScriptedHooks {
    term: u64,
    region: u16,
    grant_votes: bool,
    seen: Mutex<Vec<String>>,
}

impl ScriptedHooks {
    const fn new(term: u64, region: u16, grant_votes: bool) -> Self {
        Self {
            term,
            region,
            grant_votes,
            seen: Mutex::new(Vec::new()),
        }
    }
}

impl ElectionHooks for ScriptedHooks {
    fn self_claim(&self) -> (u64, u16) {
        (self.term, self.region)
    }

    fn observe_leader_claim(
        &self,
        term: u64,
        leader_region: u16,
        _first_seq: u64,
    ) -> Result<(), ClaimRejection> {
        self.seen
            .lock()
            .unwrap()
            .push(format!("claim:{term}:{leader_region}"));
        if term < self.term {
            return Err(ClaimRejection::Stale {
                current_term: self.term,
            });
        }
        Ok(())
    }

    fn on_heartbeat(
        &self,
        term: u64,
        leader_region: u16,
        baseline: u64,
        _prev_log: tidaldb::replication::LogPosition,
    ) -> HeartbeatExchange {
        self.seen
            .lock()
            .unwrap()
            .push(format!("hb:{term}:{leader_region}:{baseline}"));
        HeartbeatExchange {
            term: self.term,
            accepted: term >= self.term,
        }
    }

    fn on_vote(&self, rpc: VoteRpc) -> VoteReply {
        self.seen.lock().unwrap().push(format!(
            "vote:{}:{}:{}:{}",
            rpc.term, rpc.candidate.0, rpc.prevote, rpc.transfer
        ));
        VoteReply {
            term: self.term,
            granted: self.grant_votes && rpc.term > self.term,
        }
    }

    fn on_timeout_now(&self, term: u64, leader_region: u16) -> bool {
        self.seen
            .lock()
            .unwrap()
            .push(format!("timeoutnow:{term}:{leader_region}"));
        true
    }

    fn on_observed_term(&self, term: u64) {
        self.seen.lock().unwrap().push(format!("observed:{term}"));
    }

    fn report_term_acceptable(&self, reporter_term: u64) -> bool {
        reporter_term == self.term
    }
}

/// Two transports; node 1 carries scripted hooks, node 0 carries none.
fn build_pair_with_hooks(hooks: Arc<ScriptedHooks>) -> (GrpcTransport, GrpcTransport) {
    let addr0 = free_addr();
    let addr1 = free_addr();

    let t0 = GrpcTransport::new(make_config(
        ShardId(0),
        addr0,
        HashMap::from([(ShardId(1), addr1)]),
    ))
    .expect("transport 0");

    let sources = ServingSources::default();
    sources.set_election_hooks(hooks as Arc<dyn ElectionHooks>);
    let t1 = GrpcTransport::new_with_sources(
        make_config(ShardId(1), addr1, HashMap::from([(ShardId(0), addr0)])),
        sources,
    )
    .expect("transport 1");

    thread::sleep(Duration::from_millis(100));
    (t0, t1)
}

const fn vote_request(term: u64, prevote: bool) -> proto::VoteRequest {
    proto::VoteRequest {
        term,
        candidate_region: 0,
        last_log_term: 0,
        last_log_seq: 10,
        prevote,
        transfer: false,
    }
}

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

fn drain_events(rx: &mpsc::Receiver<ElectionNetEvent>, expect: usize) -> Vec<ElectionNetEvent> {
    let mut events = Vec::new();
    while events.len() < expect {
        match rx.recv_timeout(Duration::from_secs(5)) {
            Ok(e) => events.push(e),
            Err(e) => panic!("expected {expect} election events, got {events:?} ({e})"),
        }
    }
    events
}

#[test]
fn vote_roundtrip_grant_and_refusal_through_hooks() {
    let hooks = Arc::new(ScriptedHooks::new(3, 1, true));
    let (t0, _t1) = build_pair_with_hooks(Arc::clone(&hooks));

    let net = t0.election_net();
    let (tx, rx) = mpsc::channel();

    // Term 5 > voter's 3 with grant_votes: granted, reply carries voter term.
    net.fan_votes(&[ShardId(1)], vote_request(5, true), &tx);
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::VoteReply {
            from: ShardId(1),
            prevote: true,
            term: 3,
            granted: true,
        },
    );

    // Term 2 < voter's 3: refused, the higher reply term is the step-down cue.
    net.fan_votes(&[ShardId(1)], vote_request(2, false), &tx);
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::VoteReply {
            from: ShardId(1),
            prevote: false,
            term: 3,
            granted: false,
        },
    );

    let seen = hooks.seen.lock().unwrap();
    assert!(
        seen.contains(&"vote:5:0:true:false".to_string()),
        "the pre-vote reached the hooks verbatim: {seen:?}"
    );
}

#[test]
fn pre_m11p4_peer_without_hooks_counts_as_not_granted() {
    // Node 0 (no hooks) is the "old binary"; node 1 campaigns against it.
    let hooks = Arc::new(ScriptedHooks::new(0, 1, true));
    let (_t0, t1) = build_pair_with_hooks(Arc::clone(&hooks));

    let net = t1.election_net();
    let (tx, rx) = mpsc::channel();
    net.fan_votes(&[ShardId(0)], vote_request(1, false), &tx);
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::VoteUnreachable {
            from: ShardId(0),
            prevote: false,
        },
        "Unimplemented from a pre-m11p4 peer is not-granted, never a crash"
    );
}

#[test]
fn stale_term_ship_is_fenced_current_term_passes() {
    let hooks = Arc::new(ScriptedHooks::new(4, 1, false));
    let (t0, t1) = build_pair_with_hooks(Arc::clone(&hooks));

    // Stale term 2 < 4: fenced as a permanent failure (FAILED_PRECONDITION),
    // never enqueued.
    let err = t0
        .send_segment(ShardId(1), payload_with_term(2, 0, 1))
        .expect_err("a stale-term ship must be fenced");
    assert!(
        matches!(err, TransportError::Permanent { ref reason } if reason.contains("stale")),
        "got: {err:?}"
    );

    // Current term passes and the receiver sees it.
    t0.send_segment(ShardId(1), payload_with_term(4, 0, 1))
        .expect("a current-term ship must pass the fence");
    let received = t1.recv_segment().expect("the fenced node received it");
    assert_eq!(received.term, 4);
    assert_eq!(received.id.seqno, 1);

    let seen = hooks.seen.lock().unwrap();
    assert!(
        seen.contains(&"claim:2:0".to_string()) && seen.contains(&"claim:4:0".to_string()),
        "both claims reached the hooks: {seen:?}"
    );
}

#[test]
fn heartbeat_exchange_carries_term_and_acceptance() {
    let hooks = Arc::new(ScriptedHooks::new(7, 1, false));
    let (t0, _t1) = build_pair_with_hooks(Arc::clone(&hooks));

    let net = t0.election_net();
    let (tx, rx) = mpsc::channel();

    // A heartbeat from term 7 (current): accepted.
    net.fan_heartbeats(
        &[ShardId(1)],
        &proto::HeartbeatRequest {
            shard_id: 0,
            region_id: 0,
            term: 7,
            leader_region: 0,
            stream_baseline: 12,
            ..Default::default()
        },
        &tx,
    );
    // A heartbeat from a deposed term 6: refused, reply carries 7.
    net.fan_heartbeats(
        &[ShardId(1)],
        &proto::HeartbeatRequest {
            shard_id: 0,
            region_id: 0,
            term: 6,
            leader_region: 0,
            stream_baseline: 0,
            ..Default::default()
        },
        &tx,
    );
    let events = drain_events(&rx, 2);
    assert!(
        events.contains(&ElectionNetEvent::HeartbeatReply {
            from: ShardId(1),
            term: 7,
            accepted: true,
        }),
        "current-term heartbeat accepted: {events:?}"
    );
    assert!(
        events.contains(&ElectionNetEvent::HeartbeatReply {
            from: ShardId(1),
            term: 7,
            accepted: false,
        }),
        "stale-term heartbeat refused with the responder term: {events:?}"
    );

    let seen = hooks.seen.lock().unwrap();
    assert!(
        seen.contains(&"hb:7:0:12".to_string()),
        "the term's activation baseline rode the heartbeat: {seen:?}"
    );
}

#[test]
fn timeout_now_reaches_hooks_and_unreachable_peer_reports() {
    let hooks = Arc::new(ScriptedHooks::new(5, 1, false));
    let (t0, _t1) = build_pair_with_hooks(Arc::clone(&hooks));

    let net = t0.election_net();
    let (tx, rx) = mpsc::channel();
    net.send_timeout_now(
        ShardId(1),
        proto::TimeoutNowRequest {
            term: 5,
            leader_region: 0,
        },
        &tx,
    );
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::TimeoutNowDelivered {
            target: ShardId(1),
            accepted: true,
        },
    );

    // An unknown peer reports unreachable instead of hanging the driver.
    net.send_timeout_now(
        ShardId(9),
        proto::TimeoutNowRequest {
            term: 5,
            leader_region: 0,
        },
        &tx,
    );
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::TimeoutNowUnreachable { target: ShardId(9) },
    );

    assert!(
        hooks
            .seen
            .lock()
            .unwrap()
            .contains(&"timeoutnow:5:0".to_string()),
        "the transfer request reached the target's hooks"
    );
}

#[test]
fn stale_term_frontier_report_is_refused() {
    let hooks = Arc::new(ScriptedHooks::new(3, 1, false));
    let (t0, t1) = build_pair_with_hooks(Arc::clone(&hooks));
    let _ = &t1;

    // notify_applied on the hook-less node 0 reports term 0 ≠ 3: the report
    // must be REFUSED by node 1's handler (failed_precondition) — which on
    // the reporting side is just a logged streak, never a crash. We can only
    // observe the refusal indirectly: the hint fold never happens, so node
    // 1's transport hint map for shard 0 stays empty. The structural assert
    // is that nothing panicked and the claim never reached the sink-less
    // hooks as accepted traffic.
    t0.notify_applied(ShardId(1), 42);
    thread::sleep(Duration::from_millis(300));
    assert_eq!(
        t1.peer_applied_hint(ShardId(0)),
        0,
        "a stale-term report must not fold into the responder's hint map"
    );
}