From 31cf45ce35a5ad4340aa709e4a574e79cdc465b3 Mon Sep 17 00:00:00 2001 From: Edison Jwa Date: Mon, 8 Jun 2026 22:51:41 +0900 Subject: [PATCH] feat(protocol): classify poke notification strength --- crates/chanora_protocol/src/adapter.rs | 43 ++++- crates/chanora_protocol/src/dto.rs | 4 + crates/chanora_protocol/src/lib.rs | 4 +- crates/chanora_protocol/src/poke_limiter.rs | 174 ++++++++++++++++++++ 4 files changed, 216 insertions(+), 9 deletions(-) create mode 100644 crates/chanora_protocol/src/poke_limiter.rs diff --git a/crates/chanora_protocol/src/adapter.rs b/crates/chanora_protocol/src/adapter.rs index 10628fe..d5e9ee1 100644 --- a/crates/chanora_protocol/src/adapter.rs +++ b/crates/chanora_protocol/src/adapter.rs @@ -41,6 +41,7 @@ use crate::dto::{ ChannelId, ChannelInfo, ChatMessage, ClientId, ClientInfo, ClientProfile, MessageTarget, ProtocolDelta, ServerActivity, ServerSnapshot, }; +use crate::poke_limiter::PokeLimiter; use crate::ProtocolError; const SPEAKING_ACTIVITY_WINDOW: Duration = Duration::from_millis(750); @@ -713,6 +714,7 @@ async fn connection_task( // reply channel — at most 3 s of pending state per move. let mut pending_moves: PendingMoves = HashMap::new(); let mut voice_activity: HashMap = HashMap::new(); + let mut poke_limiter = PokeLimiter::new(); // Main loop: pump events, service requests, forward voice. loop { @@ -742,6 +744,7 @@ async fn connection_task( &channels.activity, &channels.delta, &mut pending_moves, + &mut poke_limiter, ), }, Ok(Some(Err(e))) => { @@ -841,6 +844,7 @@ async fn connection_task( &channels, &mut pending_moves, &mut voice_activity, + &mut poke_limiter, ) .await; let _ = reply.send(r); @@ -903,6 +907,7 @@ fn handle_non_audio_stream_item( activity_tx: &mpsc::Sender, delta_tx: &mpsc::Sender, pending_moves: &mut PendingMoves, + poke_limiter: &mut PokeLimiter, ) { match item { StreamItem::BookEvents(events) => { @@ -964,19 +969,25 @@ fn handle_non_audio_stream_item( message, } = ev { - let mapped = match target { - tsclientlib::MessageTarget::Server => MessageTarget::Server, - tsclientlib::MessageTarget::Channel => MessageTarget::Channel, + let (mapped, poke_strength) = match target { + tsclientlib::MessageTarget::Server => (MessageTarget::Server, None), + tsclientlib::MessageTarget::Channel => (MessageTarget::Channel, None), tsclientlib::MessageTarget::Client(id) => { - MessageTarget::Client(id.0 as u64) + (MessageTarget::Client(id.0 as u64), None) + } + tsclientlib::MessageTarget::Poke(id) => { + let own_client_id = + con.get_state().ok().map(|state| state.own_client.0 as u64); + let strength = poke_limiter.record(invoker.id.0 as u64, own_client_id); + (MessageTarget::Poke(id.0 as u64), Some(strength)) } - tsclientlib::MessageTarget::Poke(id) => MessageTarget::Poke(id.0 as u64), }; let _ = chat_tx.try_send(ChatMessage { sender_id: ClientId(invoker.id.0 as u64), sender_name: sanitize(&invoker.name), message: sanitize(&message), target: mapped, + poke_strength, }); } } @@ -1164,6 +1175,7 @@ async fn fetch_client_profile( channels: &EventChannels, pending_moves: &mut PendingMoves, voice_activity: &mut HashMap, + poke_limiter: &mut PokeLimiter, ) -> Result { let target_id = TsClientId(client_id as u16); @@ -1209,6 +1221,7 @@ async fn fetch_client_profile( channels, pending_moves, voice_activity, + poke_limiter, ) .await; } @@ -1219,6 +1232,7 @@ async fn fetch_client_profile( channels, pending_moves, voice_activity, + poke_limiter, ) .await; } @@ -1233,6 +1247,7 @@ async fn fetch_client_profile( channels, pending_moves, voice_activity, + poke_limiter, ) .await { @@ -1251,6 +1266,7 @@ async fn fetch_client_profile( channels, pending_moves, voice_activity, + poke_limiter, ) .await { @@ -1264,9 +1280,16 @@ async fn fetch_client_profile( } let db_info = if refresh_plan.needs_client_db_info { - request_client_db_info(con, database_id, channels, pending_moves, voice_activity) - .await - .ok() + request_client_db_info( + con, + database_id, + channels, + pending_moves, + voice_activity, + poke_limiter, + ) + .await + .ok() } else { None }; @@ -1426,6 +1449,7 @@ async fn request_messages( channels: &EventChannels, pending_moves: &mut PendingMoves, voice_activity: &mut HashMap, + poke_limiter: &mut PokeLimiter, ) -> Result, ProtocolError> { let handle = command .send_with_result(con) @@ -1468,6 +1492,7 @@ async fn request_messages( &channels.activity, &channels.delta, pending_moves, + poke_limiter, ), } } @@ -1479,6 +1504,7 @@ async fn request_client_db_info( channels: &EventChannels, pending_moves: &mut PendingMoves, voice_activity: &mut HashMap, + poke_limiter: &mut PokeLimiter, ) -> Result { let messages = request_messages( con, @@ -1486,6 +1512,7 @@ async fn request_client_db_info( channels, pending_moves, voice_activity, + poke_limiter, ) .await?; for message in messages { diff --git a/crates/chanora_protocol/src/dto.rs b/crates/chanora_protocol/src/dto.rs index fdf076b..142d1ff 100644 --- a/crates/chanora_protocol/src/dto.rs +++ b/crates/chanora_protocol/src/dto.rs @@ -3,6 +3,8 @@ use serde::{Deserialize, Serialize}; +pub use crate::poke_limiter::PokeStrength; + /// Opaque server-side channel identifier. Internal representation is /// the upstream u64 but callers must treat it as opaque. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -54,6 +56,8 @@ pub struct ChatMessage { pub message: String, /// Target scope of this message. pub target: MessageTarget, + /// Strength classification for poke notifications. + pub poke_strength: Option, } /// A server-activity notification derived from TeamSpeak bookkeeping events. diff --git a/crates/chanora_protocol/src/lib.rs b/crates/chanora_protocol/src/lib.rs index 6b6a7ef..3d5d409 100644 --- a/crates/chanora_protocol/src/lib.rs +++ b/crates/chanora_protocol/src/lib.rs @@ -35,12 +35,14 @@ mod adapter; mod dto; +pub mod poke_limiter; pub use adapter::{ConnectConfig, DisconnectReason, InboundVoice, ProtocolClient, SnapshotProbe}; pub use dto::{ ChannelId, ChannelInfo, ChatMessage, ClientId, ClientInfo, ClientProfile, MessageTarget, - ProtocolDelta, ServerActivity, ServerSnapshot, + PokeStrength, ProtocolDelta, ServerActivity, ServerSnapshot, }; +pub use poke_limiter::PokeLimiter; // Re-export the upstream voice types so chanora_audio can build outbound // voice packets without taking a direct dependency on tsclientlib / diff --git a/crates/chanora_protocol/src/poke_limiter.rs b/crates/chanora_protocol/src/poke_limiter.rs new file mode 100644 index 0000000..bb77c2c --- /dev/null +++ b/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, +} + +#[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) -> 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, + 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 + ); + } +}