fv_streams_runtime/

cluster.rs

//! THE MEMBERSHIP HANDSHAKE (D9 — zero control plane, the leader is a worker): a run's workers
//! rendezvous through the leader and open the exchange mesh. No scheduler, no etcd:
//!
//! - **Rendezvous.** The leader listens on a control port; each `--join` worker connects, registers
//!   its exchange address, and the leader — once every worker has registered — broadcasts the sorted
//!   member set and the topology fingerprint. Every worker then derives the *same* placement
//!   ([`crate::placement`]) from that set, so no worker is told its slice (D9).
//! - **Mesh.** Every worker opens exactly one [`Link`] to every peer: it *connects* to peers with a
//!   lower id and *accepts* from peers with a higher id (the peer-id preamble on
//!   [`Link::connect_as`] lets one shared listener demultiplex the higher peers). One Link per pair,
//!   full-duplex, handed to [`Graph::start_worker`](crate::dataflow::Graph::start_worker) as a
//!   [`WorkerExchange`].
//!
//! Control messages are length-prefixed postcard (§7.2 — tiny, schema-evolvable). This is the
//! transport-agnostic handshake; failure detection and re-assignment on worker loss are 4c.

use std::collections::HashMap;
use std::net::SocketAddr;

use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::runtime::Handle;

use fv_streams_exchange::Link;

use crate::dataflow::WorkerExchange;
use crate::placement::{Members, WorkerId};

/// A worker's registration with the leader: the address peers reach its exchange on.
#[derive(Debug, Serialize, Deserialize)]
struct Register {
    exchange_addr: String,
    /// The joiner's wire-format version: a leader refuses another version before anything runs.
    wire: u8,
}

/// The leader's reply to every worker: the sorted member set and the topology fingerprint every
/// worker checks its own against.
#[derive(Debug, Serialize, Deserialize)]
struct Assignment {
    workers: Vec<String>,
    fingerprint: String,
    /// The leader's wire-format version: a joiner refuses another version, symmetrically.
    wire: u8,
}

/// The one refusal both sides give: the versions, then the remedy.
fn wire_mismatch(who: &str, got: u8, mine: u8) -> std::io::Error {
    invalid(format!(
        "{who} speaks wire format v{got}, this worker v{mine} — every worker must run the same fv-streams version"
    ))
}

/// Runtime control from the leader to a worker over the persistent control connection (D9 — the
/// leader coordinates the run; barriers travel in-band on the data edges, but the *decision* to
/// issue an epoch, commit it, or stop is the leader's). `Barrier(e)` injects epoch `e` at this
/// worker's sources (aligned across all workers, so an operator sees the same epoch on every edge);
/// `Commit(e)` releases the sources that recorded a position at `e`; `Stop` begins the ordered
/// shutdown.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ToWorker {
    Barrier(u64),
    Commit(u64),
    Stop,
}

/// A worker's report back to the leader. `Acked` = every local task reported epoch `e`, this
/// worker's objects/files are durable, and here is its manifest slice (source offsets, objects and
/// files for the tasks it owns — keyed by global task id, so the leader merges every worker's slice
/// into the one manifest it writes). `Finished` = this worker's whole graph drained. The leader
/// waits for every worker before it commits an epoch or completes shutdown.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ToLeader {
    Acked {
        epoch: u64,
        contribution: fv_streams_state::Manifest,
    },
    Finished,
}

/// The leader's live control connections to the joined workers (worker id → its stream), kept past
/// the handshake so the leader can drive epochs, commits and shutdown across processes. The leader
/// is itself a worker (D9); it coordinates its own slice in-process and the joined workers here.
pub struct LeaderCtl {
    conns: Vec<(WorkerId, TcpStream)>,
}

impl LeaderCtl {
    /// Send one control message to every joined worker.
    pub async fn broadcast(&mut self, msg: ToWorker) -> std::io::Result<()> {
        for (_, s) in self.conns.iter_mut() {
            write_msg(s, &msg).await?;
        }
        Ok(())
    }

    /// Wait for the next report from every joined worker (in id order), returning them paired with
    /// the worker that sent each. Reads one message per worker.
    pub async fn collect(&mut self) -> std::io::Result<Vec<(WorkerId, ToLeader)>> {
        let mut out = Vec::with_capacity(self.conns.len());
        for (w, s) in self.conns.iter_mut() {
            out.push((*w, read_msg(s).await?));
        }
        Ok(out)
    }
}

/// A worker's live control connection to the leader.
pub struct WorkerCtl {
    conn: TcpStream,
}

