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, JoinAsk,
    JoinHooks, JoinOutcome, MemberInfo, 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,
        // m11p5: `peers` is now `host:port` strings; the fixtures hand us
        // resolved loopback `SocketAddr`s, so stringify at the boundary.
        peers: peers.into_iter().map(|(s, a)| (s, a.to_string())).collect(),
        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,
    /// Regions this node's roster lists as `Removed` (m11p5 §3.3 typed removed
    /// signal): a heartbeat/vote from one of these is answered `removed=true`.
    removed_regions: Vec<u16>,
    seen: Mutex<Vec<String>>,
}

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

    /// As [`new`](Self::new), but the named regions are `Removed` in this node's
    /// roster — every heartbeat/vote from them carries the typed removed signal.
    fn with_removed(term: u64, region: u16, grant_votes: bool, removed: Vec<u16>) -> Self {
        Self {
            removed_regions: removed,
            ..Self::new(term, region, grant_votes)
        }
    }
}

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, hb: tidal_net::HeartbeatContext) -> HeartbeatExchange {
        let tidal_net::HeartbeatContext {
            term,
            leader_region,
            stream_baseline,
            leader_last_seq,
            ..
        } = hb;
        self.seen.lock().unwrap().push(format!(
            "hb:{term}:{leader_region}:{stream_baseline}:{leader_last_seq}"
        ));
        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
    }

    fn is_removed_member(&self, region: u16) -> bool {
        self.removed_regions.contains(&region)
    }
}

/// 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,
            removed: false,
        },
    );

    // 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,
            removed: 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,
            leader_last_seq: 99,
            ..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,
            removed: false,
        }),
        "current-term heartbeat accepted: {events:?}"
    );
    assert!(
        events.contains(&ElectionNetEvent::HeartbeatReply {
            from: ShardId(1),
            term: 7,
            accepted: false,
            removed: false,
        }),
        "stale-term heartbeat refused with the responder term: {events:?}"
    );

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

/// The TYPED REMOVED SIGNAL round-trips over a real socket (m11p5 §3.3): a node
/// (region 0) whose region the responder's roster lists as `Removed` gets
/// `removed=true` on BOTH a heartbeat reply (keyed on `region_id`) and a vote
/// reply (keyed on `candidate_region`) — the out-of-band channel a node that
/// missed the `Removed` record uses to learn of its decommission. A peer NOT in
/// the removed set gets `removed=false` (proto3 default = compatible).
#[test]
fn removed_signal_round_trips_on_heartbeat_and_vote() {
    // Node 1's roster lists region 0 as Removed; node 1 accepts term 4 traffic.
    let hooks = Arc::new(ScriptedHooks::with_removed(4, 1, true, vec![0]));
    let (t0, _t1) = build_pair_with_hooks(Arc::clone(&hooks));

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

    // A heartbeat FROM region 0 (the removed node): the reply carries removed=true.
    net.fan_heartbeats(
        &[ShardId(1)],
        &proto::HeartbeatRequest {
            shard_id: 0,
            region_id: 0,
            term: 4,
            leader_region: 0,
            stream_baseline: 0,
            ..Default::default()
        },
        &tx,
    );
    let events = drain_events(&rx, 1);
    assert_eq!(
        events[0],
        ElectionNetEvent::HeartbeatReply {
            from: ShardId(1),
            term: 4,
            accepted: true,
            removed: true,
        },
        "a heartbeat from a Removed region is answered removed=true: {events:?}"
    );

    // A VOTE from candidate region 0 (the removed node): the reply carries
    // removed=true (a removed node keeps campaigning until a refusal teaches it).
    net.fan_votes(&[ShardId(1)], vote_request(5, false), &tx);
    let events = drain_events(&rx, 1);
    assert!(
        matches!(
            events[0],
            ElectionNetEvent::VoteReply {
                from: ShardId(1),
                removed: true,
                ..
            }
        ),
        "a vote from a Removed candidate is answered removed=true: {events:?}"
    );
}

/// A heartbeat/vote from a region the responder does NOT list as Removed carries
/// `removed=false` (proto3 zero-default — backward-compatible).
#[test]
fn non_removed_peer_gets_removed_false() {
    let hooks = Arc::new(ScriptedHooks::with_removed(4, 1, true, vec![9]));
    let (t0, _t1) = build_pair_with_hooks(Arc::clone(&hooks));

    let net = t0.election_net();
    let (tx, rx) = mpsc::channel();
    net.fan_heartbeats(
        &[ShardId(1)],
        &proto::HeartbeatRequest {
            shard_id: 0,
            region_id: 0,
            term: 4,
            leader_region: 0,
            stream_baseline: 0,
            ..Default::default()
        },
        &tx,
    );
    let events = drain_events(&rx, 1);
    assert!(
        matches!(
            events[0],
            ElectionNetEvent::HeartbeatReply { removed: false, .. }
        ),
        "a heartbeat from a non-Removed region is removed=false: {events:?}"
    );
}

#[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"
    );
}

// ── m11p5 §3.3: JoinCluster RPC over a real socket ───────────────────────────

/// A scriptable join-hooks impl recording every ask and answering a fixed
/// outcome (leader-accept or non-leader-refuse).
struct ScriptedJoinHooks {
    outcome: JoinOutcome,
    seen: Mutex<Vec<JoinAsk>>,
}

impl JoinHooks for ScriptedJoinHooks {
    fn join(&self, ask: JoinAsk) -> JoinOutcome {
        self.seen.lock().unwrap().push(ask);
        self.outcome.clone()
    }
}

