diff --git a/src/client.rs b/src/client.rs index 4b1e52aa0..4a1b12385 100644 --- a/src/client.rs +++ b/src/client.rs @@ -973,7 +973,10 @@ pub struct Client { pub(crate) transport_events: Arc>>>, pub(crate) transport_factory: Arc, - pub(crate) noise_socket: Arc>>>, + /// Replaced per connection, so not a `OnceLock` — but every critical section + /// is a clone or a store, so a sync lock makes holding it across an `.await` + /// a compile error on the send path rather than a review question. + pub(crate) noise_socket: Arc>>>, /// Pending IQ/ack response waiters keyed by request id. /// @@ -1027,7 +1030,9 @@ pub struct Client { pub(crate) lid_pn_cache: Arc, pub(crate) ab_props: Arc, - pub group_cache: Mutex>>, + /// Lazily built on the first group send and never replaced afterwards, so a + /// `OnceLock` keeps the read on that path down to an atomic load. + pub group_cache: std::sync::OnceLock>, pub(crate) expected_disconnect: Arc, /// Set by `reconnect()` to suppress the "Message loop exited with an error" warning. @@ -1093,7 +1098,9 @@ pub struct Client { pub(crate) needs_initial_full_sync: Arc, - pub(crate) app_state_processor: Mutex>>, + /// Built on first app-state use and never replaced: reconnect clears the + /// processor's key cache in place rather than swapping the processor. + pub(crate) app_state_processor: std::sync::OnceLock>, pub(crate) app_state_key_requests: Arc, wacore::time::Instant>>>, /// Tracks collections currently being synced to prevent duplicate sync tasks. /// Matches WA Web's in-flight tracking set in WAWebSyncdCollectionsStateMachine. @@ -1159,7 +1166,7 @@ pub struct Client { )>, >, /// Contacts with active presence subscriptions that must be re-subscribed on reconnect. - pub(crate) presence_subscriptions: Arc>>, + pub(crate) presence_subscriptions: Arc>>, /// Metrics for granular offline sync logging pub(crate) offline_sync_metrics: Arc, /// Drives the WA Web pull-batch loop for offline backlog delivery. @@ -1241,7 +1248,12 @@ pub struct Client { /// 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>>, + /// + /// Copy-on-write behind a sync lock, guarded by `chatstate_handler_count` so + /// the default (no handler registered) never takes the lock nor builds the + /// event that only a handler would read. + pub(crate) chatstate_handlers: Arc>>, + pub(crate) chatstate_handler_count: AtomicUsize, pub(crate) pdo_pending_requests: Cache, @@ -1358,6 +1370,10 @@ pub struct Client { /// Keeps retry regressions deterministic without corrupting the test database. #[cfg(test)] pub(crate) app_state_key_share_prepare_test_failures: AtomicU32, + /// Counts `ChatStateEvent` constructions, so a test can prove the + /// no-handler fast path skips the build rather than just the invoke. + #[cfg(test)] + pub(crate) chatstate_events_built: AtomicU32, /// Holds the background saver's AbortHandle so the task lifetime follows /// `Arc` ref count instead of the Bot wrapper's. Set once by diff --git a/src/client/accessors.rs b/src/client/accessors.rs index f40ca4869..e40514f8e 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -12,19 +12,15 @@ pub struct IdentityTags { } impl Client { - pub(crate) async fn get_group_cache(&self) -> Arc { - let mut guard = self.group_cache.lock().await; - if let Some(cache) = guard.as_ref() { - return cache.clone(); - } - debug!("Initializing Group Cache for the first time."); - let cache = Arc::new( - self.cache_config - .group_cache - .build_typed_ttl(self.cache_config.cache_stores.group_cache.clone(), "group"), - ); - *guard = Some(cache.clone()); - cache + pub(crate) fn get_group_cache(&self) -> &Arc { + self.group_cache.get_or_init(|| { + debug!("Initializing Group Cache for the first time."); + Arc::new( + self.cache_config + .group_cache + .build_typed_ttl(self.cache_config.cache_stores.group_cache.clone(), "group"), + ) + }) } /// Subscribe an external event handler with an explicit event filter. @@ -158,10 +154,9 @@ impl Client { .unwrap_or_else(|p| p.into_inner()) .len(); - // Only the Arc is taken under the mutex — the walk must not block - // get_group_cache(), which every group send goes through. - let group_cache_arc = self.group_cache.lock().await.clone(); - let group_cache = match group_cache_arc { + // `get()`, not `get_group_cache()`: a report must not be what builds the + // cache, so an un-warmed client still reports zero entries. + let group_cache = match self.group_cache.get() { // Arc's HeapSize already includes size_of::(). Some(cache) => { cache @@ -188,10 +183,14 @@ impl Client { // Each count read into a local so no two guards are ever held at once. let response_waiters = self.response_waiters_guard().len(); - let presence_subscriptions = self.presence_subscriptions.lock().await.len(); + let presence_subscriptions = self + .presence_subscriptions + .lock() + .unwrap_or_else(|p| p.into_inner()) + .len(); let app_state_key_requests = self.app_state_key_requests.lock().await.len(); let app_state_syncing = self.app_state_syncing.len(); - let chatstate_handlers = self.chatstate_handlers.read().await.len(); + let chatstate_handlers = self.chatstate_handler_count.load(Ordering::Acquire); let history_sync_activity = self.history_sync_activity.snapshot(); let history_sync_tasks = CollectionStats::new( history_sync_activity.tasks as u64, diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 619bc4b6d..545ef92ac 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -51,10 +51,10 @@ impl Client { } /// Get the active noise socket, or error if not connected. - pub(crate) async fn get_noise_socket(&self) -> Result, ClientError> { + pub(crate) fn get_noise_socket(&self) -> Result, ClientError> { self.noise_socket .lock() - .await + .unwrap_or_else(|p| p.into_inner()) .clone() .ok_or(ClientError::NotConnected) } diff --git a/src/client/app_state.rs b/src/client/app_state.rs index b0c5e99d4..250eb189c 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -642,18 +642,14 @@ fn finalize_app_state_key_request_peers( } impl Client { - pub(crate) async fn get_app_state_processor(&self) -> Arc { - let mut guard = self.app_state_processor.lock().await; - if let Some(proc) = guard.as_ref() { - return proc.clone(); - } - debug!("Initializing AppStateProcessor for the first time."); - let proc = Arc::new(AppStateProcessor::new( - self.persistence_manager.backend(), - self.runtime.clone(), - )); - *guard = Some(proc.clone()); - proc + pub(crate) fn get_app_state_processor(&self) -> &Arc { + self.app_state_processor.get_or_init(|| { + debug!("Initializing AppStateProcessor for the first time."); + Arc::new(AppStateProcessor::new( + self.persistence_manager.backend(), + self.runtime.clone(), + )) + }) } /// Pre-download every external blob (snapshots + patch external mutations) @@ -1698,7 +1694,7 @@ impl Client { }); } - let proc = self.get_app_state_processor().await; + let proc = self.get_app_state_processor(); // Pre-download all external blobs for all collections in the response, // concurrently (independent CDN GETs, keyed by directPath). let pre_downloaded = self.pre_download_external_blobs(&patch_lists).await; @@ -2008,7 +2004,7 @@ impl Client { debug!(target: "Client/AppState", "Parsed patch list for {:?}: has_snapshot_ref={} has_more_patches={} patches_count={}", name, pl.snapshot_ref.is_some(), pl.has_more_patches, pl.patches.len()); - let proc = self.get_app_state_processor().await; + let proc = self.get_app_state_processor(); // Pre-download all external blobs (snapshot and patch mutations), // concurrently, keyed by directPath. @@ -2394,7 +2390,7 @@ impl Client { ), None => None, }; - let proc = self.get_app_state_processor().await; + let proc = self.get_app_state_processor(); for attempt in 1..=APP_STATE_PATCH_SEND_ATTEMPTS { // Cloned per attempt because a conflict rebuilds the patch against @@ -2531,7 +2527,7 @@ impl Client { .cloned() .ok_or_else(|| anyhow::anyhow!("external blob not pre-downloaded: {path}")) }; - let proc = self.get_app_state_processor().await; + let proc = self.get_app_state_processor(); match proc.process_parsed_patch_list(list, &download, true).await { Ok((mutations, _, _)) => { wacore::telemetry::appstate_mutations(mutations.len() as u64); diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index cbd28d840..70cbf15d5 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -3493,11 +3493,7 @@ mod tests { let hashed: Jid = format!("{contact_lid}@lid").parse().expect("jid"); assert!( - client - .pending_device_sync - .take_all() - .await - .contains(&hashed), + client.pending_device_sync.take_all().contains(&hashed), "the hashed contact must be queued for a device-list refresh" ); } @@ -3522,7 +3518,7 @@ mod tests { .await; assert!( - client.pending_device_sync.take_all().await.is_empty(), + client.pending_device_sync.take_all().is_empty(), "an unresolvable hash must not refresh an unrelated contact" ); } diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index f73075e53..66c5f9b0a 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -343,7 +343,7 @@ impl Client { transport: Arc::new(Mutex::new(None)), transport_events: Arc::new(Mutex::new(None)), transport_factory, - noise_socket: Arc::new(Mutex::new(None)), + noise_socket: Arc::new(std::sync::Mutex::new(None)), response_waiters: Arc::new(std::sync::Mutex::new(ResponseWaiterMap::default())), node_waiters: std::sync::Mutex::new(Vec::new()), @@ -378,7 +378,7 @@ impl Client { cache_config.cache_stores.lid_pn_cache.clone(), )), ab_props: Arc::new(wacore::store::ab_props::AbPropsCache::new()), - group_cache: Mutex::new(None), + group_cache: std::sync::OnceLock::new(), expected_disconnect: Arc::new(AtomicBool::new(false)), intentional_reconnect: AtomicBool::new(false), @@ -421,7 +421,7 @@ impl Client { needs_initial_full_sync: Arc::new(app_state::BootstrapGate::new(false)), - app_state_processor: Mutex::new(None), + app_state_processor: std::sync::OnceLock::new(), app_state_key_requests: Arc::new(Mutex::new(HashMap::new())), app_state_syncing: app_state::SyncInFlight::new(), app_state_send_lock: Arc::new(Mutex::new(())), @@ -438,7 +438,7 @@ impl Client { outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()), delivery_receipt_queue: std::sync::OnceLock::new(), transport_ack_queue: std::sync::OnceLock::new(), - presence_subscriptions: Arc::new(Mutex::new(HashSet::new())), + presence_subscriptions: Arc::new(std::sync::Mutex::new(HashSet::new())), socket_ready_notifier: Arc::new(event_listener::Event::new()), is_ready: Arc::new(AtomicBool::new(false)), connected_notifier: Arc::new(event_listener::Event::new()), @@ -460,10 +460,13 @@ impl Client { signal_flush_test_in_attempt: AtomicU32::new(0), #[cfg(test)] app_state_key_share_prepare_test_failures: AtomicU32::new(0), + #[cfg(test)] + chatstate_events_built: AtomicU32::new(0), 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())), + chatstate_handlers: Arc::new(std::sync::RwLock::new(Arc::from([]))), + chatstate_handler_count: AtomicUsize::new(0), pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(), pdo_requested: cache_config.pdo_requested.build_with_ttl(), device_registry_cache: device_topology::DeviceRegistryCache::new( @@ -844,7 +847,7 @@ impl Client { *self.transport.lock().await = Some(transport); *self.transport_events.lock().await = Some(transport_events); - *self.noise_socket.lock().await = Some(noise_socket); + *self.noise_socket.lock().unwrap_or_else(|p| p.into_inner()) = Some(noise_socket); self.is_connected.store(true, Ordering::Release); // Notify waiters that socket is ready (before login) @@ -1177,7 +1180,7 @@ impl Client { // afterwards would strip the replacement connection instead of the one being torn down. let transport = self.transport.lock().await.take(); *self.transport_events.lock().await = None; - *self.noise_socket.lock().await = None; + *self.noise_socket.lock().unwrap_or_else(|p| p.into_inner()) = None; if let Some(transport) = transport { transport.disconnect().await; } @@ -1235,7 +1238,7 @@ impl Client { // Reset dead-socket timestamps so stale values from the previous // connection don't trigger an immediate reconnect on the next one. self.stats.reset_connection_activity(); - self.pending_device_sync.clear().await; + self.pending_device_sync.clear(); // Reset offline sync state for next connection self.offline_sync_completed.store(false, Ordering::Relaxed); self.offline_sync_finish_started @@ -1297,9 +1300,11 @@ impl Client { *self.media_conn.write().await = None; // Clear app state key cache — keys will be re-fetched from DB on demand - let processor = self.app_state_processor.lock().await.clone(); - if let Some(processor) = processor { - processor.clear_key_cache().await; + // main took the processor out of the mutex before awaiting so the guard + // did not span the clear; the write-once cell has no guard to span, so + // the borrow is the whole of it. + if let Some(proc) = self.app_state_processor.get() { + proc.clear_key_cache().await; } #[cfg(feature = "client-lifecycle")] drop(scope_close); diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 556135f16..67d826a48 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -12,7 +12,7 @@ impl Client { /// This bypasses node logging and `sent_node_waiter` resolution — use /// [`send_node`](Client::send_node) for normal stanza sending. pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { - let noise_socket = self.get_noise_socket().await?; + let noise_socket = self.get_noise_socket()?; // Wire bytes and the last-sent timestamp are recorded by the noise // sender task at the actual transport write. noise_socket @@ -53,7 +53,7 @@ impl Client { results: &mut Vec, ) -> Result<(), ClientError> { results.clear(); - let noise_socket = match self.get_noise_socket().await { + let noise_socket = match self.get_noise_socket() { Ok(socket) => socket, Err(error) => { frames.clear(); @@ -455,11 +455,19 @@ impl Client { /// Register a chatstate handler which will be invoked when a `` stanza is received. /// /// The handler receives a `ChatStateEvent` with the parsed chat state information. - pub async fn register_chatstate_handler( - &self, - handler: Arc, - ) { - self.chatstate_handlers.write().await.push(handler); + pub fn register_chatstate_handler(&self, handler: Arc) { + let mut guard = self + .chatstate_handlers + .write() + .unwrap_or_else(|p| p.into_inner()); + let mut handlers = Vec::with_capacity(guard.len() + 1); + handlers.extend(guard.iter().cloned()); + handlers.push(handler); + *guard = Arc::from(handlers); + // Published after the snapshot is in place, so a reader that sees a + // non-zero count always finds the handler behind it. + self.chatstate_handler_count + .store(guard.len(), Ordering::Release); } /// Dispatch a parsed chatstate stanza to registered handlers. @@ -512,10 +520,20 @@ impl Client { .build(), )); - // Invoke legacy callback handlers + // Invoke legacy callback handlers. Building the event is only worth it + // once something reads it, and the default registers nothing. + if self.chatstate_handler_count.load(Ordering::Acquire) == 0 { + return; + } + #[cfg(test)] + self.chatstate_events_built.fetch_add(1, Ordering::Release); let event = ChatStateEvent::from_stanza(stanza); - let handlers = self.chatstate_handlers.read().await.clone(); - for handler in handlers { + let handlers = self + .chatstate_handlers + .read() + .unwrap_or_else(|p| p.into_inner()) + .clone(); + for handler in handlers.iter().cloned() { let event_clone = event.clone(); self.runtime .spawn(Box::pin(async move { diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 3abd77169..7c3004f5b 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -163,7 +163,6 @@ impl Client { // so resolve it once instead of locking the mutex per frame. let noise_socket = self .get_noise_socket() - .await .map_err(|_| ReadLoopError::NotStarted("no noise socket"))?; // Frame decoder to parse incoming data diff --git a/src/client/tests.rs b/src/client/tests.rs index 1268cc387..337ffe454 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3358,14 +3358,14 @@ async fn test_is_connected_not_affected_by_mutex_contention() { write_key, read_key, ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.is_connected.store(true, Ordering::Release); assert!(client.is_connected(), "should report connected"); // Hold the noise_socket mutex — this used to make is_connected() return // false via try_lock() even though the socket was Some(...) - let _guard = client.noise_socket.lock().await; + let _guard = client.noise_socket.lock().unwrap(); assert!( client.is_connected(), "is_connected() must return true even while noise_socket mutex is held" @@ -3432,7 +3432,7 @@ async fn disconnect_does_not_signal_connection_cleanup_before_outbound_flush() { ); *client.transport.lock().await = Some(transport); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.is_connected.store(true, Ordering::Release); let cleanup_signal = client.connection_shutdown_signal(); @@ -3513,7 +3513,7 @@ async fn install_test_noise_socket( NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.set_connected_for_test(true); } @@ -3818,7 +3818,7 @@ async fn delivery_receipt_worker_sends_and_releases_flush() { NoiseCipher::new(&key).expect("valid key"), ); *client.transport.lock().await = Some(transport); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.is_connected.store(true, Ordering::Release); client.ack_received_message(&receipt_test_info("RCPT-WORKER-1")); @@ -3949,7 +3949,7 @@ async fn flush_waits_for_queued_delivery_receipts() { NoiseCipher::new(&key).expect("valid key"), ); *client.transport.lock().await = Some(transport); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.is_connected.store(true, Ordering::Release); client.ack_received_message(&receipt_test_info("RCPT-QUEUE-1")); @@ -4751,3 +4751,159 @@ async fn a_terminal_connect_failure_releases_a_parked_wait() { "reporting that no connection arrived" ); } + +/// The write-once cells (`group_cache`, `app_state_processor`) are read on the +/// send and app-state paths on the strength of never being rebuilt. These prove +/// the three ways that could break: a second call, a reconnect cleanup, and a +/// racing first call. +#[tokio::test] +async fn write_once_cells_return_the_same_instance_across_calls() { + let client = crate::test_utils::create_test_client().await; + + let group_cache = client.get_group_cache().clone(); + let processor = client.get_app_state_processor().clone(); + + assert!( + Arc::ptr_eq(&group_cache, client.get_group_cache()), + "the group cache must not be rebuilt on a second read" + ); + assert!( + Arc::ptr_eq(&processor, client.get_app_state_processor()), + "the app-state processor must not be rebuilt on a second read" + ); +} + +#[tokio::test] +async fn reconnect_cleanup_leaves_the_write_once_cells_installed() { + let client = crate::test_utils::create_test_client().await; + + let group_cache = client.get_group_cache().clone(); + let processor = client.get_app_state_processor().clone(); + + client.cleanup_connection_state().await; + + assert!( + Arc::ptr_eq(&group_cache, client.get_group_cache()), + "cleanup must not drop the group cache" + ); + assert!( + Arc::ptr_eq(&processor, client.get_app_state_processor()), + "cleanup clears the processor's key cache in place, it does not replace it" + ); + assert!( + client.group_cache.get().is_some() && client.app_state_processor.get().is_some(), + "and neither cell may fall back to uninitialized" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_first_readers_agree_on_one_instance() { + let client = crate::test_utils::create_test_client().await; + + // OS threads and a blocking barrier, not tasks on a worker pool: the + // getters are synchronous now, so every reader can be inside `get_or_init` + // at once instead of at most `worker_threads` of them, and none can be + // spawned late enough to find the cell already warm. + const READERS: usize = 16; + let results = tokio::task::spawn_blocking(move || { + let start = std::sync::Barrier::new(READERS); + std::thread::scope(|scope| { + let readers: Vec<_> = (0..READERS) + .map(|_| { + let client = &client; + let start = &start; + scope.spawn(move || { + start.wait(); + ( + client.get_group_cache().clone(), + client.get_app_state_processor().clone(), + ) + }) + }) + .collect(); + readers + .into_iter() + .map(|reader| reader.join().expect("reader thread should not panic")) + .collect::>() + }) + }) + .await + .expect("blocking scope should not panic"); + + let (first_cache, first_processor) = &results[0]; + for (cache, processor) in &results { + assert!( + Arc::ptr_eq(first_cache, cache), + "every racing reader must observe the same group cache" + ); + assert!( + Arc::ptr_eq(first_processor, processor), + "every racing reader must observe the same app-state processor" + ); + } +} + +fn test_chatstate_stanza() -> wacore::iq::chatstate::ChatstateStanza { + use wacore::iq::chatstate::{ChatstateSource, ChatstateStanza, ReceivedChatState}; + + ChatstateStanza { + source: ChatstateSource::User { + from: "15550001111@s.whatsapp.net".parse().expect("valid jid"), + }, + state: ReceivedChatState::Typing, + } +} + +#[tokio::test] +async fn chatstate_dispatch_skips_the_event_build_with_no_handlers() { + let client = crate::test_utils::create_test_client().await; + + client + .dispatch_chatstate_event(test_chatstate_stanza()) + .await; + + assert_eq!( + client.chatstate_events_built.load(Ordering::Acquire), + 0, + "the default registers no handler, so nothing should read the event" + ); +} + +#[tokio::test] +async fn chatstate_dispatch_reaches_every_registered_handler() { + let client = crate::test_utils::create_test_client().await; + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + + for tag in ["first", "second"] { + let seen = seen.clone(); + client.register_chatstate_handler(Arc::new(move |event| { + seen.lock() + .unwrap_or_else(|p| p.into_inner()) + .push((tag, event.chat.to_string())); + })); + } + + client + .dispatch_chatstate_event(test_chatstate_stanza()) + .await; + + crate::test_utils::poll_until("both chatstate handlers ran", || { + seen.lock().unwrap_or_else(|p| p.into_inner()).len() == 2 + }) + .await; + + let mut seen = seen.lock().unwrap_or_else(|p| p.into_inner()).clone(); + seen.sort(); + assert_eq!( + seen, + vec![ + ("first", "15550001111@s.whatsapp.net".to_string()), + ("second", "15550001111@s.whatsapp.net".to_string()), + ] + ); + assert_eq!( + client.chatstate_events_built.load(Ordering::Acquire), + 1, + "the event is built once and cloned per handler" + ); +} diff --git a/src/client/voip.rs b/src/client/voip.rs index ddf0d3e05..141966317 100644 --- a/src/client/voip.rs +++ b/src/client/voip.rs @@ -2011,7 +2011,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); (client, count) } @@ -2039,7 +2039,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client } @@ -3590,7 +3590,7 @@ mod tests { NoiseCipher::new(&[0u8; 32]).expect("valid key"), NoiseCipher::new(&[0u8; 32]).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(gated_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(gated_socket)); crate::test_utils::answer_iq( &client, &request_id, @@ -4731,7 +4731,7 @@ mod tests { NoiseCipher::new(&[0u8; 32]).expect("valid key"), NoiseCipher::new(&[0u8; 32]).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); let accept = tokio::spawn({ let client = client.clone(); @@ -4854,7 +4854,7 @@ mod tests { NoiseCipher::new(&[0u8; 32]).expect("valid key"), NoiseCipher::new(&[0u8; 32]).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(blocking_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(blocking_socket)); crate::test_utils::answer_iq( &client, &request_id, @@ -4997,7 +4997,7 @@ mod tests { NoiseCipher::new(&[0u8; 32]).expect("valid key"), NoiseCipher::new(&[0u8; 32]).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(gated_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(gated_socket)); crate::test_utils::answer_iq( &client, &request_id, diff --git a/src/features/chat_actions.rs b/src/features/chat_actions.rs index eee36ea08..bc79a5793 100644 --- a/src/features/chat_actions.rs +++ b/src/features/chat_actions.rs @@ -783,7 +783,7 @@ impl Client { use rand::Rng; use wacore::appstate::encode::encode_record; - let proc = self.get_app_state_processor().await; + let proc = self.get_app_state_processor(); let key_id = proc .backend .get_latest_sync_key_id() diff --git a/src/features/groups.rs b/src/features/groups.rs index 88285e5f9..c365c9647 100644 --- a/src/features/groups.rs +++ b/src/features/groups.rs @@ -350,13 +350,12 @@ pub(crate) struct GroupMetadataGuard<'a> { impl GroupMetadataGuard<'_> { pub(crate) async fn current(&self) -> Option> { - self.client.get_group_cache().await.get(self.jid).await + self.client.get_group_cache().get(self.jid).await } async fn cache(&self, info: Arc) { self.client .get_group_cache() - .await .insert(self.jid.clone(), info) .await; } @@ -399,11 +398,7 @@ impl GroupMetadataGuard<'_> { self.jid ); } - self.client - .get_group_cache() - .await - .invalidate(self.jid) - .await; + self.client.get_group_cache().invalidate(self.jid).await; } } @@ -445,7 +440,7 @@ impl<'a> Groups<'a> { jid: &Jid, freshness: crate::cache::Freshness, ) -> Result, GroupError> { - let cache = self.client.get_group_cache().await; + let cache = self.client.get_group_cache(); let mut cached = cache.get(jid).await; if freshness == crate::cache::Freshness::CachePreferred && let Some(cached) = cached.take() @@ -1751,7 +1746,7 @@ mod tests { ], AddressingMode::Pn, ); - let cache = client.get_group_cache().await; + let cache = client.get_group_cache(); cache.insert(group_jid.clone(), Arc::new(info)).await; let a = cache.get(&group_jid).await.expect("warm hit"); @@ -1771,7 +1766,7 @@ mod tests { vec!["12025550101@s.whatsapp.net".parse().unwrap()], AddressingMode::Pn, )); - let cache = client.get_group_cache().await; + let cache = client.get_group_cache(); cache.insert(group.clone(), Arc::clone(&previous)).await; let result = client @@ -1799,7 +1794,7 @@ mod tests { let parent: Jid = "120363000000000001@g.us".parse().unwrap(); let unrelated: Jid = "120363000000000002@g.us".parse().unwrap(); let removed: Jid = "12025550103@s.whatsapp.net".parse().unwrap(); - let cache = client.get_group_cache().await; + let cache = client.get_group_cache(); for jid in [&parent, &unrelated] { cache .insert( diff --git a/src/features/presence.rs b/src/features/presence.rs index a65a981a4..1bcb1e44f 100644 --- a/src/features/presence.rs +++ b/src/features/presence.rs @@ -140,24 +140,30 @@ impl<'a> Presence<'a> { debug!("presence subscribe: subscribing to {}", jid); let node = self.build_subscription_node(jid).await; self.client.send_node(node).await?; - self.client.track_presence_subscription(jid.clone()).await; + self.client.track_presence_subscription(jid.clone()); Ok(()) } /// Re-subscribe presence if the JID has an active subscription. /// Does not modify the tracking set. + /// + /// The check is re-read per JID rather than taken from the resubscribe + /// snapshot: an `unsubscribe` landing mid-resubscribe must not be undone. pub(crate) async fn re_subscribe_when_active(&self, jid: &Jid) -> Result<(), PresenceError> { - if !self - .client - .presence_subscriptions - .lock() - .await - .contains(jid) - { + if !self.client.is_presence_subscription_tracked(jid) { return Ok(()); } let node = self.build_subscription_node(jid).await; + // Re-read after the token lookup, which awaits. An `unsubscribe` landing + // in that window has already sent its own stanza, so subscribing now + // would leave the peer subscribed while we no longer track it. This + // narrows the window rather than closing it — `send_node` awaits too — + // but the lookup is the wide half and the re-read costs an uncontended + // lock. + if !self.client.is_presence_subscription_tracked(jid) { + return Ok(()); + } self.client.send_node(node).await?; Ok(()) } @@ -174,31 +180,38 @@ impl<'a> Presence<'a> { debug!("presence unsubscribe: unsubscribing from {}", jid); let node = self.build_unsubscription_node(jid); self.client.send_node(node).await?; - self.client.untrack_presence_subscription(jid).await; + self.client.untrack_presence_subscription(jid); Ok(()) } } impl Client { - pub(crate) async fn track_presence_subscription(&self, jid: Jid) { - self.presence_subscriptions.lock().await.insert(jid); + fn lock_presence_subscriptions( + &self, + ) -> std::sync::MutexGuard<'_, std::collections::HashSet> { + self.presence_subscriptions + .lock() + .unwrap_or_else(|p| p.into_inner()) } - pub(crate) async fn untrack_presence_subscription(&self, jid: &Jid) { - self.presence_subscriptions.lock().await.remove(jid); + pub(crate) fn track_presence_subscription(&self, jid: Jid) { + self.lock_presence_subscriptions().insert(jid); } - pub(crate) async fn tracked_presence_subscriptions(&self) -> Vec { - self.presence_subscriptions - .lock() - .await - .iter() - .cloned() - .collect() + pub(crate) fn untrack_presence_subscription(&self, jid: &Jid) { + self.lock_presence_subscriptions().remove(jid); + } + + pub(crate) fn is_presence_subscription_tracked(&self, jid: &Jid) -> bool { + self.lock_presence_subscriptions().contains(jid) + } + + pub(crate) fn tracked_presence_subscriptions(&self) -> Vec { + self.lock_presence_subscriptions().iter().cloned().collect() } pub(crate) async fn resubscribe_presence_subscriptions(&self, expected_generation: u64) { - let subscribed_jids = self.tracked_presence_subscriptions().await; + let subscribed_jids = self.tracked_presence_subscriptions(); if subscribed_jids.is_empty() { return; } @@ -413,10 +426,10 @@ mod tests { let client = bot.client(); let jid = Jid::from_str("1234567890@s.whatsapp.net").expect("valid jid"); - client.track_presence_subscription(jid.clone()).await; - client.track_presence_subscription(jid.clone()).await; + client.track_presence_subscription(jid.clone()); + client.track_presence_subscription(jid.clone()); - let tracked = client.tracked_presence_subscriptions().await; + let tracked = client.tracked_presence_subscriptions(); assert_eq!(tracked, vec![jid]); } @@ -437,11 +450,11 @@ mod tests { let client = bot.client(); let jid = Jid::from_str("1234567890@s.whatsapp.net").expect("valid jid"); - client.track_presence_subscription(jid.clone()).await; - client.untrack_presence_subscription(&jid).await; + client.track_presence_subscription(jid.clone()); + client.untrack_presence_subscription(&jid); assert!( - client.tracked_presence_subscriptions().await.is_empty(), + client.tracked_presence_subscriptions().is_empty(), "unsubscribe tracking should remove the jid" ); } @@ -475,4 +488,107 @@ mod tests { "unsubscribe stanza should not have children" ); } + + /// The resubscribe loop snapshots the tracked set, then re-checks each JID + /// before sending. This gates the send so an `unsubscribe` lands after the + /// snapshot was taken but before the loop reaches that JID. + #[tokio::test] + async fn resubscribe_skips_a_jid_unsubscribed_mid_loop() { + use crate::client::NodeFilter; + use bytes::Bytes; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + struct GatedTransport { + started: async_channel::Sender<()>, + release: async_channel::Receiver<()>, + gate_next_send: AtomicBool, + sends: Arc, + } + + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl crate::transport::Transport for GatedTransport { + async fn send(&self, _data: Bytes) -> Result<(), anyhow::Error> { + self.sends.fetch_add(1, Ordering::AcqRel); + if self.gate_next_send.swap(false, Ordering::AcqRel) { + self.started + .send(()) + .await + .map_err(|_| anyhow::anyhow!("gate observer closed"))?; + self.release + .recv() + .await + .map_err(|_| anyhow::anyhow!("gate closed"))?; + } + Ok(()) + } + + async fn disconnect(&self) {} + } + + let (client, _transport) = crate::test_utils::create_iq_test_client().await; + + let (started_tx, started_rx) = async_channel::bounded(1); + let (release_tx, release_rx) = async_channel::bounded(1); + let sends = Arc::new(AtomicUsize::new(0)); + let gated = crate::socket::NoiseSocket::new( + Arc::new(TokioRuntime), + Arc::new(GatedTransport { + started: started_tx, + release: release_rx, + gate_next_send: AtomicBool::new(true), + sends: sends.clone(), + }), + wacore::handshake::NoiseCipher::new(&[0u8; 32]).expect("valid key"), + wacore::handshake::NoiseCipher::new(&[0u8; 32]).expect("valid key"), + ); + *client.noise_socket.lock().unwrap() = Some(Arc::new(gated)); + + let first: Jid = "12025550111@s.whatsapp.net".parse().expect("valid jid"); + let second: Jid = "12025550122@s.whatsapp.net".parse().expect("valid jid"); + client.track_presence_subscription(first.clone()); + client.track_presence_subscription(second.clone()); + + // Which JID the set yields first is not fixed, so learn it from the + // stanza rather than assuming an iteration order. + let sent = client.wait_for_sent_node(NodeFilter::tag("presence")); + let generation = client.connection_generation.load(Ordering::SeqCst); + let resubscribe = { + let client = client.clone(); + tokio::spawn(async move { client.resubscribe_presence_subscriptions(generation).await }) + }; + + let node = sent.await.expect("the loop sends the first subscribe"); + let sent_to = node + .attrs + .get("to") + .cloned() + .expect("subscribe carries a target"); + started_rx.recv().await.expect("the first send is gated"); + + let unsubscribed = if sent_to == first.to_string() { + second + } else { + first + }; + client.untrack_presence_subscription(&unsubscribed); + + release_tx.send(()).await.expect("gate released"); + resubscribe + .await + .expect("resubscribe task should not panic"); + + assert_eq!( + sends.load(Ordering::Acquire), + 1, + "only the JID still tracked when the loop reached it may be re-subscribed" + ); + assert!( + !client + .tracked_presence_subscriptions() + .contains(&unsubscribed), + "and the unsubscribe must stand" + ); + } } diff --git a/src/handlers/call.rs b/src/handlers/call.rs index 6b54d937a..bda6587b8 100644 --- a/src/handlers/call.rs +++ b/src/handlers/call.rs @@ -2600,7 +2600,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); (client, sends) } @@ -2640,7 +2640,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); (client, started_rx, release_tx) } @@ -3883,7 +3883,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); let node = node_to_owned_ref(&offer_stanza()); let mut cancelled = false; @@ -4594,7 +4594,7 @@ mod tests { NoiseCipher::new(&key).expect("valid key"), NoiseCipher::new(&key).expect("valid key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); let (handler, rx) = ChannelEventHandler::new(); client.subscribe_handler(handler).detach(); diff --git a/src/message/receive.rs b/src/message/receive.rs index c43a1022e..9ee479b14 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -1410,7 +1410,7 @@ impl Client { is_offline: bool, ) { // Dedup: skip if we already have a sync pending/in-flight for this user - if !self.pending_device_sync.add(&user_jid).await { + if !self.pending_device_sync.add(&user_jid) { return; } @@ -1857,7 +1857,7 @@ mod tests { .schedule_unknown_device_sync(user.clone(), true) .await; - let pending = client.pending_device_sync.take_all().await; + let pending = client.pending_device_sync.take_all(); assert_eq!(pending, vec![user]); } } diff --git a/src/message/tests.rs b/src/message/tests.rs index 12d2b6e18..b92a7d802 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -5539,7 +5539,7 @@ async fn capturing_client( write_key, read_key, ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.set_connected_for_test(true); seed_test_pn(&client).await; // Live-path semantics by default; drain tests re-enter drain state @@ -8911,7 +8911,7 @@ async fn app_state_key_share_transport_retry_waits_for_reconnect() { // handshake and installs fresh keys and counters. Swap the socket to match, // since the one the failed frame poisoned dies with the old connection. client.connection_generation.fetch_add(1, Ordering::AcqRel); - *client.noise_socket.lock().await = Some(Arc::new(crate::socket::NoiseSocket::new( + *client.noise_socket.lock().unwrap() = Some(Arc::new(crate::socket::NoiseSocket::new( Arc::new(crate::runtime_impl::TokioRuntime), transport.clone() as Arc, wacore::handshake::NoiseCipher::new(&[0u8; 32]).expect("32-byte key"), diff --git a/src/pending_device_sync.rs b/src/pending_device_sync.rs index 82c772f4f..3ebf0476a 100644 --- a/src/pending_device_sync.rs +++ b/src/pending_device_sync.rs @@ -5,22 +5,27 @@ use std::collections::HashSet; use wacore_binary::Jid; pub(crate) struct PendingDeviceSync { - pending: async_lock::Mutex>, + /// Sync lock: no critical section here does anything but a set operation. + pending: std::sync::Mutex>, } impl PendingDeviceSync { pub(crate) fn new() -> Self { Self { - pending: async_lock::Mutex::new(HashSet::new()), + pending: std::sync::Mutex::new(HashSet::new()), } } + fn lock(&self) -> std::sync::MutexGuard<'_, HashSet> { + self.pending.lock().unwrap_or_else(|p| p.into_inner()) + } + /// Insert a user. Returns `true` if newly inserted, `false` if already present. /// /// Takes `&Jid` and clones only on a real insert, so the dedup path (the common /// case during a retry storm) does no allocation. - pub(crate) async fn add(&self, jid: &Jid) -> bool { - let mut pending = self.pending.lock().await; + pub(crate) fn add(&self, jid: &Jid) -> bool { + let mut pending = self.lock(); if pending.contains(jid) { false } else { @@ -29,11 +34,11 @@ impl PendingDeviceSync { } } - pub(crate) async fn take_all(&self) -> Vec { - self.pending.lock().await.drain().collect() + pub(crate) fn take_all(&self) -> Vec { + self.lock().drain().collect() } - pub(crate) async fn clear(&self) { - self.pending.lock().await.clear(); + pub(crate) fn clear(&self) { + self.lock().clear(); } } diff --git a/src/retry.rs b/src/retry.rs index dfaf21567..d3bf9e8dd 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -1692,7 +1692,7 @@ mod tests { NoiseCipher::new(&key).expect("write cipher"), NoiseCipher::new(&key).expect("read cipher"), ); - *client.noise_socket.lock().await = Some(Arc::new(socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(socket)); } async fn seed_retry_lease( diff --git a/src/test_utils.rs b/src/test_utils.rs index fb8b20cac..81e082dcd 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -329,7 +329,7 @@ pub(crate) async fn create_iq_test_client() -> ( NoiseCipher::new(&[0u8; 32]).expect("32-byte key"), Some(client.stats.clone()), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); client.set_connected_for_test(true); client .is_running diff --git a/src/usync.rs b/src/usync.rs index 317e48cc7..531dde9a5 100644 --- a/src/usync.rs +++ b/src/usync.rs @@ -487,7 +487,7 @@ impl Client { ) )] pub(crate) async fn flush_pending_device_sync(&self) { - let pending = self.pending_device_sync.take_all().await; + let pending = self.pending_device_sync.take_all(); if pending.is_empty() { return; } @@ -513,7 +513,7 @@ impl Client { pending.len() ); for jid in &pending { - self.pending_device_sync.add(jid).await; + self.pending_device_sync.add(jid); } } } diff --git a/src/voip/facade.rs b/src/voip/facade.rs index 043ff5c7b..c760538f8 100644 --- a/src/voip/facade.rs +++ b/src/voip/facade.rs @@ -3776,7 +3776,7 @@ mod tests { NoiseCipher::new(&key).expect("key"), NoiseCipher::new(&key).expect("key"), ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + *client.noise_socket.lock().unwrap() = Some(Arc::new(noise_socket)); } struct GatedSendTransport { @@ -3975,7 +3975,6 @@ mod tests { let group = Jid::new("120363000000000001", Server::Group); client .get_group_cache() - .await .insert( group.clone(), Arc::new(GroupInfo::new(