impl WorkerCtl {
    /// The next control message from the leader.
    pub async fn recv(&mut self) -> std::io::Result<ToWorker> {
        read_msg(&mut self.conn).await
    }

    /// Report to the leader.
    pub async fn report(&mut self, msg: ToLeader) -> std::io::Result<()> {
        write_msg(&mut self.conn, &msg).await
    }
}

/// This process's coordination role once the handshake settles: the leader holds connections to the
/// joined workers; a joiner holds its connection to the leader.
pub enum Coord {
    Leader(LeaderCtl),
    Worker(WorkerCtl),
}

fn invalid<E: std::fmt::Display>(e: E) -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
}

/// Write a length-prefixed postcard message.
async fn write_msg<T: Serialize>(s: &mut TcpStream, msg: &T) -> std::io::Result<()> {
    let bytes = postcard::to_allocvec(msg).map_err(invalid)?;
    s.write_all(&(bytes.len() as u32).to_le_bytes()).await?;
    s.write_all(&bytes).await?;
    s.flush().await
}

/// Read a length-prefixed postcard message.
async fn read_msg<T: DeserializeOwned>(s: &mut TcpStream) -> std::io::Result<T> {
    let mut len = [0u8; 4];
    s.read_exact(&mut len).await?;
    let n = u32::from_le_bytes(len) as usize;
    let mut buf = vec![0u8; n];
    s.read_exact(&mut buf).await?;
    postcard::from_bytes(&buf).map_err(invalid)
}

/// The leader side of the rendezvous: wait for `n_workers - 1` registrations on `control`, then
/// broadcast the sorted member set + `fingerprint` to every worker. Returns the leader's own
/// [`Members`] (its `me` is `my_exchange_addr`'s index in the sorted set). The leader is worker 0
/// only if its address sorts first — placement depends on the addresses, not on who leads.
pub async fn lead(
    control: &TcpListener,
    my_exchange_addr: &str,
    n_workers: usize,
    fingerprint: &str,
) -> std::io::Result<(Members, LeaderCtl)> {
    // Register every worker, keeping its address paired with its connection.
    let mut regs: Vec<(String, TcpStream)> = Vec::new();
    let mut addrs = vec![my_exchange_addr.to_string()];
    let mut foreign: Option<(String, u8)> = None;
    while addrs.len() < n_workers {
        let (mut s, _) = control.accept().await?;
        let reg: Register = read_msg(&mut s).await?;
        if reg.wire != fv_streams_exchange::WIRE_VERSION {
            foreign.get_or_insert((reg.exchange_addr.clone(), reg.wire));
        }
        addrs.push(reg.exchange_addr.clone());
        regs.push((reg.exchange_addr, s));
    }
    addrs.sort();
    addrs.dedup();
    let members = Members::new(addrs.clone(), my_exchange_addr).map_err(invalid)?;
    let assign = Assignment {
        workers: addrs.clone(),
        fingerprint: fingerprint.to_string(),
        wire: fv_streams_exchange::WIRE_VERSION,
    };
    // Reply to each worker and keep its connection, tagged with the worker id it sorted into. (The
    // reply goes out even when a worker is refused below, so it learns the leader's version too.)
    let mut conns = Vec::with_capacity(regs.len());
    for (addr, mut s) in regs {
        write_msg(&mut s, &assign).await?;
        let wid = addrs
            .iter()
            .position(|a| *a == addr)
            .expect("registered address is a member") as WorkerId;
        conns.push((wid, s));
    }
    conns.sort_by_key(|(w, _)| *w);
    if let Some((addr, got)) = foreign {
        return Err(wire_mismatch(
            &format!("the worker at {addr}"),
            got,
            fv_streams_exchange::WIRE_VERSION,
        ));
    }
    Ok((members, LeaderCtl { conns }))
}

/// How this process participates in a run: solo (no clustering), or one worker of a cluster —
/// either the leader (waits for `n_workers` to register) or a joiner (registers with a leader).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClusterCfg {
    /// The leader's control address to join, or `None` if this process is the leader.
    pub join: Option<String>,
    /// The leader's expected total worker count (the leader waits for this many).
    pub n_workers: usize,
    /// This worker's exchange listen address (peers connect here). `127.0.0.1:0` (ephemeral) for a
    /// same-box run; a routable `host:port` for multi-node.
    pub exchange_bind: String,
    /// The leader's control listen address.
    pub control_bind: String,
    /// Flows per worker pair, and the per-channel credit window in bytes.
    pub flows: usize,
    pub window: usize,
}

