feat(protocol): classify poke notification strength

crates/chanora_protocol/src/poke_limiter.rs

@@ -0,0 +1,174 @@
+//! Per-connection poke strength classification.
+
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+/// Notification strength assigned to an inbound poke.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub enum PokeStrength {
+    /// Poke should be surfaced at full strength.
+    Strong,
+    /// Poke is rate-limited but below overflow severity.
+    Suppressed,
+    /// Poke remains suppressed after repeated suppressed pokes.
+    SuppressedOverflow,
+}
+
+/// Per-connection poke limiter.
+#[derive(Debug)]
+pub struct PokeLimiter {
+    window: Duration,
+    entries: HashMap<u64, PokeEntry>,
+}
+
+#[derive(Debug)]
+struct PokeEntry {
+    tokens: u8,
+    last_refill: Instant,
+    suppressed_in_window: u8,
+}
+
+impl PokeLimiter {
+    const CAPACITY: u8 = 2;
+    const OVERFLOW_THRESHOLD: u8 = 3;
+
+    /// Create a limiter using the default five-minute refill interval.
+    pub fn new() -> Self {
+        Self {
+            window: Duration::from_secs(5 * 60),
+            entries: HashMap::new(),
+        }
+    }
+
+    /// Record a poke at the current instant.
+    pub fn record(&mut self, sender_id: u64, own_client_id: Option<u64>) -> PokeStrength {
+        self.record_at(sender_id, own_client_id, Instant::now())
+    }
+
+    /// Record a poke at an injected instant.
+    pub fn record_at(
+        &mut self,
+        sender_id: u64,
+        own_client_id: Option<u64>,
+        now: Instant,
+    ) -> PokeStrength {
+        if own_client_id == Some(sender_id) {
+            return PokeStrength::Suppressed;
+        }
+
+        let entry = self.entries.entry(sender_id).or_insert(PokeEntry {
+            tokens: Self::CAPACITY,
+            last_refill: now,
+            suppressed_in_window: 0,
+        });
+
+        if now.duration_since(entry.last_refill) >= self.window {
+            entry.tokens = Self::CAPACITY;
+            entry.last_refill = now;
+            entry.suppressed_in_window = 0;
+        }
+
+        if entry.tokens > 0 {
+            entry.tokens -= 1;
+            return PokeStrength::Strong;
+        }
+
+        entry.suppressed_in_window = entry.suppressed_in_window.saturating_add(1);
+        if entry.suppressed_in_window >= Self::OVERFLOW_THRESHOLD {
+            PokeStrength::SuppressedOverflow
+        } else {
+            PokeStrength::Suppressed
+        }
+    }
+}
+
+impl Default for PokeLimiter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn first_two_pokes_are_strong_third_is_suppressed() {
+        let mut limiter = PokeLimiter::new();
+        let now = Instant::now();
+
+        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
+            PokeStrength::Strong
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
+            PokeStrength::Suppressed
+        );
+    }
+
+    #[test]
+    fn self_poke_is_suppressed_without_consuming_token() {
+        let mut limiter = PokeLimiter::new();
+        let now = Instant::now();
+
+        assert_eq!(limiter.record_at(7, Some(7), now), PokeStrength::Suppressed);
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
+            PokeStrength::Strong
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
+            PokeStrength::Strong
+        );
+    }
+
+    #[test]
+    fn overflow_after_three_suppressed_pokes_in_five_minutes() {
+        let mut limiter = PokeLimiter::new();
+        let now = Instant::now();
+
+        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
+            PokeStrength::Strong
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
+            PokeStrength::Suppressed
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(3)),
+            PokeStrength::Suppressed
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(4)),
+            PokeStrength::SuppressedOverflow
+        );
+    }
+
+    #[test]
+    fn refill_after_interval_uses_injected_time_without_sleeping() {
+        let mut limiter = PokeLimiter::new();
+        let now = Instant::now();
+
+        assert_eq!(limiter.record_at(7, Some(1), now), PokeStrength::Strong);
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(1)),
+            PokeStrength::Strong
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(2)),
+            PokeStrength::Suppressed
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(5 * 60)),
+            PokeStrength::Strong
+        );
+        assert_eq!(
+            limiter.record_at(7, Some(1), now + Duration::from_secs(5 * 60 + 1)),
+            PokeStrength::Strong
+        );
+    }
+}