add broadcast setting to follow settings to allow followers to broadcast under certain conditions (list of actors or percent of changes received)

Author: jeromegn

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

@@ -304,7 +304,16 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age
 
                 match transport.open_bi(addr).await {
                     Ok((tx, rx)) => {
-                        match api::peer::follow::follow(&agent, tx, rx, Some(from), false).await {
+                        match api::peer::follow::follow(
+                            &agent,
+                            tx,
+                            rx,
+                            Some(from),
+                            false,
+                            follow.broadcast.as_ref(),
+                        )
+                        .await
+                        {
                             Ok(dbv) => {
                                 info!("following terminated, last db version: {dbv:?}");
                                 last_from = dbv;

crates/corro-agent/src/api/peer/follow.rs

@@ -8,11 +8,13 @@ use corro_types::{
     base::{CrsqlDbVersion, CrsqlSeq, Version},
     broadcast::{BiPayload, ChangeSource, ChangeV1, Changeset, Timestamp},
     change::ChunkedChanges,
+    config::FollowBroadcast,
     sqlite::SqlitePoolError,
 };
 use futures::{Stream, StreamExt};
 use metrics::counter;
 use quinn::{RecvStream, SendStream};
+use rand::{rngs::OsRng, Rng};
 use rusqlite::{params_from_iter, Row, ToSql};
 use speedy::{Readable, Writable};
 use tokio::{sync::mpsc, task::block_in_place};
@@ -238,6 +240,8 @@ pub async fn read_follow_msg<R: Stream<Item = std::io::Result<BytesMut>> + Unpin
 pub async fn recv_follow(
     agent: &Agent,
     mut read: FramedRead<RecvStream, LengthDelimitedCodec>,
+    local_only: bool,
+    broadcast: Option<&FollowBroadcast>,
 ) -> Result<Option<CrsqlDbVersion>, FollowError> {
     let mut last_db_version = None;
     let tx_changes = agent.tx_changes();
@@ -255,8 +259,17 @@ pub async fn recv_follow(
                         "received changeset for version(s) {:?} and db_version {db_version:?}",
                         changeset.versions()
                     );
+                    let change_src = if local_only
+                        || broadcast
+                            .map(|bcast| should_broadcast(&changeset.actor_id, bcast))
+                            .unwrap_or(false)
+                    {
+                        ChangeSource::Broadcast
+                    } else {
+                        ChangeSource::Follow
+                    };
                     tx_changes
-                        .send((changeset, ChangeSource::Follow))
+                        .send((changeset, change_src))
                         .await
                         .map_err(|_| FollowError::ChannelClosed)?;
                     if let Some(db_version) = db_version {
@@ -270,12 +283,20 @@ pub async fn recv_follow(
     Ok(last_db_version)
 }
 
+fn should_broadcast(actor_id: &ActorId, broadcast: &FollowBroadcast) -> bool {
+    match broadcast {
+        FollowBroadcast::ActorIds(set) => set.contains(actor_id),
+        FollowBroadcast::Percent(percent) => OsRng.gen_range(0..100) < *percent,
+    }
+}
+
 pub async fn follow(
     agent: &Agent,
     mut tx: SendStream,
     recv: RecvStream,
     from: Option<CrsqlDbVersion>,
     local_only: bool,
+    broadcast: Option<&FollowBroadcast>,
 ) -> Result<Option<CrsqlDbVersion>, FollowError> {
     let mut codec = LengthDelimitedCodec::builder()
         .max_frame_length(100 * 1_024 * 1_024)
@@ -302,7 +323,7 @@ pub async fn follow(
             .new_codec(),
     );
 
-    recv_follow(agent, framed).await
+    recv_follow(agent, framed, local_only, broadcast).await
 }
 
 #[cfg(test)]

crates/corro-types/src/config.rs

@@ -1,10 +1,15 @@
-use std::net::{Ipv6Addr, SocketAddr, SocketAddrV6};
+use std::{
+    collections::HashSet,
+    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
+};
 
 use camino::Utf8PathBuf;
 use corro_base_types::CrsqlDbVersion;
 use serde::{Deserialize, Serialize};
 use serde_with::{formats::PreferOne, serde_as, OneOrMany};
 
+use crate::actor::ActorId;
+
 pub const DEFAULT_GOSSIP_PORT: u16 = 4001;
 const DEFAULT_GOSSIP_IDLE_TIMEOUT: u32 = 30;
 
@@ -38,6 +43,15 @@ pub struct FollowConfig {
     pub addr: SocketAddr,
     #[serde(default)]
     pub from: FollowFrom,
+    #[serde(default)]
+    pub broadcast: Option<FollowBroadcast>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub enum FollowBroadcast {
+    ActorIds(HashSet<ActorId>),
+    Percent(u8),
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
@@ -385,8 +399,17 @@ impl ConfigBuilder {
         self
     }
 
-    pub fn follow(mut self, addr: SocketAddr, from: FollowFrom) -> Self {
-        self.follow = Some(FollowConfig { addr, from });
+    pub fn follow(
+        mut self,
+        addr: SocketAddr,
+        from: FollowFrom,
+        broadcast: Option<FollowBroadcast>,
+    ) -> Self {
+        self.follow = Some(FollowConfig {
+            addr,
+            from,
+            broadcast,
+        });
         self
     }