Review of upstream PR #985 (oxidezap/whatsapp-rust)

examples/retry_quarantine.rs

@@ -0,0 +1,265 @@
+//! Bot-side retry-receipt quarantine, built on the [`RetryAdmission`] hook.
+//!
+//! WhatsApp Web processes every inbound retry receipt (no volume throttle), and
+//! this SDK does the same by default. In a large group a cohort of members
+//! whose sessions never establish can retry on every message, driving the
+//! repair path (markForgetSenderKey, key-bundle processing, resend) at storm
+//! rate. That cost only shows up at bot send volume, so the policy lives here in
+//! the bot rather than in the SDK: the core stays WA Web compliant, and this
+//! opt-in bounds the storm for deployments that want it.
+//!
+//! This mirrors the mechanism proposed in oxidezap/whatsapp-rust#982 by
+//! @Salientekill: a token bucket keyed by (chat user, requester user), burst 2,
+//! refill 2/day. One mark repairs a healthy member (the next send carries the
+//! SKDM), so past the burst further receipts from the same pair are dropped
+//! before any repair work; the refill keeps genuine recovery possible. The
+//! device is excluded from the key on purpose: WA Web re-targets the whole user
+//! when the primary goes cold, so all devices of a broken account share one
+//! budget. Our own companion devices never reach this policy (the SDK exempts
+//! `is_peer`), so their group session can always rebuild.
+//!
+//!   cargo run --example retry_quarantine
+//!
+//! Run with `RUST_LOG=debug` to see the SDK's "dropped by RetryAdmission policy"
+//! line when a pair is quarantined.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::Ordering;
+
+use log::{error, info};
+use portable_atomic::AtomicU64;
+use wacore::time::Instant;
+use whatsapp_rust::RetryAdmission;
+use whatsapp_rust::prelude::*;
+
+/// Lazy monotonic token bucket. Pure given `now`/`burst`/`refill_per_sec`, so a
+/// lowered burst clamps on the next access and an idle pair never accrues an
+/// unbounded reserve.
+struct TokenBucket {
+    tokens: f64,
+    last_refill: Instant,
+}
+
+impl TokenBucket {
+    fn new(initial: f64, now: Instant) -> Self {
+        Self {
+            tokens: initial,
+            last_refill: now,
+        }
+    }
+
+    fn try_take(&mut self, now: Instant, burst: f64, refill_per_sec: f64) -> bool {
+        // saturating: two receipts in the same instant give elapsed 0, never a panic.
+        let elapsed = now
+            .saturating_duration_since(self.last_refill)
+            .as_secs_f64();
+        self.tokens = (self.tokens + elapsed * refill_per_sec).min(burst);
+        self.last_refill = now;
+        if self.tokens >= 1.0 {
+            self.tokens -= 1.0;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+/// Per-(chat user, requester user) quarantine. A `burst` of 0 disables it.
+struct RetryQuarantine {
+    buckets: Mutex<HashMap<(String, String), TokenBucket>>,
+    burst: f64,
+    refill_per_sec: f64,
+    /// Capacity guard: the keyspace is O(groups x broken members), so bound it.
+    /// A production policy would use an LRU / moka cache here (as #982 does with
+    /// a dedicated capacity); this example just refuses new pairs past the cap,
+    /// which fails open (admits) rather than growing without limit.
+    max_pairs: usize,
+    quarantined_total: AtomicU64,
+}
+
+impl RetryQuarantine {
+    fn new(burst: u32, refill_per_day: u32, max_pairs: usize) -> Self {
+        Self {
+            buckets: Mutex::new(HashMap::new()),
+            burst: burst as f64,
+            refill_per_sec: refill_per_day as f64 / 86_400.0,
+            max_pairs,
+            quarantined_total: AtomicU64::new(0),
+        }
+    }
+
+    /// Receipts dropped since start (your own observability, no SDK coupling).
+    fn quarantined_total(&self) -> u64 {
+        self.quarantined_total.load(Ordering::Relaxed)
+    }
+}
+
+impl RetryAdmission for RetryQuarantine {
+    fn admit(&self, chat: &Jid, requester: &Jid, _retry_count: u8) -> bool {
+        if self.burst == 0.0 {
+            return true;
+        }
+        let now = Instant::now();
+        let key = (chat.user.to_string(), requester.user.to_string());
+
+        // Synchronous, so the receive path never awaits a policy: the critical
+        // section is a map lookup plus the bucket arithmetic.
+        let mut buckets = self.buckets.lock().expect("quarantine mutex poisoned");
+        // Fail open if we would exceed the cap with a brand-new pair, so the
+        // guard never blocks repair for a pair it isn't already tracking.
+        if !buckets.contains_key(&key) && buckets.len() >= self.max_pairs {
+            return true;
+        }
+        let bucket = buckets
+            .entry(key)
+            .or_insert_with(|| TokenBucket::new(self.burst, now));
+        let allowed = bucket.try_take(now, self.burst, self.refill_per_sec);
+        drop(buckets);
+
+        if !allowed {
+            self.quarantined_total.fetch_add(1, Ordering::Relaxed);
+        }
+        allowed
+    }
+}
+
+fn main() {
+    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
+
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("failed to build tokio runtime");
+
+    rt.block_on(async {
+        let store = match SqliteStore::new("whatsapp.db").await {
+            Ok(store) => store,
+            Err(e) => {
+                error!("failed to create SQLite backend: {e}");
+                return;
+            }
+        };
+
+        let bot = match Bot::builder()
+            .with_backend(store)
+            .on_qr_code(|code, _timeout| async move {
+                info!("scan to pair:\n{code}");
+            })
+            .on_connected(|_client| async {
+                info!("connected; inbound group retry receipts now pass through the quarantine");
+            })
+            .build()
+            .await
+        {
+            Ok(bot) => bot,
+            Err(e) => {
+                error!("failed to build bot: {e}");
+                return;
+            }
+        };
+
+        // Opt in before connecting. Without this call the client keeps its
+        // default behavior (every retry receipt admitted, matching WA Web).
+        // burst 2, refill 2/day, up to 32768 tracked pairs.
+        let quarantine = Arc::new(RetryQuarantine::new(2, 2, 32_768));
+        if !bot.client().set_retry_admission(quarantine.clone()) {
+            error!("retry admission policy was already set");
+            return;
+        }
+
+        bot.run().await;
+        info!(
+            "quarantined {} retry receipt(s) this session",
+            quarantine.quarantined_total()
+        );
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Duration;
+
+    #[test]
+    fn bounds_per_pair_and_isolates_pairs() {
+        // burst 2, refill 0: two attempts, then quarantined.
+        let q = RetryQuarantine::new(2, 0, 100);
+        let g: Jid = "123-456@g.us".parse().unwrap();
+        let a: Jid = "111@lid".parse().unwrap();
+        let b: Jid = "222@lid".parse().unwrap();
+        assert!(q.admit(&g, &a, 1));
+        assert!(q.admit(&g, &a, 1));
+        assert!(!q.admit(&g, &a, 1), "third receipt quarantined");
+        // A different requester in the same chat has its own budget.
+        assert!(q.admit(&g, &b, 1));
+        assert_eq!(q.quarantined_total(), 1);
+    }
+
+    #[test]
+    fn burst_zero_disables() {
+        let q = RetryQuarantine::new(0, 0, 100);
+        let g: Jid = "123-456@g.us".parse().unwrap();
+        let a: Jid = "111@lid".parse().unwrap();
+        for _ in 0..10 {
+            assert!(q.admit(&g, &a, 1));
+        }
+        assert_eq!(q.quarantined_total(), 0);
+    }
+
+    #[test]
+    fn fails_open_past_capacity_for_new_pairs() {
+        // Cap 1: the first pair is tracked and bounded, a second brand-new pair
+        // is admitted rather than evicting/blocking.
+        let q = RetryQuarantine::new(1, 0, 1);
+        let g: Jid = "123-456@g.us".parse().unwrap();
+        let a: Jid = "111@lid".parse().unwrap();
+        let b: Jid = "222@lid".parse().unwrap();
+        assert!(q.admit(&g, &a, 1));
+        assert!(!q.admit(&g, &a, 1), "tracked pair still bounded");
+        assert!(q.admit(&g, &b, 1), "new pair fails open past cap");
+    }
+
+    // The bucket is the seam where time matters; `admit` reads the real clock,
+    // so the refill/clamp claims are driven here with explicit `now` values.
+
+    #[test]
+    fn refill_restores_admission_gradually() {
+        // A quarantined pair must recover, but one token at a time — a refill is
+        // not a reset, so a still-broken member cannot immediately storm again.
+        let t0 = Instant::now();
+        let (burst, refill_per_sec) = (2.0, 1.0); // 1 token/sec
+        let mut b = TokenBucket::new(burst, t0);
+
+        assert!(b.try_take(t0, burst, refill_per_sec));
+        assert!(b.try_take(t0, burst, refill_per_sec));
+        assert!(!b.try_take(t0, burst, refill_per_sec), "burst exhausted");
+
+        // One refill period later exactly one token is back.
+        let t1 = t0 + Duration::from_secs(1);
+        assert!(b.try_take(t1, burst, refill_per_sec), "one token refilled");
+        assert!(
+            !b.try_take(t1, burst, refill_per_sec),
+            "only one token, not a full reset"
+        );
+    }
+
+    #[test]
+    fn idle_does_not_accrue_beyond_burst() {
+        // An idle pair must not bank an unbounded reserve: after a long quiet
+        // spell it still gets only `burst` immediate attempts, so a member that
+        // was silent for hours cannot burst 3600 retries at once.
+        let t0 = Instant::now();
+        let (burst, refill_per_sec) = (2.0, 1.0);
+        let mut b = TokenBucket::new(burst, t0);
+
+        let much_later = t0 + Duration::from_secs(3600);
+        assert!(b.try_take(much_later, burst, refill_per_sec));
+        assert!(b.try_take(much_later, burst, refill_per_sec));
+        assert!(
+            !b.try_take(much_later, burst, refill_per_sec),
+            "tokens clamp at burst regardless of idle time"
+        );
+    }
+}

src/client.rs

@@ -724,6 +724,13 @@ pub struct Client {
     pub(crate) inbound_durability_hook:
         std::sync::OnceLock<Arc<dyn crate::types::durability_hook::InboundDurabilityHook>>,
 
+    /// Optional retry-receipt admission policy (see
+    /// [`crate::types::retry_admission::RetryAdmission`]): an operator opt-in to
+    /// drop some group/status retries. `None` (default) keeps WA Web behavior
+    /// with a single lock-free `OnceLock::get` on the receive path.
+    pub(crate) retry_admission:
+        std::sync::OnceLock<Arc<dyn crate::types::retry_admission::RetryAdmission>>,
+
     /// Chat state (typing indicator) handlers registered by external consumers.
     /// Each handler receives a `ChatStateEvent` describing the chat, optional participant and state.
     pub(crate) chatstate_handlers: Arc<RwLock<Vec<ChatStateHandler>>>,
@@ -768,13 +775,16 @@ pub struct Client {
     pub(crate) group_distribution_locks: Cache<Jid, Arc<async_lock::Mutex<()>>>,
 
     /// Last `(devices, sender-key-device map)` Arc pair with an empty `needs_skdm`,
-    /// so a warm repeat send skips `filter_skdm_targets`. `Weak` keeps the pointer
-    /// comparison ABA-safe, matching `GroupDevicesMemo`.
+    /// plus the map's generation at that point, so a warm repeat send skips
+    /// `filter_skdm_targets`. `Weak` keeps the pointer comparison ABA-safe
+    /// (matching `GroupDevicesMemo`); the generation catches an in-place cold
+    /// flip that leaves the `Arc` pointer unchanged.
     pub(crate) skdm_warm_memo: Cache<
         Jid,
         (
             std::sync::Weak<wacore::send::ResolvedGroupDevices>,
             std::sync::Weak<crate::sender_key_device_cache::SenderKeyDeviceMap>,
+            u64,
         ),
     >,
 

src/client/accessors.rs

@@ -83,6 +83,24 @@ impl Client {
         self.resend_rate_limiter.set_rate(burst, refill_per_min);
     }
 
+    /// Register a [`RetryAdmission`] policy: an opt-in gate that can drop inbound
+    /// group/status retry receipts from other accounts before any repair work
+    /// runs. Unset (the default) admits every receipt, matching WhatsApp Web,
+    /// with zero overhead on the receive path.
+    ///
+    /// Set once, before connecting; a later call is ignored and returns `false`
+    /// (the already-registered policy stays in effect). Live tuning belongs
+    /// inside the policy itself (e.g. atomics), not in re-registration. See
+    /// `examples/retry_quarantine.rs`.
+    ///
+    /// [`RetryAdmission`]: crate::types::retry_admission::RetryAdmission
+    pub fn set_retry_admission(
+        &self,
+        policy: Arc<dyn crate::types::retry_admission::RetryAdmission>,
+    ) -> bool {
+        self.retry_admission.set(policy).is_ok()
+    }
+
     /// Cumulative wire I/O and activity counters for this client session.
     ///
     /// Always available, no feature gate: recording costs one relaxed atomic

src/client/lifecycle.rs

@@ -237,6 +237,7 @@ impl Client {
             passkey_opening: AtomicBool::new(false),
             custom_enc_handlers: std::sync::OnceLock::new(),
             inbound_durability_hook: std::sync::OnceLock::new(),
+            retry_admission: std::sync::OnceLock::new(),
             chatstate_handlers: Arc::new(RwLock::new(Vec::new())),
             pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(),
             pdo_requested: cache_config.pdo_requested.build_with_ttl(),