impl ClusterCfg {
    /// Parse a cluster config from the runner's settings; `None` is a solo run (no `--join` and at
    /// most one worker) — the pre-Phase-4 single-process path. A worker (`join` set) or a leader
    /// (`n_workers > 1`) is clustered.
    pub fn parse(
        join: &str,
        n_workers: usize,
        exchange_bind: &str,
        control_bind: &str,
        flows: usize,
        window: usize,
    ) -> Option<ClusterCfg> {
        if join.is_empty() && n_workers <= 1 {
            return None;
        }
        Some(ClusterCfg {
            join: (!join.is_empty()).then(|| join.to_string()),
            n_workers: n_workers.max(2),
            exchange_bind: if exchange_bind.is_empty() {
                "127.0.0.1:0".into()
            } else {
                exchange_bind.into()
            },
            control_bind: if control_bind.is_empty() {
                "0.0.0.0:7100".into()
            } else {
                control_bind.into()
            },
            flows: flows.max(1),
            window: window.max(1),
        })
    }
}

/// Settle this worker's membership: bind its exchange listener, then lead or join. `base_fp` is the
/// topology fingerprint *without* the member clause — the leader broadcasts it and every joiner
/// checks its own against it, so a worker whose steps/vnodes differ is rejected before it runs.
/// Returns the settled [`Members`], the exchange listener (for [`mesh`]), and this process's
/// coordination role ([`Coord`]) with its live control connection(s).
pub async fn connect(cfg: &ClusterCfg, base_fp: &str) -> std::io::Result<(Members, TcpListener, Coord)> {
    let exchange = TcpListener::bind(&cfg.exchange_bind).await?;
    let my_addr = exchange.local_addr()?.to_string();
    let (members, coord) = match &cfg.join {
        Some(leader) => {
            eprintln!("fv-streams: joining the leader at {leader} (exchange on {my_addr})…");
            let addr: SocketAddr = leader.parse().map_err(invalid)?;
            let (m, leader_fp, ctl) = join(addr, &my_addr).await?;
            if leader_fp != base_fp {
                return Err(invalid(
                    "topology fingerprint mismatch with the leader — this worker's steps/inputs/vnodes differ; \
                     every worker must run the same pipeline.toml and STREAM_VNODES",
                ));
            }
            eprintln!("fv-streams: joined — worker {} of {}", m.me(), m.len());
            (m, Coord::Worker(ctl))
        }
        None => {
            let control = TcpListener::bind(&cfg.control_bind).await?;
            let addr = control.local_addr()?;
            eprintln!("fv-streams: leader ready on {addr} (exchange {my_addr}).");
            eprintln!("fv-streams: workers join with:  fv-streams run <pipeline.toml> --join {addr}");
            eprintln!("fv-streams: waiting for {} worker(s) to join…", cfg.n_workers - 1);
            let (m, ctl) = lead(&control, &my_addr, cfg.n_workers, base_fp).await?;
            eprintln!("fv-streams: all {} workers joined — starting.", m.len());
            (m, Coord::Leader(ctl))
        }
    };
    Ok((members, exchange, coord))
}

/// The worker side: connect to the leader's control address, register `my_exchange_addr`, and
/// receive the membership. Returns this worker's [`Members`] and the leader-declared fingerprint
/// (the caller checks it against its own local fingerprint — a mismatch is a topology disagreement).
pub async fn join(leader_control: SocketAddr, my_exchange_addr: &str) -> std::io::Result<(Members, String, WorkerCtl)> {
    join_as(leader_control, my_exchange_addr, fv_streams_exchange::WIRE_VERSION).await
}

/// `join` announcing wire version `wire` (the real one — or, in a test, a foreign one).
async fn join_as(
    leader_control: SocketAddr,
    my_exchange_addr: &str,
    wire: u8,
) -> std::io::Result<(Members, String, WorkerCtl)> {
    let mut s = TcpStream::connect(leader_control).await?;
    write_msg(
        &mut s,
        &Register {
            exchange_addr: my_exchange_addr.to_string(),
            wire,
        },
    )
    .await?;
    let assign: Assignment = read_msg(&mut s).await?;
    if assign.wire != wire {
        // `wire` is this joiner's own version (what it announced); a foreign leader is refused
        return Err(wire_mismatch("the leader", assign.wire, wire));
    }
    let members = Members::new(assign.workers, my_exchange_addr).map_err(invalid)?;
    Ok((members, assign.fingerprint, WorkerCtl { conn: s }))
}

