diff --git a/src/client.rs b/src/client.rs index b7ebb12c7..ca1836273 100644 --- a/src/client.rs +++ b/src/client.rs @@ -3,6 +3,7 @@ mod adapters; mod app_state; mod context_impl; mod device_registry; +pub(crate) mod device_topology; mod iq_ops; mod lid_pn; mod lifecycle; @@ -526,8 +527,24 @@ pub struct Client { /// LRU cache for device registry (matches WhatsApp Web's 5000 entry limit). /// Maps user ID to DeviceListRecord for fast device existence checks. /// Backed by persistent storage. - pub(crate) device_registry_cache: - TypedCache>, + /// Device registry fused with its topology tracker: every write records + /// the change by construction, so the group-devices memo below can never + /// be left stale by a forgotten bump. + pub(crate) device_registry_cache: crate::client::device_topology::DeviceRegistryCache, + /// Shared topology tracker (generation + changed-users log). LidPnCache + /// records mapping changes into it; the memo validates against it. + pub(crate) device_topology: Arc, + /// Whether the group-devices memo may be used: false when the registry + /// or LID-PN caches are store-backed (a shared external store can be + /// written by other processes, which the in-process topology tracker + /// cannot observe). + pub(crate) group_devices_memo_enabled: bool, + /// Per-group memo of the fully resolved (LID-converted) device list, + /// validated by GroupInfo identity + the device topology. Serves the + /// per-send full-set resolution in `resolve_skdm_targets` so a warm + /// repeat send skips the per-member cache fan-out. + pub(crate) group_devices_memo: + Cache>, /// Router for dispatching stanzas to their appropriate handlers pub(crate) stanza_router: crate::handlers::router::StanzaRouter, diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index 32593368c..ca1c09d65 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -10,6 +10,23 @@ use wacore_binary::Jid; use super::Client; +/// Per-group device-list snapshot for `resolve_group_devices_memoized`. +/// Valid while the producing `GroupInfo` Arc is still the cached one AND the +/// device-topology generation is unchanged. +pub(crate) struct GroupDevicesMemo { + /// Weak identity of the producing GroupInfo: pointer equality is ABA-safe + /// because the Weak keeps the allocation alive, while the heavy data + /// (participants, maps) is freed as soon as the metadata cache drops its + /// Arc — the memo retains a struct-sized header, not the whole GroupInfo. + pub(crate) group_info: std::sync::Weak, + pub(crate) generation: u64, + /// Member identifiers in BOTH namespaces (participant users, their mapped + /// counterparts, resolved device users): the scoped-invalidation check + /// tests the topology log's touched users against this set. + pub(crate) members: Arc>, + pub(crate) devices: Arc>, +} + /// Result of resolving a user identifier to lookup keys. /// This makes the LID/PN relationship explicit instead of using magic indices. #[derive(Debug, Clone)] @@ -61,6 +78,162 @@ impl Client { .to_string() } + /// Resolve a group's full (LID-converted) device list, memoized per group. + /// + /// The input set is a pure function of `group_info` (participants + LID + /// normalization), so the memo is valid exactly while BOTH hold: + /// the same `GroupInfo` snapshot (`Arc` identity — any metadata refresh or + /// membership change produces a new `Arc`) and an unchanged + /// `device_topology_generation` (any registry/mapping write bumps it). + /// On a warm repeat send this turns the per-member cache fan-out + /// (2 lookups per participant) into one memo hit. + pub(crate) async fn resolve_group_devices_memoized( + &self, + group: &Jid, + group_info: &Arc, + own_sending_jid: &Jid, + ) -> Result>, anyhow::Error> { + // Store-backed registry or mapping caches can be written by OTHER + // processes (e.g. shared Redis across pods), which this process's + // topology tracker cannot observe; the memo's freshness contract + // doesn't hold there, so it is disabled and every send resolves. + if !self.group_devices_memo_enabled { + return Ok(Arc::new( + self.resolve_group_devices_uncached(group_info, own_sending_jid) + .await?, + )); + } + // Load the generation BEFORE resolving (do NOT move this after + // get_user_devices): a write racing the resolve bumps it afterwards, + // so the memo we store is already stale by its own stamp and the next + // read revalidates. Loading after would stamp racing writes as seen + // and serve their effects stale. + let generation = self.device_topology.current(); + + if let Some(memo) = self.group_devices_memo.get(group).await + && std::ptr::eq(memo.group_info.as_ptr(), Arc::as_ptr(group_info)) + { + if memo.generation == generation { + // Refcount bump: the snapshot is immutable, so a hit shares + // it instead of cloning the device Vec. + return Ok(Arc::clone(&memo.devices)); + } + // Stale stamp: when every change since it touched only users + // outside this group, re-stamp instead of recomputing, so write + // storms on unrelated groups don't tank the hit rate. Any doubt + // (log overflow, member touched) falls through to the recompute. + if self + .device_topology + .unchanged_for(memo.generation, |user| memo.members.contains(user)) + { + self.group_devices_memo + .insert( + group.clone(), + Arc::new(GroupDevicesMemo { + group_info: memo.group_info.clone(), + generation, + members: Arc::clone(&memo.members), + devices: Arc::clone(&memo.devices), + }), + ) + .await; + return Ok(Arc::clone(&memo.devices)); + } + } + + let devices = self + .resolve_group_devices_uncached(group_info, own_sending_jid) + .await?; + + // Member identifiers in both namespaces, so the scoped-invalidation + // check can match however a write was keyed: writes record every + // resolved lookup alias (see DeviceRegistryCache::insert callers), and + // this set carries each member's group-facing identity (participant + // user + mapped counterpart) plus the namespace the resolved device + // JIDs ended up in. + let mut members = std::collections::HashSet::with_capacity( + group_info.participants.len() * 2 + devices.len() + 2, + ); + for participant in &group_info.participants { + members.insert(participant.user.clone()); + if participant.is_lid() + && let Some(pn) = group_info.phone_jid_for_lid_user(&participant.user) + { + members.insert(pn.user.clone()); + } else if let Some(lid) = group_info.lid_user_for_phone_user(&participant.user) { + members.insert(lid.clone()); + } + } + members.insert(own_sending_jid.user.clone()); + for device in &devices { + members.insert(device.user.clone()); + } + + let devices = Arc::new(devices); + self.group_devices_memo + .insert( + group.clone(), + Arc::new(GroupDevicesMemo { + group_info: Arc::downgrade(group_info), + generation, + members: Arc::new(members), + devices: Arc::clone(&devices), + }), + ) + .await; + Ok(devices) + } + + /// The memo's recompute body: derive the resolve set from `group_info` + /// (participants + LID normalization, appending self when the server + /// snapshot omitted it — mirroring `ensure_self_in_group`, so keying the + /// memo off the pre-ensure Arc stays equivalent) and resolve it. + async fn resolve_group_devices_uncached( + &self, + group_info: &Arc, + own_sending_jid: &Jid, + ) -> Result, anyhow::Error> { + let is_lid_mode = group_info.addressing_mode == wacore::types::message::AddressingMode::Lid; + let mut jids_to_resolve: Vec = group_info + .participants + .iter() + .map(|jid| { + if is_lid_mode + && jid.is_lid() + && let Some(pn) = group_info.phone_jid_for_lid_user(&jid.user) + { + return pn.to_non_ad(); + } + jid.to_non_ad() + }) + .collect(); + if !group_info + .participants + .iter() + .any(|participant| wacore_binary::JidExt::is_same_user_as(participant, own_sending_jid)) + { + let own = if is_lid_mode + && own_sending_jid.is_lid() + && let Some(pn) = group_info.phone_jid_for_lid_user(&own_sending_jid.user) + { + pn.to_non_ad() + } else { + own_sending_jid.to_non_ad() + }; + jids_to_resolve.push(own); + } + + let mut devices = self.get_user_devices(&jids_to_resolve).await?; + if is_lid_mode { + // WA Web expects LID addressing in SKDM nodes for LID groups. + devices = devices + .into_iter() + .map(|d| group_info.phone_device_jid_into_lid(d)) + .collect(); + } + Ok(devices) + } + /// Resolve a user identifier to its lookup keys with type information. /// /// Returns a `UserLookupKeys` enum that explicitly represents: @@ -165,7 +338,7 @@ impl Client { // Cache under the record's actual stored key, not our guessed one, // to keep the cache and backend consistent. self.device_registry_cache - .insert(record.user.clone(), Arc::new(record)) + .promote(record.user.clone(), Arc::new(record)) .await; return has_device; } @@ -205,8 +378,18 @@ impl Client { let record_for_cache = record.clone(); // Use canonical_key directly as cache key (no extra clone) + // Record every lookup alias, not just canonical+original: a LID-keyed + // update must also touch the mapped PN, or a PN-addressed group's memo + // (whose member set only knows the PN side) would re-stamp stale. self.device_registry_cache - .insert(canonical_key.clone(), Arc::new(record_for_cache)) + .insert( + canonical_key.clone(), + Arc::new(record_for_cache), + lookup + .all_keys() + .into_iter() + .chain(std::iter::once(original_user.as_str())), + ) .await; let backend = self.persistence_manager.backend(); @@ -264,8 +447,16 @@ impl Client { record.user.clone_from(&canonical_key); let record_for_cache = record.clone(); + // Same alias rule as update_device_list: record every lookup key. self.device_registry_cache - .insert(canonical_key.clone(), Arc::new(record_for_cache)) + .insert( + canonical_key.clone(), + Arc::new(record_for_cache), + lookup + .all_keys() + .into_iter() + .chain(std::iter::once(original_user.as_str())), + ) .await; if canonical_key != original_user { @@ -343,6 +534,11 @@ impl Client { if let Err(e) = self.persistence_manager.backend().delete_devices(key).await { warn!("Failed to delete device registry from DB for {key}: {e}"); } + // Invalidate again after the delete: a concurrent reader that read + // the doomed DB row can promote() it back between the first + // invalidate and the delete commit (same guard as the canonical + // flip path in update_device_list). + self.device_registry_cache.invalidate(key).await; } debug!("Invalidated device cache for user: {} ({:?})", user, lookup); @@ -660,7 +856,7 @@ impl Client { match backend.get_devices(key).await { Ok(Some(record)) => { self.device_registry_cache - .insert(record.user.clone(), Arc::new(record.clone())) + .promote(record.user.clone(), Arc::new(record.clone())) .await; return Some(record); } @@ -715,7 +911,7 @@ impl Client { continue; } self.device_registry_cache - .insert(record.user.clone(), Arc::new(record)) + .promote(record.user.clone(), Arc::new(record)) .await; return Some(devices); } @@ -793,12 +989,16 @@ impl Client { record.user = lid.to_string(); if let Err(e) = backend.update_device_list(record.clone()).await { + // The backend row may have changed even on error, so the + // change is recorded before the early return; the success + // path records once via the fused cache insert below. + self.device_topology.record([pn, lid]); warn!("Failed to migrate device registry to LID: {}", e); return; } self.device_registry_cache - .insert(lid.to_string(), Arc::new(record)) + .insert(lid.to_string(), Arc::new(record), [lid, pn]) .await; // Drop the PN-keyed row in both cache and DB. Invalidate @@ -853,7 +1053,7 @@ mod tests { }; client .device_registry_cache - .insert(user.into(), Arc::new(record)) + .raw_insert_for_tests(user.into(), Arc::new(record)) .await; } @@ -894,6 +1094,386 @@ mod tests { ); } + /// Locks the three validity gates of the group-devices memo: a repeat + /// resolve with the same GroupInfo Arc + generation is a memo hit (proved + /// by serving a raw cache change STALE), any topology bump recomputes, + /// and a refreshed GroupInfo (new Arc, same content) recomputes. + #[tokio::test] + async fn group_devices_memo_hits_and_invalidates() { + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + + let client = create_test_client().await; + let group: Jid = "120363000000000042@g.us".parse().expect("group jid"); + let user_a = "5511999990001"; + let user_b = "5511999990002"; + setup_device_record(&client, user_a, &[0, 5]).await; + setup_device_record(&client, user_b, &[0]).await; + + let group_info = Arc::new(GroupInfo::new( + vec![Jid::pn(user_a), Jid::pn(user_b)], + AddressingMode::Pn, + )); + + let first = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve should succeed"); + assert_eq!(first.len(), 3, "0+5 for A, 0 for B"); + + // Raw cache write WITHOUT a topology bump: the memo must keep serving + // the snapshot (this is what proves the repeat call was a hit and not + // a silent recompute). + setup_device_record(&client, user_a, &[0]).await; + let stale = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve should succeed"); + assert_eq!( + stale, first, + "same Arc + same generation must be a memo hit" + ); + + // A topology change touching a MEMBER invalidates and the recompute + // sees the new record. + client.device_topology.record([user_a]); + let fresh = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve should succeed"); + assert_eq!(fresh.len(), 2, "post-bump resolve must see the raw change"); + + // A refreshed GroupInfo (new Arc, identical content) must recompute + // even with an unchanged generation. + setup_device_record(&client, user_b, &[0, 9]).await; + let refreshed_info = Arc::new(GroupInfo::new( + vec![Jid::pn(user_a), Jid::pn(user_b)], + AddressingMode::Pn, + )); + let after_refresh = client + .resolve_group_devices_memoized( + &group, + &refreshed_info, + &refreshed_info.participants[0], + ) + .await + .expect("resolve should succeed"); + assert_eq!( + after_refresh.len(), + 3, + "a new GroupInfo Arc must invalidate the memo by identity" + ); + } + + /// Locks the scoped invalidation: changes touching only OTHER groups' + /// users re-stamp the memo (still a hit), a member's change recomputes, + /// and the doubt fallbacks (global event, log overflow) recompute. + #[tokio::test] + async fn group_devices_memo_scoped_invalidation() { + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + + let client = create_test_client().await; + let group: Jid = "120363000000000077@g.us".parse().expect("group jid"); + let user_a = "5511999990011"; + setup_device_record(&client, user_a, &[0, 5]).await; + let group_info = Arc::new(GroupInfo::new(vec![Jid::pn(user_a)], AddressingMode::Pn)); + + let first = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(first.len(), 2); + + // Raw change (not recorded) + changes touching only a NON-member: + // the memo must re-stamp and keep serving the snapshot. + setup_device_record(&client, user_a, &[0]).await; + client.device_topology.record(["5511000000001"]); + client.device_topology.record(["5511000000002"]); + let stale = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!( + stale, first, + "non-member changes must re-stamp, not recompute" + ); + + // A member's change recomputes and sees the raw change. + client.device_topology.record([user_a]); + let fresh = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(fresh.len(), 1, "member change must recompute"); + + // Global events (mapping cache clear, warm-up) poison the fast path. + setup_device_record(&client, user_a, &[0, 5, 9]).await; + client.device_topology.record_global(); + let after_global = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(after_global.len(), 3, "global event must recompute"); + + // Log overflow past the memo's stamp: cannot prove cleanliness, + // must recompute. + setup_device_record(&client, user_a, &[0]).await; + for _ in 0..300 { + client.device_topology.record(["5511000000003"]); + } + let after_overflow = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(after_overflow.len(), 1, "log overflow must recompute"); + } + + /// A mapping add for a member (logged under BOTH its LID and PN keys) + /// must invalidate even when the group only knows one namespace. + #[tokio::test] + async fn group_devices_memo_invalidated_by_member_mapping_change() { + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + + let client = create_test_client().await; + let group: Jid = "120363000000000078@g.us".parse().expect("group jid"); + let pn = "5511999990012"; + setup_device_record(&client, pn, &[0]).await; + let group_info = Arc::new(GroupInfo::new(vec![Jid::pn(pn)], AddressingMode::Pn)); + + let first = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(first.len(), 1); + + // Raw change, then learn a LID mapping for the member: the add logs + // (lid, pn) and the memo's member set carries the PN, so it must + // recompute even though the group never saw the LID. + setup_device_record(&client, pn, &[0, 7]).await; + client + .add_lid_pn_mapping( + "100000000000077", + pn, + crate::lid_pn_cache::LearningSource::Usync, + ) + .await + .expect("mapping"); + let fresh = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!( + fresh.len(), + 2, + "a member's mapping change must invalidate the memo" + ); + } + + /// Review fix: a server group snapshot that omits self used to be + /// rebuilt by ensure_self_in_group on every send (fresh Arc), making the + /// memo permanently miss. Keying off the pre-ensure Arc and appending + /// self inside the derivation keeps the identity stable. + #[tokio::test] + async fn memo_hits_when_self_missing_from_group_snapshot() { + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + + let client = create_test_client().await; + let group: Jid = "120363000000000080@g.us".parse().expect("group jid"); + let member = "5511999990014"; + let own = Jid::pn("5511999990015"); + setup_device_record(&client, member, &[0]).await; + setup_device_record(&client, "5511999990015", &[0, 3]).await; + + // Self deliberately absent from the snapshot. + let group_info = Arc::new(GroupInfo::new(vec![Jid::pn(member)], AddressingMode::Pn)); + + let first = client + .resolve_group_devices_memoized(&group, &group_info, &own) + .await + .expect("resolve"); + assert_eq!(first.len(), 3, "member device + own's two devices"); + + let second = client + .resolve_group_devices_memoized(&group, &group_info, &own) + .await + .expect("resolve"); + assert!( + Arc::ptr_eq(&first, &second), + "a self-missing group snapshot must still produce memo hits" + ); + } + + /// Codex P2 regression: a PN-addressed group's memo only knows the PN + /// side of a member when the cached GroupInfo carries no LID map, but a + /// later usync update can arrive keyed by the LID (canonical == original). + /// The write must record every lookup alias so the memo recomputes + /// instead of re-stamping stale. + #[tokio::test] + async fn lid_keyed_update_invalidates_pn_group_memo() { + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + + let client = create_test_client().await; + let group: Jid = "120363000000000079@g.us".parse().expect("group jid"); + let pn = "5511999990013"; + let lid = "100000000000079"; + + // Mapping known BEFORE the memo: the canonical record lives under the + // LID, while the group only references the member by PN. + client + .add_lid_pn_mapping(lid, pn, crate::lid_pn_cache::LearningSource::Usync) + .await + .expect("mapping"); + client + .update_device_list(wacore::store::traits::DeviceListRecord { + user: pn.into(), + devices: vec![wacore::store::traits::DeviceInfo { + device_id: 0, + key_index: None, + }], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .expect("seed record"); + + let group_info = Arc::new(GroupInfo::new(vec![Jid::pn(pn)], AddressingMode::Pn)); + let first = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!(first.len(), 1); + + // The update arrives keyed by the LID: canonical == original == LID, + // so without the alias rule only the LID would be recorded and the + // PN-only member set would re-stamp the stale snapshot. + client + .update_device_list(wacore::store::traits::DeviceListRecord { + user: lid.into(), + devices: vec![ + wacore::store::traits::DeviceInfo { + device_id: 0, + key_index: None, + }, + wacore::store::traits::DeviceInfo { + device_id: 11, + key_index: None, + }, + ], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .expect("LID-keyed update"); + + let fresh = client + .resolve_group_devices_memoized(&group, &group_info, &group_info.participants[0]) + .await + .expect("resolve"); + assert_eq!( + fresh.len(), + 2, + "a LID-keyed update for a member must invalidate the PN group's memo" + ); + } + + /// Locks the invariant that every device-topology write path bumps the + /// generation. patch_device_add/patch_device_remove funnel their writes + /// through update_device_list, so the funnel is what is asserted. + #[tokio::test] + async fn topology_mutators_bump_the_generation() { + let client = create_test_client().await; + let current_gen = |c: &Arc| c.device_topology.current(); + + let before = current_gen(&client); + client + .update_device_list(wacore::store::traits::DeviceListRecord { + user: "5511999990003".into(), + devices: vec![wacore::store::traits::DeviceInfo { + device_id: 0, + key_index: None, + }], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .expect("update_device_list"); + assert!( + current_gen(&client) > before, + "update_device_list must bump" + ); + + let before = current_gen(&client); + client + .update_device_lists(vec![wacore::store::traits::DeviceListRecord { + user: "5511999990004".into(), + devices: vec![wacore::store::traits::DeviceInfo { + device_id: 0, + key_index: None, + }], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }]) + .await + .expect("update_device_lists"); + assert!( + current_gen(&client) > before, + "update_device_lists must bump" + ); + + let before = current_gen(&client); + client.invalidate_device_cache("5511999990003").await; + assert!( + current_gen(&client) > before, + "invalidate_device_cache must bump" + ); + + let before = current_gen(&client); + client + .add_lid_pn_mapping( + "100000000000042", + "5511999990004", + crate::lid_pn_cache::LearningSource::Usync, + ) + .await + .expect("mapping should persist"); + assert!( + current_gen(&client) > before, + "add_lid_pn_mapping must bump" + ); + + // Fresh pair: the record must live under its PN key (no mapping yet) + // for the migration to find and move it. + client + .update_device_list(wacore::store::traits::DeviceListRecord { + user: "5511999990005".into(), + devices: vec![wacore::store::traits::DeviceInfo { + device_id: 0, + key_index: None, + }], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .expect("seed PN-keyed record"); + let before = current_gen(&client); + client + .migrate_device_registry_on_lid_discovery("5511999990005", "100000000000043") + .await; + assert!( + current_gen(&client) > before, + "migrate_device_registry_on_lid_discovery must bump" + ); + } + #[tokio::test] async fn warm_registry_hit_shares_arc_not_deep_clone() { let client = create_test_client().await; @@ -1209,7 +1789,7 @@ mod tests { }; client .device_registry_cache - .insert("15551234567".to_string(), Arc::new(record)) + .raw_insert_for_tests("15551234567".to_string(), Arc::new(record)) .await; // Patch: update device 3 key_index to 5 @@ -1300,7 +1880,7 @@ mod tests { client .device_registry_cache - .insert( + .raw_insert_for_tests( "15551234567".to_string(), Arc::new(record_with_raw_id("15551234567", &[0, 5], 1)), ) @@ -1349,7 +1929,7 @@ mod tests { client .device_registry_cache - .insert( + .raw_insert_for_tests( "15551234567".to_string(), Arc::new(record_with_raw_id("15551234567", &[0, 5], 1)), ) @@ -1666,7 +2246,7 @@ mod tests { }; client .device_registry_cache - .insert("15551234567".into(), Arc::new(record)) + .raw_insert_for_tests("15551234567".into(), Arc::new(record)) .await; // Warm the sender key device cache @@ -1921,7 +2501,7 @@ mod tests { // the mapping was learned. client .device_registry_cache - .insert(pn.into(), Arc::new(legacy)) + .raw_insert_for_tests(pn.into(), Arc::new(legacy)) .await; setup_lid_pn(&client, lid, pn).await; diff --git a/src/client/device_topology.rs b/src/client/device_topology.rs new file mode 100644 index 000000000..842b86977 --- /dev/null +++ b/src/client/device_topology.rs @@ -0,0 +1,178 @@ +//! Device-topology change tracking for the per-group device-list memo. +//! +//! "Topology" here means anything that can change a device-list answer: +//! registry record writes/invalidations and LID-PN mapping changes. Instead of +//! trusting every write path to remember a manual generation bump, the bump +//! lives INSIDE the write chokepoints ([`DeviceRegistryCache`] and +//! `LidPnCache::add`), so a writer cannot forget it by construction. +//! +//! Each change also logs WHICH canonical users it touched (both namespaces), +//! so a memo whose generation went stale can prove "none of the changed users +//! are in my group" and re-stamp itself instead of recomputing. Every doubtful +//! case (log overflow, global events) degrades to a recompute, never to +//! serving stale data. + +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use portable_atomic::AtomicU64; +use wacore_binary::CompactString; + +/// Bounded log capacity. Sized so a burst (e.g. a usync response for a large +/// group) still fits; overflow just disables the scoped-revalidation fast +/// path until affected memos recompute once. +const TOPOLOGY_LOG_CAPACITY: usize = 256; + +struct TopologyLog { + /// (generation that the change produced, canonical user touched). + entries: VecDeque<(u64, CompactString)>, + /// Highest generation evicted from `entries` (0 = nothing evicted). + /// A memo older than this cannot be proven clean and must recompute. + floor: u64, +} + +/// Shared tracker: a monotonic generation plus the bounded changed-users log. +pub(crate) struct DeviceTopology { + generation: AtomicU64, + log: std::sync::Mutex, +} + +impl DeviceTopology { + pub(crate) fn new() -> Arc { + Arc::new(Self { + generation: AtomicU64::new(0), + log: std::sync::Mutex::new(TopologyLog { + entries: VecDeque::with_capacity(TOPOLOGY_LOG_CAPACITY), + floor: 0, + }), + }) + } + + pub(crate) fn current(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + /// Record one topology change touching the given users (pass BOTH + /// namespaces of an identity when known: a mapping change alters which + /// canonical record either key resolves to). + pub(crate) fn record<'a>(&self, users: impl IntoIterator) { + let mut log = self.log.lock().unwrap_or_else(|p| p.into_inner()); + let generation = self.generation.load(Ordering::Acquire) + 1; + for user in users { + if log.entries.len() == TOPOLOGY_LOG_CAPACITY + && let Some((evicted_gen, _)) = log.entries.pop_front() + { + log.floor = evicted_gen; + } + log.entries + .push_back((generation, CompactString::from(user))); + } + // Publish the generation only after the log holds the users, so a + // reader that observes the new generation can always find (or rule + // out) the corresponding entries. + self.generation.store(generation, Ordering::Release); + } + + /// Record a change whose blast radius is unknown (bulk warm-up, cache + /// clear): bumps and poisons the scoped fast path so every memo + /// recomputes once. + pub(crate) fn record_global(&self) { + let mut log = self.log.lock().unwrap_or_else(|p| p.into_inner()); + let generation = self.generation.load(Ordering::Acquire) + 1; + log.entries.clear(); + log.floor = generation; + self.generation.store(generation, Ordering::Release); + } + + /// Whether every change after `since` only touched users for which + /// `is_member` returns false. `false` on any doubt (log overflow past + /// `since`), so callers recompute. + pub(crate) fn unchanged_for(&self, since: u64, is_member: impl Fn(&str) -> bool) -> bool { + let log = self.log.lock().unwrap_or_else(|p| p.into_inner()); + if log.floor > since { + return false; + } + log.entries + .iter() + .filter(|(generation, _)| *generation > since) + .all(|(_, user)| !is_member(user)) + } +} + +/// The device registry cache plus its topology tracker, fused so every write +/// records the change. Reads are pass-through; the only write entry points +/// are [`insert`](Self::insert), [`invalidate`](Self::invalidate) and the +/// non-recording [`promote`](Self::promote) (whose data is by definition what +/// the DB fallback already answered). +pub(crate) struct DeviceRegistryCache { + cache: crate::cache_store::TypedCache>, + topology: Arc, +} + +impl DeviceRegistryCache { + pub(crate) fn new( + cache: crate::cache_store::TypedCache>, + topology: Arc, + ) -> Self { + Self { cache, topology } + } + + pub(crate) async fn get( + &self, + key: &str, + ) -> Option> { + self.cache.get(key).await + } + + /// Write a record and log the touched users. `touched` carries the keys + /// whose answers change (canonical key, plus the original alias when the + /// canonical flipped). + pub(crate) async fn insert<'a>( + &self, + key: String, + record: Arc, + touched: impl IntoIterator, + ) { + self.cache.insert(key, record).await; + self.topology.record(touched); + } + + pub(crate) async fn invalidate(&self, key: &str) { + self.cache.invalidate(key).await; + self.topology.record([key]); + } + + /// Cache-fill from the DB row the fallback path would have returned: the + /// answer is unchanged, so no topology change is recorded. + pub(crate) async fn promote( + &self, + key: String, + record: Arc, + ) { + self.cache.insert(key, record).await; + } + + #[cfg(feature = "debug-diagnostics")] + pub(crate) fn entry_count(&self) -> u64 { + self.cache.entry_count() + } + + /// Test-only passthrough for moka maintenance flushes. + #[cfg(test)] + pub(crate) async fn run_pending_tasks(&self) { + self.cache.run_pending_tasks().await; + } + + /// Test-only raw write that bypasses topology recording, for fixture + /// seeding and for proving that memo hits really are hits (a raw change + /// must be served stale). + #[cfg(test)] + pub(crate) async fn raw_insert_for_tests( + &self, + key: String, + record: Arc, + ) { + self.cache.insert(key, record).await; + } +} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 8bee66724..f5400590a 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -2,6 +2,10 @@ use super::*; +/// Max groups with a cached resolved-device snapshot. LRU eviction covers +/// accounts in more groups; an evicted entry just recomputes on next send. +const GROUP_DEVICES_MEMO_CAPACITY: u64 = 64; + impl Client { pub fn shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { self.shutdown_notifier.subscribe() @@ -107,6 +111,7 @@ impl Client { let (tx, rx) = async_channel::bounded(32); + let device_topology = crate::client::device_topology::DeviceTopology::new(); let this = Self { runtime: runtime.clone(), core, @@ -215,10 +220,19 @@ impl Client { custom_enc_handlers: std::sync::OnceLock::new(), chatstate_handlers: Arc::new(RwLock::new(Vec::new())), pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(), - device_registry_cache: cache_config.device_registry_cache.build_typed_ttl( - cache_config.cache_stores.device_registry_cache.clone(), - "device_registry", + device_registry_cache: crate::client::device_topology::DeviceRegistryCache::new( + cache_config.device_registry_cache.build_typed_ttl( + cache_config.cache_stores.device_registry_cache.clone(), + "device_registry", + ), + Arc::clone(&device_topology), ), + device_topology, + group_devices_memo_enabled: cache_config.cache_stores.device_registry_cache.is_none() + && cache_config.cache_stores.lid_pn_cache.is_none(), + group_devices_memo: Cache::builder() + .max_capacity(GROUP_DEVICES_MEMO_CAPACITY) + .build(), stanza_router: Self::create_stanza_router(), synchronous_ack: false, http_client, @@ -232,6 +246,10 @@ impl Client { }; let arc = Arc::new(this); + // Mapping changes alter which canonical record a device lookup + // resolves to, so LidPnCache records into the same topology tracker. + arc.lid_pn_cache + .attach_topology(Arc::clone(&arc.device_topology)); let _ = arc.self_weak.set(Arc::downgrade(&arc)); // Warm up the LID-PN cache from persistent storage diff --git a/src/handlers/notification/mod.rs b/src/handlers/notification/mod.rs index 04d706c94..e56cb0946 100644 --- a/src/handlers/notification/mod.rs +++ b/src/handlers/notification/mod.rs @@ -815,7 +815,7 @@ mod tests { }; client .device_registry_cache - .insert("5511999999999".into(), Arc::new(record)) + .raw_insert_for_tests("5511999999999".into(), Arc::new(record)) .await; // Seed a stored identity so the had-prior-identity gate runs the full reset @@ -1118,7 +1118,7 @@ mod tests { // cleanup has something to do, but deliberately do NOT seed an identity. client .device_registry_cache - .insert( + .raw_insert_for_tests( "5511666666666".into(), Arc::new(wacore::store::traits::DeviceListRecord { user: "5511666666666".into(), diff --git a/src/lid_pn_cache.rs b/src/lid_pn_cache.rs index dde9035dd..c5687b6e7 100644 --- a/src/lid_pn_cache.rs +++ b/src/lid_pn_cache.rs @@ -46,6 +46,11 @@ pub struct LidPnCache { lid_to_entry: TypedCache, Arc>, /// Phone number -> Entry mapping (stores the most recent LID for that PN) pn_to_entry: TypedCache, Arc>, + /// Device-topology tracker (attached by Client construction): a mapping + /// change alters which canonical record either key resolves to, so adds + /// record both identifiers. Recording lives here, at the write + /// chokepoint, so callers cannot forget it. + topology: std::sync::OnceLock>, /// PN -> the LID this process durably persisted for it. Lets the learn hot /// path skip a re-persist without swallowing the first live persist of a /// mapping an offline replay only warmed in memory. Keyed by the pair so a @@ -85,15 +90,27 @@ impl LidPnCache { // Always in-memory: tracks per-process persist state, never the // mapping itself, so it must not go through the shared store. persisted: TypedCache::from_moka(config.build_with_tti()), + topology: std::sync::OnceLock::new(), }, None => Self { lid_to_entry: TypedCache::from_moka(config.build_with_tti()), pn_to_entry: TypedCache::from_moka(config.build_with_tti()), persisted: TypedCache::from_moka(config.build_with_tti()), + topology: std::sync::OnceLock::new(), }, } } + /// Attach the device-topology tracker. Mapping writes before the attach + /// (none in practice: Client construction attaches before warm-up) are + /// simply not scoped. + pub(crate) fn attach_topology( + &self, + topology: Arc, + ) { + let _ = self.topology.set(topology); + } + /// Returns approximate entry counts for the LID and PN maps. #[cfg(feature = "debug-diagnostics")] pub fn entry_counts(&self) -> (u64, u64) { @@ -185,6 +202,9 @@ impl LidPnCache { .insert(shared.phone_number.clone(), shared) .await; } + if let Some(topology) = self.topology.get() { + topology.record([&*entry.lid, &*entry.phone_number]); + } } /// Whether this process has durably persisted exactly `phone -> lid`. @@ -231,6 +251,9 @@ impl LidPnCache { self.lid_to_entry.clear().await; self.pn_to_entry.clear().await; self.persisted.clear().await; + if let Some(topology) = self.topology.get() { + topology.record_global(); + } } /// Get the number of LID entries in the cache. diff --git a/src/send.rs b/src/send.rs index f9118f938..eeba34663 100644 --- a/src/send.rs +++ b/src/send.rs @@ -781,21 +781,18 @@ impl Client { /// For LID mode, uses `group_info.phone_jid_for_lid_user` to query devices /// via PN when available (LID usync is unreliable for own JID), then /// converts the result back to LID. Same fallback as `prepare_group_stanza`. - #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.resolve_skdm_targets", level = "debug", skip_all, fields(group = %wacore_binary::jid::observe_str(group_jid))))] - async fn resolve_skdm_targets( + /// Load (or lazily build) the per-group sender-key device map. + /// + /// Atomic get-or-init: if another task invalidated the cache during our + /// DB read, get_or_init's single-flight guarantee means the stale data + /// won't be inserted — the invalidation wins and the next caller re-inits. + async fn skdm_device_map( &self, group_jid: &str, - group_info: &wacore::client::context::GroupInfo, - own_sending_jid: &Jid, - ) -> Option<(Vec, Vec)> { + ) -> std::sync::Arc { use crate::sender_key_device_cache::SenderKeyDeviceMap; - - // Atomic get-or-init: if another task invalidated the cache during our - // DB read, get_or_init's single-flight guarantee means the stale data - // won't be inserted — the invalidation wins and the next caller re-inits. let pm = self.persistence_manager.clone(); - let cached_map = self - .sender_key_device_cache + self.sender_key_device_cache .get_or_init(group_jid, async { let db_rows = pm .get_sender_key_devices(group_jid) @@ -810,10 +807,58 @@ impl Client { }); std::sync::Arc::new(SenderKeyDeviceMap::from_db_rows(&db_rows)) }) - .await; + .await + } + + /// Filter the resolved device set down to the subset still needing SKDM. + /// + /// No empty-cache early-exit: WA Web iterates an empty `senderKey` Map + /// as `false` per participant, so the filter must run unconditionally. + fn filter_skdm_targets( + &self, + group_jid: &str, + all_devices: &[Jid], + cached_map: &crate::sender_key_device_cache::SenderKeyDeviceMap, + own_sending_jid: &Jid, + ) -> Vec { + let needs_skdm: Vec = all_devices + .iter() + .filter(|device| { + if device.is_hosted() { + return false; + } + if device.user == own_sending_jid.user && device.device == own_sending_jid.device { + return false; + } + // O(1) lookups into pre-indexed cache + !cached_map + .device_has_key(&device.user, device.device) + .unwrap_or(false) + || cached_map.is_user_forgotten(&device.user) + }) + .cloned() + .collect(); + + log::debug!( + "Resolved {} devices ({} need SKDM) for {}", + all_devices.len(), + needs_skdm.len(), + group_jid + ); + needs_skdm + } + + /// SKDM target resolution for the status path, whose `GroupInfo` is built + /// fresh per send (no stable identity to memoize against). + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.resolve_skdm_targets", level = "debug", skip_all, fields(group = %wacore_binary::jid::observe_str(group_jid))))] + async fn resolve_skdm_targets( + &self, + group_jid: &str, + group_info: &wacore::client::context::GroupInfo, + own_sending_jid: &Jid, + ) -> Option<(std::sync::Arc>, Vec)> { + let cached_map = self.skdm_device_map(group_jid).await; - // No empty-cache early-exit: WA Web iterates an empty `senderKey` Map - // as `false` per participant, so the filter below must run unconditionally. let is_lid_mode = group_info.addressing_mode == wacore::types::message::AddressingMode::Lid; let jids_to_resolve: Vec = group_info .participants @@ -839,36 +884,41 @@ impl Client { } else { all_devices }; - - // Borrow for the filter so `all_devices` survives to feed the - // phash (the full set), while `needs_skdm` is just the subset - // still missing the key. - let needs_skdm: Vec = all_devices - .iter() - .filter(|device| { - if device.is_hosted() { - return false; - } - if device.user == own_sending_jid.user - && device.device == own_sending_jid.device - { - return false; - } - // O(1) lookups into pre-indexed cache - !cached_map - .device_has_key(&device.user, device.device) - .unwrap_or(false) - || cached_map.is_user_forgotten(&device.user) - }) - .cloned() - .collect(); - - log::debug!( - "Resolved {} devices ({} need SKDM) for {}", - all_devices.len(), - needs_skdm.len(), - group_jid + let all_devices = std::sync::Arc::new(all_devices); + let needs_skdm = + self.filter_skdm_targets(group_jid, &all_devices, &cached_map, own_sending_jid); + Some((all_devices, needs_skdm)) + } + Err(e) => { + log::warn!( + "Failed to resolve devices for SKDM check in {}: {:?}", + group_jid, + e ); + None + } + } + } + + /// SKDM target resolution for cached-group sends: the full device set + /// comes from the per-group memo (`resolve_group_devices_memoized`), so a + /// warm repeat send skips the per-member registry fan-out entirely. + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.resolve_skdm_targets_memoized", level = "debug", skip_all, fields(group = %group_jid)))] + async fn resolve_skdm_targets_memoized( + &self, + group: &Jid, + group_jid: &str, + group_info: &std::sync::Arc, + own_sending_jid: &Jid, + ) -> Option<(std::sync::Arc>, Vec)> { + let cached_map = self.skdm_device_map(group_jid).await; + match self + .resolve_group_devices_memoized(group, group_info, own_sending_jid) + .await + { + Ok(all_devices) => { + let needs_skdm = + self.filter_skdm_targets(group_jid, &all_devices, &cached_map, own_sending_jid); Some((all_devices, needs_skdm)) } Err(e) => { @@ -1309,6 +1359,11 @@ impl Client { crate::types::message::AddressingMode::Pn => (own_jid.clone(), "pn"), }; + // Memo identity must be the CACHED Arc: ensure_self_in_group clones + // a fresh GroupInfo whenever self is absent from the snapshot, which + // would make the memo miss on every send to such groups. The memoized + // resolver applies the same self-append internally. + let group_info_for_memo = std::sync::Arc::clone(&group_info); // resolve_skdm_targets and prepare_group_stanza both read the // participant list and expect self to be present. let group_info = ensure_self_in_group(group_info, &own_sending_jid); @@ -1369,18 +1424,25 @@ impl Client { // phash on every group send); `skdm_target_devices` is the subset // still missing the key. On the cold/`force_skdm` path both are // `None` and `prepare_group_stanza` resolves the set itself. - let (all_devices_for_phash, skdm_target_devices): (Option>, Option>) = - if force_skdm { - (None, None) - } else { - match self - .resolve_skdm_targets(&to_str, &group_info, &own_sending_jid) - .await - { - Some((all, needs)) => (Some(all), Some(needs)), - None => (None, None), - } - }; + let (all_devices_for_phash, skdm_target_devices): ( + Option>>, + Option>, + ) = if force_skdm { + (None, None) + } else { + match self + .resolve_skdm_targets_memoized( + &to, + &to_str, + &group_info_for_memo, + &own_sending_jid, + ) + .await + { + Some((all, needs)) => (Some(all), Some(needs)), + None => (None, None), + } + }; match wacore::send::prepare_group_stanza( &*self.runtime, @@ -2867,7 +2929,7 @@ mod tests { }; client .device_registry_cache - .insert((*user).into(), Arc::new(record)) + .raw_insert_for_tests((*user).into(), Arc::new(record)) .await; } diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index a12067420..6b97fa4d4 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -723,7 +723,8 @@ fn run_group_send(d: &mut GrpSendData) { // only emits a phash if it gets the full device set. Mirror the real // warm-send caller by passing it; the cold/force_skdm path resolves the set // itself and keeps None. - let all_devices_for_phash = (!d.force_skdm).then(|| d.participants.clone()); + let all_devices_for_phash = + (!d.force_skdm).then(|| std::sync::Arc::new(d.participants.clone())); let mut group_info = GroupInfo::new(std::mem::take(&mut d.participants), AddressingMode::Pn); let own_base = own_jid.to_non_ad(); if !group_info diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs index 1d604aa2d..16dcf9dc0 100644 --- a/wacore/src/send/group.rs +++ b/wacore/src/send/group.rs @@ -139,7 +139,7 @@ pub async fn prepare_group_stanza< // sends so the phash covers every device + self even when no SKDM is sent; // `None` on the cold `force_skdm` path (the set is resolved here) and for // status broadcasts (which keep the prior phash behavior). - all_devices_for_phash: Option>, + all_devices_for_phash: Option>>, edit: Option, extra_stanza_nodes: &[Node], ) -> Result { @@ -307,6 +307,7 @@ pub async fn prepare_group_stanza< // already holds the full resolved set. if let Some(src) = all_devices_for_phash .as_deref() + .map(Vec::as_slice) .or(distribution_list.as_deref()) { let phash_set = build_group_phash_set(src, &own_sending_jid);