diff --git a/src/client.rs b/src/client.rs index 25216e009..9557c8d10 100644 --- a/src/client.rs +++ b/src/client.rs @@ -177,7 +177,24 @@ type ChatStateHandler = Arc; #[derive(Clone)] pub(crate) struct ChatLane { pub enqueue_lock: Arc>, - pub queue_tx: async_channel::Sender>, + pub queue_tx: async_channel::Sender, +} + +impl ChatLane { + pub(crate) fn try_enqueue( + &self, + node: Arc, + ) -> Result<(), async_channel::TrySendError> { + self.queue_tx.try_send(QueuedChatMessage { + node, + lane_liveness: Arc::clone(&self.enqueue_lock), + }) + } +} + +pub(crate) struct QueuedChatMessage { + pub node: Arc, + pub lane_liveness: Arc>, } const APP_STATE_RETRY_MAX_ATTEMPTS: u32 = 6; diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index b515062a5..0ed4bcc2d 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -166,6 +166,7 @@ impl Client { .build(), chat_lanes: Cache::builder() .max_capacity(cache_config.chat_lanes_capacity.max(1)) + .evict_guard(|lane: &ChatLane| Arc::strong_count(&lane.enqueue_lock) <= 1) .build(), lid_pn_cache: Arc::new(LidPnCache::with_config( &cache_config.lid_pn_cache, diff --git a/src/client/tests.rs b/src/client/tests.rs index b59319476..809e32d79 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -2824,6 +2824,84 @@ async fn held_group_distribution_lane_survives_capacity_pressure() { assert!(first_again.try_lock().is_some()); } +#[tokio::test] +async fn active_chat_lane_survives_capacity_pressure() { + fn test_lane() -> (ChatLane, async_channel::Receiver) { + let (queue_tx, queue_rx) = async_channel::unbounded(); + ( + ChatLane { + enqueue_lock: Arc::new(async_lock::Mutex::new(())), + queue_tx, + }, + queue_rx, + ) + } + + let config = crate::cache_config::CacheConfig { + chat_lanes_capacity: 1, + ..Default::default() + }; + let client = crate::test_utils::create_test_client_with_config( + "chat_lane_eviction", + Arc::new(MockHttpClient), + config, + ) + .await; + + let first: Jid = "120363000000000021@g.us".parse().unwrap(); + let second: Jid = "120363000000000022@g.us".parse().unwrap(); + let (first_lane, first_rx) = test_lane(); + let first_tx_probe = first_lane.queue_tx.clone(); + client.chat_lanes.insert(first.clone(), first_lane).await; + + let first_lane = client.chat_lanes.get(&first).await.unwrap(); + let node = NodeBuilder::new("message") + .attr("from", first.clone()) + .attr("id", "ACTIVE-LANE-1") + .build(); + first_lane.try_enqueue(node_to_owned_ref(node)).unwrap(); + drop(first_lane); + let active_message = first_rx.recv().await.unwrap(); + + let (second_lane, _second_rx) = test_lane(); + client.chat_lanes.insert(second, second_lane).await; + + let first_again = client + .chat_lanes + .get(&first) + .await + .expect("an active lane must remain cached"); + assert!(first_again.queue_tx.same_channel(&first_tx_probe)); + + let next_node = NodeBuilder::new("message") + .attr("from", first.clone()) + .attr("id", "ACTIVE-LANE-2") + .build(); + first_again + .try_enqueue(node_to_owned_ref(next_node)) + .unwrap(); + drop(first_again); + drop(active_message); + let next_active_message = first_rx.recv().await.unwrap(); + + let third: Jid = "120363000000000023@g.us".parse().unwrap(); + let (third_lane, _third_rx) = test_lane(); + client.chat_lanes.insert(third, third_lane).await; + assert!( + client.chat_lanes.get(&first).await.is_some(), + "a lane with an in-flight message must not be evicted" + ); + + drop(next_active_message); + let fourth: Jid = "120363000000000024@g.us".parse().unwrap(); + let (fourth_lane, _fourth_rx) = test_lane(); + client.chat_lanes.insert(fourth, fourth_lane).await; + assert!( + client.chat_lanes.get(&first).await.is_none(), + "an idle lane must become evictable again" + ); +} + /// Proves that `is_connected()` no longer gives false negatives under mutex /// contention. Before the fix, `try_lock()` would fail when another task held /// the noise_socket mutex, causing `is_connected()` to return `false` even diff --git a/src/handlers/message.rs b/src/handlers/message.rs index 597f9ae8b..bfb474bc5 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -1,5 +1,5 @@ use super::traits::StanzaHandler; -use crate::client::{ChatLane, Client}; +use crate::client::{ChatLane, Client, QueuedChatMessage}; use async_trait::async_trait; use log::warn; use std::sync::Arc; @@ -45,7 +45,7 @@ impl MessageHandler { // Lock serializes enqueue order for this chat let _guard = lane.enqueue_lock.lock().await; - if let Err(e) = lane.queue_tx.try_send(node) { + if let Err(e) = lane.try_enqueue(node) { warn!("Failed to enqueue message for processing: {e}"); // Cancel ack so server redelivers *cancelled = true; @@ -75,7 +75,7 @@ impl StanzaHandler for MessageHandler { /// Construct a ChatLane with a spawned worker task. Extracted to keep the /// init closure passed to `get_with_by_ref` small. fn create_chat_lane(client: &Arc) -> ChatLane { - let (tx, rx) = async_channel::unbounded::>(); + let (tx, rx) = async_channel::unbounded::(); let client_for_worker = client.clone(); let spawn_generation = client @@ -85,7 +85,11 @@ fn create_chat_lane(client: &Arc) -> ChatLane { client .runtime .spawn(Box::pin(async move { - while let Ok(msg_node) = rx.recv().await { + while let Ok(QueuedChatMessage { + node: msg_node, + lane_liveness, // Prevents capacity eviction until processing finishes. + }) = rx.recv().await + { if client_for_worker .connection_generation .load(std::sync::atomic::Ordering::Acquire) @@ -111,6 +115,7 @@ fn create_chat_lane(client: &Arc) -> ChatLane { MAX_MESSAGE_DELAY_MS / 1000 ); } + drop(lane_liveness); } })) .detach();