introduce new type

Author: somtochiama

Cargo.lock

@@ -24,7 +24,7 @@ use corro_types::{
     actor::{Actor, ActorId},
     agent::{Agent, Bookie, SplitPool},
     base::CrsqlSeq,
-    broadcast::{BroadcastInput, BroadcastV1, ChangeSource, ChangeV1, FocaInput},
+    broadcast::{BroadcastInput, BroadcastV1, BroadcastV2, ChangeSource, ChangeV1, FocaInput},
     channel::CorroReceiver,
     members::MemberAddedResult,
     sync::generate_sync,
@@ -690,7 +690,7 @@ pub async fn handle_changes(
             continue;
         }
 
-        let src_str: &'static str = src.into();
+        let src_str: &'static str = (&src).into();
         let recv_lag = change.ts().and_then(|ts| {
             let mut our_ts = Timestamp::from(agent.clock().new_timestamp());
             if ts > our_ts {
@@ -705,7 +705,10 @@ pub async fn handle_changes(
             Some((our_ts.0 - ts.0).to_duration())
         });
 
-        if matches!(src, ChangeSource::Broadcast) {
+        if matches!(
+            src,
+            ChangeSource::Broadcast | ChangeSource::BroadcastV2(_, _)
+        ) {
             counter!("corro.broadcast.recv.count", "kind" => "change").increment(1);
         }
 
@@ -766,20 +769,32 @@ pub async fn handle_changes(
             }
         }
 
-        assert_sometimes!(
-            matches!(src, ChangeSource::Sync),
-            "Corrosion receives changes through sync"
-        );
-        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(
+        if !change.is_empty() {
+            let bcast = match src.clone() {
+                ChangeSource::Broadcast => {
+                    assert_sometimes!(true, "Corrosion rebroadcasts changes");
+                    Some(BroadcastInput::Rebroadcast(BroadcastV1::Change(
                         change.clone(),
                     )))
-            {
-                debug!("broadcasts are full or done!");
+                }
+                ChangeSource::BroadcastV2(set, num_broadcasts) => {
+                    assert_sometimes!(true, "Corrosion rebroadcasts changes");
+                    Some(BroadcastInput::RebroadcastV2(BroadcastV2 {
+                        change: BroadcastV1::Change(change.clone()),
+                        set,
+                        num_broadcasts,
+                    }))
+                }
+                ChangeSource::Sync => {
+                    assert_sometimes!(true, "Corrosion receives changes through sync");
+                    None
+                }
+            };
+
+            if let Some(bcast) = bcast {
+                if let Err(_e) = agent.tx_bcast().try_send(bcast) {
+                    debug!("broadcasts are full or done!");
+                }
             }
         }
 

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

@@ -1,6 +1,6 @@
 use corro_types::{
     actor::ClusterId,
-    broadcast::{BroadcastV1, ChangeSource, ChangeV1, UniPayload, UniPayloadV1},
+    broadcast::{BroadcastV1, BroadcastV2, ChangeSource, ChangeV1, UniPayload, UniPayloadV1},
     channel::CorroSender,
 };
 use metrics::counter;
@@ -66,16 +66,38 @@ pub fn spawn_unipayload_handler(
 
                                             match payload {
                                                 UniPayload::V1 {
-                                                    data:
-                                                        UniPayloadV1::Broadcast(BroadcastV1::Change(
-                                                            change,
-                                                        )),
+                                                    data: payload_data,
                                                     cluster_id: payload_cluster_id,
                                                 } => {
                                                     if cluster_id != payload_cluster_id {
                                                         continue;
                                                     }
-                                                    changes.push((change, ChangeSource::Broadcast));
+
+                                                    match payload_data {
+                                                        UniPayloadV1::Broadcast(
+                                                            BroadcastV1::Change(change),
+                                                        ) => {
+                                                            changes.push((
+                                                                change,
+                                                                ChangeSource::Broadcast,
+                                                            ));
+                                                        }
+                                                        UniPayloadV1::BroadcastV2(
+                                                            BroadcastV2 {
+                                                                change: BroadcastV1::Change(change),
+                                                                set,
+                                                                num_broadcasts,
+                                                            },
+                                                        ) => {
+                                                            changes.push((
+                                                                change,
+                                                                ChangeSource::BroadcastV2(
+                                                                    set,
+                                                                    num_broadcasts,
+                                                                ),
+                                                            ));
+                                                        }
+                                                    }
                                                 }
                                             }
                                         }

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

@@ -38,7 +38,10 @@ use tripwire::Tripwire;
 use corro_types::{
     actor::{Actor, ActorId},
     agent::Agent,
-    broadcast::{BroadcastInput, DispatchRuntime, FocaCmd, FocaInput, UniPayload, UniPayloadV1},
+    broadcast::{
+        BroadcastChangeV1, BroadcastInput, BroadcastV1, BroadcastV2, DispatchRuntime, FocaCmd,
+        FocaInput, UniPayload, UniPayloadV1,
+    },
     channel::{bounded, CorroReceiver, CorroSender},
 };
 
@@ -523,6 +526,10 @@ async fn handle_broadcasts(
                 let (bcast, is_local) = match input {
                     BroadcastInput::Rebroadcast(bcast) => (bcast, false),
                     BroadcastInput::AddBroadcast(bcast) => (bcast, true),
+                    BroadcastInput::RebroadcastV2(bcast) => {
+                        let BroadcastV2::Change(BroadcastChangeV1 { change, .. }) = bcast;
+                        (BroadcastV1::Change(change), false)
+                    }
                 };
                 trace!("adding broadcast: {bcast:?}, local? {is_local}");
 

crates/corro-agent/src/broadcast/mod.rs

@@ -19,6 +19,7 @@ config = { workspace = true }
 consul-client = { version = "0.1.0-alpha.0", path = "../consul-client" }
 corro-api-types = { version = "0.1.0-alpha.1", path = "../corro-api-types" }
 corro-base-types = { version = "0.1.0-alpha.1", path = "../corro-base-types" }
+probabilistic-set = { version = "0.1.0-alpha.1", path = "../probabilistic-set" }
 deadpool = { workspace = true }
 enquote = { workspace = true }
 fallible-iterator = { workspace = true }

crates/corro-types/Cargo.toml

@@ -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 probabilistic_set::ProbabilisticSet;
 use rusqlite::{
     types::{FromSql, FromSqlError},
     ToSql,
@@ -48,6 +49,7 @@ pub enum UniPayload {
 #[derive(Debug, Clone, Readable, Writable)]
 pub enum UniPayloadV1 {
     Broadcast(BroadcastV1),
+    BroadcastV2(BroadcastV2),
 }
 
 #[derive(Debug, Clone, Readable, Writable)]
@@ -94,6 +96,13 @@ pub enum BroadcastV1 {
     Change(ChangeV1),
 }
 
+#[derive(Clone, Debug, Readable, Writable)]
+pub struct BroadcastV2 {
+    pub change: BroadcastV1,
+    pub set: ProbSet,
+    pub num_broadcasts: u8,
+}
+
 #[derive(Debug, Clone, PartialEq, Readable, Writable)]
 pub struct ColumnChange {
     pub cid: ColumnName,
@@ -103,11 +112,12 @@ pub struct ColumnChange {
     pub cl: i64,
 }
 
-#[derive(Debug, Clone, Copy, strum::IntoStaticStr)]
+#[derive(Debug, Clone, strum::IntoStaticStr)]
 #[strum(serialize_all = "snake_case")]
 pub enum ChangeSource {
     Broadcast,
     Sync,
+    BroadcastV2(ProbabilisticSet, u64),
 }
 
 #[derive(Debug, Clone, PartialEq, Readable, Writable)]
@@ -586,6 +596,7 @@ pub enum BroadcastDecodeError {
 pub enum BroadcastInput {
     Rebroadcast(BroadcastV1),
     AddBroadcast(BroadcastV1),
+    RebroadcastV2(BroadcastV2),
 }
 
 pub struct DispatchRuntime<T> {

crates/corro-types/src/broadcast.rs

@@ -0,0 +1,12 @@
+[package]
+name = "probabilistic-set"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+rand = { workspace = true }
+serde = { workspace = true }
+speedy = { workspace = true }
+
+[lints]
+workspace = true

crates/probabilistic-set/Cargo.toml

@@ -0,0 +1,77 @@
+use rand::Rng;
+use serde::{Deserialize, Serialize};
+use speedy::{Readable, Writable};
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+
+#[derive(Serialize, Debug, Deserialize, Clone, Readable, Writable)]
+pub struct ProbabilisticSet {
+    // todo: consider making bits and array?
+    bits: Vec<u8>,
+    size_bits: usize,
+    seed: u64, // Store the random seed
+}
+
+impl ProbabilisticSet {
+    pub fn new(expected_items: usize, bits_per_item: usize) -> Self {
+        let size_bits = expected_items * bits_per_item;
+        let seed = rand::thread_rng().r#gen(); // Generate random seed
+
+        ProbabilisticSet {
+            bits: vec![0; (size_bits + 7) / 8],
+            size_bits,
+            seed,
+        }
+    }
+
+    // Create with a specific seed (useful for testing or controlled randomness)
+    pub fn with_seed(expected_items: usize, bits_per_item: usize, seed: u64) -> Self {
+        let size_bits: usize = expected_items * bits_per_item;
+        ProbabilisticSet {
+            bits: vec![0; (size_bits + 7) / 8],
+            size_bits,
+            seed,
+        }
+    }
+
+    pub fn insert(&mut self, item: u128) {
+        let idx = self.hash(item) % self.size_bits;
+        let byte_idx = idx / 8;
+        let bit_idx = idx % 8;
+        self.bits[byte_idx] |= 1 << bit_idx;
+    }
+
+    pub fn contains(&self, item: u128) -> bool {
+        let idx = self.hash(item) % self.size_bits;
+        let byte_idx = idx / 8;
+        let bit_idx = idx % 8;
+        self.bits[byte_idx] & (1 << bit_idx) != 0
+    }
+
+    fn hash(&self, item: u128) -> usize {
+        let mut hasher = DefaultHasher::new();
+        // Hash both the seed and the item
+        self.seed.hash(&mut hasher);
+        item.hash(&mut hasher);
+        hasher.finish() as usize
+    }
+
+    // Get current seed (useful for debugging)
+    pub fn seed(&self) -> u64 {
+        self.seed
+    }
+
+    // Change the seed (and clear the bits since the hash function effectively changes)
+    pub fn reseed(&mut self) {
+        self.seed = rand::thread_rng().r#gen();
+        self.bits.fill(0);
+    }
+
+    pub fn size_bytes(&self) -> usize {
+        // Include size of seed in total
+        self.bits.len() + std::mem::size_of::<u64>()
+    }
+}
+
+#[cfg(test)]
+mod tests {}