diff --git a/Cargo.lock b/Cargo.lock index f10766d77..f30938e7d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -896,6 +896,7 @@ version = "0.1.0" dependencies = [ "antithesis_sdk", "arc-swap", + "async-trait", "axum", "axum-extra", "backoff", @@ -926,6 +927,7 @@ dependencies = [ "opentelemetry", "parking_lot", "pin-project-lite", + "plum-foca", "quinn", "quinn-plaintext", "quinn-proto", @@ -1078,6 +1080,8 @@ dependencies = [ name = "corro-tests" version = "0.1.0" dependencies = [ + "async-trait", + "bytes", "corro-agent", "corro-client", "corro-types", @@ -1148,6 +1152,7 @@ dependencies = [ "opentelemetry", "papaya", "parking_lot", + "plum-foca", "rand 0.9.2", "rangemap", "rcgen", @@ -3687,6 +3692,20 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "plum-foca" +version = "0.1.0" +dependencies = [ + "indexmap 2.11.0", + "rand 0.9.2", + "serde", + "speedy", + "strum", + "thiserror 1.0.69", + "tracing", + "tracing-subscriber", +] + [[package]] name = "portable-atomic" version = "1.11.1" diff --git a/Cargo.toml b/Cargo.toml index 54e34b065..a55068e41 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] members = [ ".antithesis/client/rust-load-generator", "crates/*", - "integration-tests" + "integration-tests", ] resolver = "2" diff --git a/crates/corro-admin/src/lib.rs b/crates/corro-admin/src/lib.rs index 0758ba537..5526e9226 100644 --- a/crates/corro-admin/src/lib.rs +++ b/crates/corro-admin/src/lib.rs @@ -10,7 +10,8 @@ use corro_types::{ actor::{ActorId, ClusterId}, agent::{Agent, BookedVersions, Bookie, WriteConn}, base::{CrsqlDbVersion, CrsqlSeq}, - broadcast::{FocaCmd, FocaInput}, + broadcast::{FocaCmd, FocaInput, PlumtreeInput}, + config::BroadcastMethod, sqlite::SqlitePoolError, sync::generate_sync, updates::Handle, @@ -111,9 +112,15 @@ pub enum Command { Actor(ActorCommand), Subs(SubsCommand), Log(LogCommand), + Plumtree(PlumtreeCommand), Reload, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum PlumtreeCommand { + Stats, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum SyncCommand { Generate, @@ -627,6 +634,49 @@ async fn handle_conn( } } }, + Command::Plumtree(cmd) => { + if !matches!(agent.broadcast_method(), BroadcastMethod::Plumtree) { + send_error( + &mut stream, + "plumtree broadcast method is not enabled on this node", + ) + .await; + continue; + } + + match cmd { + PlumtreeCommand::Stats => { + let (tx, rx) = oneshot::channel(); + if let Err(e) = agent + .tx_plumtree() + .send(PlumtreeInput::QueryStats(tx)) + .await + { + send_error(&mut stream, e).await; + continue; + } + + let stats = match rx.await { + Ok(stats) => stats, + Err(e) => { + send_error(&mut stream, e).await; + continue; + } + }; + + let value = match serde_json::to_value(&stats) { + Ok(json) => json, + Err(e) => { + send_error(&mut stream, e).await; + continue; + } + }; + + send(&mut stream, Response::Json(value)).await; + send_success(&mut stream).await; + } + } + } }, Ok(None) => { debug!("done with admin conn"); diff --git a/crates/corro-agent/Cargo.toml b/crates/corro-agent/Cargo.toml index 6a215249f..73675bbe8 100644 --- a/crates/corro-agent/Cargo.toml +++ b/crates/corro-agent/Cargo.toml @@ -10,6 +10,7 @@ repository.workspace = true workspace = true [dependencies] +async-trait = { workspace = true } antithesis_sdk = { workspace = true } arc-swap = { workspace = true } axum = { workspace = true } @@ -37,6 +38,7 @@ metrics = { workspace = true } opentelemetry = { workspace = true } parking_lot = { workspace = true } pin-project-lite.workspace = true +plum-foca = { path = "../plum-foca" } quinn = { workspace = true } quinn-proto = { workspace = true } quinn-plaintext = { workspace = true } diff --git a/crates/corro-agent/src/agent/handlers.rs b/crates/corro-agent/src/agent/handlers.rs index c032cfbb5..2cdb361e1 100644 --- a/crates/corro-agent/src/agent/handlers.rs +++ b/crates/corro-agent/src/agent/handlers.rs @@ -24,8 +24,9 @@ use corro_types::{ actor::{Actor, ActorId}, agent::{Agent, Bookie, SplitPool}, base::CrsqlSeq, - broadcast::{BroadcastInput, BroadcastV1, ChangeSource, ChangeV1, FocaInput}, + broadcast::{ChangeSource, ChangeV1, FocaInput, PlumtreeUpdates}, channel::CorroReceiver, + config::BroadcastMethod, members::MemberAddedResult, sqlite::log_slow_inflight_queries, sync::generate_sync, @@ -133,12 +134,7 @@ pub fn spawn_incoming_connection_handlers( // Spawn handler tasks for this connection spawn_foca_handler(&agent, &tripwire, &conn); - uni::spawn_unipayload_handler( - &tripwire, - &conn, - agent.cluster_id(), - agent.tx_changes().clone(), - ); + uni::spawn_unipayload_handler(&tripwire, &conn, agent.clone()); bi::spawn_bipayload_handler(&agent, &bookie, &tripwire, &conn); }); } @@ -338,7 +334,7 @@ pub async fn handle_notifications( info!("Member Up {actor:?} (result: {member_added_res:?})"); match member_added_res { - MemberAddedResult::NewMember | MemberAddedResult::Removed => { + MemberAddedResult::NewMember(_) | MemberAddedResult::Removed => { if matches!(member_added_res, MemberAddedResult::Removed) { debug!("Member Removed {actor:?} due to member id mismatch"); counter!( @@ -362,10 +358,13 @@ pub async fn handle_notifications( error!("could not send new foca cluster size: {e}"); } } + + send_plumtree_member_update(&agent, &actor, member_added_res).await; } - MemberAddedResult::Updated => { + MemberAddedResult::Updated(_) => { debug!("Member Updated {actor:?}"); // anything else to do here? + // TODO: do we want to send rtt updates to plumtree } MemberAddedResult::Ignored => { // TODO: it's unclear if this is needed or @@ -402,14 +401,26 @@ pub async fn handle_notifications( error!("could not send new foca cluster size: {e}"); } } + + send_plumtree_member_update(&agent, &actor, MemberAddedResult::Removed).await; } counter!("corro.swim.notification", "type" => "memberdown").increment(1); } OwnedNotification::Rename(a, b) => { - let mut lock = agent.members().write(); - let del_res = lock.remove_member(&a); - let add_res = lock.add_member(&b); - info!("Member Rename {a:?} to {b:?} (del_res: {del_res:?}, add_res: {add_res:?})"); + let (del_res, add_res) = { + let mut lock = agent.members().write(); + let del_res = lock.remove_member(&a); + let add_res = lock.add_member(&b); + info!( + "Member Rename {a:?} to {b:?} (del_res: {del_res:?}, add_res: {add_res:?})" + ); + (del_res, add_res) + }; + + if del_res { + send_plumtree_member_update(&agent, &a, MemberAddedResult::Removed).await; + } + send_plumtree_member_update(&agent, &b, add_res).await; } OwnedNotification::Active => { info!("Current node is considered ACTIVE"); @@ -432,6 +443,33 @@ pub async fn handle_notifications( } } +async fn send_plumtree_member_update(agent: &Agent, actor: &Actor, result: MemberAddedResult) { + if agent.broadcast_method() != BroadcastMethod::Plumtree { + return; + } + + let update = { + match result { + MemberAddedResult::Removed => Some(PlumtreeUpdates::MemberDown(actor.id())), + MemberAddedResult::NewMember(state) | MemberAddedResult::Updated(state) => { + Some(PlumtreeUpdates::MemberUp { + actor_id: actor.id(), + addr: state.addr, + ring: state.ring, + // rtt_ms: members.avg_rtt_ms(&actor.id()).unwrap_or(u64::MAX), + }) + } + MemberAddedResult::Ignored => None, + } + }; + + if let Some(update) = update { + if let Err(e) = agent.tx_plumtree_updates().send(update).await { + error!("could not forward plumtree update: {e}"); + } + } +} + /// We keep a write-ahead-log, which under write-pressure can grow to /// multiple gigabytes and needs periodic truncation. We don't want /// to schedule this task too often since it locks the whole DB. @@ -1074,18 +1112,11 @@ pub async fn handle_changes( matches!(src, ChangeSource::Sync), "Corrosion receives changes through sync" ); - // Rebroadcast changes received from broadcast + + // Rebroadcast changes received from broadcast (gossip method). if matches!(src, ChangeSource::Broadcast) && !change.is_empty() { assert_sometimes!(true, "Corrosion rebroadcasts changes"); - if let Err(_e) = - agent - .tx_bcast() - .try_send(BroadcastInput::Rebroadcast(BroadcastV1::Change( - change.clone(), - ))) - { - debug!("broadcasts are full or done!"); - } + agent.broadcaster().rebroadcast(&change); } // Handle the new change - queue it and potentially spawn a batch diff --git a/crates/corro-agent/src/agent/run_root.rs b/crates/corro-agent/src/agent/run_root.rs index 920fdef56..591cddd24 100644 --- a/crates/corro-agent/src/agent/run_root.rs +++ b/crates/corro-agent/src/agent/run_root.rs @@ -10,7 +10,8 @@ use crate::{ reaper::spawn_reaper, setup, util, AgentOptions, }, - broadcast::runtime_loop, + broadcast::{handle_broadcasts, runtime_loop, BroadcastOpts}, + plumtree::spawn_plumtree_loop, transport::Transport, }; @@ -18,7 +19,7 @@ use corro_types::{ agent::{Agent, Bookie}, base::CrsqlSeq, channel::bounded, - config::{Config, PerfConfig}, + config::{BroadcastMethod, Config, PerfConfig}, }; use futures::FutureExt; @@ -58,6 +59,8 @@ async fn run( rx_clear_buf, rx_changes, rx_foca, + rx_plumtree, + rx_plumtree_updates, subs_manager, subs_bcast_cache, updates_bcast_cache, @@ -84,20 +87,22 @@ async fn run( let (notifications_tx, notifications_rx) = bounded(pconf.notifications_channel_len, "notifications"); - let member_states = util::load_member_states(&agent).await; + let loaded_member_states = util::load_member_states(&agent).await; + let member_states: Vec<_> = loaded_member_states + .iter() + .map(|(address, member, _)| (*address, member.clone())) + .collect(); //// Start the main SWIM runtime loop - runtime_loop( + let foca_config = runtime_loop( // here the agent already has the current cluster_id, we don't need to pass one agent.actor(None, agent.config().gossip.member_id), agent.clone(), - transport.clone(), rx_foca, - rx_bcast, to_send_tx, notifications_tx, tripwire.clone(), - member_states.clone(), + member_states, ); //// Update member connection RTTs @@ -106,7 +111,29 @@ async fn run( handlers::spawn_swim_announcer(&agent, gossip_addr, tripwire.clone()); // Load existing cluster members into the SWIM runtime - util::initialise_foca(&agent, member_states).await; + util::initialise_foca(&agent, loaded_member_states).await; + + match agent.broadcast_method() { + BroadcastMethod::Gossip => spawn_counted(handle_broadcasts( + agent.clone(), + rx_bcast, + transport.clone(), + foca_config, + tripwire.clone(), + BroadcastOpts::default(), + )), + BroadcastMethod::Plumtree => spawn_counted( + spawn_plumtree_loop( + agent.clone(), + transport.clone(), + rx_plumtree, + rx_plumtree_updates, + agent.tx_changes().clone(), + tripwire.clone(), + ) + .inspect(|_| info!("plumtree loop is done")), + ), + }; // Load schema from paths if let Err(e) = execute_schema_from_paths(&agent).await { diff --git a/crates/corro-agent/src/agent/setup.rs b/crates/corro-agent/src/agent/setup.rs index 30ffdbdfc..1db590272 100644 --- a/crates/corro-agent/src/agent/setup.rs +++ b/crates/corro-agent/src/agent/setup.rs @@ -33,7 +33,9 @@ use corro_types::{ actor::ActorId, agent::{migrate, Agent, AgentConfig, BookedVersions, Bookie, SplitPool}, base::{CrsqlDbVersion, CrsqlDbVersionRange}, - broadcast::{BroadcastInput, ChangeSource, ChangeV1, FocaInput}, + broadcast::{ + BroadcastInput, ChangeSource, ChangeV1, FocaInput, PlumtreeInput, PlumtreeUpdates, + }, channel::{bounded, CorroReceiver}, config::Config, members::Members, @@ -54,6 +56,9 @@ pub struct AgentOptions { pub rx_clear_buf: CorroReceiver<(ActorId, CrsqlDbVersionRange)>, pub rx_changes: CorroReceiver<(ChangeV1, ChangeSource)>, pub rx_foca: CorroReceiver, + // pub tx_plumtree: CorroSender, + pub rx_plumtree: CorroReceiver, + pub rx_plumtree_updates: CorroReceiver, pub rtt_rx: TokioReceiver<(SocketAddr, Duration)>, pub subs_manager: SubsManager, pub subs_bcast_cache: SharedMatcherBroadcastCache, @@ -156,6 +161,9 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age let (tx_bcast, rx_bcast) = bounded(conf.perf.bcast_channel_len, "bcast"); let (tx_changes, rx_changes) = bounded(conf.perf.changes_channel_len, "changes"); let (tx_foca, rx_foca) = bounded(conf.perf.foca_channel_len, "foca"); + let (tx_plumtree, rx_plumtree) = bounded(conf.perf.bcast_channel_len, "plumtree"); + let (tx_plumtree_updates, rx_plumtree_updates) = + bounded(conf.perf.foca_channel_len, "plumtree_updates"); // Load all actors' bookie state synchronously. let start = Instant::now(); @@ -179,6 +187,8 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age rx_clear_buf, rx_changes, rx_foca, + rx_plumtree, + rx_plumtree_updates, rtt_rx, subs_manager: subs_manager.clone(), subs_bcast_cache, @@ -202,6 +212,8 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age tx_clear_buf, tx_changes, tx_foca, + tx_plumtree, + tx_plumtree_updates, write_sema, schema: RwLock::new(schema), cluster_id, diff --git a/crates/corro-agent/src/agent/tests.rs b/crates/corro-agent/src/agent/tests.rs index 1d123977c..e3db68d0e 100644 --- a/crates/corro-agent/src/agent/tests.rs +++ b/crates/corro-agent/src/agent/tests.rs @@ -42,6 +42,7 @@ use corro_types::{ agent::Agent, api::{ColumnName, TableName}, change::row_to_change, + config::BroadcastMethod, pubsub::pack_columns, }; @@ -126,15 +127,18 @@ async fn http_api_requested_endpoint_name_header_enforced() -> eyre::Result<()> Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] - -async fn insert_rows_and_gossip() -> eyre::Result<()> { +async fn insert_rows_and_broadcast(method: BroadcastMethod) -> eyre::Result<()> { _ = tracing_subscriber::fmt::try_init(); let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple(); - let ta1 = launch_test_agent(|conf| conf.build(), tripwire.clone()).await?; + let ta1 = launch_test_agent( + |conf| conf.broadcast_method(method).build(), + tripwire.clone(), + ) + .await?; let ta2 = launch_test_agent( |conf| { - conf.bootstrap(vec![ta1.agent.gossip_addr().to_string()]) + conf.broadcast_method(method) + .bootstrap(vec![ta1.agent.gossip_addr().to_string()]) .build() }, tripwire.clone(), @@ -329,20 +333,41 @@ async fn insert_rows_and_gossip() -> eyre::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn insert_rows_and_gossip() -> eyre::Result<()> { + insert_rows_and_broadcast(BroadcastMethod::Gossip).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn insert_rows_and_plumtree() -> eyre::Result<()> { + insert_rows_and_broadcast(BroadcastMethod::Plumtree).await +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn chill_test() -> eyre::Result<()> { - configurable_stress_test(2, 1, 4).await + configurable_stress_test(2, 1, 4, BroadcastMethod::Gossip).await } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn stress_test() -> eyre::Result<()> { - configurable_stress_test(30, 10, 200).await + configurable_stress_test(30, 10, 200, BroadcastMethod::Gossip).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn stress_plumtree_test() -> eyre::Result<()> { + configurable_stress_test(30, 10, 200, BroadcastMethod::Plumtree).await } #[ignore] #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn stresser_test() -> eyre::Result<()> { - configurable_stress_test(45, 15, 1500).await + configurable_stress_test(45, 15, 1500, BroadcastMethod::Gossip).await +} + +#[ignore] +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn stresser_plumtree_test() -> eyre::Result<()> { + configurable_stress_test(45, 15, 1500, BroadcastMethod::Plumtree).await } /// Default parameters: @@ -357,6 +382,7 @@ pub async fn configurable_stress_test( num_nodes: usize, connectivity: usize, input_count: usize, + method: BroadcastMethod, ) -> eyre::Result<()> { _ = tracing_subscriber::fmt::try_init(); let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple(); @@ -381,6 +407,7 @@ pub async fn configurable_stress_test( launch_test_agent( |conf| { conf.gossip_addr(gossip_addr) + .broadcast_method(method) .bootstrap( bootstrap .iter() diff --git a/crates/corro-agent/src/agent/uni.rs b/crates/corro-agent/src/agent/uni.rs index 369925a4f..06e4b5857 100644 --- a/crates/corro-agent/src/agent/uni.rs +++ b/crates/corro-agent/src/agent/uni.rs @@ -1,23 +1,27 @@ use corro_types::{ - actor::ClusterId, - broadcast::{BroadcastV1, ChangeSource, ChangeV1, UniPayload, UniPayloadV1}, - channel::CorroSender, + agent::Agent, + broadcast::{ + BroadcastV1, ChangeSource, ChangeV1, PlumtreeInput, PlumtreeMsgV1, PlumtreeWire, + UniPayload, UniPayloadV1, + }, + config::BroadcastMethod, }; use metrics::counter; +use plum_foca::GossipMsg; use speedy::Readable; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LengthDelimitedCodec}; -use tracing::{debug, error, trace}; +use tracing::{debug, error, trace, warn}; use tripwire::Tripwire; /// Spawn a task that accepts unidirectional broadcast streams, then /// spawns another task for each incoming stream to handle. -pub fn spawn_unipayload_handler( - tripwire: &Tripwire, - conn: &quinn::Connection, - cluster_id: ClusterId, - tx_changes: CorroSender<(ChangeV1, ChangeSource)>, -) { +pub fn spawn_unipayload_handler(tripwire: &Tripwire, conn: &quinn::Connection, agent: Agent) { + let cluster_id = agent.cluster_id(); + let broadcast_method = agent.broadcast_method(); + let tx_changes = agent.tx_changes().clone(); + let tx_plumtree = agent.tx_plumtree().clone(); + tokio::spawn({ let conn = conn.clone(); let mut tripwire = tripwire.clone(); @@ -46,6 +50,8 @@ pub fn spawn_unipayload_handler( tokio::spawn({ let tx_changes = tx_changes.clone(); + let tx_plumtree = tx_plumtree.clone(); + let broadcast_method = broadcast_method; async move { let mut framed = FramedRead::new( rx, @@ -54,7 +60,7 @@ pub fn spawn_unipayload_handler( .new_codec(), ); - let mut changes = vec![]; + let mut changes: Vec = vec![]; loop { match StreamExt::next(&mut framed).await { Some(Ok(b)) => { @@ -75,7 +81,41 @@ pub fn spawn_unipayload_handler( if cluster_id != payload_cluster_id { continue; } - changes.push((change, ChangeSource::Broadcast)); + changes.push(change); + } + UniPayload::V1 { + data: + UniPayloadV1::Plumtree(PlumtreeWire::V1 { + data: wire_msg, + }), + cluster_id: payload_cluster_id, + } => { + if cluster_id != payload_cluster_id { + continue; + } + let tx = tx_plumtree.clone(); + + match broadcast_method { + // route gossips if we are using broadcast method + BroadcastMethod::Gossip => { + if let PlumtreeMsgV1::Gossip(msg) = + wire_msg + { + warn!("broadcast algorithm set to gossip but node receieved plumtree message"); + changes.push(msg.payload); + } + } + BroadcastMethod::Plumtree => { + if let Err(e) = tx + .send(PlumtreeInput::Wire(wire_msg)) + .await + { + error!( + "could not route Plumtree msg: {e}" + ); + } + } + } } } } @@ -92,10 +132,34 @@ pub fn spawn_unipayload_handler( } } - for change in changes.into_iter().rev() { - if let Err(e) = tx_changes.send(change).await { - error!("could not send change for processing: {e}"); - return; + match broadcast_method { + BroadcastMethod::Plumtree => { + warn!("broadcast algorithm set to plumtree but node receieved gossip message."); + for change in changes.into_iter().rev() { + let wire = PlumtreeMsgV1::Gossip(GossipMsg { + round: 1, + sender: change.actor_id, + payload: change, + }); + if let Err(e) = + tx_plumtree.send(PlumtreeInput::Wire(wire)).await + { + error!( + "could not route legacy gossip change to plumtree: {e}" + ); + return; + } + } + } + BroadcastMethod::Gossip => { + for change in changes.into_iter().rev() { + if let Err(e) = + tx_changes.send((change, ChangeSource::Broadcast)).await + { + error!("could not send change for processing: {e}"); + return; + } + } } } } diff --git a/crates/corro-agent/src/agent/util.rs b/crates/corro-agent/src/agent/util.rs index b41a986d7..6c018ad80 100644 --- a/crates/corro-agent/src/agent/util.rs +++ b/crates/corro-agent/src/agent/util.rs @@ -74,15 +74,22 @@ use tripwire::{PreemptibleFutureExt, Tripwire}; const REQUESTED_ENDPOINT_NAME_HEADER: &str = "x-corrosion-requested-endpoint-name"; -pub async fn initialise_foca(agent: &Agent, states: Vec<(SocketAddr, Member)>) { +pub async fn initialise_foca(agent: &Agent, states: Vec<(SocketAddr, Member, Option)>) { if !states.is_empty() { let mut foca_states = BTreeMap::>::new(); { // block to drop the members write lock let mut members = agent.members().write(); - for (address, foca_state) in states { - members.by_addr.insert(address, foca_state.id().id()); + for (address, foca_state, rtt_min) in states { + if matches!(foca_state.state(), foca::State::Alive) { + let actor = foca_state.id(); + members.add_member(actor); + if let Some(rtt_ms) = rtt_min { + members.add_rtt(address, Duration::from_millis(rtt_ms)); + } + } + if matches!(foca_state.state(), foca::State::Suspect) { continue; } @@ -139,28 +146,29 @@ pub async fn initialise_foca(agent: &Agent, states: Vec<(SocketAddr, Member Vec<(SocketAddr, Member)> { +/// Load the existing known member state, addresses, and persisted RTT samples. +pub async fn load_member_states(agent: &Agent) -> Vec<(SocketAddr, Member, Option)> { match agent.pool().read().await { Ok(conn) => block_in_place(|| { - match conn.prepare("SELECT address,foca_state FROM __corro_members") { + match conn.prepare("SELECT address, foca_state, rtt_min FROM __corro_members") { Ok(mut prepped) => { match prepped .query_map([], |row| Ok(( row.get::<_, String>(0)?.parse().map_err(|e| rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(e)))?, - row.get::<_, String>(1)? + row.get::<_, String>(1)?, + row.get::<_, Option>(2)?, )) ) - .and_then(|rows| rows.collect::>>()) + .and_then(|rows| rows.collect::)>>>()) { Ok(members) => { - members.into_iter().filter_map(|(address, state)| match serde_json::from_str::>(state.as_str()) { - Ok(fs) => Some((address, fs)), + members.into_iter().filter_map(|(address, state, rtt_min)| match serde_json::from_str::>(state.as_str()) { + Ok(fs) => Some((address, fs, rtt_min.map(|rtt| rtt as u64))), Err(e) => { error!("could not deserialize foca member state: {e} (json: {state})"); None } - }).collect::)>>() + }).collect::, Option)>>() } Err(e) => { error!("could not query for foca member states: {e}"); diff --git a/crates/corro-agent/src/api/peer/mod.rs b/crates/corro-agent/src/api/peer/mod.rs index 9951dff3a..a2326e0f7 100644 --- a/crates/corro-agent/src/api/peer/mod.rs +++ b/crates/corro-agent/src/api/peer/mod.rs @@ -1675,7 +1675,7 @@ mod tests { use corro_types::{ api::{ColumnName, TableName}, broadcast::ChangesetPerTable, - config::{Config, TlsClientConfig, TlsConfig, DEFAULT_GOSSIP_CLIENT_ADDR}, + config::{BroadcastConfig, Config, TlsClientConfig, TlsConfig, DEFAULT_GOSSIP_CLIENT_ADDR}, pubsub::pack_columns, tls::{generate_ca, generate_client_cert, generate_server_cert}, }; @@ -2421,6 +2421,7 @@ mod tests { max_mtu: None, disable_gso: false, member_id: None, + broadcast: BroadcastConfig::Gossip, }; let server = gossip_server_endpoint(&gossip_config).await?; diff --git a/crates/corro-agent/src/api/public/mod.rs b/crates/corro-agent/src/api/public/mod.rs index 2f129fe6c..3c2ec054a 100644 --- a/crates/corro-agent/src/api/public/mod.rs +++ b/crates/corro-agent/src/api/public/mod.rs @@ -667,8 +667,8 @@ mod tests { use corro_types::{ api::RowId, base::CrsqlDbVersion, - broadcast::{BroadcastInput, BroadcastV1, ChangeV1, Changeset}, - config::Config, + broadcast::{ChangeV1, Changeset, PlumtreeInput}, + config::{BroadcastMethod, Config}, schema::SqliteType, }; use futures::StreamExt; @@ -693,12 +693,14 @@ mod tests { .db_path(dir.path().join("corrosion.db").display().to_string()) .gossip_addr("127.0.0.1:0".parse()?) .api_addr("127.0.0.1:0".parse()?) + .broadcast_method(BroadcastMethod::Plumtree) .build()?, tripwire, ) .await?; let rx_bcast = &mut agent_options.rx_bcast; + let rx_plumtree = &mut agent_options.rx_plumtree; execute_schema(&agent, vec![corro_tests::TEST_SCHEMA.to_owned()]).await?; @@ -718,20 +720,20 @@ mod tests { assert!(body.0.results.len() == 1); - let msg = rx_bcast + let msg = rx_plumtree .recv() .await .expect("not msg received on bcast channel"); assert!(matches!( msg, - BroadcastInput::AddBroadcast(BroadcastV1::Change(ChangeV1 { + PlumtreeInput::Broadcast(ChangeV1 { changeset: Changeset::FullV2 { version: CrsqlDbVersion(1), .. }, .. - })) + }) )); assert_eq!(agent.booked().read().last(), Some(CrsqlDbVersion(1))); diff --git a/crates/corro-agent/src/broadcast/mod.rs b/crates/corro-agent/src/broadcast/mod.rs index 85c683b73..5b1a29426 100644 --- a/crates/corro-agent/src/broadcast/mod.rs +++ b/crates/corro-agent/src/broadcast/mod.rs @@ -43,16 +43,16 @@ use corro_types::{ sqlite::unnest_param, }; -use crate::{agent::util::log_at_pow_10, transport::Transport}; +use crate::{agent::util::log_at_pow_10, transport::TransportExt}; #[derive(Clone)] -struct TimerSpawner { - send: mpsc::UnboundedSender<(Duration, Timer)>, +pub struct TimerSpawner { + send: mpsc::UnboundedSender<(Duration, T)>, } -impl TimerSpawner { - pub fn new(timer_tx: mpsc::Sender<(Timer, Instant)>) -> Self { - let (send, mut recv) = mpsc::unbounded_channel(); +impl TimerSpawner { + pub fn new(timer_tx: mpsc::Sender<(T, Instant)>) -> Self { + let (send, mut recv) = mpsc::unbounded_channel::<(Duration, T)>(); let rt = tokio::runtime::Builder::new_current_thread() .enable_all() @@ -86,7 +86,7 @@ impl TimerSpawner { Self { send } } - pub fn spawn(&self, task: (Duration, Timer)) { + pub fn spawn(&self, task: (Duration, T)) { self.send .send(task) .expect("Thread with LocalSet has shut down."); @@ -122,14 +122,12 @@ fn handle_timer( pub fn runtime_loop( actor: Actor, agent: Agent, - transport: Transport, mut rx_foca: CorroReceiver, - rx_bcast: CorroReceiver, to_send_tx: CorroSender<(Actor, Bytes)>, notifications_tx: CorroSender>, tripwire: Tripwire, member_states: Vec<(SocketAddr, Member)>, -) { +) -> Arc> { debug!("starting runtime loop for actor: {actor:?}"); let rng = StdRng::from_os_rng(); @@ -391,17 +389,10 @@ pub fn runtime_loop( } }); - spawn_counted(handle_broadcasts( - agent, - rx_bcast, - transport, - config, - tripwire, - Default::default(), - )); + config } -type BroadcastRateLimiter = RateLimiter< +pub(crate) type TransmitRateLimiter = RateLimiter< governor::state::NotKeyed, governor::state::InMemoryState, governor::clock::QuantaClock, @@ -423,10 +414,10 @@ impl Default for BroadcastOpts { } } -async fn handle_broadcasts( +pub(crate) async fn handle_broadcasts( agent: Agent, mut rx_bcast: CorroReceiver, - transport: Transport, + transport: impl TransportExt + Clone + Send + 'static, config: Arc>, mut tripwire: Tripwire, opts: BroadcastOpts, @@ -473,7 +464,7 @@ async fn handle_broadcasts( let mut limited_log_count = 0; - let bytes_per_sec: BroadcastRateLimiter = RateLimiter::direct(Quota::per_second(unsafe { + let bytes_per_sec: TransmitRateLimiter = RateLimiter::direct(Quota::per_second(unsafe { NonZeroU32::new_unchecked(10 * 1024 * 1024) })) .with_middleware(); @@ -623,12 +614,7 @@ async fn handle_broadcasts( ring0_count += 1; ring0.insert(addr); - match try_transmit_broadcast( - &bytes_per_sec, - payload.clone(), - transport.clone(), - addr, - ) { + match try_transmit_uni(&bytes_per_sec, payload.clone(), transport.clone(), addr) { Err(e) => { log_at_pow_10( "could not spawn broadcast transmission: {e}", @@ -732,7 +718,7 @@ async fn handle_broadcasts( let mut spawn_count = 0; trace!("broadcasting to: {:?}", broadcast_to); for addr in broadcast_to { - match try_transmit_broadcast( + match try_transmit_uni( &bytes_per_sec, pending.payload.clone(), transport.clone(), @@ -1063,7 +1049,7 @@ impl PendingBroadcast { } #[derive(Debug, thiserror::Error)] -enum TransmitError { +pub(crate) enum TransmitError { #[error("payload > u32::MAX: {0}")] TooBig(usize), #[error(transparent)] @@ -1073,10 +1059,10 @@ enum TransmitError { } #[tracing::instrument(skip(payload, transport), fields(buf_size = payload.len()), level = "debug")] -fn try_transmit_broadcast( - bytes_per_sec: &BroadcastRateLimiter, +pub(crate) fn try_transmit_uni( + bytes_per_sec: &TransmitRateLimiter, payload: Bytes, - transport: Transport, + transport: impl TransportExt + Send + 'static, addr: SocketAddr, ) -> Result + Send>>, TransmitError> { trace!("singly broadcasting to {addr}"); @@ -1119,12 +1105,15 @@ fn try_transmit_broadcast( #[cfg(test)] mod tests { use super::*; - use crate::agent::spawn_unipayload_handler; - use corro_tests::launch_test_agent; + use crate::agent::{setup, spawn_unipayload_handler}; + use crate::transport::Transport; + use corro_tests::test_config; use corro_types::{ base::{dbsr, CrsqlDbVersion, CrsqlSeq}, broadcast::{BroadcastV1, ChangeV1, Changeset}, }; + // use plum_foca::PlumtreeMsg; + use std::collections::HashSet; use uuid::Uuid; #[test] @@ -1183,13 +1172,15 @@ mod tests { .with_max_level(tracing::Level::DEBUG) .try_init(); let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple(); - let ta1 = launch_test_agent(|conf| conf.build(), tripwire.clone()).await?; + let (_tmpdir, conf) = test_config(|c| c.build())?; + let (ta1_agent, opts) = setup(conf, tripwire.clone()).await?; + let mut rx_changes = opts.rx_changes; let (tx_bcast, rx_bcast) = bounded(100, "bcast"); let (tx_rtt, _) = mpsc::channel(100); let config = Arc::new(RwLock::new(make_foca_config(1.try_into().unwrap(), None))); - let transport = Transport::new(&ta1.config.gossip, tx_rtt).await?; + let transport = Transport::new(&ta1_agent.config().gossip, tx_rtt).await?; let server_config = quinn_plaintext::server_config(); let endpoint = quinn::Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap())?; @@ -1200,13 +1191,13 @@ mod tests { ActorId(Uuid::new_v4()), ta2_gossip_addr, Default::default(), - ta1.agent.cluster_id(), + ta1_agent.cluster_id(), None, ); - ta1.agent.members().write().add_member(&ta2_actor); + ta1_agent.members().write().add_member(&ta2_actor); let bcast = BroadcastV1::Change(ChangeV1 { - actor_id: ta1.agent.actor_id(), + actor_id: ta1_agent.actor_id(), changeset: Changeset::Full { version: CrsqlDbVersion(0), changes: vec![], @@ -1218,13 +1209,13 @@ mod tests { let mut ser_buf = BytesMut::new(); UniPayload::V1 { data: UniPayloadV1::Broadcast(bcast), - cluster_id: ta1.agent.cluster_id(), + cluster_id: ta1_agent.cluster_id(), } .write_to_stream((&mut ser_buf).writer())?; let estimated_size = ser_buf.len(); tokio::spawn(handle_broadcasts( - ta1.agent.clone(), + ta1_agent.clone(), rx_bcast, transport, config, @@ -1235,7 +1226,7 @@ mod tests { }, )); - let actor_id = ta1.agent.actor_id(); + let actor_id = ta1_agent.actor_id(); for i in 0..5 { tx_bcast .send(BroadcastInput::Rebroadcast(BroadcastV1::Change(ChangeV1 { @@ -1255,8 +1246,7 @@ mod tests { info!("accepting connection"); let conn = conn.await.unwrap(); - let (tx_changes, mut rx_changes) = bounded(100, "changes"); - spawn_unipayload_handler(&tripwire, &conn, ta1.agent.cluster_id(), tx_changes); + spawn_unipayload_handler(&tripwire, &conn, ta1_agent.clone()); // we should receive five items starting from the biggest version for i in (0..5).rev() { diff --git a/crates/corro-agent/src/lib.rs b/crates/corro-agent/src/lib.rs index 70a778959..7e5ef2d3d 100644 --- a/crates/corro-agent/src/lib.rs +++ b/crates/corro-agent/src/lib.rs @@ -1,4 +1,5 @@ pub mod agent; pub mod api; pub mod broadcast; +pub mod plumtree; pub mod transport; diff --git a/crates/corro-agent/src/plumtree/mod.rs b/crates/corro-agent/src/plumtree/mod.rs new file mode 100644 index 000000000..26d7a3f5a --- /dev/null +++ b/crates/corro-agent/src/plumtree/mod.rs @@ -0,0 +1,1151 @@ +use std::{ + collections::{HashMap, VecDeque}, + net::SocketAddr, + num::NonZeroU32, + ops::RangeInclusive, + time::{Duration, Instant}, +}; + +use bytes::{BufMut, Bytes, BytesMut}; +use corro_types::{ + actor::{ActorId, ClusterId}, + agent::{Agent, Bookie}, + base::{CrsqlDbVersion, CrsqlSeq}, + broadcast::{ + ChangeId, ChangeSource, ChangeV1, ChangesetId, PlumtreeInput, PlumtreeMsgV1, PlumtreeStats, + PlumtreeUpdates, PlumtreeWire, UniPayload, UniPayloadV1, + }, + channel::{bounded, CorroReceiver, CorroSender}, +}; +use governor::{Quota, RateLimiter}; +use indexmap::IndexMap; +use metrics::{counter, gauge}; +use plum_foca::{Payload, PlumPrio, PlumtreeState, RttInfo, SeenStore, Timer}; +use rangemap::RangeInclusiveSet; +use speedy::Writable; +use tokio::{sync::mpsc, task::JoinSet, time::interval}; +use tokio_util::codec::{Encoder, LengthDelimitedCodec}; +use tracing::{debug, error, info, trace, warn}; +use tripwire::Tripwire; + +use crate::{ + agent::util::log_at_pow_10, + broadcast::{try_transmit_uni, TimerSpawner, TransmitError, TransmitRateLimiter}, + transport::TransportExt, +}; + +#[derive(Debug)] +struct SeenEntry { + seqs: Option>, + last_seq: Option, + #[allow(unused)] + round: plum_foca::Round, + duplicate_count: u32, +} + +struct ChangeSeenStore { + entries: IndexMap<(ActorId, CrsqlDbVersion), SeenEntry>, + max_entries: usize, + bookie: Bookie, +} + +impl ChangeSeenStore { + fn new(max_entries: usize, bookie: Bookie) -> Self { + Self { + entries: IndexMap::new(), + max_entries, + bookie, + } + } +} + +impl SeenStore for ChangeSeenStore { + fn evict_if_needed(&mut self) { + if self.entries.len() > self.max_entries { + self.entries.drain(0..self.entries.len() - self.max_entries); + } + } + + fn contains(&self, id: &ChangeId) -> bool { + self.contains_local(id) || self.contains_booked(id) + } + + fn size(&self) -> usize { + self.entries.len() + } + + fn observe(&mut self, id: ChangeId, round: plum_foca::Round) -> Option { + if self.contains_booked(&id) { + // we already received this change through a sync or dropped it during a prune + counter!("corro.plumtree.change.synced").increment(1); + return Some(0); + } + let actor_id = id.actor_id; + match &id.changeset_id { + ChangesetId::Full { + version, + seqs, + last_seq, + } => match self.entries.entry((actor_id, *version)) { + indexmap::map::Entry::Vacant(e) => { + let incoming = RangeInclusiveSet::from_iter([RangeInclusive::from(*seqs)]); + e.insert(SeenEntry { + seqs: Some(incoming), + last_seq: Some(*last_seq), + round, + duplicate_count: 0, + }); + + None + } + indexmap::map::Entry::Occupied(mut e) => { + let entry = e.get_mut(); + + if entry.seqs.is_none() { + entry.duplicate_count += 1; + return Some(entry.duplicate_count); + } + + let incoming = RangeInclusive::from(*seqs); + let stored = entry.seqs.as_mut().unwrap(); + + if stored.gaps(&incoming).next().is_none() { + // we already seen this seqs, increase duplicate count if change is complete + let full = CrsqlSeq(0)..=*last_seq; + let is_complete = stored.gaps(&full).next().is_none(); + + // if we seen this partial change, but we are still incomplete, return zero duplicates. + // We don't want partials to cause a prune while we are in the middle of receiving the change. + if !is_complete { + return Some(0); + } + + entry.duplicate_count += 1; + return Some(entry.duplicate_count); + } + + stored.insert(incoming); + None + } + }, + // Empty changesets are actually not used rn + ChangesetId::Empty { versions } => { + let min_seen = versions + .map(|version| { + if let Some(entry) = self.entries.get_mut(&(actor_id, version)) { + entry.last_seq = None; + entry.duplicate_count += 1; + return entry.duplicate_count; + } + + self.entries.insert( + (actor_id, version), + SeenEntry { + seqs: None, + last_seq: None, + round, + duplicate_count: 0, + }, + ); + + 0 + }) + .min() + .unwrap_or(0); + + // there's at least one version are seeing for the first time + if min_seen == 0 { + None + } else { + Some(min_seen) + } + } + } + } +} + +impl ChangeSeenStore { + fn contains_booked(&self, id: &ChangeId) -> bool { + let actor_id = id.actor_id; + if let Some(bookie) = self.bookie.get(&actor_id) { + match &id.changeset_id { + ChangesetId::Full { version, seqs, .. } => { + bookie.read().contains(*version, Some(*seqs)) + } + ChangesetId::Empty { versions, .. } => versions + .clone() + .all(|version| bookie.read().contains(version, None)), + } + } else { + false + } + } + + fn contains_local(&self, id: &ChangeId) -> bool { + let actor_id = id.actor_id; + match &id.changeset_id { + ChangesetId::Full { version, seqs, .. } => { + let entry = self.entries.get(&(actor_id, *version)); + let Some(entry) = entry else { + return false; + }; + entry + .seqs + .as_ref() + .map(|old_seqs| { + let incoming = RangeInclusive::from(*seqs); + old_seqs.gaps(&incoming).count() == 0 + }) + .unwrap_or(false) + } + ChangesetId::Empty { versions, .. } => versions.clone().all(|version| { + self.entries + .get(&(actor_id, version)) + .map(|e| e.seqs.is_none()) + .unwrap_or(false) + }), + } + } +} + +/// Implements `plum_foca::Runtime` for Corrosion, bridging the generic protocol +/// to Corrosion's transport, change processing, and timer infrastructure. +struct CorrosionPlumtreeRuntime { + tx_changes: CorroSender<(ChangeV1, ChangeSource)>, + timer_spawner: TimerSpawner>, + tx_msgs: CorroSender<(PlumPrio, Vec, PlumtreeMsgV1)>, +} + +impl CorrosionPlumtreeRuntime { + fn new( + tx_changes: CorroSender<(ChangeV1, ChangeSource)>, + timer_spawner: TimerSpawner>, + tx_msgs: CorroSender<(PlumPrio, Vec, PlumtreeMsgV1)>, + ) -> Self { + Self { + tx_changes, + timer_spawner, + tx_msgs, + } + } +} + +impl plum_foca::Runtime for CorrosionPlumtreeRuntime { + fn send_all( + &mut self, + peers: Vec, + msg: plum_foca::PlumtreeMsg, + priority: PlumPrio, + ) { + if let Err(e) = self.tx_msgs.try_send((priority, peers, msg)) { + error!("plumtree: could not send message: {e}"); + } + } + + fn send(&mut self, to: ActorId, msg: PlumtreeMsgV1, prio: PlumPrio) { + if let Err(e) = self.tx_msgs.try_send((prio, vec![to], msg)) { + error!("plumtree: could not send message: {e}"); + } + } + + fn deliver(&mut self, payload: ChangeV1) { + let tx = self.tx_changes.clone(); + tokio::spawn(async move { + match tokio::time::timeout( + Duration::from_secs(1), + tx.send((payload, ChangeSource::Broadcast)), + ) + .await + { + Ok(Err(e)) => error!("plumtree: could not deliver change: {e}"), + Err(_) => error!("plumtree: timed out delivering change after 1s"), + Ok(Ok(())) => {} + } + }); + } + + fn schedule(&mut self, timer: plum_foca::Timer, after: Duration) { + self.timer_spawner.spawn((after, timer)); + } + + fn notify(&mut self, notification: plum_foca::Notification<'_, ChangeId, ActorId>) { + trace!("plumtree notification: {notification:?}"); + match notification { + plum_foca::Notification::PeerMovedToEager(_) => { + counter!("corro.plumtree.peer_to_eager").increment(1); + } + plum_foca::Notification::PeerMovedToLazy(_) => { + counter!("corro.plumtree.peer_to_lazy").increment(1); + } + plum_foca::Notification::PeerDroppedFromEager(_) => { + counter!("corro.plumtree.peer_dropped_from_eager").increment(1); + } + plum_foca::Notification::PeerEvictedFromLazy(_) => { + counter!("corro.plumtree.peer_evicted").increment(1); + } + plum_foca::Notification::DuplicateMessage(_) => { + counter!("corro.plumtree.duplicate_message").increment(1); + } + plum_foca::Notification::PayloadNotCached(_) => { + counter!("corro.plumtree.payload_not_cached").increment(1); + } + plum_foca::Notification::MessageMissing(count) => { + counter!("corro.plumtree.message_missing").increment(count as u64); + } + plum_foca::Notification::PruneSuppressed(_) => { + counter!("corro.plumtree.prune_suppressed").increment(1); + } + plum_foca::Notification::Rebalance => { + counter!("corro.plumtree.rebalance.total").increment(1); + } + } + } + + fn now(&self) -> Instant { + Instant::now() + } +} + +pub async fn spawn_plumtree_loop( + agent: Agent, + transport: T, + rx_plumtree: CorroReceiver, + rx_plumtree_updates: CorroReceiver, + tx_changes: CorroSender<(ChangeV1, ChangeSource)>, + tripwire: Tripwire, +) { + let plumtree_config = agent + .config() + .gossip + .plumtree() + .cloned() + .unwrap_or_default(); + + let config = plum_foca::Config { + ihave_timeout: Duration::from_millis(150), + optimization_threshold: Some(plumtree_config.optimization_threshold), + max_cached_payloads: 5000, + num_eager: None, + min_lazy: None, + max_lazy: None, + prune_threshold: plumtree_config.prune_threshold, + max_received_entries: 10000, + prune_throttle: plumtree_config.prune_throttle_secs.map(Duration::from_secs), + eager_ratios: plumtree_config.eager_ratios, + }; + + plumtree_loop( + agent, + transport, + rx_plumtree, + rx_plumtree_updates, + tx_changes, + config, + tripwire, + ) + .await; +} + +pub async fn plumtree_loop( + agent: Agent, + transport: T, + mut rx_plumtree: CorroReceiver, + mut rx_plumtree_updates: CorroReceiver, + tx_changes: CorroSender<(ChangeV1, ChangeSource)>, + config: plum_foca::Config, + mut tripwire: Tripwire, +) { + let seen = ChangeSeenStore::new(config.max_received_entries, agent.bookie().clone()); + let mut state: PlumtreeState = + PlumtreeState::new_with_store(agent.actor_id(), config, seen); + + let (plumtree_timer_tx, mut plumtree_timer_rx) = mpsc::channel(10); + let timer_spawner = TimerSpawner::new(plumtree_timer_tx); + + let (tx_msgs, rx_msgs) = bounded(agent.config().perf.bcast_channel_len, "plumtree_msgs"); + + let send_agent = agent.clone(); + let send_transport = transport.clone(); + let send_msgs_handle = tokio::spawn(send_messages_loop(send_agent, send_transport, rx_msgs)); + + // send out ihave digests to lazy peers + let mut ihave_tick_interval = interval(Duration::from_millis(150)); + let mut maintenance_interval = interval(Duration::from_secs(60)); + + let mut rt = CorrosionPlumtreeRuntime::new(tx_changes, timer_spawner, tx_msgs); + + let peers: Vec<_> = plumtree_topology_map(&agent).into_iter().collect(); + info!("added {} peers to plumtree from members", peers.len()); + state.add_peers_bulk_with_rtt(peers, &mut rt); + + enum Branch { + Input(PlumtreeInput), + Updates(PlumtreeUpdates), + HandleTimer(Timer), + IHaveTick, + MaintenanceTick, + } + + loop { + let branch = tokio::select! { + biased; + _ = &mut tripwire => { + info!("plumtree_loop: tripwire fired, sending shutdown prunes"); + state.handle_shutdown(&mut rt); + break; + }, + updates = rx_plumtree_updates.recv() => match updates { + Some(updates) => Branch::Updates(updates), + None => { + warn!("plumtree_loop: updates channel closed"); + break; + } + }, + input = rx_plumtree.recv() => match input { + Some(input) => Branch::Input(input), + None => { + warn!("plumtree_loop: input channel closed"); + break; + } + }, + Some((timer, _seq)) = plumtree_timer_rx.recv() => { + Branch::HandleTimer(timer) + } + _ = ihave_tick_interval.tick() => { + Branch::IHaveTick + } + _ = maintenance_interval.tick() => { + Branch::MaintenanceTick + } + }; + + match branch { + Branch::Input(input) => match input { + PlumtreeInput::Wire(msg) => { + let msg_type: &'static str = (&msg).into(); + trace!("plumtree: received {msg_type} message"); + counter!("corro.plumtree.messages", "msg_type" => msg_type).increment(1); + match msg { + PlumtreeMsgV1::Gossip(g) => { + trace!("handling plumtree gossip"); + state.handle_gossip(g, &mut rt); + } + PlumtreeMsgV1::IHave(ih) => { + trace!("handling plumtree ihave"); + state.handle_ihave(ih, &mut rt); + } + PlumtreeMsgV1::Graft(g) => { + trace!("handling plumtree graft"); + if !g.send { + counter!("corro.plumtree.graft.optimization").increment(1); + } + state.handle_graft(g, &mut rt); + } + PlumtreeMsgV1::Prune(p) => { + trace!("handling plumtree prune"); + state.handle_prune(p, &mut rt); + } + } + } + PlumtreeInput::Broadcast(change) => { + let id = change.message_id(); + trace!("plumtree: broadcasting change: {id:?}"); + state.broadcast(id, change, &mut rt); + } + PlumtreeInput::QueryStats(reply) => { + let stats = PlumtreeStats { + eager_peers: state.eager_peers().iter().copied().collect(), + lazy_peers: state.lazy_peers().iter().copied().collect(), + ring_locked_peers: state.ring_locked_peers().iter().copied().collect(), + known_peers: state.known_peers().len(), + num_eager_target: state.num_eager(), + min_lazy_target: state.min_lazy(), + max_lazy_target: state.max_lazy(), + lazy_queue_len: state.lazy_queue().len(), + seen_cache_size: state.seen_cache_size(), + payload_cache_size: state.payload_cache_size(), + }; + let _ = reply.send(stats); + } + }, + Branch::Updates(updates) => match updates { + PlumtreeUpdates::MemberUp { + actor_id, + addr: _, + ring, + } => { + info!("plumtree: receieved member up: {actor_id}, ring: {ring:?}"); + state.peer_up(actor_id, Some(RttInfo { ring }), &mut rt); + } + PlumtreeUpdates::MemberDown(actor_id) => { + info!("plumtree: receieved member down: {actor_id}"); + state.peer_down(&actor_id, &mut rt); + } + }, + Branch::HandleTimer(timer) => { + state.timer_fired(timer, &mut rt); + } + Branch::IHaveTick => { + trace!("plumtree: sending out ihave digests"); + state.tick(&mut rt); + } + Branch::MaintenanceTick => { + trace!("plumtree: updating peer topology"); + state.update_peer_topology(plumtree_topology_map(&agent), &mut rt); + state.cache_evict_if_needed(&mut rt); + + gauge!("corro.plumtree.eager_peers").set(state.eager_peers().len() as f64); + gauge!("corro.plumtree.lazy_peers").set(state.lazy_peers().len() as f64); + gauge!("corro.plumtree.ring_locked_peers") + .set(state.ring_locked_peers().len() as f64); + gauge!("corro.plumtree.known_peers").set(state.known_peers().len() as f64); + gauge!("corro.plumtree.lazy_queue").set(state.lazy_queue().len() as f64); + + gauge!("corro.plumtree.payload_cache_size").set(state.payload_cache_size() as f64); + gauge!("corro.plumtree.seen_cache_size").set(state.seen_cache_size() as f64); + } + } + } + + drop(rt); + if let Err(e) = send_msgs_handle.await { + error!("plumtree send loop task failed to join: {e}"); + } +} + +#[derive(Debug)] +struct PendingPlumtreeSend { + peers: Vec, + payload: Bytes, +} + +struct P1GossipBatch { + peers: Vec, + buf: BytesMut, +} + +async fn send_messages_loop( + agent: Agent, + transport: T, + mut rx_msgs: CorroReceiver<(PlumPrio, Vec, PlumtreeMsgV1)>, +) { + const MAX_INFLIGHT: usize = 500; + const P1_GOSSIP_BATCH_INTERVAL: Duration = Duration::from_millis(10); + const P1_GOSSIP_BATCH_CUTOFF: usize = 1024 * 1024; + + let cluster_id = agent.cluster_id(); + let max_queue_len = agent.config().perf.processing_queue_len; + let batch_gossip = agent + .config() + .gossip + .plumtree() + .map(|p| p.batch_gossip) + .unwrap_or(false); + + let mut codec = LengthDelimitedCodec::builder() + .max_frame_length(10 * 1_024 * 1_024) + .new_codec(); + let mut ser_buf = BytesMut::new(); + let mut frame_buf = BytesMut::new(); + + let mut p0_queue: VecDeque = VecDeque::new(); + let mut p1_queue: VecDeque = VecDeque::new(); + let mut p1_gossip_batch = P1GossipBatch { + peers: Vec::new(), + buf: BytesMut::new(), + }; + let mut gossip_batch_interval = interval(P1_GOSSIP_BATCH_INTERVAL); + let mut join_set = JoinSet::new(); + let mut limited_log_count = 0; + let mut drop_log_count = 0; + + let bytes_per_sec: TransmitRateLimiter = RateLimiter::direct(Quota::per_second(unsafe { + NonZeroU32::new_unchecked(10 * 1024 * 1024) + })) + .with_middleware(); + + #[allow(clippy::large_enum_variant)] + enum Branch { + Msg((PlumPrio, Vec, PlumtreeMsgV1)), + GossipBatchDeadline, + } + + loop { + let branch = tokio::select! { + biased; + _ = join_set.join_next(), if !join_set.is_empty() => { + continue; + }, + msg = rx_msgs.recv() => match msg { + Some(msg) => Branch::Msg(msg), + None => { + warn!("plumtree send loop: message channel closed"); + break; + } + }, + _ = gossip_batch_interval.tick() => Branch::GossipBatchDeadline, + }; + + let mut rate_limited = false; + + match branch { + Branch::GossipBatchDeadline => { + if !p1_gossip_batch.buf.is_empty() { + p1_queue.push_back(PendingPlumtreeSend { + peers: std::mem::take(&mut p1_gossip_batch.peers), + payload: p1_gossip_batch.buf.split().freeze(), + }); + } + } + Branch::Msg((prio, peers, msg)) => { + debug!("plumtree: msg: {msg:?}, peers: {peers:?}"); + let p1_gossip = matches!(&msg, PlumtreeMsgV1::Gossip(_)); + let payload = match encode_plumtree_wire( + cluster_id, + &mut codec, + &mut ser_buf, + &mut frame_buf, + msg, + ) { + Ok(payload) => payload, + Err(()) => continue, + }; + + if batch_gossip && p1_gossip { + // gossip is sent to latest eager peers + p1_gossip_batch.peers = peers; + p1_gossip_batch.buf.extend_from_slice(&payload); + if p1_gossip_batch.buf.len() >= P1_GOSSIP_BATCH_CUTOFF { + p1_queue.push_back(PendingPlumtreeSend { + peers: std::mem::take(&mut p1_gossip_batch.peers), + payload: p1_gossip_batch.buf.split().freeze(), + }); + } + } else { + let pending = PendingPlumtreeSend { peers, payload }; + match prio { + PlumPrio::P0 => p0_queue.push_back(pending), + PlumPrio::P1 => p1_queue.push_back(pending), + } + } + } + } + + drain_plumtree_queue( + &agent, + &transport, + &bytes_per_sec, + &mut join_set, + &mut p0_queue, + MAX_INFLIGHT, + &mut rate_limited, + &mut limited_log_count, + ); + if !rate_limited { + drain_plumtree_queue( + &agent, + &transport, + &bytes_per_sec, + &mut join_set, + &mut p1_queue, + MAX_INFLIGHT, + &mut rate_limited, + &mut limited_log_count, + ); + } + + if drop_oldest_plumtree_send(&mut p0_queue, &mut p1_queue, max_queue_len).is_some() { + log_at_pow_10( + "dropped old plumtree message from send queue", + &mut drop_log_count, + ); + counter!("corro.plumtree.send.dropped").increment(1); + } + } + + info!("plumtree send loop is done"); +} + +fn encode_plumtree_wire( + cluster_id: ClusterId, + codec: &mut LengthDelimitedCodec, + ser_buf: &mut BytesMut, + frame_buf: &mut BytesMut, + msg: PlumtreeMsgV1, +) -> Result { + ser_buf.clear(); + if let Err(e) = (UniPayload::V1 { + data: UniPayloadV1::Plumtree(PlumtreeWire::V1 { data: msg }), + cluster_id, + }) + .write_to_stream(ser_buf.writer()) + { + error!("plumtree: failed to serialize wire msg: {e}"); + return Err(()); + } + + frame_buf.clear(); + if let Err(e) = codec.encode(ser_buf.split().freeze(), frame_buf) { + error!("plumtree: failed to frame wire msg: {e}"); + return Err(()); + } + + Ok(frame_buf.split().freeze()) +} + +fn resolve_peer_addrs(agent: &Agent, peers: &[ActorId]) -> Vec { + let members = agent.members().read(); + peers + .iter() + .filter_map(|id| members.states.get(id).map(|st| st.addr)) + .collect() +} + +#[allow(clippy::too_many_arguments)] +fn drain_plumtree_queue( + agent: &Agent, + transport: &T, + bytes_per_sec: &TransmitRateLimiter, + join_set: &mut JoinSet<()>, + queue: &mut VecDeque, + max_inflight: usize, + rate_limited: &mut bool, + limited_log_count: &mut u64, +) { + while !queue.is_empty() && join_set.len() < max_inflight { + let pending = queue.pop_front().unwrap(); + let addrs = resolve_peer_addrs(agent, &pending.peers); + if addrs.is_empty() { + warn!( + peers = ?pending.peers, + "plumtree: no addresses for peers, dropping message" + ); + continue; + } + + debug!("plumtree: sending plumtree msg to {addrs:?}"); + let mut spawn_count = 0; + let addr_count = addrs.len(); + for addr in addrs { + if join_set.len() >= max_inflight { + break; + } + + match try_transmit_uni( + bytes_per_sec, + pending.payload.clone(), + transport.clone(), + addr, + ) { + Err(e) => match e { + TransmitError::TooBig(_) | TransmitError::InsufficientCapacity(_) => { + error!("plumtree: could not spawn transmission: {e}"); + } + TransmitError::QuotaExceeded(_) => { + *rate_limited = true; + counter!("corro.plumtree.send.rate_limited").increment(1); + log_at_pow_10("plumtree broadcasts rate limited", limited_log_count); + break; + } + }, + Ok(fut) => { + join_set.spawn(async move { + fut.await; + counter!("corro.plumtree.send.total").increment(1); + }); + spawn_count += 1; + } + } + } + + if *rate_limited && spawn_count == 0 && addr_count > 0 { + queue.push_front(pending); + break; + } + + counter!("corro.plumtree.send.spawn").increment(spawn_count); + } +} + +fn drop_oldest_plumtree_send( + p0_queue: &mut VecDeque, + p1_queue: &mut VecDeque, + max: usize, +) -> Option { + if p0_queue.len() + p1_queue.len() <= max { + return None; + } + // drop from low-priority queue first + p1_queue.pop_back().or_else(|| p0_queue.pop_back()) +} + +/// Ring + RTT snapshot from in-memory [`Members`]. +fn plumtree_topology_map(agent: &Agent) -> HashMap { + let members = agent.members().read(); + members + .states + .iter() + .map(|(id, st)| (*id, RttInfo { ring: st.ring })) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::setup; + use crate::transport::{TransportError, TransportExt}; + use async_trait::async_trait; + use bytes::Bytes; + use corro_tests::test_config; + use corro_types::{ + actor::Actor, + base::{dbsr, CrsqlDbVersion, CrsqlSeq}, + broadcast::{ChangeV1, Changeset}, + members::Members, + }; + use parking_lot::RwLock; + use plum_foca::EagerRatios; + + use std::{collections::HashMap, net::SocketAddr, sync::Arc}; + // use plum_foca::PlumtreeMsg; + use rand::seq::IndexedRandom; + use rangemap::RangeInclusiveSet; + use speedy::Readable; + use tokio::{task::JoinSet, time::Duration}; + use tokio_stream::StreamExt; + use tokio_util::codec::{FramedRead, LengthDelimitedCodec}; + use tokio_util::sync::CancellationToken; + + use rand::{rngs::StdRng, SeedableRng}; + + #[derive(Clone, Debug)] + pub struct TestTransport { + nodes: Arc>>>, + } + + impl TestTransport { + pub fn new() -> Self { + Self { + nodes: Arc::new(RwLock::new(HashMap::new())), + } + } + + pub async fn add_node(&self, addr: SocketAddr) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(50000); + self.nodes.write().insert(addr, tx); + rx + } + } + + #[async_trait] + impl TransportExt for TestTransport { + async fn send_datagram(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError> { + let tx = self.nodes.write().get(&addr).unwrap().clone(); + tokio::spawn(async move { + let _ = tx.send(data).await; + }); + Ok(()) + } + + async fn send_uni(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError> { + let tx = self.nodes.write().get(&addr).unwrap().clone(); + tokio::spawn(async move { + let _ = tx.send(data).await; + }); + Ok(()) + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn test_plumtree_broadcast_spread() -> eyre::Result<()> { + let _ = tracing_subscriber::fmt::try_init(); + let (tripwire, _, _) = Tripwire::new_simple(); + let num_nodes = 10; + let num_changes = 200; + let transport = TestTransport::new(); + let config = Arc::new(RwLock::new(foca::Config::new_wan( + num_nodes.try_into().unwrap(), + ))); + + println!("config: {:?}", config.read().max_transmissions); + let mut processed_join_set = JoinSet::new(); + let mut plumtree_join_set = JoinSet::new(); + + let mut tas: Vec<_> = Vec::new(); + let mut send_tas: Vec<_> = Vec::new(); + let mut members = Members::default(); + for _ in 0..num_nodes { + let (_, test_conf) = test_config(|conf| conf.build())?; + let (agent, opts) = setup(test_conf.clone(), tripwire.clone()).await?; + members.add_member(&Actor::new( + agent.actor_id(), + agent.gossip_addr(), + agent.clock().new_timestamp().into(), + agent.cluster_id(), + None, + )); + send_tas.push((agent.actor_id(), agent.tx_plumtree().clone())); + tas.push((agent, opts)); + } + + let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple(); + // let mut tx_bcasts: Vec<_> = Vec::new(); + let cancel_token = CancellationToken::new(); + + for (agent, mut opts) in tas.into_iter() { + let agent_clone = agent.clone(); + agent_clone.members().write().states = members.states.clone(); + agent_clone.members().write().by_addr = members.by_addr.clone(); + agent_clone.members().write().rtts = members.rtts.clone(); + + let cancel = cancel_token.clone(); + + let mut transport_rx: mpsc::Receiver = + transport.add_node(agent_clone.gossip_addr()).await; + + let tx_plumtree = agent_clone.tx_plumtree().clone(); + let transport_clone = transport.clone(); + let tripwire_clone = tripwire.clone(); + tokio::spawn(async move { + let config = plum_foca::Config { + ihave_timeout: Duration::from_millis(500), + optimization_threshold: Some(5), + max_cached_payloads: 4096, + num_eager: Some(3), + min_lazy: Some(5), + max_lazy: Some(5), + prune_threshold: 3, + max_received_entries: 10000, + prune_throttle: None, + eager_ratios: EagerRatios::default(), + }; + plumtree_loop( + agent_clone.clone(), + transport_clone, + opts.rx_plumtree, + opts.rx_plumtree_updates, + agent_clone.tx_changes().clone(), + config, + tripwire_clone, + ) + .await; + }); + + // one sec sleep so plumtree applies membershiip update + tokio::time::sleep(Duration::from_secs(1)).await; + println!("spawned plumtree loop"); + + #[derive(Default)] + struct PlumStats { + gossip: u64, + ihave: u64, + graft: u64, + prune: u64, + } + + let cancel_clone = cancel.clone(); + plumtree_join_set.spawn(async move { + let mut stats = PlumStats::default(); + loop { + tokio::select! { + biased; + _ = cancel_clone.cancelled() => break, + Some(b) = transport_rx.recv() => { + let mut framed = FramedRead::new( + b.as_ref(), + LengthDelimitedCodec::builder() + .max_frame_length(100 * 1_024 * 1_024) + .new_codec(), + ); + while let Some(Ok(frame)) = framed.next().await { + if let Ok(UniPayload::V1 { + data: UniPayloadV1::Plumtree(msg), + .. + }) = UniPayload::read_from_buffer(&frame) + { + let PlumtreeWire::V1 { data } = msg; + + let msg_type: &'static str = (&data).into(); + match msg_type { + "gossip" => stats.gossip += 1, + "i_have" => stats.ihave += 1, + "graft" => stats.graft += 1, + "prune" => stats.prune += 1, + _ => warn!("unexpected message type: {msg_type}"), + } + tx_plumtree + .send(PlumtreeInput::Wire(data)) + .await + .ok(); + } + } + + } + } + } + stats + }); + + // let agent_clone: Agent = agent.clone(); + processed_join_set.spawn(async move { + let actor_id = agent.actor_id(); + let mut seen_map: RangeInclusiveSet = RangeInclusiveSet::new(); + let mut duplicate = 0; + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => break, + Some(changes) = opts.rx_changes.recv() => { + match changes { + ( + ChangeV1 { + actor_id: _, + changeset, + }, + ChangeSource::Broadcast, + ) => { + if seen_map.contains(&changeset.versions().start()) { + duplicate += 1; + continue; + } + seen_map.insert(changeset.versions().start()..=changeset.versions().end()); + } + _ => { + warn!("unexpected change source: {:?}", changes); + } + } + } + } + } + + (actor_id, duplicate, seen_map) + }); + } + + println!("done spawning processing threads"); + let mut rng = StdRng::from_os_rng(); + let chunk_size = 10u64; + let chunk_pause = Duration::from_millis(200); + let mut all_changes = Vec::new(); + for chunk_start in (0..num_changes).step_by(chunk_size as usize) { + let chunk_end = (chunk_start + chunk_size).min(num_changes); + for i in chunk_start..chunk_end { + let (actor_id, tx_plumtree) = send_tas.choose(&mut rng).unwrap(); + let change = ChangeV1 { + actor_id: *actor_id, + changeset: Changeset::Full { + version: CrsqlDbVersion(i), + changes: vec![], + seqs: dbsr!(0, 0), + last_seq: CrsqlSeq(0), + ts: Default::default(), + }, + }; + all_changes.push(change.clone()); + tx_plumtree + .send(PlumtreeInput::Broadcast(change)) + .await + .unwrap(); + } + + tokio::time::sleep(chunk_pause).await; + } + + println!("done sending changes"); + tokio::time::sleep(Duration::from_secs(10)).await; + cancel_token.cancel(); + println!("done cancelling"); + drop(send_tas); + + println!("sending cancel signal"); + tripwire_tx.send(()).await.unwrap(); + tripwire_worker.await; + // spawn::wait_for_all_pending_handles().await; + + let results = processed_join_set.join_all().await; + let plumtree_results = plumtree_join_set.join_all().await; + + let mut total_map = results + .into_iter() + .map(|(actor_id, duplicate, seen_map)| (actor_id, (duplicate, seen_map))) + .collect::>(); + + let plumstats_map = plumtree_results.into_iter().collect::>(); + + for change in all_changes { + let (_, seen_map) = total_map + .entry(change.actor_id) + .or_insert((0, RangeInclusiveSet::new())); + seen_map + .insert(change.changeset.versions().start()..=change.changeset.versions().end()); + } + + let total_expected = num_nodes as u64 * num_changes; + let total_seen: u64 = total_map + .values() + .map(|(_, seen_map)| { + seen_map + .iter() + .map(|v| v.end().0 - v.start().0 + 1) + .sum::() + }) + .sum(); + let extra_recvs: u64 = total_map + .values() + .map(|(duplicate, _)| *duplicate as u64) + .sum(); + + let (total_gossip, total_ihave, total_graft, total_prune) = plumstats_map.iter().fold( + (0, 0, 0, 0), + |(acc_gossip, acc_ihave, acc_graft, acc_prune), stats| { + ( + acc_gossip + stats.gossip, + acc_ihave + stats.ihave, + acc_graft + stats.graft, + acc_prune + stats.prune, + ) + }, + ); + let total_control = total_ihave + total_graft + total_prune; + + println!(); + println!("--- Plumtree spread results (async) ---"); + println!("nodes: {num_nodes}, changes: {num_changes}"); + println!(); + + for (actor_id, (duplicate, seen_map)) in total_map { + println!("actor_id: {actor_id}, duplicate: {duplicate}, seen: {seen_map:?}"); + } + + println!("delivery: {total_seen} / {total_expected}"); + println!( + "delivery %: {:.2}", + total_seen as f64 / total_expected as f64 * 100.0 + ); + println!("duplicate deliveries: {extra_recvs}"); + println!( + "duplicate %: {:.4}", + extra_recvs as f64 / total_expected as f64 * 100.0 + ); + println!(); + println!("gossip messages: {total_gossip}"); + println!( + "gossip duplicates: {:.1}", + total_gossip as f64 / total_expected as f64 + ); + println!("ihave messages: {total_ihave}"); + println!("graft messages: {total_graft}"); + println!("prune messages: {total_prune}"); + println!("total control: {total_control}"); + println!( + "control per change: {:.1}", + total_control as f64 / num_changes as f64 + ); + + assert!( + total_seen as f64 / total_expected as f64 > 0.99, + "delivery rate too low: {total_seen}/{total_expected}" + ); + Ok(()) + } +} diff --git a/crates/corro-agent/src/transport.rs b/crates/corro-agent/src/transport.rs index b224b91e8..0861c9183 100644 --- a/crates/corro-agent/src/transport.rs +++ b/crates/corro-agent/src/transport.rs @@ -6,6 +6,7 @@ use std::{ time::{Duration, Instant}, }; +use async_trait::async_trait; use bytes::Bytes; use corro_types::config::GossipConfig; use metrics::{counter, gauge, histogram}; @@ -22,6 +23,12 @@ use tracing::{debug, debug_span, info, trace, warn, Instrument}; use crate::api::peer::gossip_client_endpoint; +#[async_trait] +pub trait TransportExt { + async fn send_datagram(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError>; + async fn send_uni(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError>; +} + #[derive(Debug, Clone)] pub struct Transport(Arc); @@ -565,3 +572,79 @@ fn test_conn(conn: &Connection) -> bool { } } } + +#[async_trait] +impl TransportExt for Transport { + #[tracing::instrument(skip(self, data), fields(buf_size = data.len()), level = "debug", err)] + async fn send_datagram(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError> { + let conn = self.connect(addr, TrafficClass::Foca).await?; + debug!("connected to {addr}"); + + match conn.send_datagram(data.clone()) { + Ok(send) => { + debug!("sent datagram to {addr}"); + return Ok(send); + } + Err(SendDatagramError::ConnectionLost(e)) => { + debug!("retryable error attempting to send datagram: {e}"); + } + Err(e) => { + counter!("corro.transport.send_datagram.errors", "addr" => addr.to_string(), "error" => e.to_string()).increment(1); + if matches!(e, SendDatagramError::TooLarge) { + warn!(%addr, "attempted to send a larger-than-PMTU datagram. len: {}, pmtu: {:?}", data.len(), conn.max_datagram_size()); + } + return Err(e.into()); + } + } + + let conn = self.connect(addr, TrafficClass::Foca).await?; + debug!("re-connected to {addr}"); + Ok(conn.send_datagram(data)?) + } + + #[tracing::instrument(skip(self, data), fields(buf_size = data.len()), level = "debug", err)] + async fn send_uni(&self, addr: SocketAddr, data: Bytes) -> Result<(), TransportError> { + let len = data.len(); + let conn = self.connect(addr, TrafficClass::Broadcast).await?; + + let mut stream = match conn + .open_uni() + .instrument(debug_span!("quic_open_uni")) + .await + { + Ok(stream) => stream, + Err(e @ ConnectionError::VersionMismatch) => { + return Err(e.into()); + } + Err(e) => { + debug!("retryable error attempting to open unidirectional stream: {e}"); + let conn = self.connect(addr, TrafficClass::Broadcast).await?; + conn.open_uni() + .instrument(debug_span!("quic_open_uni")) + .await? + } + }; + + stream + .write_chunk(data) + .instrument(debug_span!("quic_write_chunk")) + .await?; + + stream + .finish() + .expect("unreachable, the stream does not leave this method"); + + stream + .stopped() + .instrument(debug_span!("quic_stopped")) + .await?; + + counter!( + "corro.transport.tx.bytes.v2.total", + "traffic" => TrafficClass::Broadcast.as_str() + ) + .increment(len as u64); + + Ok(()) + } +} diff --git a/crates/corro-tests/Cargo.toml b/crates/corro-tests/Cargo.toml index 72f7f7a8e..92d07a187 100644 --- a/crates/corro-tests/Cargo.toml +++ b/crates/corro-tests/Cargo.toml @@ -7,6 +7,8 @@ homepage.workspace = true repository.workspace = true [dependencies] +async-trait = { workspace = true } +bytes = { workspace = true } corro-agent = { path = "../corro-agent" } corro-client = { path = "../corro-client" } corro-types = { path = "../corro-types" } diff --git a/crates/corro-tests/src/lib.rs b/crates/corro-tests/src/lib.rs index e34df588c..acfc010fd 100644 --- a/crates/corro-tests/src/lib.rs +++ b/crates/corro-tests/src/lib.rs @@ -63,10 +63,9 @@ pub struct TestAgent { pub config: Config, } -pub async fn launch_test_agent Result>( +pub fn test_config Result>( f: F, - tripwire: Tripwire, -) -> eyre::Result { +) -> eyre::Result<(TempDir, Config)> { let tmpdir = TempDir::new(tempfile::tempdir()?); let schema_path = tmpdir.path().join("schema"); @@ -77,12 +76,17 @@ pub async fn launch_test_agent Result Result>( + f: F, + tripwire: Tripwire, +) -> eyre::Result { + let (tmpdir, conf) = test_config(f)?; let (agent, bookie, _, _) = start_with_config(conf.clone(), tripwire).await?; Ok(TestAgent { diff --git a/crates/corro-types/Cargo.toml b/crates/corro-types/Cargo.toml index 7b83acad1..536d2d837 100644 --- a/crates/corro-types/Cargo.toml +++ b/crates/corro-types/Cargo.toml @@ -36,6 +36,7 @@ metrics = { workspace = true } once_cell = { workspace = true } opentelemetry = { workspace = true } parking_lot = { workspace = true } +plum-foca = { path = "../plum-foca" } rand = { workspace = true } rangemap = { workspace = true } rcgen = { workspace = true } diff --git a/crates/corro-types/src/agent.rs b/crates/corro-types/src/agent.rs index 34c93b76e..aa3f502d8 100644 --- a/crates/corro-types/src/agent.rs +++ b/crates/corro-types/src/agent.rs @@ -33,9 +33,12 @@ use tripwire::Tripwire; use crate::{ actor::{Actor, ActorId, ClusterId, MemberId}, base::{CrsqlDbVersion, CrsqlDbVersionRange, CrsqlSeq}, - broadcast::{BroadcastInput, ChangeSource, ChangeV1, FocaInput, Timestamp}, + broadcast::{ + BroadcastInput, BroadcastV1, ChangeSource, ChangeV1, FocaInput, PlumtreeInput, + PlumtreeUpdates, Timestamp, + }, channel::{bounded, CorroSender}, - config::Config, + config::{BroadcastMethod, Config}, metrics_tracker::MetricsTracker, pubsub::SubsManager, schema::Schema, @@ -71,7 +74,8 @@ pub struct AgentConfig { pub tx_clear_buf: CorroSender<(ActorId, CrsqlDbVersionRange)>, pub tx_changes: CorroSender<(ChangeV1, ChangeSource)>, pub tx_foca: CorroSender, - + pub tx_plumtree: CorroSender, + pub tx_plumtree_updates: CorroSender, pub write_sema: Arc, pub schema: RwLock, @@ -106,6 +110,9 @@ pub struct AgentInner { tx_clear_buf: CorroSender<(ActorId, CrsqlDbVersionRange)>, tx_changes: CorroSender<(ChangeV1, ChangeSource)>, tx_foca: CorroSender, + tx_plumtree: CorroSender, + tx_plumtree_updates: CorroSender, + broadcaster: Broadcaster, write_sema: Arc, schema: RwLock, cluster_id: ArcSwap, @@ -123,6 +130,10 @@ pub struct Limits { impl Agent { pub fn new(config: AgentConfig) -> Self { + let broadcaster = match config.config.load().gossip.broadcast.method() { + BroadcastMethod::Gossip => Broadcaster::Gossip(config.tx_bcast.clone()), + BroadcastMethod::Plumtree => Broadcaster::Plumtree(config.tx_plumtree.clone()), + }; Self(Arc::new(AgentInner { actor_id: config.actor_id, pool: config.pool, @@ -140,6 +151,9 @@ impl Agent { tx_clear_buf: config.tx_clear_buf, tx_changes: config.tx_changes, tx_foca: config.tx_foca, + tx_plumtree: config.tx_plumtree, + tx_plumtree_updates: config.tx_plumtree_updates, + broadcaster, write_sema: config.write_sema, schema: config.schema, cluster_id: ArcSwap::from_pointee(config.cluster_id), @@ -212,6 +226,14 @@ impl Agent { &self.0.tx_foca } + pub fn tx_plumtree(&self) -> &CorroSender { + &self.0.tx_plumtree + } + + pub fn tx_plumtree_updates(&self) -> &CorroSender { + &self.0.tx_plumtree_updates + } + pub fn write_sema(&self) -> &Arc { &self.0.write_sema } @@ -252,6 +274,16 @@ impl Agent { self.0.config.load() } + pub fn broadcast_method(&self) -> BroadcastMethod { + self.config().gossip.broadcast_method() + } + + /// Route changes to the active dissemination method, resolved once from + /// [`BroadcastMethod`] at construction. + pub fn broadcaster(&self) -> &Broadcaster { + &self.0.broadcaster + } + pub fn set_config(&self, new_conf: Config) { self.0.config.store(Arc::new(new_conf)) } @@ -306,6 +338,53 @@ impl Agent { } } +/// Routes changes based on configured broadcast method. +pub enum Broadcaster { + Gossip(CorroSender), + Plumtree(CorroSender), +} + +impl Broadcaster { + pub fn broadcast_local(&self, change: ChangeV1) { + match self { + Broadcaster::Gossip(tx) => { + let tx = tx.clone(); + tokio::spawn(async move { + if let Err(e) = tx + .send(BroadcastInput::AddBroadcast(BroadcastV1::Change(change))) + .await + { + debug!("could not queue change for gossip broadcast: {e}"); + } + }); + } + Broadcaster::Plumtree(tx) => { + let tx = tx.clone(); + tokio::spawn(async move { + if let Err(e) = tx.send(PlumtreeInput::Broadcast(change)).await { + error!("could not send change for plumtree broadcast: {e}"); + } + }); + } + } + } + + pub fn rebroadcast(&self, change: &ChangeV1) { + match self { + Broadcaster::Gossip(tx) => { + let input = BroadcastInput::Rebroadcast(BroadcastV1::Change(change.clone())); + let tx = tx.clone(); + tokio::spawn(async move { + if let Err(e) = tx.send(input).await { + debug!("could not queue change for gossip rebroadcast: {e}"); + } + }); + } + Broadcaster::Plumtree(_) => {} + } + } +} + pub fn migrate(clock: Arc, conn: &mut Connection) -> rusqlite::Result<()> { let migrations: Vec> = vec![ Box::new(init_migration as fn(&Transaction) -> rusqlite::Result<()>), diff --git a/crates/corro-types/src/broadcast.rs b/crates/corro-types/src/broadcast.rs index e63e02218..e9e140222 100644 --- a/crates/corro-types/src/broadcast.rs +++ b/crates/corro-types/src/broadcast.rs @@ -1,6 +1,6 @@ use std::{ - cmp, collections::HashMap, fmt, io, num::NonZeroU32, num::ParseIntError, ops::Deref, - time::Duration, + cmp, collections::HashMap, fmt, io, net::SocketAddr, num::NonZeroU32, num::ParseIntError, + ops::Deref, time::Duration, }; use antithesis_sdk::assert_sometimes; @@ -10,6 +10,7 @@ use corro_base_types::{CrsqlDbVersionRange, CrsqlSeqRange}; use foca::{Identity, Member, Notification, Runtime, Timer}; use indexmap::{map::Entry, IndexMap}; use metrics::counter; +use plum_foca::Payload; use rusqlite::{ types::{FromSql, FromSqlError}, ToSql, @@ -48,6 +49,33 @@ pub enum UniPayload { #[derive(Debug, Clone, Readable, Writable)] pub enum UniPayloadV1 { Broadcast(BroadcastV1), + Plumtree(PlumtreeWire), +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Readable, Writable)] +pub struct ChangeId { + pub actor_id: ActorId, + pub changeset_id: ChangesetId, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Readable, Writable)] +pub enum ChangesetId { + Full { + version: CrsqlDbVersion, + seqs: CrsqlSeqRange, + last_seq: CrsqlSeq, + }, + Empty { + versions: CrsqlDbVersionRange, + }, +} + +/// Concrete Plumtree message type used on the wire. +pub type PlumtreeMsgV1 = plum_foca::PlumtreeMsg; + +#[derive(Debug, Clone, Readable, Writable)] +pub enum PlumtreeWire { + V1 { data: PlumtreeMsgV1 }, } #[derive(Debug, Clone, Readable, Writable)] @@ -89,6 +117,51 @@ pub enum AuthzV1 { Token(String), } +#[derive(Debug)] +pub enum PlumtreeInput { + /// A wire message received from a remote peer (deserialized from uni-stream). + Wire(PlumtreeMsgV1), + /// Application wants to originate a new broadcast. + Broadcast(ChangeV1), + /// Diagnostics: request a snapshot of the current plumtree state. + QueryStats(oneshot::Sender), +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct PlumtreeStats { + /// Peers we push GOSSIP to directly (the spanning-tree branches). + pub eager_peers: Vec, + /// Peers we only send IHAVE digests to. + pub lazy_peers: Vec, + /// Peers pinned to eager because of their ring/latency. + pub ring_locked_peers: Vec, + /// Total number of peers known to the overlay. + pub known_peers: usize, + /// Target number of eager peers (fanout). + pub num_eager_target: usize, + /// Target minimum number of lazy peers. + pub min_lazy_target: usize, + /// Target maximum number of lazy peers. + pub max_lazy_target: usize, + /// Number of pending IHAVE digests queued for lazy peers. + pub lazy_queue_len: usize, + /// Number of message ids currently tracked in the seen cache. + pub seen_cache_size: usize, + /// Number of cached payloads available to answer GRAFTs. + pub payload_cache_size: usize, +} + +#[derive(Debug)] +pub enum PlumtreeUpdates { + MemberUp { + actor_id: ActorId, + addr: SocketAddr, + ring: Option, + // rtt_ms: u64, + }, + MemberDown(ActorId), +} + #[derive(Clone, Debug, Readable, Writable)] pub enum BroadcastV1 { Change(ChangeV1), @@ -124,6 +197,22 @@ impl Deref for ChangeV1 { } } +impl Payload for ChangeV1 { + type MessageId = ChangeId; + type NodeId = ActorId; + + fn message_id(&self) -> ChangeId { + ChangeId { + actor_id: self.actor_id, + changeset_id: self.changeset.id(), + } + } + + fn origin(&self) -> Self::NodeId { + self.actor_id + } +} + #[derive(Debug, Clone, PartialEq, Readable, Writable)] pub enum Changeset { Empty { @@ -432,6 +521,37 @@ impl Changeset { Changeset::Empty { .. } | Changeset::EmptySet { .. } => Box::new(std::iter::empty()), } } + + pub fn id(&self) -> ChangesetId { + match self { + Changeset::Full { + version, + seqs, + last_seq, + .. + } => ChangesetId::Full { + version: *version, + seqs: *seqs, + last_seq: *last_seq, + }, + Changeset::FullV2 { + version, + seqs, + last_seq, + .. + } => ChangesetId::Full { + version: *version, + seqs: *seqs, + last_seq: *last_seq, + }, + Changeset::Empty { versions, .. } => ChangesetId::Empty { + versions: *versions, + }, + Changeset::EmptySet { .. } => { + todo!() + } + } + } } #[derive(Debug, thiserror::Error)] @@ -703,20 +823,10 @@ pub async fn broadcast_changes( match_changes(agent.subs_manager(), &changeset, db_version); match_changes(agent.updates_manager(), &changeset, db_version); - let tx_bcast = agent.tx_bcast().clone(); assert_sometimes!(true, "Corrosion broadcasts changes"); - tokio::spawn(async move { - if let Err(e) = tx_bcast - .send(BroadcastInput::AddBroadcast(BroadcastV1::Change( - ChangeV1 { - actor_id, - changeset, - }, - ))) - .await - { - error!("could not send change message for broadcast: {e}"); - } + agent.broadcaster().broadcast_local(ChangeV1 { + actor_id, + changeset, }); } Err(e) => { diff --git a/crates/corro-types/src/channel.rs b/crates/corro-types/src/channel.rs index a51ce7d8e..023d4b84e 100644 --- a/crates/corro-types/src/channel.rs +++ b/crates/corro-types/src/channel.rs @@ -73,17 +73,18 @@ pub fn bounded( let (tx, rx) = channel(capacity); let threshold = (capacity as f64 * 0.9) as usize; - let inner_channel = tx.clone(); + // weeak sender so metrics loop won't keep channel open + let weak_channel = tx.downgrade(); tokio::spawn(async move { let mut ticks_since_report = 0; let mut tick = interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut prev = inner_channel.capacity(); + let mut prev = capacity; loop { tick.tick().await; - if inner_channel.is_closed() { + let Some(inner_channel) = weak_channel.upgrade() else { break; - } + }; let current = inner_channel.capacity(); if ticks_since_report >= 30 || (prev != current && (current < threshold || prev < threshold)) diff --git a/crates/corro-types/src/config.rs b/crates/corro-types/src/config.rs index cc8513056..3be717b77 100755 --- a/crates/corro-types/src/config.rs +++ b/crates/corro-types/src/config.rs @@ -211,6 +211,93 @@ pub enum AuthzConfig { BearerToken(String), } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum BroadcastMethod { + Gossip, + Plumtree, +} + +fn default_broadcast_config() -> BroadcastConfig { + BroadcastConfig::Gossip +} + +pub fn default_plumtree_prune_threshold() -> u32 { + 5 +} + +pub fn default_plumtree_optimization_threshold() -> u32 { + 7 +} + +pub fn default_plumtree_prune_throttle_secs() -> Option { + Some(1) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PlumtreeConfig { + #[serde(default = "default_plumtree_prune_threshold")] + pub prune_threshold: u32, + #[serde(default = "default_plumtree_optimization_threshold")] + pub optimization_threshold: u32, + #[serde(default)] + pub batch_gossip: bool, + /// Suppress repeat PRUNEs to the same peer within this window. `None` disables. + #[serde(default = "default_plumtree_prune_throttle_secs")] + pub prune_throttle_secs: Option, + /// Near/Mid/Far split when selecting eager and lazy peers. + #[serde(default)] + pub eager_ratios: plum_foca::EagerRatios, +} + +impl Default for PlumtreeConfig { + fn default() -> Self { + Self { + prune_threshold: default_plumtree_prune_threshold(), + optimization_threshold: default_plumtree_optimization_threshold(), + batch_gossip: false, + prune_throttle_secs: default_plumtree_prune_throttle_secs(), + eager_ratios: plum_foca::EagerRatios::default(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum BroadcastConfig { + Gossip, + Plumtree(PlumtreeConfig), +} + +impl Default for BroadcastConfig { + fn default() -> Self { + Self::Gossip + } +} + +impl BroadcastConfig { + pub fn method(&self) -> BroadcastMethod { + match self { + Self::Gossip => BroadcastMethod::Gossip, + Self::Plumtree(_) => BroadcastMethod::Plumtree, + } + } + + pub fn plumtree(&self) -> Option<&PlumtreeConfig> { + match self { + Self::Plumtree(cfg) => Some(cfg), + Self::Gossip => None, + } + } + + pub fn from_method(method: BroadcastMethod) -> Self { + match method { + BroadcastMethod::Gossip => Self::Gossip, + BroadcastMethod::Plumtree => Self::Plumtree(PlumtreeConfig::default()), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GossipConfig { #[serde(alias = "addr")] @@ -232,6 +319,18 @@ pub struct GossipConfig { pub disable_gso: bool, #[serde(default)] pub member_id: Option, + #[serde(default = "default_broadcast_config")] + pub broadcast: BroadcastConfig, +} + +impl GossipConfig { + pub fn broadcast_method(&self) -> BroadcastMethod { + self.broadcast.method() + } + + pub fn plumtree(&self) -> Option<&PlumtreeConfig> { + self.broadcast.plumtree() + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -370,6 +469,8 @@ pub enum ConfigError { Config(#[from] config::ConfigError), #[error("gossip.max_mtu value {value} is below the QUIC minimum of 1200 (RFC 9000)")] InvalidMaxMtu { value: u16 }, + #[error("{0}")] + InvalidEagerRatios(#[from] plum_foca::EagerRatiosError), } impl Config { @@ -394,6 +495,9 @@ impl Config { return Err(ConfigError::InvalidMaxMtu { value: mtu }); } } + if let Some(plumtree) = self.gossip.plumtree() { + plumtree.eager_ratios.validate()?; + } Ok(()) } } @@ -418,6 +522,7 @@ pub struct ConfigBuilder { member_id: Option, max_mtu: Option, disable_gso: bool, + broadcast: Option, } impl ConfigBuilder { @@ -497,6 +602,21 @@ impl ConfigBuilder { self } + pub fn broadcast_method(mut self, method: BroadcastMethod) -> Self { + self.broadcast = Some(BroadcastConfig::from_method(method)); + self + } + + pub fn broadcast(mut self, broadcast: BroadcastConfig) -> Self { + self.broadcast = Some(broadcast); + self + } + + pub fn plumtree(mut self, plumtree: PlumtreeConfig) -> Self { + self.broadcast = Some(BroadcastConfig::Plumtree(plumtree)); + self + } + /// Disable Generic Segmentation Offload (GSO) for the QUIC gossip transport. pub fn disable_gso(mut self, disable: bool) -> Self { self.disable_gso = disable; @@ -543,6 +663,7 @@ impl ConfigBuilder { max_mtu: self.max_mtu, disable_gso: self.disable_gso, member_id: self.member_id, + broadcast: self.broadcast.unwrap_or_else(default_broadcast_config), }, perf: self.perf.unwrap_or_default(), admin: AdminConfig { @@ -609,3 +730,44 @@ pub struct TableReapConfig { #[serde(default)] pub match_filter: Option, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn broadcast_config_defaults_to_gossip() { + let cfg: GossipConfig = serde_json::from_value(serde_json::json!({ + "bind_addr": "127.0.0.1:4001", + })) + .unwrap(); + assert_eq!(cfg.broadcast.method(), BroadcastMethod::Gossip); + assert!(cfg.plumtree().is_none()); + } + + #[test] + fn broadcast_config_plumtree() { + let cfg: GossipConfig = serde_json::from_value(serde_json::json!({ + "bind_addr": "127.0.0.1:4001", + "broadcast": { + "plumtree": { + "prune_threshold": 7 + } + } + })) + .unwrap(); + assert_eq!(cfg.broadcast.method(), BroadcastMethod::Plumtree); + assert_eq!(cfg.plumtree().unwrap().prune_threshold, 7); + } + + #[test] + fn plumtree_eager_ratios_must_sum_to_100() { + let ratios = plum_foca::EagerRatios { + near_pct: 60, + mid_pct: 30, + far_pct: 20, + }; + assert!(ratios.validate().is_err()); + assert!(plum_foca::EagerRatios::default().validate().is_ok()); + } +} diff --git a/crates/corro-types/src/members.rs b/crates/corro-types/src/members.rs index 3eadcfd26..ee4bdc213 100644 --- a/crates/corro-types/src/members.rs +++ b/crates/corro-types/src/members.rs @@ -48,9 +48,82 @@ impl MemberState { const RING_BUCKETS: [Range; 6] = [0..6, 6..15, 15..50, 50..100, 100..200, 200..300]; +/// Number of recent RTT samples retained per member and used to compute the +/// median RTT for ring assignment. +const RTT_WINDOW: usize = 20; + +/// Hysteresis dead-band as a percentage of the shared boundary value, applied +/// only to moves between *adjacent* ring buckets. Prevents ring flapping when +/// the average RTT sits right on a bucket boundary. Non-adjacent jumps are not +/// affected and switch immediately. +const RING_HYSTERESIS_PCT: u64 = 20; + +/// Index of the `RING_BUCKETS` range containing `avg`. Values at or above the +/// last bucket (>= 300ms) are clamped into the top bucket so a ring is always +/// assigned. +fn bucket_for(avg: u64) -> u8 { + RING_BUCKETS + .iter() + .position(|r| r.contains(&avg)) + .unwrap_or(RING_BUCKETS.len() - 1) as u8 +} + +/// Compute the ring for `avg`, applying hysteresis only when new ring bucket +/// is adjacent to prevent small rtt jumps from moving rings. +fn ring_with_hysteresis(current: Option, avg: u64) -> u8 { + let target = bucket_for(avg); + let Some(current) = current else { + return target; + }; + if target == current || target.abs_diff(current) != 1 { + return target; + } + if target > current { + // Moving up: exceed the end-of-current boundary by the margin. + let boundary = RING_BUCKETS[current as usize].end; + let margin = boundary * RING_HYSTERESIS_PCT / 100; + if avg >= boundary + margin { + target + } else { + current + } + } else { + // Moving down: drop below the start-of-current boundary by the margin. + let boundary = RING_BUCKETS[current as usize].start; + let margin = boundary * RING_HYSTERESIS_PCT / 100; + if avg + margin <= boundary { + target + } else { + current + } + } +} + +fn median_rtt(rtt: &Rtt) -> Option { + let len = rtt.buf.len(); + if len == 0 { + return None; + } + + let mut scratch = [0u64; RTT_WINDOW]; + let (head, tail) = rtt.buf.as_slices(); + scratch[..head.len()].copy_from_slice(head); + scratch[head.len()..len].copy_from_slice(tail); + + let values = &mut scratch[..len]; + values.sort_unstable(); + + let mid = len / 2; + Some(if len % 2 == 0 { + (values[mid - 1] + values[mid]) / 2 + } else { + values[mid] + }) +} + #[derive(Debug, Default, Clone)] pub struct Rtt { - pub buf: CircularBuffer<20, u64>, + pub buf: CircularBuffer, } #[derive(Default)] @@ -61,11 +134,11 @@ pub struct Members { pub rtts: BTreeMap, } -#[derive(Debug, PartialEq)] +#[derive(Debug)] pub enum MemberAddedResult { - NewMember, + NewMember(MemberState), Removed, - Updated, + Updated(MemberState), Ignored, } @@ -87,6 +160,13 @@ impl Members { } } + /// Median RTT in milliseconds for this member (same statistic as + /// [`Self::recalculate_rings`]), or `None` if there are no samples yet. + pub fn avg_rtt_ms(&self, actor_id: &ActorId) -> Option { + let addr = self.states.get(actor_id)?.addr; + median_rtt(self.rtts.get(&addr)?) + } + // A result of `true` means that the effective list of // cluster member addresses has changed pub fn add_member(&mut self, actor: &Actor) -> MemberAddedResult { @@ -108,8 +188,8 @@ impl Members { }; } + let is_new = !self.states.contains_key(&actor_id); let member = self.states.entry(actor_id).or_insert_with(|| { - ret = MemberAddedResult::NewMember; MemberState::new( actor.addr(), actor.ts(), @@ -118,6 +198,10 @@ impl Members { ) }); + if is_new { + ret = MemberAddedResult::NewMember(member.clone()); + } + trace!("member: {member:?}"); // The received timestamp is older than the previously known @@ -137,12 +221,12 @@ impl Members { member.ts = actor.ts(); member.cluster_id = actor.cluster_id(); member.member_id = actor.member_id(); - ret = MemberAddedResult::Updated; + ret = MemberAddedResult::Updated(member.clone()); } // If we just inserted, add the actor to the by_addr set and // recalculate the RTT rings. - if ret == MemberAddedResult::NewMember { + if matches!(ret, MemberAddedResult::NewMember(_)) { self.by_addr.insert(actor.addr(), actor.id()); self.recalculate_rings(actor.addr()); } @@ -178,31 +262,27 @@ impl Members { self.recalculate_rings(addr) } - /// For a given member, calculate the average RTT and update - /// `self.ring` with the index of the corresponding bucket in - /// `RING_BUCKETS`. + /// For a given member, calculate the median RTT and update `self.ring` with + /// the index of the corresponding bucket in `RING_BUCKETS`, applying + /// hysteresis on moves between adjacent buckets (see [`ring_with_hysteresis`]) + /// to avoid flapping when the RTT sits near a bucket boundary. fn recalculate_rings(&mut self, addr: SocketAddr) { if let Some(actor_id) = self.by_addr.get(&addr) { - if let Some(avg) = self.rtts.get(&addr).and_then(|rtt| { - // If the ring buffer isn't empty - (!rtt.buf.is_empty()).then(|| { - // We can only access the ring buffer via two - // slices, so we sum both of them together - (rtt.buf.as_slices().0.iter().sum::() - + rtt.buf.as_slices().1.iter().sum::()) - // Then average over the full size of the ring - // buffer for the average of recent RTTs - / rtt.buf.len() as u64 - }) - }) { - if let Some(state) = self.states.get_mut(actor_id) { - // We check which range-bucket the RTT is - // contained in, then update the stored index - for (ring, n) in RING_BUCKETS.iter().enumerate() { - if n.contains(&avg) { - state.ring = Some(ring as u8); - break; + if let Some(rtt) = self.rtts.get(&addr) { + let median = median_rtt(rtt); + + if let Some(median) = median { + if let Some(state) = self.states.get_mut(actor_id) { + let new_ring = ring_with_hysteresis(state.ring, median); + if state.ring != Some(new_ring) { + info!( + "actor: {actor_id}, rtt: {:?}{:?}, old ring: {:?}, new ring: {new_ring}, median: {median}", + rtt.buf.as_slices().0, + rtt.buf.as_slices().1, + state.ring, + ); } + state.ring = Some(new_ring); } } } diff --git a/crates/corrosion/src/main.rs b/crates/corrosion/src/main.rs index 1660b60a2..a8637b17d 100644 --- a/crates/corrosion/src/main.rs +++ b/crates/corrosion/src/main.rs @@ -587,6 +587,13 @@ async fn process_cli(cli: Cli) -> eyre::Result<()> { conn.send_command(corro_admin::Command::Log(corro_admin::LogCommand::Reset)) .await?; } + Command::Plumtree(PlumtreeCommand::Stats) => { + let mut conn = AdminConn::connect(cli.admin_path()).await?; + conn.send_command(corro_admin::Command::Plumtree( + corro_admin::PlumtreeCommand::Stats, + )) + .await?; + } } Ok(()) @@ -767,6 +774,16 @@ enum Command { /// Log related commands #[command(subcommand)] Log(LogCommand), + + /// Plumtree broadcast overlay commands + #[command(subcommand)] + Plumtree(PlumtreeCommand), +} + +#[derive(Subcommand)] +enum PlumtreeCommand { + /// Dump key plumtree overlay stats + Stats, } #[derive(Subcommand)] diff --git a/crates/plum-foca/Cargo.toml b/crates/plum-foca/Cargo.toml new file mode 100644 index 000000000..c0d5e03d4 --- /dev/null +++ b/crates/plum-foca/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "plum-foca" +version = "0.1.0" +edition = "2024" + +[dependencies] +indexmap = { workspace = true } +serde = { workspace = true, features = ["derive"] } +thiserror = { workspace = true } +speedy = { workspace = true } +strum = { workspace = true } +tracing = { workspace = true } +rand = { workspace = true } + +[dev-dependencies] +tracing-subscriber = { workspace = true } + +[lints] +workspace = true diff --git a/crates/plum-foca/src/lib.rs b/crates/plum-foca/src/lib.rs new file mode 100644 index 000000000..7b2a7d3b6 --- /dev/null +++ b/crates/plum-foca/src/lib.rs @@ -0,0 +1,2349 @@ +use indexmap::{IndexMap, IndexSet}; +use rand::rngs::SmallRng; +use rand::seq::{IteratorRandom, SliceRandom}; +use rand::{Rng, SeedableRng}; +use serde::{Deserialize, Serialize}; +use speedy::{Readable, Writable}; +use std::collections::{HashMap, HashSet}; +use std::fmt::Debug; +use std::hash::Hash; +use std::time::{Duration, Instant}; +use thiserror::Error; +use tracing::{debug, info, trace, warn}; + +pub trait MessageId: Clone + Eq + Hash + Debug + Send + 'static {} + +pub trait Payload: Clone + Debug + Send + 'static { + type MessageId: MessageId; + type NodeId: NodeId; + fn message_id(&self) -> Self::MessageId; + fn origin(&self) -> Self::NodeId; +} + +pub trait NodeId: Copy + Eq + Hash + Ord + Debug + Send {} + +impl MessageId for T where T: Clone + Eq + Hash + Debug + Send + 'static {} +impl NodeId for T where T: Copy + Eq + Hash + Ord + Debug + Send + 'static {} + +pub type Round = u32; + +#[derive(Clone, Copy)] +pub enum PlumPrio { + P0, + P1, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct RttInfo { + pub ring: Option, + // pub rtt_ms: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RingBucket { + Near, + Mid, + Far, +} + +impl RingBucket { + fn of(info: RttInfo) -> Self { + match info.ring { + Some(0 | 1) => Self::Near, + Some(2 | 3) => Self::Mid, + Some(_) | None => Self::Far, + } + } +} + +/// Split eager (and lazy) peers across Near/Mid/Far RTT buckets. +/// Values are percentages of the target fanout; default values are (~50% near, ~30% mid, ~20% far). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct EagerRatios { + pub near_pct: u8, + pub mid_pct: u8, + pub far_pct: u8, +} + +impl Default for EagerRatios { + fn default() -> Self { + Self { + near_pct: 50, + mid_pct: 30, + far_pct: 20, + } + } +} + +impl EagerRatios { + pub fn validate(&self) -> Result<(), EagerRatiosError> { + let sum = self.near_pct as u16 + self.mid_pct as u16 + self.far_pct as u16; + if sum != 100 { + return Err(EagerRatiosError { + near_pct: self.near_pct, + mid_pct: self.mid_pct, + far_pct: self.far_pct, + sum, + }); + } + Ok(()) + } +} + +#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)] +#[error( + "eager ratios must sum to 100, got near={near_pct}, mid={mid_pct}, far={far_pct} (sum={sum})" +)] +pub struct EagerRatiosError { + pub near_pct: u8, + pub mid_pct: u8, + pub far_pct: u8, + pub sum: u16, +} + +/// Tunable parameters for the Plumtree protocol. +#[derive(Debug, Clone)] +pub struct Config { + /// How long to wait after receiving an IHave before sending a GRAFT + /// if the full message hasn't arrived yet. + pub ihave_timeout: Duration, + /// Optimization threshold in rounds. This is the minimum amount by which a lazy peer's + /// IHAVE round must beat our eager-delivery round before we attempt to graft the lazy peer. + pub optimization_threshold: Option, + /// Number of eager peers (fanout). `None` derives it from cluster size: + /// `round(log10(known_peers + 1) * 3)` (min 3). + pub num_eager: Option, + /// Minimum number of lazy peers. `None` defaults to `num_eager * 1.5`. + pub min_lazy: Option, + /// Maximum number of lazy peers. `None` defaults to `num_eager * 2`. + pub max_lazy: Option, + /// Maximum number of times a message can be received before pruning sender. + /// We trade some duplication for tree stability. + pub prune_threshold: u32, + /// cap on seen entries; oldest are evicted when exceeded. + pub max_received_entries: usize, + /// cap on cached payload used to respond to GRAFT requests. + pub max_cached_payloads: usize, + /// If set, suppress repeat PRUNEs to the same peer within this window. + pub prune_throttle: Option, + /// Near/Mid/Far split when selecting eager and lazy peers during rebalance. + pub eager_ratios: EagerRatios, +} + +impl Default for Config { + fn default() -> Self { + Self { + ihave_timeout: Duration::from_secs(1), + optimization_threshold: Some(3), + num_eager: None, + min_lazy: None, + max_lazy: None, + prune_threshold: 1, + max_received_entries: 10000, + max_cached_payloads: 8192, + prune_throttle: Some(Duration::from_secs(1)), + eager_ratios: EagerRatios::default(), + } + } +} + +/// Resolved eager/lazy peer targets used by the protocol. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct FanoutTargets { + num_eager: usize, + min_lazy: usize, + max_lazy: usize, +} + +fn resolve_fanout(known_peers: usize, config: &Config) -> FanoutTargets { + let num_eager = config.num_eager.unwrap_or_else(|| { + let cluster = (known_peers + 1) as f64; + ((cluster.log10() * 3.0).round() as usize).max(3) + }); + FanoutTargets { + num_eager, + min_lazy: config.min_lazy.unwrap_or(num_eager * 3 / 2), + max_lazy: config.max_lazy.unwrap_or(num_eager * 2), + } +} + +pub trait SeenStore { + fn evict_if_needed(&mut self); + fn contains(&self, id: &I) -> bool; + fn observe(&mut self, id: I, round: Round) -> Option; + fn size(&self) -> usize; +} + +/// Timers produced by the protocol, these are handled with the `schedule` method. +/// The runtime schedules these externally and calls `PlumtreeState::timer_fired` when they expire. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum Timer { + /// Fires after `config.ihave_timeout`. For each id still missing, the node + /// sends a GRAFT request to senders. + IHaveTimeoutBatch { + ids: Vec, + retries: u32, + senders: Vec, + }, +} + +/// Observable protocol events for metrics, logging, or diagnostics. +#[derive(Debug, PartialEq, Eq)] +pub enum Notification<'a, I: MessageId, N: NodeId> { + PeerMovedToEager(&'a N), + PeerDroppedFromEager(&'a N), + PeerMovedToLazy(&'a N), + PeerEvictedFromLazy(&'a N), + DuplicateMessage(&'a I), + PayloadNotCached(&'a I), + MessageMissing(usize), + PruneSuppressed(&'a N), + Rebalance, +} + +pub trait Runtime, N: NodeId> { + /// Send a protocol message to a specific peer. + fn send(&mut self, to: N, msg: PlumtreeMsg, prio: PlumPrio); + + // Send a message to a group of peers + fn send_all(&mut self, peers: Vec, msg: PlumtreeMsg, prio: PlumPrio); + + /// Deliver a received message to the application layer. + fn deliver(&mut self, payload: P); + + /// Request the caller to schedule a timer. When the duration elapses, + /// the caller must invoke `PlumtreeState::timer_fired`. + fn schedule(&mut self, timer: Timer, after: Duration); + + /// Observable protocol event. + fn notify(&mut self, notification: Notification<'_, I, N>); + + /// Current time. The state machine is sans-IO(sans-io.readthedocs.io) and owns no clock, so the + /// runtime supplies one. + fn now(&self) -> Instant; +} + +/// A protocol message exchanged between Plumtree peers. +/// +/// Each peer keeps two overlays over the same neighbor set: an eager set that +/// forms a spanning tree and receives full payloads (push), and a lazy set +/// that only receives message_id digests and pulls the payload on demand. +/// `Graft` and `Prune` move a link between the two sets to repair and optimize +/// the tree. +#[derive(Debug, Clone, Readable, Writable, strum::IntoStaticStr)] +#[strum(serialize_all = "snake_case")] +pub enum PlumtreeMsg +where + I: MessageId, + P: Payload, + N: NodeId, +{ + /// Full payload pushed to an eager-set peer. This is how messages actually + /// propagate along the spanning tree. + Gossip(GossipMsg), + /// Lazy-push digest: contains batch of recently-seen message ids that is sent to lazy + /// peers on a tick. Lets a peer notice a message it never received eagerly + /// and `Graft` to fetch it. + IHave(IHaveMsg), + /// Ask the recipient to add us to its eager set (and, when `send` is set, + /// to reply with the full payload). Sent when an `IHave` reveals a message + /// that never arrived eagerly (tree repair) or to adopt a shorter path + /// (optimization). + Graft(GraftMsg), + /// Ask the recipient to remove us from its eager set, demoting the link to + /// lazy. Sent after a duplicate `Gossip` arrives once the duplicate count + /// crosses `Config::prune_threshold`, which signals a redundant tree edge. + Prune(PruneMsg), +} + +/// Full payload sent immediately to eager peers. +#[derive(Debug, Clone, Readable, Writable)] +pub struct GossipMsg +where + I: MessageId, + P: Payload, + N: NodeId, +{ + pub round: Round, + pub sender: N, + pub payload: P, +} + +/// Digest-only batch sent to lazy peers on each tick. +#[derive(Debug, Clone, Readable, Writable)] +pub struct IHaveMsg +where + I: MessageId, + N: NodeId, +{ + pub sender: N, + pub digests: Vec>, +} + +#[derive(Debug, Clone, Readable, Writable)] +pub struct IHaveDigest { + pub id: I, + pub round: Round, +} + +#[derive(Debug, Clone, Readable, Writable)] +pub struct GraftRequest { + pub id: I, + pub round: Round, +} + +#[derive(Debug, Clone, Readable, Writable)] +pub struct GraftMsg +where + I: MessageId, + N: NodeId, +{ + pub sender: N, + /// When true, respond with GOSSIP for each cached request id. + pub send: bool, + pub requests: Vec>, +} + +#[derive(Debug, Clone)] +pub struct PruneMsg +where + I: MessageId, + N: NodeId, +{ + pub sender: N, + pub triggered_by: Option, +} + +// NOTE: speedy's derive can't express the bounds needed to (de)serialize an +// `Option` field over a generic `I`, so we implement the traits by hand. +impl<'a, C, I, N> speedy::Readable<'a, C> for PruneMsg +where + C: speedy::Context, + I: MessageId + speedy::Readable<'a, C>, + N: NodeId + speedy::Readable<'a, C>, +{ + fn read_from>(reader: &mut R) -> Result { + Ok(Self { + sender: reader.read_value()?, + triggered_by: reader.read_value()?, + }) + } + + fn minimum_bytes_needed() -> usize { + >::minimum_bytes_needed() + + as speedy::Readable<'a, C>>::minimum_bytes_needed() + } +} + +impl speedy::Writable for PruneMsg +where + C: speedy::Context, + I: MessageId + speedy::Writable, + N: NodeId + speedy::Writable, +{ + fn write_to>(&self, writer: &mut T) -> Result<(), C::Error> { + writer.write_value(&self.sender)?; + writer.write_value(&self.triggered_by)?; + Ok(()) + } +} + +#[derive(Debug, Clone)] +struct MissingEntry { + ihave_sender: N, + round: Round, +} + +#[derive(Debug)] +struct PayloadCache { + entries: IndexMap, + max_size: usize, +} + +impl PayloadCache { + fn new(max_size: usize) -> Self { + Self { + entries: IndexMap::with_capacity(max_size), + max_size, + } + } + + fn insert(&mut self, id: I, payload: P, round: Round) { + self.entries.insert(id, (payload, round)); + } + + fn get(&self, id: &I) -> Option<&(P, Round)> { + self.entries.get(id) + } + + fn evict_if_needed(&mut self) { + if self.entries.len() > self.max_size { + self.entries.drain(0..self.entries.len() - self.max_size); + } + } + + fn size(&self) -> usize { + self.entries.len() + } +} + +// Mostly a dupe of ThrottleMap where we provide the time +#[derive(Debug)] +struct PruneThrottle { + until: HashMap, +} + +impl Default for PruneThrottle { + fn default() -> Self { + Self { + until: HashMap::new(), + } + } +} + +impl PruneThrottle { + fn throttled(&self, peer: &N, now: Instant) -> bool { + self.until.get(peer).is_some_and(|t| now < *t) + } + + fn record(&mut self, peer: N, now: Instant, ttl: Duration) { + self.until.insert(peer, now + ttl); + } + + fn clear_expired(&mut self, now: Instant) { + self.until.retain(|_, t| *t > now); + } +} + +const RING_EXTRA_CONFIRMATIONS: u32 = 5; + +/// Full Plumtree protocol state for one local node. +#[derive(Debug)] +pub struct PlumtreeState< + I: MessageId, + P: Payload, + N: NodeId, + S: SeenStore, +> { + /// This node's own id, excluded from every peer set. + local_id: N, + /// Tunables (fanout sizes, thresholds, throttles). + config: Config, + + /// Peers on the spanning tree: they get full `Gossip` payloads (push). + eager_peers: HashSet, + /// Lazy peers only get `IHave` digests (lazy push). + lazy_peers: IndexSet, + /// Every peer we know about, eager or lazy; drives. + known_peers: HashSet, + /// Eager peers pinned to the tree because they are our immediate ring + /// neighbors; they are never pruned, ensuring a node is never isolated. + ring_locked: HashSet, + /// Per-per rtt information. + peer_topology: HashMap, + /// Candidate topology change awaiting confirmation + pending_topology: HashMap, + + /// Digests buffered since the last tick, flushed to lazy peers as `IHave`. + lazy_queue: Vec>, + /// Messages announced via `IHave` but not yet received + missing: HashMap>, + + /// Dedup store: message ids we've already delivered, to drop duplicates. + seen: S, + /// Recently-delivered payloads, kept so we can answer inbound `Graft`s. + cache: PayloadCache, + /// PRNG for random peer selection (seedable for deterministic tests). + rng: SmallRng, + /// Rate-limits outbound `Prune`s per peer to avoid flapping links. + prune_throttle: PruneThrottle, + /// set on membership updates so we'd rebalance peers on next tick. + needs_rebalance: bool, + // target number of eager and lazy peers, calculated based on config + // or cluster size if config is not set. + fanout: FanoutTargets, +} + +impl, N: NodeId, S: SeenStore> + PlumtreeState +{ + pub fn new_with_store(local_id: N, config: Config, seen: S) -> Self { + Self::new_with_store_seeded(local_id, config, seen, rand::rng().random()) + } + + /// Like [`Self::new_with_store`] but with a deterministic RNG seed. + /// + /// All randomized decisions (eager/lazy selection, lazy eviction, graft + /// sender choice) draw from this RNG, so simulations and tests are + /// reproducible given the same seed and event order. + pub fn new_with_store_seeded(local_id: N, config: Config, seen: S, seed: u64) -> Self { + let cache_size = config.max_cached_payloads; + let fanout = resolve_fanout(0, &config); + Self { + local_id, + config, + eager_peers: HashSet::new(), + lazy_peers: IndexSet::new(), + known_peers: HashSet::new(), + ring_locked: HashSet::new(), + peer_topology: HashMap::new(), + pending_topology: HashMap::new(), + lazy_queue: Vec::new(), + missing: HashMap::new(), + seen, + cache: PayloadCache::new(cache_size), + rng: SmallRng::seed_from_u64(seed), + prune_throttle: PruneThrottle::default(), + needs_rebalance: false, + fanout, + } + } + + pub fn ring_locked_peers(&self) -> &HashSet { + &self.ring_locked + } + + pub fn has_message(&self, id: &I) -> bool { + self.seen.contains(id) + } + + pub fn eager_peers(&self) -> &HashSet { + &self.eager_peers + } + + pub fn lazy_peers(&self) -> &IndexSet { + &self.lazy_peers + } + + pub fn known_peers(&self) -> &HashSet { + &self.known_peers + } + + pub fn lazy_queue(&self) -> &Vec> { + &self.lazy_queue + } + + pub fn config(&self) -> &Config { + &self.config + } + + pub fn num_eager(&self) -> usize { + self.fanout.num_eager + } + + pub fn min_lazy(&self) -> usize { + self.fanout.min_lazy + } + + pub fn max_lazy(&self) -> usize { + self.fanout.max_lazy + } + + pub fn payload_cache_size(&self) -> usize { + self.cache.size() + } + + pub fn seen_cache_size(&self) -> usize { + self.seen.size() + } + /// Recompute cached fanout when `config.num_eager` is unset and cluster + /// size produces a new target. Returns `true` if any effective value changed. + fn maybe_recompute_fanout(&mut self) -> bool { + if self.config.num_eager.is_some() { + return false; + } + let new = resolve_fanout(self.known_peers.len(), &self.config); + if new == self.fanout { + return false; + } + self.fanout = new; + true + } + + // used only in sim tests + pub fn force_eager(&mut self, peer: N) { + if self.known_peers.contains(&peer) { + self.lazy_peers.swap_remove(&peer); + self.eager_peers.insert(peer); + } + } + + pub fn local_id(&self) -> &N { + &self.local_id + } + + // --- Protocol methods --- + + /// Handle an incoming GOSSIP message carrying a full payload. + /// + /// Unlike the paper, we do NOT promote the sender to eager on receipt. + /// In a multi-sender network, the peer that forwarded sender A's + /// message fast may be a poor path for sender B. Eager promotion + /// only happens through intentional GRAFT (IHave timeout or + /// optimization path). + pub fn handle_gossip(&mut self, msg: GossipMsg, rt: &mut impl Runtime) { + let GossipMsg { + round, + sender, + payload, + } = msg; + let id = payload.message_id(); + + let self_actor_id = self.local_id; + if let Some(duplicates) = self.seen.observe(id.clone(), round) { + if duplicates > self.config.prune_threshold && !self.ring_locked.contains(&sender) { + // TODO: suppress prunes for recently grafted peers? this would help + // with tree stability + let suppressed = self + .config + .prune_throttle + .is_some_and(|_| self.prune_throttle.throttled(&sender, rt.now())); + if suppressed { + rt.notify(Notification::PruneSuppressed(&sender)); + } else { + trace!( + ?self_actor_id, + ?sender, + "sending PRUNE due to duplicate gossip, triggered_by: {id:?}" + ); + + if let Some(ttl) = self.config.prune_throttle { + self.prune_throttle.record(sender, rt.now(), ttl); + } + + rt.send( + sender, + PlumtreeMsg::Prune(PruneMsg { + sender: self.local_id, + triggered_by: Some(id.clone()), + }), + PlumPrio::P1, + ); + self.move_to_lazy(&sender, rt); + } + } + rt.notify(Notification::DuplicateMessage(&id)); + return; + } + + self.cache.insert(id.clone(), payload.clone(), round); + rt.deliver(payload.clone()); + + let next_round = round + 1; + let peers: Vec = self + .eager_peers + .iter() + .filter(|p| **p != sender && **p != self.local_id && payload.origin() != **p) + .copied() + .collect(); + trace!( + ?self_actor_id, + ?sender, + "forwarding gossip to {} eager peers (id: {id:?}, round: {next_round}, peers: {peers:?})", + peers.len(), + ); + if !peers.is_empty() { + rt.send_all( + peers, + PlumtreeMsg::Gossip(GossipMsg { + round: next_round, + sender: self.local_id, + payload, + }), + PlumPrio::P1, + ); + } + + // if peer is not eager, ensure they are in lazy + if !self.eager_peers.contains(&sender) { + self.ensure_in_lazy(&sender, rt); + } + + if let Some(entry) = self.missing.remove(&id) + && let Some(optimization_threshold) = self.config.optimization_threshold + && entry.round + optimization_threshold < round + && entry.ihave_sender != sender + { + let sender = &entry.ihave_sender; + debug!( + ?self_actor_id, + "sending graft to {sender:?} (optimization from id {id:?} with round {round})" + ); + rt.send( + *sender, + PlumtreeMsg::Graft(GraftMsg { + sender: self.local_id, + send: false, + requests: vec![GraftRequest { + id: id.clone(), + round: entry.round, + }], + }), + PlumPrio::P0, + ); + self.move_to_eager(&entry.ihave_sender, rt); + + // paper has a prune here but we might prune a good path for different sender, + // possibly need more info to determine if we should prune. + } + + self.enqueue_ihave(id, round); + } + + /// Handle an incoming IHave digest batch. + /// + /// For each digest we haven't already received, record it in the + /// `missing` set and schedule one `IHaveTimeoutBatch` timer. If the full + /// GOSSIP doesn't arrive before the timer fires, we'll GRAFT. + pub fn handle_ihave(&mut self, msg: IHaveMsg, rt: &mut impl Runtime) { + let IHaveMsg { sender, digests } = msg; + + let mut senders = vec![sender]; + senders.extend(self.random_eager_peers(1, &sender)); + + let mut new_ids = Vec::new(); + for digest in digests { + if self.seen.contains(&digest.id) { + continue; + } + if self.missing.contains_key(&digest.id) { + continue; + } + + new_ids.push(digest.id.clone()); + self.missing.insert( + digest.id, + MissingEntry { + ihave_sender: sender, + round: digest.round, + }, + ); + } + + let self_actor_id = &self.local_id; + if !new_ids.is_empty() { + trace!( + ?self_actor_id, + ?sender, + "handle_ihave, scheduling graft timeout for {}", + new_ids.len(), + ); + rt.schedule( + Timer::IHaveTimeoutBatch { + ids: new_ids, + retries: 0, + senders, + }, + self.config.ihave_timeout, + ); + } + } + + /// Handle an incoming GRAFT request. + /// + /// The sender is asking us to add them back to our eager set and + /// (re)send the full payload for each requested message. + pub fn handle_graft(&mut self, msg: GraftMsg, rt: &mut impl Runtime) { + let GraftMsg { + sender, + send, + requests, + } = msg; + + self.move_to_eager(&sender, rt); + + if !send { + return; + } + + let self_actor_id = &self.local_id; + for req in requests { + if !self.seen.contains(&req.id) { + continue; + } + if let Some((payload, round)) = self.cache.get(&req.id).cloned() { + debug!( + ?self_actor_id, + ?sender, + "sending cached payloads requested through graft ({:?})", + req.id + ); + rt.send( + sender, + PlumtreeMsg::Gossip(GossipMsg { + round, + sender: self.local_id, + payload, + }), + PlumPrio::P0, + ); + } else { + debug!( + ?self_actor_id, + ?sender, + "requested payload no longer cached ({:?})", + req.id, + ); + rt.notify(Notification::PayloadNotCached(&req.id)); + } + } + } + + /// Handle an incoming PRUNE. + /// + pub fn handle_prune(&mut self, msg: PruneMsg, rt: &mut impl Runtime) { + debug!(self_actor_id = ?self.local_id, sender = ?msg.sender, "received prune, triggered by {:?}", msg.triggered_by); + if !self.ring_locked.contains(&msg.sender) { + self.move_to_lazy(&msg.sender, rt); + } + } + + /// Handle a graceful shutdown. + /// + /// Sends a PRUNE to every eager peer so they stop treating us as part of + /// the spanning tree and re-route around us. + pub fn handle_shutdown(&mut self, rt: &mut impl Runtime) { + info!(self_actor_id = ?self.local_id, "sending prunes to eager peers due to shutdown"); + let peers = self.eager_peers.iter().copied().collect::>(); + rt.send_all( + peers, + PlumtreeMsg::Prune(PruneMsg { + sender: self.local_id, + triggered_by: None, + }), + PlumPrio::P1, + ); + } + + /// Originate a new message from this node. + /// + /// Marks it as received, caches the payload, sends GOSSIP to all + /// eager peers (round 0), and enqueues IHave for lazy peers. + pub fn broadcast(&mut self, id: I, payload: P, rt: &mut impl Runtime) { + debug!( + self_actor_id = ?self.local_id, + "broadcasting message to {} eager and {} lazy", + self.eager_peers.len(), + self.lazy_peers.len() + ); + + self.seen.observe(id.clone(), 0); + self.cache.insert(id.clone(), payload.clone(), 0); + + let peers = self + .eager_peers + .iter() + .filter(|p| **p != self.local_id) + .copied() + .collect::>(); + rt.send_all( + peers, + PlumtreeMsg::Gossip(GossipMsg { + round: 0, + sender: self.local_id, + payload, + }), + PlumPrio::P0, + ); + + self.enqueue_ihave(id, 0); + } + + /// Handle a fired timer. + /// + /// `IHaveTimeoutBatch`: for each id still missing, send GRAFT and promote + /// the target peer to eager; schedule retry at a shorter timeout. + pub fn timer_fired(&mut self, timer: Timer, rt: &mut impl Runtime) { + match timer { + Timer::IHaveTimeoutBatch { + ids, + retries, + senders, + } => self.handle_ihave_timeout(ids, retries, senders, rt), + } + } + + fn handle_ihave_timeout( + &mut self, + ids: Vec, + retries: u32, + senders: Vec, + rt: &mut impl Runtime, + ) { + if senders.is_empty() || retries >= senders.len() as u32 { + return; + } + + let send_to = senders[retries as usize]; + let mut graft_requests = Vec::new(); + + for id in ids { + if self.seen.contains(&id) { + trace!("missing change already received, noop"); + self.missing.remove(&id); + continue; + } + + let Some(entry) = self.missing.get(&id).cloned() else { + trace!("no longer missing, noop"); + continue; + }; + + graft_requests.push(GraftRequest { + id: id.clone(), + round: entry.round, + }); + } + + if !graft_requests.is_empty() { + debug!( + self_actor_id = ?self.local_id, + "Handling ihave time out for {} requests", graft_requests.len(), + ); + // todo: maybe move this after we actually receive a response? + self.move_to_eager(&send_to, rt); + for chunk in graft_requests.chunks(10) { + rt.send( + send_to, + PlumtreeMsg::Graft(GraftMsg { + sender: self.local_id, + send: true, + requests: chunk.to_vec(), + }), + PlumPrio::P0, + ); + } + // supress prunes for a recently grafted peers + self.prune_throttle + .record(send_to, rt.now(), Duration::from_secs(3 * 60)); + + let ids = graft_requests.iter().map(|r| r.id.clone()).collect(); + // reschedule a retry if we still have more senders + if retries + 1 < senders.len() as u32 { + rt.schedule( + Timer::IHaveTimeoutBatch { + ids, + retries: retries + 1, + senders, + }, + self.config.ihave_timeout / 2, + ); + } else { + rt.notify(Notification::MessageMissing(ids.len())); + for id in ids { + self.missing.remove(&id); + } + } + } + } + + /// Periodic sends for all pending IHave digests to lazy peers. + /// + /// The caller should invoke this on a regular interval. + pub fn tick(&mut self, rt: &mut impl Runtime) { + let digests = self.drain_lazy_queue(); + if self.lazy_peers.is_empty() || digests.is_none() { + return; + } + let peers: Vec = self.lazy_peers.iter().copied().collect(); + rt.send_all( + peers, + PlumtreeMsg::IHave(IHaveMsg { + sender: self.local_id, + digests: digests.unwrap(), + }), + PlumPrio::P1, + ); + } + + fn random_eager_peers(&mut self, count: usize, exclude: &N) -> Vec { + self.eager_peers + .iter() + .filter(|p| *p != exclude) + .choose_multiple(&mut self.rng, count) + .into_iter() + .copied() + .collect() + } + + // --- Peer membership --- + + /// Add multiple peers at once (e.g. during bootstrap). + pub fn add_peers_bulk(&mut self, peers: Vec, rt: &mut impl Runtime) { + let entries = peers.into_iter().map(|p| (p, RttInfo::default())).collect(); + self.add_peers_bulk_with_rtt(entries, rt); + } + + /// Bootstrap peers together with RTT ring info for [`Self::rebalance`]. + pub fn add_peers_bulk_with_rtt( + &mut self, + peers: Vec<(N, RttInfo)>, + rt: &mut impl Runtime, + ) { + for (peer, info) in peers { + if peer != self.local_id { + self.known_peers.insert(peer); + self.peer_topology.insert(peer, info); + self.pending_topology.remove(&peer); + } + } + self.maybe_recompute_fanout(); + self.rebalance(rt); + } + + /// A new peer has come online and wants to join the overlay. + pub fn peer_up(&mut self, peer: N, rtt: Option, rt: &mut impl Runtime) { + if peer == self.local_id || self.known_peers.contains(&peer) { + return; + } + + self.peer_topology.insert(peer, rtt.unwrap_or_default()); + self.known_peers.insert(peer); + self.pending_topology.remove(&peer); + self.maybe_recompute_fanout(); + if self.eager_peers.len() < self.num_eager() { + self.move_to_eager(&peer, rt); + } else { + self.insert_into_lazy(peer, rt); + } + self.needs_rebalance = true; + } + + /// Run the deferred rebalance if one is due — either a membership + /// change flagged it (`needs_rebalance`), or an overlay peer crossed a + /// bucket boundary (Near/Mid/Far). + pub fn update_peer_topology( + &mut self, + updates: impl IntoIterator, + rt: &mut impl Runtime, + ) { + let mut topology_changed = false; + + let prev_count = self.known_peers().len(); + for (peer, info) in updates { + match self.peer_topology.get(&peer).copied() { + None => { + info!("topology added peer {:?}, new: {:?}", peer, info); + self.commit_topology(peer, info); + topology_changed = true; + } + Some(existing) => { + let old_bucket = RingBucket::of(existing); + let new_bucket = RingBucket::of(info); + + // change topology if we get a ring for the first time + if existing.ring.is_none() && info.ring.is_some() { + self.pending_topology.remove(&peer); + self.peer_topology.insert(peer, info); + if new_bucket != old_bucket { + topology_changed = true; + } + continue; + } + + if new_bucket == old_bucket { + self.pending_topology.remove(&peer); + continue; + } + + let confirmations = match self.pending_topology.get(&peer) { + Some((pending_info, count)) + if RingBucket::of(*pending_info) == new_bucket => + { + count + 1 + } + _ => 1, + }; + + if confirmations >= RING_EXTRA_CONFIRMATIONS { + info!( + "topology changed for peer {:?}, old: {:?}, new: {:?} (stable for {confirmations} runs)", + peer, existing, info + ); + self.commit_topology(peer, info); + topology_changed = true; + } else { + // Not yet stable, record the candidate and wait. + self.pending_topology.insert(peer, (info, confirmations)); + } + } + } + } + + if self.known_peers().len() != prev_count { + self.maybe_recompute_fanout(); + } + + if self.needs_rebalance || topology_changed { + info!( + "rebalancing peers: topology_changed: {topology_changed}, needs rebalance: {:?}", + self.needs_rebalance + ); + self.rebalance(rt); + self.needs_rebalance = false; + } + } + + fn commit_topology(&mut self, peer: N, info: RttInfo) { + self.peer_topology.insert(peer, info); + self.known_peers.insert(peer); + self.pending_topology.remove(&peer); + } + + pub fn peer_down(&mut self, peer: &N, _rt: &mut impl Runtime) { + let was_eager = self.eager_peers.remove(peer); + let was_lazy = self.lazy_peers.swap_remove(peer); + self.known_peers.remove(peer); + self.ring_locked.remove(peer); + self.pending_topology.remove(peer); + if was_eager || was_lazy || self.maybe_recompute_fanout() { + self.needs_rebalance = true; + } + } + + // --- Rebalance --- + + /// Clears existing lazy / eager peers and selects them again. + /// + /// Ring neighbors are always eager. Remaining eager slots are split across + /// near, mid, and far RTT bucket. Locked peers count toward the bucket targets. + /// + /// todo: maybe re-implement rebalance so we don't clear exisiting eager/lazy peers. + fn rebalance(&mut self, rt: &mut impl Runtime) { + rt.notify(Notification::Rebalance); + + self.eager_peers.clear(); + self.lazy_peers.clear(); + self.ring_locked.clear(); + + if self.known_peers.is_empty() { + return; + } + + self.set_ring_neighbors(); + if self.known_peers.len() <= self.num_eager() { + self.eager_peers.extend(self.known_peers.iter().copied()); + return; + } + + // ring-locked peers are always eager so a node is never isolated. + // count them per bucket so they count toward that bucket's target. + let mut locked_near = 0; + let mut locked_mid = 0; + let mut locked_far = 0; + for p in self.ring_locked.iter().copied() { + match self.peer_bucket(&p) { + RingBucket::Near => locked_near += 1, + RingBucket::Mid => locked_mid += 1, + RingBucket::Far => locked_far += 1, + } + self.eager_peers.insert(p); + } + + let num_eager = self.num_eager(); + + // pools exclude ring_locked peers so they aren't reselected. + let (near_pool, mid_pool, far_pool) = self.bucket_pools(); + + // Near/Mid/Far split for eager and lazy peer selection (see `EagerRatios`). + let (near_t, mid_t, far_t) = Self::bucket_targets( + num_eager, + near_pool.len() + locked_near, + mid_pool.len() + locked_mid, + far_pool.len() + locked_far, + self.config.eager_ratios, + ); + let eager_near = near_t.saturating_sub(locked_near); + let eager_mid = mid_t.saturating_sub(locked_mid); + let eager_far = far_t.saturating_sub(locked_far); + + self.eager_peers + .extend(near_pool.iter().take(eager_near).copied()); + self.eager_peers + .extend(mid_pool.iter().take(eager_mid).copied()); + self.eager_peers + .extend(far_pool.iter().take(eager_far).copied()); + + // same selection logic for lazy peers since they are eager candidates. + // pools already exclude ring_locked, so lazy can never grab a locked peer. + let num_lazy = self.min_lazy(); + let (lazy_near, lazy_mid, lazy_far) = Self::bucket_targets( + num_lazy, + near_pool.len().saturating_sub(eager_near), + mid_pool.len().saturating_sub(eager_mid), + far_pool.len().saturating_sub(eager_far), + self.config.eager_ratios, + ); + + self.lazy_peers + .extend(near_pool.iter().skip(eager_near).take(lazy_near).copied()); + self.lazy_peers + .extend(mid_pool.iter().skip(eager_mid).take(lazy_mid).copied()); + self.lazy_peers + .extend(far_pool.iter().skip(eager_far).take(lazy_far).copied()); + trace!( + self_actor_id = ?self.local_id, + eager = self.eager_peers.len(), + lazy = self.lazy_peers.len(), + ring_locked = self.ring_locked.len(), + known = self.known_peers.len(), + "rebalance complete (eager: {:?}, lazy: {:?})", self.eager_peers, self.lazy_peers + ); + } + + fn bucket_pools(&mut self) -> (Vec, Vec, Vec) { + let mut near = Vec::new(); + let mut mid = Vec::new(); + let mut far = Vec::new(); + for p in self.known_peers.iter() { + if self.eager_peers.contains(p) || self.lazy_peers.contains(p) { + continue; + } + match self.peer_bucket(p) { + RingBucket::Near => near.push(*p), + RingBucket::Mid => mid.push(*p), + RingBucket::Far => far.push(*p), + } + } + + near.shuffle(&mut self.rng); + mid.shuffle(&mut self.rng); + far.shuffle(&mut self.rng); + (near, mid, far) + } + + fn bucket_targets( + num_eager: usize, + near_cap: usize, + mid_cap: usize, + far_cap: usize, + ratios: EagerRatios, + ) -> (usize, usize, usize) { + debug_assert!(ratios.validate().is_ok()); + let near_pct = ratios.near_pct as usize; + let mid_pct = ratios.mid_pct as usize; + let near = (num_eager * near_pct).div_ceil(100).min(near_cap); + let mid = (num_eager * mid_pct).div_ceil(100).min(mid_cap); + let far = num_eager.saturating_sub(near + mid).min(far_cap); + + // update near and mid, if we still have some room + let near = near_cap.min(num_eager.saturating_sub(mid + far)); + let mid = mid_cap.min(num_eager.saturating_sub(near + far)); + (near, mid, far) + } + + fn peer_bucket(&self, p: &N) -> RingBucket { + RingBucket::of(self.peer_topology.get(p).copied().unwrap_or_default()) + } + + fn enqueue_ihave(&mut self, id: I, round: Round) { + let digest = IHaveDigest { id, round }; + self.lazy_queue.push(digest); + } + + fn drain_lazy_queue(&mut self) -> Option>> { + if self.lazy_queue.is_empty() { + return None; + } + Some(std::mem::take(&mut self.lazy_queue)) + } + + /// Recompute ring neighbors based on current known_peers. + /// + /// Clears old ring_locked set and recalculates the two peers that + /// are immediately before and after local_id in sorted order. + fn set_ring_neighbors(&mut self) { + self.ring_locked.clear(); + + let mut peers: Vec<_> = self.known_peers.iter().collect(); + peers.push(&self.local_id); + if peers.len() <= 1 { + return; + } + + peers.sort(); + if let Some(position) = peers.iter().position(|p| *p == &self.local_id) { + let len = peers.len(); + let after_idx = (position + 1) % len; + let before_idx = (position + len - 1) % len; + + self.ring_locked.insert(*peers[before_idx]); + self.ring_locked.insert(*peers[after_idx]); + } + } + + fn move_to_eager(&mut self, peer: &N, rt: &mut impl Runtime) { + if self.eager_peers.contains(peer) || !self.known_peers.contains(peer) { + return; + } + + self.eager_peers.insert(*peer); + let was_lazy = self.lazy_peers.swap_remove(peer); + trace!( + self_actor_id = ?self.local_id, + ?peer, + "peer moved to eager (was_lazy: {was_lazy})" + ); + rt.notify(Notification::PeerMovedToEager(peer)); + } + + fn move_to_lazy(&mut self, peer: &N, rt: &mut impl Runtime) { + if self.ring_locked.contains(peer) { + warn!( + self_actor_id = ?self.local_id, + ?peer, + "move_to_lazy skipped (ring-locked eager peer)" + ); + return; + } + + if self.lazy_peers.contains(peer) || !self.known_peers.contains(peer) { + return; + } + + let was_eager = self.eager_peers.remove(peer); + if was_eager { + rt.notify(Notification::PeerDroppedFromEager(peer)); + } + + self.insert_into_lazy(*peer, rt); + trace!( + self_actor_id = ?self.local_id, + ?peer, + "peer moved to lazy: (was_eager: {was_eager})" + ); + } + + /// Ensure a peer is at least in the lazy set so they receive IHave + /// digests. Does nothing if the peer is already in eager or lazy. + fn ensure_in_lazy(&mut self, peer: &N, rt: &mut impl Runtime) { + if self.eager_peers.contains(peer) + || self.lazy_peers.contains(peer) + || !self.known_peers.contains(peer) + { + return; + } + self.insert_into_lazy(*peer, rt); + } + + /// Insert a peer into the lazy set, evicting a random non-ring-locked + /// peer if at capacity. Returns true if the peer was inserted. + /// + fn insert_into_lazy(&mut self, peer: N, rt: &mut impl Runtime) -> bool { + if self.lazy_peers.contains(&peer) { + return true; + } + + if self.lazy_peers.len() < self.max_lazy() { + self.lazy_peers.insert(peer); + rt.notify(Notification::PeerMovedToLazy(&peer)); + return true; + } + + // Lazy is full — evict a random lazy peer. + let idx = self.rng.random_range(0..self.lazy_peers().len()); + let evicted = self.lazy_peers.swap_remove_index(idx); + + if let Some(evicted) = evicted { + trace!( + self_actor_id = ?self.local_id, + ?peer, + "inserted into lazy, evicted peer {evicted:?} to make room" + ); + self.lazy_peers.insert(peer); + rt.notify(Notification::PeerMovedToLazy(&peer)); + rt.notify(Notification::PeerEvictedFromLazy(&evicted)); + true + } else { + debug!( + self_actor_id = ?self.local_id, + ?peer, + "unable to insert into lazy, peer dropped" + ); + false + } + } + + pub fn cache_evict_if_needed(&mut self, rt: &mut impl Runtime) { + self.seen.evict_if_needed(); + self.cache.evict_if_needed(); + self.prune_throttle.clear_expired(rt.now()); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::iter; + + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + pub(crate) struct TestMsgId(u8); + + pub(crate) type TestNodeId = u8; + + #[derive(Debug, Clone, PartialEq)] + pub(crate) struct TestPayload(pub Vec); + + impl Payload for TestPayload { + type MessageId = TestMsgId; + type NodeId = TestNodeId; + fn message_id(&self) -> Self::MessageId { + TestMsgId(self.0[1]) + } + + fn origin(&self) -> Self::NodeId { + self.0[0] + } + } + + #[derive(Debug, Clone, Default)] + struct TestSeenStore { + entries: HashMap, + } + + impl SeenStore for TestSeenStore { + fn evict_if_needed(&mut self) { + // no-op + } + + fn size(&self) -> usize { + self.entries.len() + } + + fn contains(&self, id: &TestMsgId) -> bool { + self.entries.contains_key(id) + } + + fn observe(&mut self, id: TestMsgId, round: Round) -> Option { + let existing = self.entries.get_mut(&id); + if let Some((_, seen)) = existing { + // *existing = round; + *seen += 1; + return Some(*seen); + } + + self.entries.insert(id, (round, 1)); + None + } + } + + pub(crate) fn test_config() -> Config { + Config { + ihave_timeout: Duration::from_secs(3), + optimization_threshold: Some(3), + max_cached_payloads: 128, + num_eager: Some(5), + min_lazy: Some(10), + max_lazy: Some(15), + prune_threshold: 1, + max_received_entries: 10000, + prune_throttle: None, + eager_ratios: EagerRatios::default(), + } + } + + fn graft_msg( + sender: TestNodeId, + send: bool, + requests: Vec<(TestMsgId, Round)>, + ) -> GraftMsg { + GraftMsg { + sender, + send, + requests: requests + .into_iter() + .map(|(id, round)| GraftRequest { id, round }) + .collect(), + } + } + + fn state() -> PlumtreeState { + PlumtreeState::new_with_store(0u8, test_config(), TestSeenStore::default()) + } + + /// Accumulates all runtime calls for assertion in tests. + #[derive(Debug, Default)] + pub(crate) struct AccumulatingRuntime { + pub sent: Vec<(TestNodeId, PlumtreeMsg)>, + pub delivered: Vec, + pub scheduled: Vec<(Timer, Duration)>, + } + + impl AccumulatingRuntime { + fn clear(&mut self) { + self.sent.clear(); + self.delivered.clear(); + self.scheduled.clear(); + } + } + + impl Runtime for AccumulatingRuntime { + fn send_all( + &mut self, + peers: Vec, + msg: PlumtreeMsg, + _prio: PlumPrio, + ) { + for peer in peers { + self.send(peer, msg.clone(), PlumPrio::P0); + } + } + + fn send( + &mut self, + to: TestNodeId, + msg: PlumtreeMsg, + _prio: PlumPrio, + ) { + self.sent.push((to, msg)); + } + + fn deliver(&mut self, payload: TestPayload) { + self.delivered.push(payload); + } + + fn schedule(&mut self, timer: Timer, after: Duration) { + self.scheduled.push((timer, after)); + } + + fn notify(&mut self, _notification: Notification<'_, TestMsgId, TestNodeId>) {} + + fn now(&self) -> Instant { + Instant::now() + } + } + + fn payload(node: u8, v: u8) -> TestPayload { + TestPayload(vec![node, v]) + } + + /// Extract the inner GossipMsg from a PlumtreeMsg, panicking otherwise. + fn unwrap_gossip( + m: &PlumtreeMsg, + ) -> &GossipMsg { + match m { + PlumtreeMsg::Gossip(g) => g, + other => panic!("expected Gossip, got {other:?}"), + } + } + + fn unwrap_prune( + m: &PlumtreeMsg, + ) -> &PruneMsg { + match m { + PlumtreeMsg::Prune(p) => p, + other => panic!("expected Prune, got {other:?}"), + } + } + + fn unwrap_graft( + m: &PlumtreeMsg, + ) -> &GraftMsg { + match m { + PlumtreeMsg::Graft(g) => g, + other => panic!("expected Graft, got {other:?}"), + } + } + + fn unwrap_ihave( + m: &PlumtreeMsg, + ) -> &IHaveMsg { + match m { + PlumtreeMsg::IHave(ih) => ih, + other => panic!("expected IHave, got {other:?}"), + } + } + + #[test] + fn derived_fanout_updates_with_cluster_growth() { + let mut cfg = test_config(); + cfg.num_eager = None; + cfg.min_lazy = None; + cfg.max_lazy = None; + let mut s = PlumtreeState::new_with_store(0u8, cfg, TestSeenStore::default()); + let mut rt = AccumulatingRuntime::default(); + + assert_eq!(s.num_eager(), 3); + assert_eq!(s.min_lazy(), 4); + assert_eq!(s.max_lazy(), 6); + + // 99 known peers -> cluster size 100 -> log10(100)*3 = 6 eager. + let peers: Vec<_> = (1..=99).collect(); + s.add_peers_bulk(peers, &mut rt); + assert_eq!(s.known_peers().len(), 99); + assert_eq!(s.num_eager(), 6); + assert_eq!(s.min_lazy(), 9); + assert_eq!(s.max_lazy(), 12); + assert!(s.eager_peers().len() <= s.num_eager()); + } + + #[test] + fn test_membership_events() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + s.peer_up(1, None, &mut rt); + assert!(s.eager_peers.contains(&1)); + assert!(!s.lazy_peers.contains(&1)); + + for i in 1..=5 { + s.peer_up(i, None, &mut rt); + } + assert_eq!(s.eager_peers.len(), 5); + + // lazy sets gets peers when eager is full + s.peer_up(6, None, &mut rt); + assert!(s.known_peers.contains(&6)); + // Sixth peer may be ring-locked (must be eager) or placed in lazy. + assert!(s.eager_peers.contains(&6) || s.lazy_peers.contains(&6)); + + let prev_eager = s.eager_peers().clone(); + let prev_lazy = s.lazy_peers().clone(); + // duplicate peer up events should be ignored + s.peer_up(1, None, &mut rt); + assert_eq!(s.eager_peers().clone(), prev_eager); + assert_eq!(s.lazy_peers().clone(), prev_lazy); + + // peer_up only flags a rebalance; flush it so ring neighbors are picked. + s.update_peer_topology(iter::empty::<(TestNodeId, RttInfo)>(), &mut rt); + + // test that ring-locked peers are eager + // since id is zero, ring locked peers are 1 and 6 + assert!(s.ring_locked_peers().contains(&1)); + assert!(s.ring_locked_peers().contains(&6)); + + assert!(s.eager_peers().contains(&1)); + assert!(s.eager_peers().contains(&6)); + assert!(!s.lazy_peers().contains(&1)); + assert!(!s.lazy_peers().contains(&6)); + + // new peer that takes a locked position gets ring-locked after rebalance + s.peer_up(7, None, &mut rt); + s.update_peer_topology(iter::empty::<(TestNodeId, RttInfo)>(), &mut rt); + assert!(s.ring_locked_peers().contains(&7)); + assert!(s.eager_peers().contains(&7)); + assert!(!s.lazy_peers().contains(&7)); + assert_eq!(s.known_peers().len(), 7); + + // removal of a ring-locked peer flags a rebalance; flush selects new ones + s.peer_down(&7, &mut rt); + s.update_peer_topology(iter::empty::<(TestNodeId, RttInfo)>(), &mut rt); + assert!(!s.ring_locked_peers().contains(&7)); + assert!(!s.eager_peers().contains(&7)); + assert!(!s.lazy_peers().contains(&7)); + assert_eq!(s.known_peers().len(), 6); + + // prunes should never delete ring-locked peers + s.handle_prune( + PruneMsg { + sender: 1, + triggered_by: Some(TestMsgId(0)), + }, + &mut rt, + ); + assert!(s.eager_peers().contains(&1)); + + // prunes should move non-ring-locked peers to lazy + let rand_eager = *s + .eager_peers() + .iter() + .find(|p| !s.ring_locked_peers().contains(p)) + .unwrap(); + s.handle_prune( + PruneMsg { + sender: rand_eager, + triggered_by: Some(TestMsgId(0)), + }, + &mut rt, + ); + assert!(!s.eager_peers().contains(&rand_eager)); + assert!(s.lazy_peers().contains(&rand_eager)); + + // peer down removes member from any internal sets + s.peer_down(&6, &mut rt); + assert!(!s.eager_peers().contains(&6)); + assert!(!s.lazy_peers().contains(&6)); + assert_eq!(s.known_peers().len(), 5); + } + + #[test] + fn peer_up_both_full_ignored() { + let mut cfg = test_config(); + cfg.num_eager = Some(2); + cfg.min_lazy = Some(2); + cfg.max_lazy = Some(2); + let mut s: PlumtreeState = + PlumtreeState::new_with_store(0u8, cfg, TestSeenStore::default()); + let mut rt = AccumulatingRuntime::default(); + + s.peer_up(1, None, &mut rt); // eager + s.peer_up(2, None, &mut rt); // eager + s.peer_up(3, None, &mut rt); // lazy + s.peer_up(4, None, &mut rt); // lazy + s.peer_up(5, None, &mut rt); // one peer remains only in known_peers (caps exceeded) + + assert_eq!(s.known_peers.len(), 5); + assert_eq!(s.eager_peers.len(), 2); + assert_eq!(s.lazy_peers.len(), 2); + let orphan = s + .known_peers + .iter() + .filter(|p| !s.eager_peers.contains(*p) && !s.lazy_peers.contains(*p)) + .count(); + assert_eq!(orphan, 1); + } + + #[test] + fn prune_drops_eager_peer() { + let mut cfg = test_config(); + cfg.num_eager = Some(3); + cfg.min_lazy = Some(5); + cfg.max_lazy = Some(5); + let mut s: PlumtreeState = + PlumtreeState::new_with_store(0u8, cfg, TestSeenStore::default()); + let mut rt = AccumulatingRuntime::default(); + + s.peer_up(1, None, &mut rt); // eager + s.peer_up(2, None, &mut rt); // eager + s.peer_up(3, None, &mut rt); // eager (full) + s.peer_up(4, None, &mut rt); // lazy + s.peer_up(5, None, &mut rt); // lazy + s.update_peer_topology(iter::empty::<(TestNodeId, RttInfo)>(), &mut rt); + assert_eq!(s.eager_peers.len(), 3); + assert_eq!(s.lazy_peers.len(), 2); + + // Receive a PRUNE from the non-locked eager peer (who that is depends on shuffle). + let prune_sender = *s + .eager_peers + .iter() + .find(|p| !s.ring_locked_peers().contains(*p)) + .expect("one non-locked eager slot-filled peer"); + assert!(!s.ring_locked_peers().contains(&prune_sender)); + s.handle_prune( + PruneMsg { + sender: prune_sender, + triggered_by: Some(TestMsgId(99)), + }, + &mut rt, + ); + assert_eq!(s.eager_peers.len(), 2, "peer should be dropped from eager"); + assert_eq!(s.lazy_peers.len(), 3); + assert!(!s.eager_peers.contains(&prune_sender)); + assert!(s.lazy_peers.contains(&prune_sender)); + + // prune for ring-locked peer shouldn't succeed + s.handle_prune( + PruneMsg { + sender: 1, + triggered_by: Some(TestMsgId(99)), + }, + &mut rt, + ); + assert!(s.eager_peers().contains(&1)); + } + + #[test] + fn broadcast_sends_to_eager_and_enqueues_lazy() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + s.peer_up(1, None, &mut rt); // eager + s.peer_up(2, None, &mut rt); // eager + // manually move 3 to lazy + s.lazy_peers.insert(3); + + s.broadcast(TestMsgId(42), payload(0, 42), &mut rt); + + // Should have sent GOSSIP to eager peers 1 and 2 + assert_eq!(rt.sent.len(), 2); + for (to, m) in &rt.sent { + let g = unwrap_gossip(m); + assert_eq!(g.payload.message_id(), TestMsgId(42)); + assert_eq!(g.round, 0); + assert_eq!(g.sender, 0); // local_id + assert!(s.eager_peers.contains(to)); + } + + // Lazy queue for peer 3 should have one digest + assert_eq!(s.lazy_queue.len(), 1); + + // Message should be marked received + assert!(s.has_message(&TestMsgId(42))); + + // Receive a GOSSIP from peer 4 (not yet in our peer set) + rt.clear(); + s.handle_gossip( + GossipMsg { + round: 1, + sender: 1, + payload: payload(1, 10), + }, + &mut rt, + ); + // Delivered once + assert_eq!(rt.delivered.len(), 1); + assert_eq!(rt.delivered[0].message_id(), TestMsgId(10)); + + // Forwarded to eager peers 2, 3 (not sender 1) + let gossip_targets: Vec = rt + .sent + .iter() + .filter(|(_, m)| matches!(m, PlumtreeMsg::Gossip(_))) + .map(|(to, _)| *to) + .collect(); + assert!(gossip_targets.contains(&2)); + assert!(!gossip_targets.contains(&1)); // not back to sender + + // Forwarded gossips have round + 1 + for (_, m) in &rt.sent { + if let PlumtreeMsg::Gossip(g) = m { + assert_eq!(g.round, 2); + assert_eq!(g.sender, 0); + } + } + + assert!(!s.lazy_peers.contains(&1)); + // lazy queue + assert_eq!(s.lazy_queue.len(), 2); + assert!(s.has_message(&TestMsgId(10))); + + rt.clear(); + + // unknown peer message gets processed but isn't added to peer set + s.handle_gossip( + GossipMsg { + round: 0, + sender: 5, + payload: payload(5, 15), + }, + &mut rt, + ); + assert_eq!(s.lazy_queue.len(), 3); + assert!(s.has_message(&TestMsgId(15))); + assert!(!s.lazy_peers.contains(&5)); + + // duplicate gossip gets pruned + rt.clear(); + s.handle_gossip( + GossipMsg { + round: 1, + sender: 11, + payload: payload(5, 15), + }, + &mut rt, + ); + + // No delivery + assert!(rt.delivered.is_empty()); + + // PRUNE sent to duplicate sender 11 + assert_eq!(rt.sent.len(), 1); + let (to, m) = &rt.sent[0]; + assert_eq!(*to, 11); + let prune = unwrap_prune(m); + assert_eq!(prune.sender, 0); + assert_eq!(prune.triggered_by, Some(TestMsgId(15))); + + assert!(!s.lazy_peers.contains(&11)); + assert!(!s.eager_peers.contains(&11)); + } + + #[test] + fn handle_gossip_optimization_grafts_shorter_path() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + // optimization_threshold = 3 + + // Peer 5 tells us about msg(42) via IHave at round 1 + s.peer_up(5, None, &mut rt); + s.handle_ihave( + IHaveMsg { + sender: 5, + digests: vec![ + IHaveDigest { + id: TestMsgId(42), + round: 1, + }, + IHaveDigest { + id: TestMsgId(45), + round: 2, + }, + ], + }, + &mut rt, + ); + assert!(s.missing.contains_key(&TestMsgId(42))); + rt.clear(); + + // Now the GOSSIP arrives from peer 1 at round 10 (1 + 3 < 10) + s.peer_up(1, None, &mut rt); + s.handle_gossip( + GossipMsg { + round: 10, + sender: 1, + payload: payload(10, 42), + }, + &mut rt, + ); + + // Should have sent GRAFT to peer 5 (shorter path) + let grafts: Vec<_> = rt + .sent + .iter() + .filter(|(_, m)| matches!(m, PlumtreeMsg::Graft(_))) + .collect(); + assert_eq!(grafts.len(), 1); + let (to, m) = &grafts[0]; + assert_eq!(*to, 5); + let graft = unwrap_graft(m); + assert_eq!(graft.requests[0].round, 1); + assert!(!graft.send); + + // Peer 5 promoted to eager + assert!(s.eager_peers.contains(&5)); + + // Missing entry removed + assert!(!s.missing.contains_key(&TestMsgId(42))); + + // check that nothing gets triggered if round is very close + rt.clear(); + s.handle_gossip( + GossipMsg { + round: 3, + sender: 1, + payload: payload(10, 45), + }, + &mut rt, + ); + + let grafts: Vec<_> = rt + .sent + .iter() + .filter(|(_, m)| matches!(m, PlumtreeMsg::Graft(_))) + .collect(); + + assert!(grafts.is_empty()); + assert!(!s.missing.contains_key(&TestMsgId(45))); + } + + // ----------------------------------------------------------------------- + // handle_ihave + // ----------------------------------------------------------------------- + + #[test] + fn handle_ihave_schedules_timer() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + s.seen.observe(TestMsgId(3), 0); + s.handle_ihave( + IHaveMsg { + sender: 5, + digests: vec![ + IHaveDigest { + id: TestMsgId(1), + round: 0, + }, + IHaveDigest { + id: TestMsgId(2), + round: 1, + }, + IHaveDigest { + id: TestMsgId(3), + round: 1, + }, + ], + }, + &mut rt, + ); + + assert_eq!(s.missing.len(), 2); + assert_eq!(rt.scheduled.len(), 1); + // already received changes aren't added to missing + assert!(!s.missing.contains_key(&TestMsgId(3))); + assert_eq!( + rt.scheduled[0].0, + Timer::IHaveTimeoutBatch { + ids: vec![TestMsgId(1), TestMsgId(2)], + retries: 0, + senders: vec![5], + } + ); + assert_eq!(rt.scheduled[0].1, Duration::from_secs(3)); + } + + #[test] + fn handle_graft_sends_cached_payload() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + // First, broadcast so the payload is cached + s.broadcast(TestMsgId(1), payload(42, 1), &mut rt); + rt.clear(); + + // Peer 5 sends GRAFT + s.known_peers.insert(5); + s.handle_graft(graft_msg(5, true, vec![(TestMsgId(1), 0)]), &mut rt); + + // Peer 5 promoted to eager + assert!(s.eager_peers.contains(&5)); + + // GOSSIP sent back with cached payload + assert_eq!(rt.sent.len(), 1); + let (to, m) = &rt.sent[0]; + assert_eq!(*to, 5); + let g = unwrap_gossip(m); + assert_eq!(g.payload, payload(42, 1)); + + rt.clear(); + s.seen.observe(TestMsgId(99), 0); + s.handle_graft(graft_msg(5, true, vec![(TestMsgId(99), 0)]), &mut rt); + + // seen but not cached — graft must not send gossip back + assert!(s.has_message(&TestMsgId(99))); + assert!(rt.sent.is_empty()); + } + + #[test] + fn handle_ihave_timeout_partial_noop() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + s.peer_up(5, None, &mut rt); + s.handle_ihave( + IHaveMsg { + sender: 5, + digests: vec![ + IHaveDigest { + id: TestMsgId(1), + round: 0, + }, + IHaveDigest { + id: TestMsgId(2), + round: 0, + }, + ], + }, + &mut rt, + ); + s.seen.observe(TestMsgId(1), 0); + rt.sent.clear(); + rt.scheduled.clear(); + + s.timer_fired( + Timer::IHaveTimeoutBatch { + ids: vec![TestMsgId(1), TestMsgId(2)], + retries: 0, + senders: vec![5, 6], + }, + &mut rt, + ); + + assert_eq!(rt.sent.len(), 1); + { + let (to, m) = &rt.sent[0]; + let graft = unwrap_graft(m); + assert_eq!(*to, 5); + assert_eq!(graft.requests.len(), 1); + assert_eq!(graft.requests[0].id, TestMsgId(2)); + assert_eq!(graft.requests[0].round, 0); + } + + assert!(s.eager_peers.contains(&5)); + assert!(!s.missing.contains_key(&TestMsgId(1))); + assert!(s.missing.contains_key(&TestMsgId(2))); + assert_eq!(rt.scheduled.len(), 1); + // schedule retry + let scheduled = Timer::IHaveTimeoutBatch { + ids: vec![TestMsgId(2)], + retries: 1, + senders: vec![5, 6], + }; + assert_eq!(rt.scheduled[0].0, scheduled,); + + rt.sent.clear(); + rt.scheduled.clear(); + + s.timer_fired(scheduled, &mut rt); + + let graft = unwrap_graft(&rt.sent[0].1); + assert_eq!(graft.requests.len(), 1); + assert_eq!(rt.sent[0].0, 6); + assert!(!s.missing.contains_key(&TestMsgId(2))); + assert_eq!(rt.scheduled.len(), 0); + } + + #[test] + fn timer_fired_ihave_timeout_chunks_graft_requests() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + let digests: Vec<_> = (1..=15) + .map(|n| IHaveDigest { + id: TestMsgId(n), + round: 0, + }) + .collect(); + + s.handle_ihave(IHaveMsg { sender: 5, digests }, &mut rt); + rt.sent.clear(); + rt.scheduled.clear(); + + let ids: Vec<_> = (1..=15).map(TestMsgId).collect(); + s.timer_fired( + Timer::IHaveTimeoutBatch { + ids, + retries: 0, + senders: vec![5], + }, + &mut rt, + ); + + assert_eq!(rt.sent.len(), 2); + let g0 = unwrap_graft(&rt.sent[0].1); + let g1 = unwrap_graft(&rt.sent[1].1); + assert_eq!(g0.requests.len(), 10); + assert_eq!(g1.requests.len(), 5); + assert!(rt.sent.iter().all(|(to, _)| *to == 5)); + } + + #[test] + fn timer_fired_noop_if_already_received() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + s.handle_ihave( + IHaveMsg { + sender: 5, + digests: vec![IHaveDigest { + id: TestMsgId(1), + round: 0, + }], + }, + &mut rt, + ); + + // GOSSIP arrives before timer fires, removing from missing + s.handle_gossip( + GossipMsg { + round: 0, + sender: 2, + payload: payload(2, 1), + }, + &mut rt, + ); + rt.sent.clear(); + + // Timer fires — but missing entry already removed + s.timer_fired( + Timer::IHaveTimeoutBatch { + ids: vec![TestMsgId(1)], + retries: 0, + senders: vec![5], + }, + &mut rt, + ); + assert!(rt.sent.is_empty()); + } + + #[test] + fn tick_flushes_lazy_queues() { + let mut s = state(); + let mut rt = AccumulatingRuntime::default(); + + s.lazy_peers.insert(3); + s.lazy_peers.insert(4); + + s.enqueue_ihave(TestMsgId(1), 0); + s.enqueue_ihave(TestMsgId(2), 1); + + s.tick(&mut rt); + + // IHave sent to both lazy peers + assert_eq!(rt.sent.len(), 2); + for (to, m) in &rt.sent { + let ih = unwrap_ihave(m); + assert_eq!(ih.sender, 0); + assert_eq!(ih.digests.len(), 2); + assert!(s.lazy_peers.contains(to)); + } + + // Queue is drained + assert!(s.drain_lazy_queue().is_none()); + } + + #[test] + fn rebalance_uses_weighted_ring_buckets() { + let mut cfg = test_config(); + cfg.num_eager = Some(8); + cfg.min_lazy = Some(10); + cfg.max_lazy = Some(10); + let mut s = PlumtreeState::new_with_store(0u8, cfg, TestSeenStore::default()); + let mut rt = AccumulatingRuntime::default(); + + let peers = (1u8..=12) + .map(|peer| { + let ring = match peer { + 1..=4 => Some(0), + 5 => Some(2), + _ => Some(5), + }; + (peer, RttInfo { ring }) + }) + .collect(); + s.add_peers_bulk_with_rtt(peers, &mut rt); + + assert_eq!(s.eager_peers.len(), 8); + assert!( + s.ring_locked_peers() + .iter() + .all(|p| s.eager_peers.contains(p)) + ); + + let near = s + .eager_peers + .iter() + .filter(|p| s.peer_bucket(p) == RingBucket::Near) + .count(); + let mid = s + .eager_peers + .iter() + .filter(|p| s.peer_bucket(p) == RingBucket::Mid) + .count(); + let far = s + .eager_peers + .iter() + .filter(|p| s.peer_bucket(p) == RingBucket::Far) + .count(); + assert!(near == 4, "expected near slots to be favored: {near}"); + assert!( + mid == 1, + "expected mid bucket coverage to have just one peer: {mid}" + ); + assert!(far == 3, "expected far bucket coverage: {far}"); + + // rest should get added to lazy + assert_eq!(s.lazy_peers().len(), 4); + } + + #[test] + fn move_to_lazy_respects_max_lazy() { + let mut cfg = test_config(); + cfg.num_eager = Some(3); + // Need room for both prior lazy (6) and demoted peer (5) so 6 is not evicted from lazy. + cfg.min_lazy = Some(2); + cfg.max_lazy = Some(2); + let mut s = PlumtreeState::new_with_store(0u8, cfg, TestSeenStore::default()); + let mut rt = AccumulatingRuntime::default(); + + // Ring locks 7 and 4; pool {5,6} → one eager, one lazy (shuffle picks which). + s.peer_up(4, None, &mut rt); + s.peer_up(5, None, &mut rt); + s.peer_up(6, None, &mut rt); + s.peer_up(7, None, &mut rt); + + let dup_sender = *s + .eager_peers + .iter() + .find(|p| !s.ring_locked_peers().contains(*p)) + .expect("one non-locked eager"); + assert!(s.lazy_peers().len() == 1); + let lazy_peer = s.lazy_peers()[0]; + assert_ne!(dup_sender, lazy_peer); + + // Duplicate gossip from non-locked eager → demote; + s.seen.observe(TestMsgId(1), 0); + s.handle_gossip( + GossipMsg { + round: 0, + sender: dup_sender, + payload: payload(dup_sender, 1), + }, + &mut rt, + ); + + assert!(!s.eager_peers.contains(&dup_sender)); + assert!(s.lazy_peers.contains(&dup_sender)); + } + + #[test] + fn full_prune_graft_cycle() { + // Simulate: A broadcasts, B and C both receive from A (eager). + // B also receives a duplicate from C → B sends PRUNE to C. + // Later C has a message B missed → B GRAFTs C back. + + // Node B (id=2). Ring [1,2,4,10,30]: neighbors of 2 are 1 and 4 — C (30) is not locked. + let cfg = test_config(); + let mut b = + PlumtreeState::::new_with_store( + 2, + cfg.clone(), + TestSeenStore::default(), + ); + let mut rt_b = AccumulatingRuntime::default(); + + b.peer_up(1, None, &mut rt_b); // A + b.peer_up(4, None, &mut rt_b); + b.peer_up(10, None, &mut rt_b); + b.peer_up(30, None, &mut rt_b); // C — not successor/predecessor of B on the ring + + // B receives GOSSIP from A + b.handle_gossip( + GossipMsg { + round: 0, + sender: 1, + payload: payload(1, 1), + }, + &mut rt_b, + ); + assert!(b.has_message(&TestMsgId(1))); + rt_b.sent.clear(); + + // B receives duplicate from C → sends PRUNE to C + b.handle_gossip( + GossipMsg { + round: 1, + sender: 30, + payload: payload(1, 1), + }, + &mut rt_b, + ); + assert_eq!(rt_b.sent.len(), 1); + let (to, m) = &rt_b.sent[0]; + assert_eq!(*to, 30); + unwrap_prune(m); // verifies it's a Prune + assert!(b.lazy_peers.contains(&30)); + rt_b.sent.clear(); + + // Now C ticks and sends IHave for TestMsgId(2) to B (lazy peer) + // Simulated: B receives the IHave + b.handle_ihave( + IHaveMsg { + sender: 30, + digests: vec![IHaveDigest { + id: TestMsgId(2), + round: 0, + }], + }, + &mut rt_b, + ); + + // Timer fires — B sends GRAFT to C + b.timer_fired( + Timer::IHaveTimeoutBatch { + ids: vec![TestMsgId(2)], + retries: 0, + senders: vec![30], + }, + &mut rt_b, + ); + assert_eq!(rt_b.sent.len(), 1); + let (to, m) = &rt_b.sent[0]; + assert_eq!(*to, 30); + unwrap_graft(m); + + // C is back in eager + assert!(b.eager_peers.contains(&30)); + assert!(!b.lazy_peers.contains(&30)); + } +} diff --git a/crates/plum-foca/tests/sim.rs b/crates/plum-foca/tests/sim.rs new file mode 100644 index 000000000..4a5258f54 --- /dev/null +++ b/crates/plum-foca/tests/sim.rs @@ -0,0 +1,2320 @@ +//! Deterministic discrete-event simulation of the Plumtree protocol at +//! cluster scale. +//! +//! Instead of spawning tasks and channels, every node is a bare +//! [`PlumtreeState`] and all effects (sends, timers, ticks) are events in a +//! single priority queue ordered by virtual time. Message latency is derived +//! from a geographic model: one (lat, lon) per region, +//! RTT ≈ great-circle-km / 75 (speed of light in fiber + routing overhead), +//! validated against measured region-to-region RTTs. Number of nodes in a re +//! region is weighted. +//! +//! ```text +//! cargo test --release -p plum-foca --test sim -- --ignored --nocapture +//! ``` + +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::fmt; +use std::sync::LazyLock; +use std::time::{Duration, Instant}; + +/// Fixed epoch so the sim can map virtual µs to a monotonic `Instant` for +/// `Runtime::now()` — keeps time-windowed protocol logic deterministic. +static SIM_EPOCH: LazyLock = LazyLock::new(Instant::now); + +use indexmap::IndexMap; +use plum_foca::{ + Config, EagerRatios, Notification, Payload, PlumPrio, PlumtreeMsg, PlumtreeState, Round, + RttInfo, Runtime, SeenStore, Timer, +}; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +type NId = u32; +type MId = u64; + +#[derive(Debug, Clone)] +struct SimPayload { + origin: NId, + seq: u32, +} + +impl SimPayload { + fn id(&self) -> MId { + ((self.origin as u64) << 32) | self.seq as u64 + } +} + +/// Decode the origin node encoded in a `SimPayload::id()`/gossip `mid` +/// (both sims use the same `(origin << 32) | seq` encoding). +fn origin_of(mid: MId) -> NId { + (mid >> 32) as NId +} + +impl Payload for SimPayload { + type MessageId = MId; + type NodeId = NId; + + fn message_id(&self) -> MId { + self.id() + } + + fn origin(&self) -> NId { + self.origin + } +} + +struct SimSeenStore { + entries: IndexMap, + cap: usize, +} + +impl SimSeenStore { + fn new(cap: usize) -> Self { + Self { + entries: IndexMap::with_capacity(cap.min(64 * 1024)), + cap, + } + } +} + +impl SeenStore for SimSeenStore { + fn evict_if_needed(&mut self) { + if self.entries.len() > self.cap { + let excess = self.entries.len() - self.cap; + self.entries.drain(0..excess); + } + } + + fn contains(&self, id: &MId) -> bool { + self.entries.contains_key(id) + } + + fn observe(&mut self, id: MId, _round: Round) -> Option { + if let Some(count) = self.entries.get_mut(&id) { + *count += 1; + return Some(*count); + } + self.entries.insert(id, 1); + None + } + + fn size(&self) -> usize { + self.entries.len() + } +} + +struct Region { + #[allow(dead_code)] + name: &'static str, + lat: f64, + lon: f64, + weight: u32, +} + +#[rustfmt::skip] +const REGIONS: &[Region] = &[ + Region { name: "ams", lat: 52.31, lon: 4.76, weight: 274 }, + Region { name: "arn", lat: 59.65, lon: 17.92, weight: 20 }, + Region { name: "bom", lat: 19.09, lon: 72.87, weight: 33 }, + Region { name: "cdg", lat: 49.01, lon: 2.55, weight: 175 }, + Region { name: "dfw", lat: 32.90, lon: -97.04, weight: 528 }, + Region { name: "ewr", lat: 40.69, lon: -74.17, weight: 245 }, + Region { name: "fra", lat: 50.03, lon: 8.57, weight: 347 }, + Region { name: "gru", lat: -23.43, lon: -46.47, weight: 111 }, + Region { name: "iad", lat: 38.94, lon: -77.46, weight: 442 }, + Region { name: "jnb", lat: -26.14, lon: 28.25, weight: 16 }, + Region { name: "lax", lat: 33.94, lon: -118.41, weight: 306 }, + Region { name: "lhr", lat: 51.47, lon: -0.45, weight: 74 }, + Region { name: "nrt", lat: 35.77, lon: 140.39, weight: 76 }, + Region { name: "ord", lat: 41.97, lon: -87.91, weight: 559 }, + Region { name: "qmx", lat: 19.43, lon: -99.13, weight: 42 }, + Region { name: "sin", lat: 1.36, lon: 103.99, weight: 209 }, + Region { name: "sjc", lat: 37.36, lon: -121.93, weight: 682 }, + Region { name: "syd", lat: -33.95, lon: 151.18, weight: 200 }, + Region { name: "yyz", lat: 43.68, lon: -79.63, weight: 114 }, +]; + +fn haversine_km(a: &Region, b: &Region) -> f64 { + let (la1, la2) = (a.lat.to_radians(), b.lat.to_radians()); + let dlat = (b.lat - a.lat).to_radians(); + let dlon = (b.lon - a.lon).to_radians(); + let h = (dlat / 2.0).sin().powi(2) + la1.cos() * la2.cos() * (dlon / 2.0).sin().powi(2); + 2.0 * 6371.0 * h.sqrt().asin() +} + +/// RTT between two regions in ms: 4ms routing floor + distance at ~2/3 c. +/// Matches measured Fly region-to-region RTTs within ~15%. +fn region_rtt_ms(a: usize, b: usize) -> u64 { + if a == b { + return 1; + } + (4.0 + haversine_km(®IONS[a], ®IONS[b]) / 75.0) as u64 +} + +/// Mirrors `corro_types::members::RING_BUCKETS`. +const RING_BUCKETS: [std::ops::Range; 6] = [0..6, 6..15, 15..50, 50..100, 100..200, 200..300]; + +fn ring_from_rtt_ms(rtt_ms: u64) -> Option { + RING_BUCKETS + .iter() + .position(|range| range.contains(&rtt_ms)) + .map(|ring| ring as u8) +} + +#[derive(Default)] +struct NotificationStats { + duplicates: u64, + missing: u64, + not_cached: u64, + promotions: u64, + demotions: u64, + prune_suppressed: u64, +} + +#[derive(Default)] +struct SimRuntime { + // Collects outgoing protocol messages + // (DstNodeId, PlumtreeMsg) + outbox: Vec<(NId, PlumtreeMsg)>, + timers: Vec<(Timer, Duration)>, + // Payloads the protocol delivered this step. + inbox: Vec, + notification_stats: NotificationStats, + /// current virtual time in µs, kept in sync with `Sim::now` + now_us: u64, +} + +impl Runtime for SimRuntime { + fn send(&mut self, to: NId, msg: PlumtreeMsg, _prio: PlumPrio) { + self.outbox.push((to, msg)); + } + + fn send_all( + &mut self, + peers: Vec, + msg: PlumtreeMsg, + _prio: PlumPrio, + ) { + for peer in peers { + self.outbox.push((peer, msg.clone())); + } + } + + fn deliver(&mut self, payload: SimPayload) { + self.inbox.push(payload); + } + + fn schedule(&mut self, timer: Timer, after: Duration) { + self.timers.push((timer, after)); + } + + fn notify(&mut self, notification: Notification<'_, MId, NId>) { + match notification { + Notification::DuplicateMessage(_) => self.notification_stats.duplicates += 1, + Notification::MessageMissing(count) => self.notification_stats.missing += count as u64, + Notification::PayloadNotCached(_) => self.notification_stats.not_cached += 1, + Notification::PeerMovedToEager(_) => self.notification_stats.promotions += 1, + Notification::PeerDroppedFromEager(_) => self.notification_stats.demotions += 1, + Notification::PeerMovedToLazy(_) | Notification::PeerEvictedFromLazy(_) => {} + Notification::PruneSuppressed(_) => self.notification_stats.prune_suppressed += 1, + Notification::Rebalance => {} + } + } + + fn now(&self) -> Instant { + *SIM_EPOCH + Duration::from_micros(self.now_us) + } +} + +enum Event { + Recv { + to: NId, + msg: PlumtreeMsg, + }, + TimerFired { + node: NId, + timer: Timer, + }, + /// Flush this node's lazy IHave queue (`PlumtreeState::tick`). + Tick { + node: NId, + }, + Originate { + node: NId, + seq: u32, + }, + /// Node stops responding. `restart_after: None` is a permanent crash; + /// `Some(d)` means it rejoins after `d` with a fresh `PlumtreeState`, + /// like a real process restart. + NodeDown { + node: NId, + restart_after: Option, + }, + /// A previously-down node rejoins (only reachable via `NodeDown { + /// restart_after: Some(_) }`). + NodeUp { + node: NId, + }, + /// `observer`'s foca instance finally notices `subject` went down, after + /// a random per-observer detection delay. + MemberDown { + observer: NId, + subject: NId, + }, + /// `observer`'s foca instance finally notices `subject` came back. + MemberUp { + observer: NId, + subject: NId, + }, +} + +struct Scheduled { + at: u64, + seq: u64, + ev: E, +} + +impl PartialEq for Scheduled { + fn eq(&self, other: &Self) -> bool { + (self.at, self.seq) == (other.at, other.seq) + } +} +impl Eq for Scheduled {} +impl PartialOrd for Scheduled { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for Scheduled { + // reversed so the BinaryHeap pops the earliest event first; seq breaks + // ties deterministically. + fn cmp(&self, other: &Self) -> Ordering { + (other.at, other.seq).cmp(&(self.at, self.seq)) + } +} + +// --------------------------------------------------------------------------- +// Simulation +// --------------------------------------------------------------------------- + +#[derive(Clone)] +struct Params { + n: usize, + num_eager: usize, + msgs_per_node: u32, + /// Window over which broadcasts are staggered (uniformly at random). + broadcast_window: Duration, + /// Give nodes real RTT ring info; `false` leaves every peer in the Far + /// bucket (ring: None), i.e. topology-blind eager selection. + rtt_aware: bool, + /// Probability that any protocol message is dropped in flight. + loss: f64, + /// PRUNE throttle window; `None` = prune every over-threshold duplicate. + prune_throttle: Option, + /// Extra eager peers force-added per node after bootstrap, modelling an + /// over-full eager set (e.g. admission without trimming). 0 = none. + eager_inflate: usize, + seed: u64, + eager_ratios: EagerRatios, + /// Number of randomly chosen nodes that crash permanently partway + /// through the run (never send or receive again). + crash_count: usize, + /// Number of randomly chosen nodes that go down and later rejoin with a + /// fresh `PlumtreeState`, modeling a real process restart. + flap_count: usize, + /// How long a flapping node stays down before rejoining. + flap_downtime: Duration, + /// Range for the random per-observer delay before a `peer_down`/`peer_up` + /// is delivered, modeling staggered SWIM failure detection across the + /// cluster (real nodes don't all notice a change at the same instant). + detection_delay_min: Duration, + detection_delay_max: Duration, +} + +impl Params { + fn new(n: usize, num_eager: usize) -> Self { + Self { + n, + num_eager, + msgs_per_node: 1, + // ~200 broadcasts/s cluster-wide regardless of N + broadcast_window: Duration::from_millis(n as u64 * 5), + rtt_aware: true, + loss: 0.0, + prune_throttle: Some(SIM_PRUNE_THROTTLE), + eager_inflate: 0, + seed: 42, + eager_ratios: EagerRatios::default(), + crash_count: 0, + flap_count: 0, + flap_downtime: Duration::from_secs(5), + detection_delay_min: Duration::from_secs(2), + detection_delay_max: Duration::from_secs(8), + } + } + + fn with_eager_ratios(mut self, eager_ratios: EagerRatios) -> Self { + self.eager_ratios = eager_ratios; + self + } + + /// `crash_count` nodes crash permanently, `flap_count` nodes go down and + /// rejoin after `flap_downtime`; all are chosen randomly and scheduled + /// uniformly over `broadcast_window`, overlapping the broadcast storm. + fn with_failures( + mut self, + crash_count: usize, + flap_count: usize, + flap_downtime: Duration, + ) -> Self { + self.crash_count = crash_count; + self.flap_count = flap_count; + self.flap_downtime = flap_downtime; + self + } + + fn with_detection_delay(mut self, min: Duration, max: Duration) -> Self { + self.detection_delay_min = min; + self.detection_delay_max = max; + self + } + + /// Production config from `spawn_plumtree_loop`, with `num_eager` + /// parametrized and the seen-store sized so it never evicts mid-run. + fn config(&self) -> Config { + Config { + ihave_timeout: Duration::from_millis(200), + optimization_threshold: Some(5), + num_eager: None, + min_lazy: None, + max_lazy: None, + prune_threshold: 5, + max_received_entries: (self.n as u32 * self.msgs_per_node) as usize + 1, + max_cached_payloads: 4096, + prune_throttle: self.prune_throttle, + eager_ratios: self.eager_ratios, + } + } +} + +const TICK_INTERVAL: Duration = Duration::from_millis(200); +/// Default prune throttle in sim: shorter than production (1s) but still +/// spans a few tick intervals at sim event density. +const SIM_PRUNE_THROTTLE: Duration = Duration::from_millis(500); +/// A gossip send crossing a link slower than this counts as long-haul. +const LONGHAUL_RTT_MS: u64 = 150; + +struct MsgStat { + sent_at: u64, + gossip_sends: u64, + deliveries: u32, + last_delivery_at: u64, + /// Nodes that have received this message (Plumtree sim only; used to + /// check survivor delivery under node failures). Always empty in the + /// gossip baseline sim. + delivered_to: HashSet, +} + +#[derive(Default)] +struct Stats { + sent_gossip: u64, + sent_ihave: u64, + sent_graft: u64, + sent_prune: u64, + longhaul_gossip: u64, + dropped: u64, + duplicates: u64, + missing: u64, + not_cached: u64, + promotions: u64, + demotions: u64, + prune_suppressed: u64, + /// One entry per delivery: number of gossip hops from the origin. + hops: Vec, + /// One entry per delivery: broadcast-to-delivery latency in µs. + latencies_us: Vec, +} + +struct Sim { + params: Params, + now: u64, // virtual time, µs + next_seq: u64, + queue: BinaryHeap>, + states: Vec>, + region_of: Vec, + rtt_ms: Vec>, + rng: StdRng, + runtime: SimRuntime, + tick_pending: Vec, + /// Round of the gossip message currently being handled, for hop + /// attribution of the deliveries it produces. + gossip_round: Option, + msgs: HashMap, + stats: Stats, + /// avg eager-set size right after bootstrap+inflation, before the run. + start_avg_eager: f64, + /// Nodes currently unreachable (crashed or mid-flap). + down: Vec, + /// Nodes that crashed permanently (subset of nodes still `down` once the + /// queue fully drains). + crashed: HashSet, + /// Nodes that went down and rejoined at least once. + flapped: HashSet, + /// Ground truth, per recipient, of every message it could never have + /// received because it was down — recorded directly at the moment it's + /// true, not inferred from timing: either the recipient was already + /// down when the message originated, or a real delivery attempt reached + /// it while down and got dropped. See `failure_miss_stats`. + missed_during_downtime: HashMap>, +} + +impl Sim { + fn new(params: Params) -> Self { + let mut rng = StdRng::seed_from_u64(params.seed); + let n = params.n; + + let nregions = REGIONS.len(); + let rtt_ms: Vec> = (0..nregions) + .map(|a| (0..nregions).map(|b| region_rtt_ms(a, b)).collect()) + .collect(); + + // Assign nodes to regions, weighted by real per-region populations. + let total_weight: u32 = REGIONS.iter().map(|r| r.weight).sum(); + let region_of: Vec = (0..n) + .map(|_| { + let mut pick = rng.random_range(0..total_weight); + for (i, r) in REGIONS.iter().enumerate() { + if pick < r.weight { + return i as u8; + } + pick -= r.weight; + } + unreachable!() + }) + .collect(); + + let config = params.config(); + let mut states: Vec<_> = (0..n) + .map(|i| { + PlumtreeState::new_with_store_seeded( + i as NId, + config.clone(), + SimSeenStore::new(config.max_received_entries), + params.seed.wrapping_add(i as u64), + ) + }) + .collect(); + + // Full-membership bootstrap: every node knows every other node, the + // way corrosion clusters work today. + let mut rt = SimRuntime::default(); + for i in 0..n { + let peers: Vec<(NId, RttInfo)> = (0..n) + .filter(|&j| j != i) + .map(|j| { + let rtt = rtt_ms[region_of[i] as usize][region_of[j] as usize]; + let info = if params.rtt_aware { + RttInfo { + ring: ring_from_rtt_ms(rtt), + } + } else { + RttInfo::default() + }; + (j as NId, info) + }) + .collect(); + states[i].add_peers_bulk_with_rtt(peers, &mut rt); + } + // bootstrap produces no traffic + assert!(rt.outbox.is_empty() && rt.timers.is_empty()); + + let mut sim = Self { + now: 0, + next_seq: 0, + queue: BinaryHeap::new(), + states, + region_of, + rtt_ms, + rng, + runtime: rt, + tick_pending: vec![false; n], + gossip_round: None, + msgs: HashMap::new(), + stats: Stats::default(), + params, + start_avg_eager: 0.0, + down: vec![false; n], + crashed: HashSet::new(), + flapped: HashSet::new(), + missed_during_downtime: HashMap::new(), + }; + + // Optionally over-fill each node's eager set (model admission without + // trimming); prune should converge it back over the run. + if sim.params.eager_inflate > 0 { + for i in 0..n { + let mut added = 0; + let mut attempts = 0; + while added < sim.params.eager_inflate && attempts < sim.params.eager_inflate * 50 { + attempts += 1; + let p = sim.rng.random_range(0..n as NId); + if p as usize == i || sim.states[i].eager_peers().contains(&p) { + continue; + } + sim.states[i].force_eager(p); + added += 1; + } + } + } + sim.start_avg_eager = sim + .states + .iter() + .map(|s| s.eager_peers().len()) + .sum::() as f64 + / n as f64; + + // Every node broadcasts, staggered over the window. + let window_us = sim.params.broadcast_window.as_micros() as u64; + for node in 0..n as NId { + for seq in 0..sim.params.msgs_per_node { + let at = sim.rng.random_range(0..window_us.max(1)); + sim.push(at, Event::Originate { node, seq }); + } + } + + // Inject node failures (crashes and down/up flaps), scheduled + // uniformly over the same window so churn overlaps the broadcast + // storm, the way a real failure would. + let fail_count = sim.params.crash_count + sim.params.flap_count; + debug_assert!( + fail_count <= n, + "not enough nodes for the requested failure counts" + ); + if fail_count > 0 { + let mut candidates: Vec = (0..n as NId).collect(); + candidates.shuffle(&mut sim.rng); + for &node in &candidates[..sim.params.crash_count] { + let at = sim.rng.random_range(0..window_us.max(1)); + sim.push( + at, + Event::NodeDown { + node, + restart_after: None, + }, + ); + } + for &node in &candidates[sim.params.crash_count..fail_count] { + let at = sim.rng.random_range(0..window_us.max(1)); + sim.push( + at, + Event::NodeDown { + node, + restart_after: Some(sim.params.flap_downtime), + }, + ); + } + } + + sim + } + + fn push(&mut self, at: u64, ev: Event) { + let seq = self.next_seq; + self.next_seq += 1; + self.queue.push(Scheduled { at, seq, ev }); + } + + /// One-way flight time in µs between two nodes, with jitter. + fn flight_us(&mut self, from: NId, to: NId) -> u64 { + let rtt = self.link_rtt_ms(from, to); + let jitter = self.rng.random_range(1.0..1.15); + ((rtt * 1000) as f64 / 2.0 * jitter) as u64 + } + + fn link_rtt_ms(&self, from: NId, to: NId) -> u64 { + self.rtt_ms[self.region_of[from as usize] as usize][self.region_of[to as usize] as usize] + } + + fn rtt_info(&self, from: NId, to: NId) -> RttInfo { + if self.params.rtt_aware { + RttInfo { + ring: ring_from_rtt_ms(self.link_rtt_ms(from, to)), + } + } else { + RttInfo::default() + } + } + + /// Random per-observer delay before a `peer_down`/`peer_up` is + /// delivered, modeling staggered SWIM failure detection. + fn random_detection_delay(&mut self) -> u64 { + let min = self.params.detection_delay_min.as_micros() as u64; + let max = self.params.detection_delay_max.as_micros() as u64; + if max <= min { + return min; + } + self.rng.random_range(min..max) + } + + /// Run until the event queue drains (the protocol quiesces). + fn run(mut self) -> Report { + self.drain_queue_to_quiescence(); + self.check_invariants(); + self.report() + } + + /// Run the event queue to quiescence. Split out from `run` so failure + /// tests can inspect `Sim`'s private state afterward instead of only + /// the aggregate `Report`. + fn drain_queue_to_quiescence(&mut self) { + while let Some(Scheduled { at, ev, .. }) = self.queue.pop() { + debug_assert!(at >= self.now); + self.now = at; + self.runtime.now_us = at; + match ev { + Event::Originate { node, seq } => { + if self.down[node as usize] { + continue; + } + let payload = SimPayload { origin: node, seq }; + let id = payload.id(); + self.msgs.insert( + id, + MsgStat { + sent_at: self.now, + gossip_sends: 0, + deliveries: 0, + last_delivery_at: 0, + delivered_to: HashSet::new(), + }, + ); + // Ground truth: any node already down when this message + // is born could never receive it, full stop. + for r in 0..self.params.n as NId { + if r != node && self.down[r as usize] { + self.missed_during_downtime.entry(r).or_default().insert(id); + } + } + self.states[node as usize].broadcast(id, payload, &mut self.runtime); + self.drain(node); + } + Event::Recv { to, msg } => { + if self.down[to as usize] { + self.stats.dropped += 1; + // Ground truth: this specific broadcast really did + // try to reach `to` and got dropped because it was + // down — only `Gossip` carries a broadcast payload; + // IHave/Graft/Prune are protocol control traffic. + if let PlumtreeMsg::Gossip(g) = &msg { + self.missed_during_downtime + .entry(to) + .or_default() + .insert(g.payload.id()); + } + continue; + } + let state = &mut self.states[to as usize]; + match msg { + PlumtreeMsg::Gossip(g) => { + self.gossip_round = Some(g.round); + state.handle_gossip(g, &mut self.runtime); + } + PlumtreeMsg::IHave(ih) => state.handle_ihave(ih, &mut self.runtime), + PlumtreeMsg::Graft(g) => state.handle_graft(g, &mut self.runtime), + PlumtreeMsg::Prune(p) => state.handle_prune(p, &mut self.runtime), + } + self.drain(to); + self.gossip_round = None; + } + Event::TimerFired { node, timer } => { + if self.down[node as usize] { + continue; + } + self.states[node as usize].timer_fired(timer, &mut self.runtime); + self.drain(node); + } + Event::Tick { node } => { + self.tick_pending[node as usize] = false; + if self.down[node as usize] { + continue; + } + let state = &mut self.states[node as usize]; + state.cache_evict_if_needed(&mut self.runtime); + state.tick(&mut self.runtime); + self.drain(node); + } + Event::NodeDown { + node, + restart_after, + } => { + self.down[node as usize] = true; + if restart_after.is_none() { + self.crashed.insert(node); + } + let n = self.params.n as NId; + for observer in 0..n { + if observer == node || self.down[observer as usize] { + continue; + } + let delay = self.random_detection_delay(); + let at = self.now + delay; + self.push( + at, + Event::MemberDown { + observer, + subject: node, + }, + ); + } + if let Some(downtime) = restart_after { + let at = self.now + downtime.as_micros() as u64; + self.push(at, Event::NodeUp { node }); + } + } + Event::NodeUp { node } => { + self.down[node as usize] = false; + self.flapped.insert(node); + let config = self.params.config(); + let seed = self.rng.random(); + self.states[node as usize] = PlumtreeState::new_with_store_seeded( + node, + config.clone(), + SimSeenStore::new(config.max_received_entries), + seed, + ); + self.tick_pending[node as usize] = false; + + // Instant rejoin bootstrap: the restarted node learns the + // current alive membership immediately (like joining + // fresh off the gossip seed addrs). + let n: u32 = self.params.n as NId; + let peers: Vec<(NId, RttInfo)> = (0..n) + .filter(|&j| j != node && !self.down[j as usize]) + .map(|j| (j, self.rtt_info(node, j))) + .collect(); + self.states[node as usize].add_peers_bulk_with_rtt(peers, &mut self.runtime); + self.drain(node); + + // Every other alive node's foca instance notices the + // rejoin after its own random detection delay. + for observer in 0..n { + if observer == node || self.down[observer as usize] { + continue; + } + let delay = self.random_detection_delay(); + let at = self.now + delay; + self.push( + at, + Event::MemberUp { + observer, + subject: node, + }, + ); + } + } + Event::MemberDown { observer, subject } => { + if !self.down[observer as usize] { + self.states[observer as usize].peer_down(&subject, &mut self.runtime); + self.drain(observer); + } + } + Event::MemberUp { observer, subject } => { + if !self.down[observer as usize] && !self.down[subject as usize] { + let info = self.rtt_info(observer, subject); + self.states[observer as usize].peer_up( + subject, + Some(info), + &mut self.runtime, + ); + self.drain(observer); + } + } + } + } + } + + /// Apply everything the node just asked the runtime to do. + fn drain(&mut self, from: NId) { + for payload in std::mem::take(&mut self.runtime.inbox) { + let stat = self + .msgs + .get_mut(&payload.id()) + .expect("delivered unknown msg"); + stat.deliveries += 1; + stat.last_delivery_at = self.now; + stat.delivered_to.insert(from); + self.stats.latencies_us.push(self.now - stat.sent_at); + // round 0 = received directly from the origin = 1 hop + let round = self.gossip_round.expect("delivery outside handle_gossip"); + self.stats.hops.push(round + 1); + } + + let notification_stats = std::mem::take(&mut self.runtime.notification_stats); + self.stats.duplicates += notification_stats.duplicates; + self.stats.missing += notification_stats.missing; + self.stats.not_cached += notification_stats.not_cached; + self.stats.promotions += notification_stats.promotions; + self.stats.demotions += notification_stats.demotions; + self.stats.prune_suppressed += notification_stats.prune_suppressed; + + for (timer, after) in std::mem::take(&mut self.runtime.timers) { + let at = self.now + after.as_micros() as u64; + self.push(at, Event::TimerFired { node: from, timer }); + } + + for (to, msg) in std::mem::take(&mut self.runtime.outbox) { + match &msg { + PlumtreeMsg::Gossip(g) => { + self.stats.sent_gossip += 1; + if self.link_rtt_ms(from, to) > LONGHAUL_RTT_MS { + self.stats.longhaul_gossip += 1; + } + if let Some(stat) = self.msgs.get_mut(&g.payload.id()) { + stat.gossip_sends += 1; + } + } + PlumtreeMsg::IHave(_) => self.stats.sent_ihave += 1, + PlumtreeMsg::Graft(_) => self.stats.sent_graft += 1, + PlumtreeMsg::Prune(_) => self.stats.sent_prune += 1, + } + if self.params.loss > 0.0 && self.rng.random_bool(self.params.loss) { + self.stats.dropped += 1; + continue; + } + let at = self.now + self.flight_us(from, to); + self.push(at, Event::Recv { to, msg }); + } + + // Mirror the agent's periodic IHave flush, but only schedule a tick + // when there is something to flush so the queue can drain. + if !self.tick_pending[from as usize] && !self.states[from as usize].lazy_queue().is_empty() + { + self.tick_pending[from as usize] = true; + let at = self.now + TICK_INTERVAL.as_micros() as u64; + self.push(at, Event::Tick { node: from }); + } + } + + fn check_invariants(&self) { + // By the time the queue fully drains, every flap has resolved (its + // `NodeUp` always fires eventually), so anything still `down` here + // must be a permanent crash. + for (i, node) in (0..self.params.n as NId).enumerate() { + if self.down[i] { + assert!( + self.crashed.contains(&node), + "node {node} still down at drain but not a permanent crash" + ); + continue; + } + let state = &self.states[i]; + let eager = state.eager_peers(); + let lazy = state.lazy_peers(); + let known = state.known_peers(); + for p in state.ring_locked_peers() { + assert!(eager.contains(p), "ring-locked peer not eager"); + } + for p in eager { + assert!(!lazy.contains(p), "peer in both eager and lazy"); + assert!(known.contains(p), "eager peer not known"); + } + for p in lazy { + assert!(known.contains(p), "lazy peer not known"); + } + assert!(lazy.len() <= state.max_lazy(), "lazy set exceeds max_lazy"); + assert!(state.lazy_queue().is_empty(), "lazy queue not drained"); + assert_eq!( + known.len(), + self.params.n - 1 - self.crashed.len(), + "full membership not reconverged for node {node}" + ); + } + } + + fn report(&self) -> Report { + // Not `n * msgs_per_node`: under failure injection a node's own + // origin is skipped if it's already down when its broadcast fires + // (see `Event::Originate`), so the true count of originated messages + // can be less. Identical to `n * msgs_per_node` whenever nothing + // gets skipped (i.e. every non-failure run). + let expected = self.msgs.len() as u64 * (self.params.n as u64 - 1); + let delivered = self.stats.latencies_us.len() as u64; + + let mut hops = self.stats.hops.clone(); + hops.sort_unstable(); + let mut lat = self.stats.latencies_us.clone(); + lat.sort_unstable(); + + // time for a fully-delivered message to reach the whole cluster + let mut full: Vec = self + .msgs + .values() + .filter(|m| m.deliveries as u64 == self.params.n as u64 - 1) + .map(|m| m.last_delivery_at - m.sent_at) + .collect(); + full.sort_unstable(); + + // duplicate copies delivered per message, in broadcast order, to show + // tree convergence. 0 = perfect spanning tree (no duplicates); >0 = extra + // copies each node receives on average. + let mut by_time: Vec<&MsgStat> = self.msgs.values().collect(); + by_time.sort_unstable_by_key(|m| m.sent_at); + let dups_per_msg = |msgs: &[&MsgStat]| -> f64 { + if msgs.is_empty() { + return 0.0; + } + let sends: u64 = msgs.iter().map(|m| m.gossip_sends).sum(); + sends as f64 / (msgs.len() as f64 * (self.params.n - 1) as f64) - 1.0 + }; + let decile = (by_time.len() / 10).max(1); + let dups_per_msg_first = dups_per_msg(&by_time[..decile]); + let dups_per_msg_last = dups_per_msg(&by_time[by_time.len() - decile..]); + + let avg_eager = self + .states + .iter() + .map(|s| s.eager_peers().len()) + .sum::() as f64 + / self.params.n as f64; + + let (missed_due_to_failure, unexplained_missing) = + failure_miss_stats(self.params.n, &self.msgs, &self.missed_during_downtime); + + Report { + label: format!( + "n={} num_eager={} eager={}/{}/{} rtt_aware={} loss={:.1}% seed={}", + self.params.n, + self.params.num_eager, + self.params.eager_ratios.near_pct, + self.params.eager_ratios.mid_pct, + self.params.eager_ratios.far_pct, + self.params.rtt_aware, + self.params.loss * 100.0, + self.params.seed, + ), + expected, + delivered, + missing: self.stats.missing, + hops_p50: pct(&hops, 0.50), + hops_p99: pct(&hops, 0.99), + hops_max: hops.last().copied().unwrap_or(0), + lat_p50_ms: pct(&lat, 0.50) / 1000, + lat_p99_ms: pct(&lat, 0.99) / 1000, + full_p50_ms: pct(&full, 0.50) / 1000, + full_p99_ms: pct(&full, 0.99) / 1000, + full_max_ms: full.last().copied().unwrap_or(0) / 1000, + dups_per_msg_first, + dups_per_msg_last, + sent_gossip: self.stats.sent_gossip, + sent_ihave: self.stats.sent_ihave, + sent_graft: self.stats.sent_graft, + sent_prune: self.stats.sent_prune, + prune_suppressed: self.stats.prune_suppressed, + longhaul_frac: self.stats.longhaul_gossip as f64 / self.stats.sent_gossip.max(1) as f64, + duplicates: self.stats.duplicates, + dropped: self.stats.dropped, + not_cached: self.stats.not_cached, + avg_eager, + start_avg_eager: self.start_avg_eager, + crashed: self.crashed.len(), + flapped: self.flapped.len(), + missed_due_to_failure, + unexplained_missing, + } + } +} + +fn pct(sorted: &[T], p: f64) -> T { + if sorted.is_empty() { + return T::default(); + } + let idx = ((sorted.len() as f64 * p) as usize).min(sorted.len() - 1); + sorted[idx] +} + +/// For every message and every would-be recipient (all nodes except the +/// origin) that never appears in `delivered_to`, classify the miss using +/// `missed_during_downtime` — ground truth recorded live during the run (see +/// its doc comment), not inferred from timing after the fact: +/// - `missed_due_to_failure`: the recipient's entry in `missed_during_downtime` +/// names this message — it was already down when the message was born, or +/// a real delivery attempt reached it while down and got dropped. +/// - `unexplained_missing`: no such record exists — nothing ever even +/// tried to reach this recipient while it was reachable, or it was +/// never down at all. A real signal, expected to be 0 in a healthy run +/// (modulo gossip's own inherent small residual non-convergence, which +/// is unrelated to any failure and pre-exists this metric). +/// Short-circuits to `(0, 0)` when nothing was ever recorded, so non-failure +/// runs (including 2k-scale ones) pay nothing for this. +fn failure_miss_stats( + n: usize, + msgs: &HashMap, + missed_during_downtime: &HashMap>, +) -> (u64, u64) { + if missed_during_downtime.is_empty() { + return (0, 0); + } + let mut missed_due_to_failure = 0u64; + let mut unexplained_missing = 0u64; + for (&mid, stat) in msgs.iter() { + let origin = origin_of(mid); + for r in 0..n as NId { + if r == origin || stat.delivered_to.contains(&r) { + continue; + } + let excused = missed_during_downtime + .get(&r) + .is_some_and(|mids| mids.contains(&mid)); + if excused { + missed_due_to_failure += 1; + } else { + unexplained_missing += 1; + } + } + } + (missed_due_to_failure, unexplained_missing) +} + +struct Report { + label: String, + expected: u64, + delivered: u64, + missing: u64, + hops_p50: u32, + hops_p99: u32, + hops_max: u32, + lat_p50_ms: u64, + lat_p99_ms: u64, + full_p50_ms: u64, + full_p99_ms: u64, + full_max_ms: u64, + dups_per_msg_first: f64, + dups_per_msg_last: f64, + sent_gossip: u64, + sent_ihave: u64, + sent_graft: u64, + sent_prune: u64, + prune_suppressed: u64, + longhaul_frac: f64, + duplicates: u64, + dropped: u64, + not_cached: u64, + avg_eager: f64, + start_avg_eager: f64, + /// Nodes that crashed permanently. + crashed: usize, + /// Nodes that went down and rejoined at least once. + flapped: usize, + /// Non-deliveries fully explained by a recipient being down (crashed, or + /// not yet (re)joined) when the message was sent — see `failure_miss_stats`. + missed_due_to_failure: u64, + /// Non-deliveries that aren't explained by any known failure — should be + /// 0 in a healthy run; a nonzero value signals a real protocol miss. + unexplained_missing: u64, +} + +impl Report { + fn delivery_pct(&self) -> f64 { + self.delivered as f64 / self.expected as f64 * 100.0 + } + + /// Delivery percentage after excluding the gaps that failures fully + /// explain — i.e. delivery among the recipients that could actually have + /// received each message. + fn adjusted_delivery_pct(&self) -> f64 { + self.delivered as f64 / (self.expected - self.missed_due_to_failure).max(1) as f64 * 100.0 + } +} + +impl fmt::Display for Report { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "--- {}", self.label)?; + writeln!( + f, + " delivery {:>10.4}% ({}/{}) missing={} not_cached={}", + self.delivery_pct(), + self.delivered, + self.expected, + self.missing, + self.not_cached, + )?; + writeln!( + f, + " failures crashed={} flapped={} missed_due_to_failure={} ({:.2}%) \ + unexplained_missing={} adjusted_delivery={:.4}%", + self.crashed, + self.flapped, + self.missed_due_to_failure, + self.missed_due_to_failure as f64 / self.expected.max(1) as f64 * 100.0, + self.unexplained_missing, + self.adjusted_delivery_pct(), + )?; + writeln!( + f, + " hops p50={} p99={} max={}", + self.hops_p50, self.hops_p99, self.hops_max + )?; + writeln!( + f, + " latency p50={}ms p99={}ms full-cluster: p50={}ms p99={}ms max={}ms", + self.lat_p50_ms, self.lat_p99_ms, self.full_p50_ms, self.full_p99_ms, self.full_max_ms + )?; + writeln!( + f, + " duplicates per msg first-decile={:.2} last-decile={:.2} duplicates={}", + self.dups_per_msg_first, self.dups_per_msg_last, self.duplicates + )?; + writeln!( + f, + " traffic gossip={} ihave={} graft={} prune={} prune_suppressed={} dropped={} longhaul-gossip={:.1}%", + self.sent_gossip, + self.sent_ihave, + self.sent_graft, + self.sent_prune, + self.prune_suppressed, + self.dropped, + self.longhaul_frac * 100.0, + )?; + writeln!( + f, + " topology avg_eager start={:.1} end={:.1}", + self.start_avg_eager, self.avg_eager + ) + } +} + +// --------------------------------------------------------------------------- +// Gossip (foca) broadcast baseline +// --------------------------------------------------------------------------- +// +// Recreates the dissemination logic of `corro-agent`'s `handle_broadcasts` +// (the `BroadcastMethod::Gossip` path that plumtree is meant to replace) so we +// have concrete numbers to compare plumtree against on the same topology. +// +// The model (see crates/corro-agent/src/broadcast/mod.rs): +// * Every node originates one change, exactly like the plumtree sim, so the +// two are measured on an identical workload. We deliberately do NOT model +// the byte-batching / broadcast_cutoff buffering — that only packs several +// changes into one datagram, it doesn't change *who* receives a change — so +// messages disseminate independently of one another. That independence lets +// `run` process originators in memory-bounded chunks: identical per-message +// results, bounded peak RAM at 2k scale (all-at-once would be >1GB). +// * Originator (`AddBroadcast`, is_local): floods ALL of its ring0 peers once +// (the "local broadcast"), plus a global pending that EXCLUDES ring0. +// * Any node, on first receipt of a change that isn't its own +// (`Rebroadcast`, is_local=false): a global pending only; ring0 is NOT +// excluded and is NOT specially flooded. +// * A global pending sends to `choose_count` freshly-chosen random peers +// (deduped via `sent_to`), then re-queues itself up to `max_transmissions` +// times with `100ms * send_count` backoff. +// choose_count = max(num_indirect_probes=3, (N - ring0)/(max_tx*10)) +// max_transmissions = foca Config::compute_max_tx(cluster_size) +// * We don't model the 10MB/s rate limiter (messages are tiny) or the +// anti-entropy SYNC path that mops up gossip stragglers — so the reported +// delivery% is the broadcast layer alone, and may sit just under 100%. + +const NUM_INDIRECT_PROBES: usize = 3; +/// `sleep_ms_base` between successive transmissions of a pending broadcast. +const GOSSIP_RESEND_BASE: Duration = Duration::from_millis(100); + +/// foca `Config::compute_max_tx`: log10(cluster_size+1) * 4, clamped to [1,255]. +fn compute_max_tx(cluster_size: u32) -> u8 { + let max_tx = f64::from(cluster_size.saturating_add(1)).log10() * 4.0; + if max_tx <= 1.0 { + 1 + } else if max_tx >= 255.0 { + 255 + } else { + max_tx as u8 + } +} + +#[derive(Clone)] +struct GossipParams { + n: usize, + /// With RTT info each node has a ring0 (peers < 6ms) to flood; blind means + /// no ring assignment, so ring0 is empty and everything is random gossip. + rtt_aware: bool, + loss: f64, + seed: u64, + /// Window over which origins are staggered by `run_staggered` (mirrors + /// `Params::broadcast_window`; unused by the chunked `run`). + broadcast_window: Duration, + /// See `Params::crash_count`/`flap_count`/`flap_downtime`. + crash_count: usize, + flap_count: usize, + flap_downtime: Duration, + /// See `Params::detection_delay_min`/`detection_delay_max`. + detection_delay_min: Duration, + detection_delay_max: Duration, +} + +impl GossipParams { + fn new(n: usize) -> Self { + Self { + n, + rtt_aware: true, + loss: 0.0, + seed: 42, + broadcast_window: Duration::from_millis(n as u64 * 5), + crash_count: 0, + flap_count: 0, + flap_downtime: Duration::from_secs(5), + detection_delay_min: Duration::from_secs(2), + detection_delay_max: Duration::from_secs(8), + } + } + + /// See `Params::with_failures`. + fn with_failures( + mut self, + crash_count: usize, + flap_count: usize, + flap_downtime: Duration, + ) -> Self { + self.crash_count = crash_count; + self.flap_count = flap_count; + self.flap_downtime = flap_downtime; + self + } + + fn with_detection_delay(mut self, min: Duration, max: Duration) -> Self { + self.detection_delay_min = min; + self.detection_delay_max = max; + self + } +} + +enum GEvent { + Originate { + node: NId, + mid: MId, + }, + /// One transmission round of a pending broadcast on `node`. + Send { + node: NId, + mid: MId, + send_count: u8, + sent_to: std::collections::HashSet, + hop: u32, + }, + Recv { + to: NId, + mid: MId, + hop: u32, + }, + /// See `Event::NodeDown`. + NodeDown { + node: NId, + restart_after: Option, + }, + /// See `Event::NodeUp`. + NodeUp { + node: NId, + }, + /// See `Event::MemberDown`. + MemberDown { + observer: NId, + subject: NId, + }, + /// See `Event::MemberUp`. + MemberUp { + observer: NId, + subject: NId, + }, +} + +struct GossipSim { + params: GossipParams, + now: u64, + next_seq: u64, + queue: BinaryHeap>, + region_of: Vec, + rtt_ms: Vec>, + /// ring0 peers (RTT ring 0) per node; empty when topology-blind. + ring0: Vec>, + seen: Vec>, + rng: StdRng, + max_tx: u8, + msgs: HashMap, + stats: Stats, + /// Nodes currently unreachable (crashed or mid-flap). See `Sim::down`. + down: Vec, + /// Nodes that crashed permanently. + crashed: HashSet, + /// Nodes that went down and rejoined at least once. + flapped: HashSet, + /// See `Sim::missed_during_downtime`. + missed_during_downtime: HashMap>, + /// Per-node view of which other nodes it currently believes are down — + /// gossip's equivalent of Plumtree's adaptive `known_peers`, since gossip + /// otherwise always samples targets from the full `0..n` regardless of + /// membership. Updated via delayed `MemberDown`/`MemberUp`, same as + /// `Sim`'s `peer_down`/`peer_up`. + believed_down: Vec>, +} + +impl GossipSim { + fn new(params: GossipParams) -> Self { + let mut rng = StdRng::seed_from_u64(params.seed); + let n = params.n; + + let nregions = REGIONS.len(); + let rtt_ms: Vec> = (0..nregions) + .map(|a| (0..nregions).map(|b| region_rtt_ms(a, b)).collect()) + .collect(); + + let total_weight: u32 = REGIONS.iter().map(|r| r.weight).sum(); + let region_of: Vec = (0..n) + .map(|_| { + let mut pick = rng.random_range(0..total_weight); + for (i, r) in REGIONS.iter().enumerate() { + if pick < r.weight { + return i as u8; + } + pick -= r.weight; + } + unreachable!() + }) + .collect(); + + // ring0(i) = peers within ring 0 (RTT < 6ms). Blind => no rings => empty. + let ring0: Vec> = (0..n) + .map(|i| { + if !params.rtt_aware { + return Vec::new(); + } + (0..n) + .filter(|&j| j != i) + .filter(|&j| { + let rtt = rtt_ms[region_of[i] as usize][region_of[j] as usize]; + ring_from_rtt_ms(rtt) == Some(0) + }) + .map(|j| j as NId) + .collect() + }) + .collect(); + + let mut sim = Self { + now: 0, + next_seq: 0, + queue: BinaryHeap::new(), + region_of, + rtt_ms, + ring0, + seen: vec![std::collections::HashSet::new(); n], + rng, + max_tx: compute_max_tx(n as u32), + msgs: HashMap::new(), + stats: Stats::default(), + down: vec![false; n], + crashed: HashSet::new(), + flapped: HashSet::new(), + missed_during_downtime: HashMap::new(), + believed_down: vec![HashSet::new(); n], + params, + }; + + // Inject the same failures as `Sim::new`, scheduled uniformly over + // `broadcast_window` so both sims can be compared under an identical + // failure timeline regardless of which `run*` method is used. + let window_us = sim.params.broadcast_window.as_micros() as u64; + let fail_count = sim.params.crash_count + sim.params.flap_count; + debug_assert!( + fail_count <= n, + "not enough nodes for the requested failure counts" + ); + if fail_count > 0 { + let mut candidates: Vec = (0..n as NId).collect(); + candidates.shuffle(&mut sim.rng); + for &node in &candidates[..sim.params.crash_count] { + let at = sim.rng.random_range(0..window_us.max(1)); + sim.push( + at, + GEvent::NodeDown { + node, + restart_after: None, + }, + ); + } + for &node in &candidates[sim.params.crash_count..fail_count] { + let at = sim.rng.random_range(0..window_us.max(1)); + sim.push( + at, + GEvent::NodeDown { + node, + restart_after: Some(sim.params.flap_downtime), + }, + ); + } + } + + sim + } + + fn push(&mut self, at: u64, ev: GEvent) { + let seq = self.next_seq; + self.next_seq += 1; + self.queue.push(Scheduled { at, seq, ev }); + } + + fn link_rtt_ms(&self, from: NId, to: NId) -> u64 { + self.rtt_ms[self.region_of[from as usize] as usize][self.region_of[to as usize] as usize] + } + + fn flight_us(&mut self, from: NId, to: NId) -> u64 { + let rtt = self.link_rtt_ms(from, to); + let jitter = self.rng.random_range(1.0..1.15); + ((rtt * 1000) as f64 / 2.0 * jitter) as u64 + } + + /// Transmit one copy of `mid` from->to, applying loss; schedules the Recv. + fn send_one(&mut self, from: NId, to: NId, mid: MId, hop: u32) { + self.stats.sent_gossip += 1; + if self.link_rtt_ms(from, to) > LONGHAUL_RTT_MS { + self.stats.longhaul_gossip += 1; + } + if let Some(s) = self.msgs.get_mut(&mid) { + s.gossip_sends += 1; + } + if self.params.loss > 0.0 && self.rng.random_bool(self.params.loss) { + self.stats.dropped += 1; + return; + } + let at = self.now + self.flight_us(from, to); + self.push(at, GEvent::Recv { to, mid, hop }); + } + + /// `choose_count` random peers not already in `exclude`. + fn sample_excluding( + &mut self, + exclude: &std::collections::HashSet, + take: usize, + ) -> Vec { + let mut out = Vec::with_capacity(take); + let mut picked = std::collections::HashSet::new(); + let n = self.params.n as NId; + let mut attempts = 0usize; + // rejection sampling: sent_to stays far below N, so collisions are rare. + while out.len() < take && attempts < take.saturating_mul(50).max(1) { + attempts += 1; + let c = self.rng.random_range(0..n); + if exclude.contains(&c) || !picked.insert(c) { + continue; + } + out.push(c); + } + out + } + + /// See `Sim::random_detection_delay`. + fn random_detection_delay(&mut self) -> u64 { + let min = self.params.detection_delay_min.as_micros() as u64; + let max = self.params.detection_delay_max.as_micros() as u64; + if max <= min { + return min; + } + self.rng.random_range(min..max) + } + + fn run(mut self) -> Report { + // Every node originates one change (matching the plumtree sim). Because + // gossip messages disseminate independently (no shared tree, no modeled + // bandwidth contention), originators are processed in memory-bounded + // chunks: identical per-message results, bounded peak RAM at 2k scale. + const CHUNK: usize = 128; + let mut origins: Vec = (0..self.params.n as NId).collect(); + // shuffle so a chunk isn't correlated with region/node-id ordering. + origins.shuffle(&mut self.rng); + + for chunk in origins.chunks(CHUNK) { + for &node in chunk { + let mid = (node as u64) << 32; + let at = self.now; + self.push(at, GEvent::Originate { node, mid }); + } + self.drain_queue(); + // messages are independent across chunks; clear per-node seen so + // memory doesn't grow with the number of chunks processed. + for s in &mut self.seen { + s.clear(); + } + } + self.report() + } + + /// Like `run`, but stages one origin per node across `broadcast_window` + /// — the same continuous-time model `Sim::run` uses — instead of firing + /// a whole chunk at the same instant. This is what makes the failure + /// timeline (`GossipParams::with_failures`, scheduled in `new` over the + /// same window) meaningfully comparable to `Sim`'s. Small-scale only: + /// unlike `run`, there's no chunking, so `seen` grows for every node — + /// don't use this at 2k+ node scale. + fn run_staggered(mut self) -> Report { + let n = self.params.n as NId; + let window_us = self.params.broadcast_window.as_micros() as u64; + for node in 0..n { + let mid = (node as u64) << 32; + let at = self.rng.random_range(0..window_us.max(1)); + self.push(at, GEvent::Originate { node, mid }); + } + self.drain_queue(); + self.report() + } + + /// Run the event queue to quiescence for the currently-pending messages. + fn drain_queue(&mut self) { + while let Some(Scheduled { at, ev, .. }) = self.queue.pop() { + debug_assert!(at >= self.now); + self.now = at; + match ev { + GEvent::Originate { node, mid } => { + if self.down[node as usize] { + continue; + } + self.seen[node as usize].insert(mid); + self.msgs.insert( + mid, + MsgStat { + sent_at: self.now, + gossip_sends: 0, + deliveries: 0, + last_delivery_at: 0, + delivered_to: HashSet::new(), + }, + ); + // Ground truth: any node already down when this message + // is born could never receive it, full stop. + for r in 0..self.params.n as NId { + if r != node && self.down[r as usize] { + self.missed_during_downtime + .entry(r) + .or_default() + .insert(mid); + } + } + // local broadcast: flood every ring0 peer once (hop 1), + // skipping any peer this node currently believes is down. + // `sent_to` still excludes the *full* ring0 (believed-down + // or not) from the random pool below, so a peer we've + // stopped flooding isn't picked again at random either. + let r0 = self.ring0[node as usize].clone(); + for &peer in &r0 { + if !self.believed_down[node as usize].contains(&peer) { + self.send_one(node, peer, mid, 1); + } + } + // global pending excluding ring0 (is_local=true). + let mut sent_to: std::collections::HashSet = r0.into_iter().collect(); + sent_to.insert(node); + self.do_send(node, mid, 0, sent_to, 1); + } + GEvent::Send { + node, + mid, + send_count, + sent_to, + hop, + } => { + self.do_send(node, mid, send_count, sent_to, hop); + } + GEvent::Recv { to, mid, hop, .. } => { + if self.down[to as usize] { + self.stats.dropped += 1; + // Ground truth: this send really did try to reach + // `to` and got dropped because it was down. + self.missed_during_downtime.entry(to).or_default().insert(mid); + continue; + } + if self.seen[to as usize].contains(&mid) { + self.stats.duplicates += 1; + continue; + } + self.seen[to as usize].insert(mid); + let stat = self.msgs.get_mut(&mid).expect("recv unknown msg"); + stat.deliveries += 1; + stat.last_delivery_at = self.now; + stat.delivered_to.insert(to); + self.stats.latencies_us.push(self.now - stat.sent_at); + self.stats.hops.push(hop); + // rebroadcast (is_local=false): global pending, ring0 NOT excluded. + let mut sent_to = std::collections::HashSet::new(); + sent_to.insert(to); + self.do_send(to, mid, 0, sent_to, hop + 1); + } + GEvent::NodeDown { + node, + restart_after, + } => { + self.down[node as usize] = true; + if restart_after.is_none() { + self.crashed.insert(node); + } + let n = self.params.n as NId; + for observer in 0..n { + if observer == node || self.down[observer as usize] { + continue; + } + let delay = self.random_detection_delay(); + let at = self.now + delay; + self.push( + at, + GEvent::MemberDown { + observer, + subject: node, + }, + ); + } + if let Some(downtime) = restart_after { + let at = self.now + downtime.as_micros() as u64; + self.push(at, GEvent::NodeUp { node }); + } + } + GEvent::NodeUp { node } => { + self.down[node as usize] = false; + self.flapped.insert(node); + // Fresh process: empty dedup cache, and an instant + // rejoin-bootstrap belief that mirrors currently-down + // membership (see `Sim`'s `NodeUp` handler). + self.seen[node as usize].clear(); + let n = self.params.n as NId; + self.believed_down[node as usize] = (0..n) + .filter(|&j| j != node && self.down[j as usize]) + .collect(); + + for observer in 0..n { + if observer == node || self.down[observer as usize] { + continue; + } + let delay = self.random_detection_delay(); + let at = self.now + delay; + self.push( + at, + GEvent::MemberUp { + observer, + subject: node, + }, + ); + } + } + GEvent::MemberDown { observer, subject } => { + if !self.down[observer as usize] { + self.believed_down[observer as usize].insert(subject); + } + } + GEvent::MemberUp { observer, subject } => { + if !self.down[observer as usize] && !self.down[subject as usize] { + self.believed_down[observer as usize].remove(&subject); + } + } + } + } + } + + /// One transmission round of a pending broadcast, then re-queue with backoff. + fn do_send( + &mut self, + node: NId, + mid: MId, + send_count: u8, + mut sent_to: std::collections::HashSet, + hop: u32, + ) { + if self.down[node as usize] { + // Dead sender: no more retransmissions. + return; + } + let ring0_count = self.ring0[node as usize].len(); + let dynamic = self.params.n.saturating_sub(ring0_count) / (self.max_tx as usize * 10); + let choose_count = NUM_INDIRECT_PROBES.max(dynamic); + let take = choose_count.min(self.params.n.saturating_sub(sent_to.len())); + // Don't waste a pick on a peer we already believe is down. + let mut exclude = sent_to.clone(); + exclude.extend(self.believed_down[node as usize].iter().copied()); + for p in self.sample_excluding(&exclude, take) { + sent_to.insert(p); + self.send_one(node, p, mid, hop); + } + + let next = send_count + 1; + if next < self.max_tx { + let at = self.now + GOSSIP_RESEND_BASE.as_micros() as u64 * next as u64; + self.push( + at, + GEvent::Send { + node, + mid, + send_count: next, + sent_to, + hop, + }, + ); + } + } + + fn report(&self) -> Report { + // Every node originates once, so expected deliveries = M*(N-1) where + // M is how many actually got to originate. Not `params.n`: under + // failure injection a node's own origin is skipped if it's already + // down when its broadcast fires (see `GEvent::Originate`), so M can + // be less than N. Identical to `n` whenever nothing gets skipped. + let expected = self.msgs.len() as u64 * (self.params.n as u64 - 1); + let delivered = self.stats.latencies_us.len() as u64; + + let mut hops = self.stats.hops.clone(); + hops.sort_unstable(); + let mut lat = self.stats.latencies_us.clone(); + lat.sort_unstable(); + + let mut full: Vec = self + .msgs + .values() + .filter(|m| m.deliveries as u64 == self.params.n as u64 - 1) + .map(|m| m.last_delivery_at - m.sent_at) + .collect(); + full.sort_unstable(); + + // measure duplicate sends over all messages. + let dups_per_msg = if self.msgs.is_empty() { + 0.0 + } else { + let total_sends: u64 = self.msgs.values().map(|m| m.gossip_sends).sum(); + total_sends as f64 / (self.msgs.len() as f64 * (self.params.n - 1) as f64) - 1.0 + }; + + let avg_ring0 = + self.ring0.iter().map(|r| r.len()).sum::() as f64 / self.params.n.max(1) as f64; + + let (missed_due_to_failure, unexplained_missing) = + failure_miss_stats(self.params.n, &self.msgs, &self.missed_during_downtime); + + Report { + label: format!( + "GOSSIP n={} all-broadcast rtt_aware={} loss={:.1}% seed={} max_tx={}", + self.params.n, + self.params.rtt_aware, + self.params.loss * 100.0, + self.params.seed, + self.max_tx, + ), + expected, + delivered, + missing: 0, + hops_p50: pct(&hops, 0.50), + hops_p99: pct(&hops, 0.99), + hops_max: hops.last().copied().unwrap_or(0), + lat_p50_ms: pct(&lat, 0.50) / 1000, + lat_p99_ms: pct(&lat, 0.99) / 1000, + full_p50_ms: pct(&full, 0.50) / 1000, + full_p99_ms: pct(&full, 0.99) / 1000, + full_max_ms: full.last().copied().unwrap_or(0) / 1000, + // gossip has no tree to converge; report the single value in both slots. + dups_per_msg_first: dups_per_msg, + dups_per_msg_last: dups_per_msg, + sent_gossip: self.stats.sent_gossip, + sent_ihave: 0, + sent_graft: 0, + sent_prune: 0, + prune_suppressed: 0, + longhaul_frac: self.stats.longhaul_gossip as f64 / self.stats.sent_gossip.max(1) as f64, + duplicates: self.stats.duplicates, + dropped: self.stats.dropped, + not_cached: 0, + // no eager set in gossip; report the avg ring0 (guaranteed-flood) size. + avg_eager: avg_ring0, + start_avg_eager: avg_ring0, + crashed: self.crashed.len(), + flapped: self.flapped.len(), + missed_due_to_failure, + unexplained_missing, + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// Install a tracing subscriber driven by `RUST_LOG` (no-op if RUST_LOG is +/// unset or a subscriber is already installed). The plumtree state machine +/// logs every graft/prune and eager<->lazy move via `trace!`, independent of +/// the sim's runtime, so this surfaces them. Each line carries `self_actor_id` +/// and `peer`, so repeated churn on the same pair is visible. +fn init_trace_logging() { + use tracing_subscriber::EnvFilter; + let _ = tracing_subscriber::fmt() + .with_ansi(false) + .with_env_filter(EnvFilter::from_default_env()) + .with_target(false) + // virtual time, not wall-clock — line order is the sim's causal order. + .without_time() + .try_init(); +} + +/// Emit plum-foca trace logs for a single 2k plumtree run so tree stability — +/// whether we keep grafting/pruning the same node pairs — can be inspected. +/// Full `plum_foca=trace` is ~18M lines (dominated by per-forward gossip), so +/// filter to the topology-churn events. Run with: +/// +/// ```text +/// RUST_LOG=plum_foca=trace cargo test --release -p plum-foca --test sim \ +/// sim_2k_tree_stability -- --ignored --nocapture 2>/dev/null \ +/// | grep -E 'moved to eager|moved to lazy|PRUNE|graft|evicted' \ +/// > /tmp/plumtree_2k_stability.log +/// ``` +#[test] +#[ignore = "emits trace logs; run in release with RUST_LOG=plum_foca=trace"] +fn sim_2k_tree_stability() { + init_trace_logging(); + let report = Sim::new(Params::new(2000, 8)).run(); + eprintln!("{report}"); +} + +/// Every node broadcasts once; every message must reach every other node. +#[test] +fn sim_small_full_delivery() { + let report = Sim::new(Params::new(200, 5)).run(); + println!("{report}"); + assert_eq!(report.delivered, report.expected, "incomplete delivery"); + assert_eq!(report.missing, 0); +} + +/// With message loss, the IHave -> timeout -> GRAFT path must recover. +#[test] +fn sim_small_loss_recovery() { + let mut params = Params::new(200, 5); + params.loss = 0.02; + let report = Sim::new(params).run(); + println!("{report}"); + assert!(report.dropped > 0, "loss knob had no effect"); + assert!( + report.delivery_pct() > 99.9, + "delivery {:.4}% under 2% loss", + report.delivery_pct() + ); + assert!(report.sent_graft > 0, "recovery never used the graft path"); +} + +/// RTT-aware eager selection should keep most gossip off long-haul links +/// compared to topology-blind selection. +#[test] +fn sim_small_rtt_awareness_pays_off() { + let aware = Sim::new(Params::new(300, 8)).run(); + let mut blind_params = Params::new(300, 8); + blind_params.rtt_aware = false; + let blind = Sim::new(blind_params).run(); + println!("rtt-aware:\n{aware}"); + println!("topology-blind:\n{blind}"); + + assert_eq!(aware.delivered, aware.expected); + assert_eq!(blind.delivered, blind.expected); + assert!( + aware.longhaul_frac < blind.longhaul_frac, + "rtt-aware long-haul fraction {:.3} not lower than blind {:.3}", + aware.longhaul_frac, + blind.longhaul_frac + ); +} + +/// 2k nodes, all broadcasting, across a sweep of num_eager values. +/// Run with: cargo test --release -p plum-foca --test sim -- --ignored --nocapture +#[test] +#[ignore = "slow; run in release mode"] +fn sim_2k_num_eager_sweep() { + let mut hops_p50 = Vec::new(); + for num_eager in [3, 5, 8, 12] { + let report = Sim::new(Params::new(2000, num_eager)).run(); + println!("{report}"); + assert_eq!( + report.delivered, report.expected, + "incomplete delivery at num_eager={num_eager}" + ); + assert_eq!(report.missing, 0); + hops_p50.push(report.hops_p50); + } + // more eager peers -> wider tree -> fewer hops + assert!( + hops_p50.last() <= hops_p50.first(), + "hops did not improve with fanout: {hops_p50:?}" + ); +} + +/// 2k nodes: compare near/mid/far eager ratio splits. +/// Run with: cargo test --release -p plum-foca --test sim sim_2k_eager_ratios_sweep -- --ignored --nocapture +#[test] +#[ignore = "slow; run in release mode"] +fn sim_2k_eager_ratios_sweep() { + const N: usize = 2000; + let cases: [(&str, EagerRatios); 4] = [ + ( + "near-heavy (default)", + EagerRatios { + near_pct: 50, + mid_pct: 30, + far_pct: 20, + }, + ), + ( + "balanced", + EagerRatios { + near_pct: 40, + mid_pct: 30, + far_pct: 30, + }, + ), + ( + "far-heavy", + EagerRatios { + near_pct: 30, + mid_pct: 35, + far_pct: 35, + }, + ), + ( + "near small", + EagerRatios { + near_pct: 20, + mid_pct: 40, + far_pct: 40, + }, + ), + ]; + + let mut reports = Vec::new(); + for (name, ratios) in cases { + let report = Sim::new(Params::new(N, 8).with_eager_ratios(ratios)).run(); + println!("{name}:\n{report}"); + assert_eq!( + report.delivered, report.expected, + "incomplete delivery for {name}" + ); + assert_eq!(report.missing, 0, "missing messages for {name}"); + reports.push((name, report)); + } +} + +/// Baseline: foca gossip broadcast at small scale. Characterizes the redundancy +/// (duplicate copies per message) and delivery of the current production path on +/// the same topology. +#[test] +fn sim_gossip_small() { + let report = GossipSim::new(GossipParams::new(200)).run(); + println!("{report}"); + // epidemic gossip reaches ~everyone, though the broadcast layer alone + // (no anti-entropy sync) can leave a tiny straggler tail. + assert!( + report.delivery_pct() > 99.0, + "gossip delivery {:.2}% unexpectedly low", + report.delivery_pct() + ); + // each relaying node sprays a large fraction of the cluster: dups/msg is high. + assert!( + report.dups_per_msg_first > 1.0, + "gossip dups/msg {:.2} unexpectedly low", + report.dups_per_msg_first + ); +} + +/// Head-to-head: the current foca gossip broadcast vs plumtree at 2k, on the +/// same topology, printed as one comparison table. This is the "what does +/// plumtree buy us" test. dups/msg is the key number (redundant copies delivered +/// per message); `sends/msg` = (dups/msg + 1)·(N-1) is the absolute per-change cost. +/// Run with: cargo test --release -p plum-foca --test sim -- --ignored --nocapture +#[test] +#[ignore = "slow; run in release mode"] +fn sim_gossip_vs_plumtree_2k() { + const N: usize = 2000; + let gossip = GossipSim::new(GossipParams::new(N)).run(); + let plum = Sim::new(Params::new(N, 8)).run(); + + println!("gossip (current production path):\n{gossip}"); + println!("plumtree (num_eager=8):\n{plum}"); + + let row = |name: &str, g: String, p: String| println!(" {name:<22}{g:>16}{p:>16}"); + let sends_per_msg = |r: &Report| (r.dups_per_msg_first + 1.0) * (N as f64 - 1.0); + println!("\n=== gossip vs plumtree (n={N}, full membership) ==="); + row("metric", "gossip".into(), "plumtree".into()); + row( + "delivery %", + format!("{:.3}", gossip.delivery_pct()), + format!("{:.3}", plum.delivery_pct()), + ); + row( + "hops p50/p99", + format!("{}/{}", gossip.hops_p50, gossip.hops_p99), + format!("{}/{}", plum.hops_p50, plum.hops_p99), + ); + row( + "full-cluster p50 ms", + gossip.full_p50_ms.to_string(), + plum.full_p50_ms.to_string(), + ); + row( + "full-cluster p99 ms", + gossip.full_p99_ms.to_string(), + plum.full_p99_ms.to_string(), + ); + row( + "dups/msg (redundancy)", + format!("{:.1}", gossip.dups_per_msg_first), + format!("{:.1}", plum.dups_per_msg_last), + ); + row( + "sends/change", + format!("{:.0}", sends_per_msg(&gossip)), + format!("{:.0}", sends_per_msg(&plum)), + ); + row( + "longhaul gossip %", + format!("{:.1}", gossip.longhaul_frac * 100.0), + format!("{:.1}", plum.longhaul_frac * 100.0), + ); + println!( + " => plumtree uses ~{:.0}x fewer transmissions per change", + sends_per_msg(&gossip) / sends_per_msg(&plum).max(1.0) + ); + + // The whole point of plumtree: drastically lower redundancy at comparable + // convergence. Guard the headline result so a regression is caught. + assert!( + plum.dups_per_msg_last < gossip.dups_per_msg_first / 5.0, + "plumtree dups/msg {:.1} not dramatically below gossip dups/msg {:.1}", + plum.dups_per_msg_last, + gossip.dups_per_msg_first + ); +} + +/// 2k-node gossip baseline, to compare directly against the plumtree numbers +/// (`sim_2k_num_eager_sweep` / `sim_2k_rtt_awareness`). Every node broadcasts +/// once, the same workload as the plumtree 2k tests. +/// Run with: cargo test --release -p plum-foca --test sim -- --ignored --nocapture +#[test] +#[ignore = "slow; run in release mode"] +fn sim_gossip_2k() { + let aware = GossipSim::new(GossipParams::new(2000)).run(); + let mut blind_params = GossipParams::new(2000); + blind_params.rtt_aware = false; + let blind = GossipSim::new(blind_params).run(); + println!("rtt-aware (ring0 flood):\n{aware}"); + println!("topology-blind (pure random gossip):\n{blind}"); +} + +/// 2k nodes: topology-aware vs topology-blind eager selection. +#[test] +#[ignore = "slow; run in release mode"] +fn sim_2k_rtt_awareness() { + init_trace_logging(); + let aware = Sim::new(Params::new(2000, 8)).run(); + let mut blind_params = Params::new(2000, 8); + blind_params.rtt_aware = false; + let blind = Sim::new(blind_params).run(); + println!("rtt-aware:\n{aware}"); + println!("topology-blind:\n{blind}"); + + assert_eq!(aware.delivered, aware.expected); + assert_eq!(blind.delivered, blind.expected); + assert!(aware.longhaul_frac < blind.longhaul_frac); +} + +/// Measure the prune-throttle (B2): same 2k run with throttling off vs the +/// sim default (500ms) per-peer PRUNE window. Expect prune traffic to drop +/// and full-cluster latency preserved (the in-flight duplicate bursts no longer +/// each draw a PRUNE). Run with: +/// cargo test --release -p plum-foca --test sim sim_2k_prune_throttle -- --ignored --nocapture +#[test] +#[ignore = "slow; run in release mode"] +fn sim_2k_prune_throttle() { + let mut baseline_params = Params::new(2000, 8); + baseline_params.prune_throttle = None; + let baseline = Sim::new(baseline_params).run(); + let throttled = Sim::new(Params::new(2000, 8)).run(); + println!("prune-throttle OFF:\n{baseline}"); + println!("prune-throttle 500ms:\n{throttled}"); + + // delivery must be unharmed + assert_eq!( + throttled.delivered, throttled.expected, + "throttle broke delivery" + ); + assert_eq!(throttled.missing, 0); + // prune traffic should fall substantially + assert!( + throttled.sent_prune < baseline.sent_prune / 2, + "throttle didn't cut prune traffic: {} -> {}", + baseline.sent_prune, + throttled.sent_prune + ); + assert!(throttled.prune_suppressed > 0, "nothing was suppressed"); + // convergence shouldn't regress (allow a small margin for timing shifts) + assert!( + throttled.full_p99_ms <= baseline.full_p99_ms + baseline.full_p99_ms / 10, + "full-cluster p99 regressed: {}ms -> {}ms", + baseline.full_p99_ms, + throttled.full_p99_ms + ); +} + +/// Some nodes crash permanently mid-run (stop sending/receiving for good). +/// Observers only learn of the crash after a random per-observer delay +/// (`MemberDown`), mirroring staggered SWIM detection. The tree must +/// self-heal around the dead nodes: since the sim always runs to full +/// quiescence, every node that never crashes must eventually receive every +/// message, and every surviving node must fully reconverge on membership +/// (checked by `check_invariants`, which already accounts for `crashed`). +#[test] +fn sim_small_crash_resilience() { + const N: usize = 200; + let params = Params::new(N, 5).with_failures(10, 0, Duration::from_secs(5)); + let mut sim = Sim::new(params); + sim.drain_queue_to_quiescence(); + sim.check_invariants(); + let report = sim.report(); + println!("{report}"); + + assert_eq!(report.crashed, 10, "crash count knob had no effect"); + // Every gap must be explained by a recipient having crashed — anything + // left over would mean the tree failed to route around a dead node. + assert_eq!( + report.unexplained_missing, 0, + "some non-delivery isn't explained by a crash — real protocol miss" + ); +} + +/// Some nodes go down and later rejoin with a fresh `PlumtreeState` (a real +/// process restart), again with a randomly-delayed `MemberDown`/`MemberUp` +/// per observer. Plumtree never replays already-delivered messages to a +/// late (re)joiner, so we can't assert full history delivery — instead, we +/// wait for every flap + detection delay to resolve, then fire a fresh +/// broadcast wave and require it reach the *entire* cluster, including the +/// nodes that just came back. +#[test] +fn sim_small_flap_recovery() { + const N: usize = 200; + let flap_downtime = Duration::from_secs(5); + let params = Params::new(N, 5) + .with_failures(0, 10, flap_downtime) + .with_detection_delay(Duration::from_secs(2), Duration::from_secs(4)); + let broadcast_window = params.broadcast_window; + let detection_delay_max = params.detection_delay_max; + let mut sim = Sim::new(params); + + // Worst case: a flap is scheduled right at the end of the window, then + // takes `flap_downtime` to come back, then up to `detection_delay_max` + // for every observer to notice. Double that last term for margin. + let settle_at = broadcast_window.as_micros() as u64 + + flap_downtime.as_micros() as u64 + + detection_delay_max.as_micros() as u64 * 2; + + // Post-recovery wave, using `seq: 1` (distinct from the `seq: 0` wave + // `Sim::new` already scheduled from the default `msgs_per_node: 1`) so + // these message ids don't collide with — and get treated as duplicates + // of — the earlier ones. + for node in 0..N as NId { + sim.push(settle_at, Event::Originate { node, seq: 1 }); + } + + sim.drain_queue_to_quiescence(); + sim.check_invariants(); + let report = sim.report(); + println!("{report}"); + + assert_eq!(report.flapped, 10, "flap count knob had no effect"); + // Every gap must be explained by a recipient not having (re)joined yet + // when a message was sent — in particular, every node is alive for the + // post-recovery wave above, so any gap there would show up here as a + // real protocol miss. + assert_eq!( + report.unexplained_missing, 0, + "some non-delivery isn't explained by a node still being down — real protocol miss" + ); +} + +/// Head-to-head: plumtree vs the gossip baseline under an *identical* crash +/// timeline (same crash count, seed, detection-delay range, broadcast +/// window) — the no-failure comparison is `sim_gossip_vs_plumtree_2k`; this +/// is its churn counterpart. Both must fully explain their delivery gap as +/// crash-caused, not as an unexplained protocol miss. +#[test] +fn sim_gossip_vs_plumtree_failures() { + const N: usize = 500; + const CRASH_COUNT: usize = 20; + let flap_downtime = Duration::from_secs(4); + let detection_delay = (Duration::from_secs(2), Duration::from_secs(5)); + + let mut plum = Sim::new( + Params::new(N, 5) + .with_failures(CRASH_COUNT, 0, flap_downtime) + .with_detection_delay(detection_delay.0, detection_delay.1), + ); + plum.drain_queue_to_quiescence(); + plum.check_invariants(); + let plum_report = plum.report(); + + let gossip_report = GossipSim::new( + GossipParams::new(N) + .with_failures(CRASH_COUNT, 0, flap_downtime) + .with_detection_delay(detection_delay.0, detection_delay.1), + ) + .run_staggered(); + + println!("gossip under failures:\n{gossip_report}"); + println!("plumtree under failures:\n{plum_report}"); + + let row = |name: &str, g: String, p: String| println!(" {name:<24}{g:>16}{p:>16}"); + println!("\n=== gossip vs plumtree under failures (n={N}, crash={CRASH_COUNT}) ==="); + row("metric", "gossip".into(), "plumtree".into()); + row( + "delivery %", + format!("{:.3}", gossip_report.delivery_pct()), + format!("{:.3}", plum_report.delivery_pct()), + ); + row( + "delivered/expected", + format!("{}/{}", gossip_report.delivered, gossip_report.expected), + format!("{}/{}", plum_report.delivered, plum_report.expected), + ); + row( + "adjusted delivery %", + format!("{:.3}", gossip_report.adjusted_delivery_pct()), + format!("{:.3}", plum_report.adjusted_delivery_pct()), + ); + row( + "duplicates", + gossip_report.duplicates.to_string(), + plum_report.duplicates.to_string(), + ); + row( + "dups/msg", + format!("{:.2}", gossip_report.dups_per_msg_first), + format!("{:.2}", plum_report.dups_per_msg_last), + ); + row( + "latency p50/p99 ms", + format!("{}/{}", gossip_report.lat_p50_ms, gossip_report.lat_p99_ms), + format!("{}/{}", plum_report.lat_p50_ms, plum_report.lat_p99_ms), + ); + row( + "missed_due_to_failure", + gossip_report.missed_due_to_failure.to_string(), + plum_report.missed_due_to_failure.to_string(), + ); + row( + "unexplained_missing", + gossip_report.unexplained_missing.to_string(), + plum_report.unexplained_missing.to_string(), + ); + + assert_eq!(gossip_report.crashed, CRASH_COUNT); + assert_eq!(plum_report.crashed, CRASH_COUNT); + assert_eq!( + gossip_report.unexplained_missing, 0, + "gossip has non-delivery not explained by the injected crashes" + ); + assert_eq!( + plum_report.unexplained_missing, 0, + "plumtree has non-delivery not explained by the injected crashes" + ); +} diff --git a/doc/config/gossip.md b/doc/config/gossip.md index 309bf1242..9fdb7a29b 100644 --- a/doc/config/gossip.md +++ b/doc/config/gossip.md @@ -62,6 +62,41 @@ This should be your "effective" MTU: `network interface's MTU - IP header size - Certain environments don't support GSO (Generic Segmentation Offload). This is detected by the QUIC implementation, but it's possible to pre-emptively disable it to avoid re-trying the initial packets without GSO as it is detected as unavailable. +#### `gossip.broadcast` + +Selects the algorithm used to disseminate changes across the cluster. Please note that all nodes in a cluster need to use the same broadcast +algorithm. + +Two methods are available: + +- `gossip` (**default**): SWIM-style gossip. Each change is re-broadcast to a random subset of peers. Robust and simple, but every hop carries the full payload, so bandwidth grows with cluster size. +- `plumtree`: [Plumtree](https://asc.di.fct.unl.pt/~jleitao/pdf/srds07-leitao.pdf) epidemic broadcast. Nodes build a spanning tree of "eager" peers that forward full payloads, while the remaining "lazy" peers only exchange lightweight `IHave` announcements and pull (`GRAFT`) the payload on demand. This cuts redundant payload traffic on larger clusters, at the cost of some protocol complexity. + +If unset, Corrosion uses `gossip`. To enable Plumtree, add a `[gossip.broadcast.plumtree]` table. All of its fields are optional and fall back to the defaults below: + +```toml +[gossip.broadcast.plumtree] +prune_threshold = 5 # optional +optimization_threshold = 7 # optional +batch_gossip = false # optional +``` + +##### `prune_threshold` + +Number of times a message may be received from the same eager peer before that peer is pruned (moved from the eager set to the lazy set with a `PRUNE`). Duplicate deliveries mean the tree has a redundant edge, so pruning trims it. + +Higher values tolerate more duplication in exchange for a more stable tree; lower values prune more aggressively, reducing redundant payload traffic but causing more tree churn. Defaults to `5`. + +##### `optimization_threshold` + +Threshold, in gossip rounds, for switching to a shorter path. When a node learns via a lazy `IHave` that a peer could have delivered a message at least this many rounds sooner than the eager peer that actually delivered it, it grafts the shorter path and prunes the longer one. + +Lower values optimize the tree more aggressively (more `GRAFT`/`PRUNE` churn); higher values leave the tree alone unless the improvement is large. Defaults to `7`. + +##### `batch_gossip` + +When `true`, coalesce multiple outgoing Plumtree gossip and membership messages instead of sending each one individually. Reduces per-packet overhead on busy clusters at the cost of a small amount of added latency. Defaults to `false`. + #### `gossip.tls` Strong encryption is highly recommended for any non-development usage of Corrosion.