/// Open the exchange mesh from a settled membership: one [`Link`] per peer — connect to lower ids,
/// accept from higher ids on `exchange` — returning the [`WorkerExchange`] `start_worker` consumes.
/// Accept and connect run concurrently so no worker deadlocks waiting on a peer that is waiting on
/// it.
pub async fn mesh(
    members: &Members,
    exchange: &TcpListener,
    handle: Handle,
    flows: usize,
    window: usize,
) -> std::io::Result<WorkerExchange> {
    let me = members.me();
    let higher: Vec<WorkerId> = members.peers().filter(|&p| p > me).collect();
    let lower: Vec<WorkerId> = members.peers().filter(|&p| p < me).collect();

    let accept_fut = Link::accept_mesh(exchange, &higher, flows, window);
    let connect_fut = async {
        let mut m: HashMap<WorkerId, Link> = HashMap::new();
        for p in &lower {
            let addr: SocketAddr = members.addr(*p).expect("peer address").parse().map_err(invalid)?;
            m.insert(*p, Link::connect_as(addr, me, flows, window).await?);
        }
        Ok::<_, std::io::Error>(m)
    };
    let (accepted, connected) = tokio::join!(accept_fut, connect_fut);
    let links = accepted?.into_iter().chain(connected?);

    let mut senders = HashMap::new();
    let mut receivers = Vec::new();
    for (peer, link) in links {
        let (s, r) = link.into_split();
        senders.insert(peer, s);
        receivers.push(r);
    }
    Ok(WorkerExchange {
        handle,
        senders,
        receivers,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use fv_streams_exchange::FrameKind;

    fn rt() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .enable_all()
            .build()
            .unwrap()
    }

    #[test]
    fn cluster_cfg_parse_is_solo_unless_joined_or_multi_worker() {
        // No join, one worker → solo (the pre-Phase-4 path).
        assert_eq!(ClusterCfg::parse("", 1, "", "", 4, 1 << 20), None);
        assert_eq!(ClusterCfg::parse("", 0, "", "", 4, 1 << 20), None);
        // A joiner is clustered even with the default worker count.
        let w = ClusterCfg::parse("10.0.0.1:7100", 1, "", "", 4, 1 << 20).unwrap();
        assert_eq!(w.join.as_deref(), Some("10.0.0.1:7100"));
        assert_eq!(w.n_workers, 2, "a joiner implies at least two workers");
        assert_eq!(w.exchange_bind, "127.0.0.1:0", "ephemeral exchange bind by default");
        // A leader is clustered when it expects more than one worker.
        let l = ClusterCfg::parse("", 3, "1.2.3.4:9000", "0.0.0.0:7100", 8, 1 << 22).unwrap();
        assert_eq!(l.join, None);
        assert_eq!((l.n_workers, l.flows, l.window), (3, 8, 1 << 22));
        assert_eq!(l.exchange_bind, "1.2.3.4:9000");
    }

    #[test]
    fn a_worker_on_another_wire_version_is_refused_at_join_by_both_sides() {
        // The fence: a joiner announcing a foreign wire version is refused by the leader, and the
        // joiner refuses the leader's reply — both naming the versions and the remedy, before any
        // task runs (the topology fingerprint fence's shape).
        let rt = rt();
        rt.block_on(async {
            let lx = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let leader_x = lx.local_addr().unwrap().to_string();
            let control = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let caddr = control.local_addr().unwrap();
            let wx = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let worker_x = wx.local_addr().unwrap().to_string();
            let leader = tokio::spawn({
                let leader_x = leader_x.clone();
                async move { lead(&control, &leader_x, 2, "fp-1").await }
            });
            let foreign = fv_streams_exchange::WIRE_VERSION + 1;
            let worker_err = join_as(caddr, &worker_x, foreign)
                .await
                .err()
                .expect("the joiner refuses");
            let leader_err = leader.await.unwrap().err().expect("the leader refuses");
            for (who, e) in [("joiner", worker_err.to_string()), ("leader", leader_err.to_string())] {
                assert!(
                    e.contains("wire format v") && e.contains("same fv-streams version"),
                    "{who}: {e}"
                );
            }
            assert!(leader_err.to_string().contains(&format!("v{foreign}")), "{leader_err}");
        });
    }

    #[test]
    fn rendezvous_settles_the_same_membership_on_every_worker() {
        let rt = rt();
        rt.block_on(async {
            // Each worker binds an exchange listener to get an address; the leader also binds a
            // control listener the workers join on.
            let lx = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let leader_x = lx.local_addr().unwrap().to_string();
            let control = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let caddr = control.local_addr().unwrap();
            let wx = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let worker_x = wx.local_addr().unwrap().to_string();

            let leader = tokio::spawn({
                let leader_x = leader_x.clone();
                async move { lead(&control, &leader_x, 2, "fp-1").await.unwrap() }
            });
            let (worker_members, fp, mut worker_ctl) = join(caddr, &worker_x).await.unwrap();
            let (leader_members, mut leader_ctl) = leader.await.unwrap();

            // Both settle on the same sorted set of two workers and agree on the fingerprint.
            assert_eq!(leader_members.len(), 2);
            assert_eq!(worker_members.len(), 2);
            assert_eq!(fp, "fp-1");
            assert_ne!(leader_members.me(), worker_members.me(), "distinct worker ids");
            assert_eq!(leader_members.fingerprint_clause(), worker_members.fingerprint_clause());

            // The control connections stay live: a runtime round-trip over them — leader drives an
            // epoch + commit, the worker acks and reports finished.
            leader_ctl.broadcast(ToWorker::Barrier(7)).await.unwrap();
            assert_eq!(worker_ctl.recv().await.unwrap(), ToWorker::Barrier(7));
            let contribution = fv_streams_state::Manifest {
                v: 1,
                epoch: 7,
                fingerprint: "fp-1".into(),
                sources: Default::default(),
                objects: vec!["s0-op3".into()],
                at_ms: 0,
                files: Default::default(),
            };
            worker_ctl
                .report(ToLeader::Acked {
                    epoch: 7,
                    contribution: contribution.clone(),
                })
                .await
                .unwrap();
            assert_eq!(
                leader_ctl.collect().await.unwrap(),
                vec![(worker_members.me(), ToLeader::Acked { epoch: 7, contribution })]
            );
            leader_ctl.broadcast(ToWorker::Commit(7)).await.unwrap();
            assert_eq!(worker_ctl.recv().await.unwrap(), ToWorker::Commit(7));
            leader_ctl.broadcast(ToWorker::Stop).await.unwrap();
            assert_eq!(worker_ctl.recv().await.unwrap(), ToWorker::Stop);
            worker_ctl.report(ToLeader::Finished).await.unwrap();
            assert_eq!(
                leader_ctl.collect().await.unwrap(),
                vec![(worker_members.me(), ToLeader::Finished)]
            );
        });
    }

    #[test]
    fn a_three_worker_mesh_links_every_pair_and_carries_frames() {
        let rt = rt();
        let handle = rt.handle().clone();
        rt.block_on(async {
            // Three workers, addresses fixed so the sorted order (hence ids) is deterministic.
            let listeners: Vec<TcpListener> = {
                let mut v = Vec::new();
                for _ in 0..3 {
                    v.push(TcpListener::bind("127.0.0.1:0").await.unwrap());
                }
                v
            };
            let addrs: Vec<String> = listeners.iter().map(|l| l.local_addr().unwrap().to_string()).collect();
            let members: Vec<Members> = addrs.iter().map(|a| Members::new(addrs.clone(), a).unwrap()).collect();

            // Every worker opens its mesh concurrently.
            let (flows, window) = (2usize, 1usize << 20);
            let mut tasks = Vec::new();
            for (m, l) in members.into_iter().zip(listeners) {
                let h = handle.clone();
                tasks.push(tokio::spawn(async move {
                    let ex = mesh(&m, &l, h, flows, window).await.unwrap();
                    (m.me(), ex)
                }));
            }
            let mut exchanges: HashMap<WorkerId, WorkerExchange> = HashMap::new();
            for t in tasks {
                let (me, ex) = t.await.unwrap();
                // Every worker links to the other two.
                assert_eq!(ex.senders.len(), 2, "worker {me} has a sender to each peer");
                assert_eq!(ex.receivers.len(), 2, "worker {me} has a receiver from each peer");
                exchanges.insert(me, ex);
            }

            // Worker 0 sends a frame to worker 2; worker 2 receives it on one of its links.
            let ex0 = exchanges.get(&0).unwrap();
            ex0.senders[&2]
                .send(5, FrameKind::Data, bytes::Bytes::from_static(b"hi-2"))
                .await
                .unwrap();
            let ex2 = exchanges.get_mut(&2).unwrap();
            let mut saw = None;
            for r in ex2.receivers.iter_mut() {
                if let Ok(Some(f)) = tokio::time::timeout(std::time::Duration::from_secs(2), r.recv()).await {
                    saw = Some(f);
                    break;
                }
            }
            let f = saw.expect("worker 2 received worker 0's frame");
            assert_eq!((f.channel, &f.body[..]), (5, &b"hi-2"[..]));
        });
    }
}