Introduce plumtree broadcast method in corrosion.

Cargo.toml

@@ -1,7 +1,7 @@
 [workspace]
 members = [ ".antithesis/client/rust-load-generator",
 	"crates/*",
-	"integration-tests"
+	"integration-tests",
 ]
 resolver = "2"
 

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");

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 }

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

crates/corro-agent/src/agent/run_root.rs

@@ -10,15 +10,16 @@ use crate::{
         reaper::spawn_reaper,
         setup, util, AgentOptions,
     },
-    broadcast::runtime_loop,
+    broadcast::{handle_broadcasts, runtime_loop, BroadcastOpts},
+    plumtree::spawn_plumtree_loop,
     transport::Transport,
 };
 
 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 {