/// Build a single transport (node 1) carrying scripted JOIN hooks, plus a
/// hook-less transport (node 0) the joiner dials to see `Unimplemented`.
fn build_pair_with_join_hooks(hooks: Arc<ScriptedJoinHooks>) -> (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_join_hooks(hooks as Arc<dyn JoinHooks>);
    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)
}

#[test]
fn join_cluster_roundtrips_the_roster_through_hooks() {
    let outcome = JoinOutcome {
        accepted: true,
        refusal_reason: String::new(),
        assigned_id: 3,
        term: 7,
        leader_region: "us-east".into(),
        leader_grpc_addr: "tidaldb-0.peers.svc:9500".into(),
        leader_http_addr: "http://tidaldb-0.peers.svc:9501".into(),
        members: vec![
            MemberInfo {
                id: 0,
                name: "us-east".into(),
                grpc_addr: "tidaldb-0.peers.svc:9500".into(),
                http_addr: "http://tidaldb-0.peers.svc:9501".into(),
                role: 0, // Voter
            },
            MemberInfo {
                id: 3,
                name: "joiner".into(),
                grpc_addr: "tidaldb-3.peers.svc:9500".into(),
                http_addr: "http://tidaldb-3.peers.svc:9501".into(),
                role: 1, // Learner
            },
        ],
        membership_version: 4,
    };
    let hooks = Arc::new(ScriptedJoinHooks {
        outcome,
        seen: Mutex::new(Vec::new()),
    });
    // Node 0 dials node 1 (the leader carrying the join hooks).
    let (t0, _t1) = build_pair_with_join_hooks(Arc::clone(&hooks));

    let resp = t0
        .join_cluster(
            ShardId(1),
            proto::JoinRequest {
                name: "joiner".into(),
                grpc_addr: "tidaldb-3.peers.svc:9500".into(),
                http_addr: "http://tidaldb-3.peers.svc:9501".into(),
                capabilities: tidal_net::CAP_KIND4_MEMBERSHIP,
            },
        )
        .expect("join RPC reached the leader hooks");

    assert!(resp.accepted);
    assert_eq!(resp.assigned_id, 3);
    assert_eq!(resp.term, 7);
    assert_eq!(resp.membership_version, 4);
    assert_eq!(resp.leader_region, "us-east");
    assert_eq!(resp.members.len(), 2);
    assert_eq!(resp.members[1].name, "joiner");
    assert_eq!(resp.members[1].role, 1, "the joiner came back a Learner");

    // The ask reached the hooks verbatim, capabilities and all.
    let seen = hooks.seen.lock().unwrap();
    assert_eq!(seen.len(), 1);
    assert_eq!(seen[0].name, "joiner");
    assert_eq!(seen[0].capabilities, tidal_net::CAP_KIND4_MEMBERSHIP);
}

#[test]
fn join_cluster_non_leader_refusal_is_a_response_not_an_error() {
    // A non-leader seed answers accepted=false WITH a leader hint — a
    // successful RPC the joiner re-targets, never a transport error.
    let outcome = JoinOutcome {
        accepted: false,
        refusal_reason: "not the leader; re-target the hint".into(),
        assigned_id: 0,
        term: 9,
        leader_region: "eu-west".into(),
        leader_grpc_addr: "tidaldb-1.peers.svc:9500".into(),
        leader_http_addr: "http://tidaldb-1.peers.svc:9501".into(),
        members: Vec::new(),
        membership_version: 0,
    };
    let hooks = Arc::new(ScriptedJoinHooks {
        outcome,
        seen: Mutex::new(Vec::new()),
    });
    let (t0, _t1) = build_pair_with_join_hooks(Arc::clone(&hooks));

    let resp = t0
        .join_cluster(
            ShardId(1),
            proto::JoinRequest {
                name: "joiner".into(),
                grpc_addr: "j:9500".into(),
                http_addr: "http://j:9501".into(),
                capabilities: tidal_net::CAP_KIND4_MEMBERSHIP,
            },
        )
        .expect("a non-leader refusal is still a successful RPC");
    assert!(!resp.accepted);
    assert_eq!(resp.leader_region, "eu-west");
    assert_eq!(
        resp.leader_grpc_addr, "tidaldb-1.peers.svc:9500",
        "the joiner re-targets the hinted leader"
    );
}

#[test]
fn join_cluster_without_hooks_is_unimplemented() {
    // Node 0 has no join hooks (a pre-m11p5 binary, or a bare transport): the
    // RPC answers Unimplemented and the joiner retries the next seed (§3.7).
    let outcome = JoinOutcome {
        accepted: true,
        refusal_reason: String::new(),
        assigned_id: 1,
        term: 1,
        leader_region: "x".into(),
        leader_grpc_addr: "x:1".into(),
        leader_http_addr: "http://x:2".into(),
        members: Vec::new(),
        membership_version: 1,
    };
    let hooks = Arc::new(ScriptedJoinHooks {
        outcome,
        seen: Mutex::new(Vec::new()),
    });
    // node 1 carries hooks; node 0 does NOT. Dial node 0 from node 1.
    let (_t0, t1) = build_pair_with_join_hooks(Arc::clone(&hooks));
    let err = t1
        .join_cluster(
            ShardId(0),
            proto::JoinRequest {
                name: "joiner".into(),
                grpc_addr: "j:9500".into(),
                http_addr: "http://j:9501".into(),
                capabilities: 0,
            },
        )
        .expect_err("a hook-less node answers Unimplemented");
    let msg = format!("{err:?}");
    assert!(
        msg.contains("Unimplemented") || msg.to_lowercase().contains("unimplemented"),
        "pre-m11p5 join is Unimplemented, never a crash: {msg}"
    );
}