From 6181e2b601ec905af306afb3053fc3efe28effc7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Mon, 1 Jun 2026 13:00:59 -0300 Subject: [PATCH] perf(lid-pn): skip PN->LID session migration for peers with no PN state migrate_signal_sessions_on_lid_discovery scanned MIGRATION_DEVICE_RANGE (100) device slots per newly-learned mapping, taking two per-address locks plus a session+identity lookup each. For a freshly-resolved peer (every member of a large group on first send) all 100 slots are empty, so the whole scan is wasted. Add SignalStore::has_signal_state_for_user (sqlite EXISTS, in-memory prefix scan; default true so unimplemented backends keep the full scan) and a SignalStoreCache::has_state_for_user that checks the in-memory cache then the backend. Guard the migration loop on it: skip entirely when the PN side has no session/identity. Behavior-identical for peers with state (the loop already migrated nothing when empty), and the full migration still runs otherwise. dhat (group-send 800 fresh members, cold first send): migrate frame 10.04 GB -> 0, total run 10.07 GB -> 0.02 GB (-99.8%); 12/12 group replies delivered. --- src/client/lid_pn.rs | 70 +++++++++++++++++++++ storages/sqlite-storage/src/sqlite_store.rs | 41 ++++++++++++ wacore/src/store/in_memory.rs | 40 ++++++++++++ wacore/src/store/signal_cache.rs | 26 ++++++++ wacore/src/store/traits.rs | 10 +++ 5 files changed, 187 insertions(+) diff --git a/src/client/lid_pn.rs b/src/client/lid_pn.rs index 5c32c6847..f47770383 100644 --- a/src/client/lid_pn.rs +++ b/src/client/lid_pn.rs @@ -404,6 +404,18 @@ impl Client { let backend = self.persistence_manager.backend(); + // Nothing to migrate unless the PN side has Signal state. For a freshly + // resolved peer (e.g. every member of a large group on first send) this + // skips MIGRATION_DEVICE_RANGE lock+lookup iterations that would all + // find nothing. On a lookup error, fall through to the full scan. + if let Ok(false) = self + .signal_cache + .has_state_for_user(pn, backend.as_ref()) + .await + { + return; + } + for device_id in 0..MIGRATION_DEVICE_RANGE { // `&str` → `CompactString` is inline for ≤24-byte user parts // (all PN/LID identifiers fit), so no String intermediate. @@ -1100,6 +1112,49 @@ mod tests { ); } + /// A freshly-resolved peer (no prior PN Signal state) must short-circuit the + /// per-device migration scan: nothing to move, so no LID session appears and + /// the MIGRATION_DEVICE_RANGE lock/lookup loop is skipped. + #[tokio::test] + async fn migrate_skips_when_no_pn_signal_state() { + use wacore::types::jid::JidExt as _; + + let client: Arc = create_test_client().await; + let pn = "5500000000777"; + let lid = "222222222222222"; + client + .add_lid_pn_mapping(lid, pn, LearningSource::PeerPnMessage) + .await + .unwrap(); + let backend = client.persistence_manager.backend(); + + // Fresh peer: no PN session or identity anywhere, so the guard skips. + assert!( + !client + .signal_cache + .has_state_for_user(pn, backend.as_ref()) + .await + .unwrap(), + "fresh peer should have no PN Signal state" + ); + + client + .migrate_signal_sessions_on_lid_discovery(pn, lid) + .await; + + // No LID session was materialized (nothing was migrated). + let lid_addr = Jid::lid_device(lid.to_string(), 0).to_protocol_address(); + assert!( + client + .signal_cache + .get_session(&lid_addr, backend.as_ref()) + .await + .unwrap() + .is_none(), + "migration of a stateless peer must not create a LID session" + ); + } + /// Migration must hold the same per-address session locks that /// encrypt/decrypt take. Otherwise a concurrent `message_encrypt` /// on the LID slot can clobber the just-migrated session (or read @@ -1118,6 +1173,21 @@ mod tests { .await .unwrap(); + // Seed a PN session so the migration actually enters its per-device + // loop. The existence guard skips when there is nothing to migrate, and + // this test is about the lock the loop takes when migrating real state. + let pn_addr = Jid::pn_device(pn.to_string(), 0).to_protocol_address(); + client + .signal_cache + .put_session( + &pn_addr, + wacore::libsignal::protocol::SessionRecord::deserialize(&tagged_session_blob( + 0xDEAD_BEEF, + )) + .expect("seed PN blob deserializes"), + ) + .await; + let lid_addr = Jid::lid_device(lid.to_string(), 0).to_protocol_address(); let lid_lock = client.session_lock_for(lid_addr.as_str()).await; let held = lid_lock.lock().await; diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 78b9a9eaf..4b85943a2 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -1267,6 +1267,47 @@ impl SignalStore for SqliteStore { .await } + async fn has_signal_state_for_user(&self, user: &str) -> Result { + let pool = self.pool.clone(); + let device_id = self.device_id; + // Address is `user@server` (device 0) or `user:dev@server`; `user` is a + // numeric PN/LID so it carries no LIKE wildcards. + let pat_at = format!("{user}@%"); + let pat_dev = format!("{user}:%"); + self.with_semaphore(move || -> Result { + let mut conn = pool + .get() + .map_err(|e| StoreError::Connection(Box::new(e)))?; + let has_session = diesel::select(diesel::dsl::exists( + sessions::table + .filter(sessions::device_id.eq(device_id)) + .filter( + sessions::address + .like(&pat_at) + .or(sessions::address.like(&pat_dev)), + ), + )) + .get_result::(&mut conn) + .map_err(|e| StoreError::Database(Box::new(e)))?; + if has_session { + return Ok(true); + } + let has_identity = diesel::select(diesel::dsl::exists( + identities::table + .filter(identities::device_id.eq(device_id)) + .filter( + identities::address + .like(&pat_at) + .or(identities::address.like(&pat_dev)), + ), + )) + .get_result::(&mut conn) + .map_err(|e| StoreError::Database(Box::new(e)))?; + Ok(has_identity) + }) + .await + } + async fn put_session(&self, address: &str, session: &[u8]) -> Result<()> { self.put_session_for_device(address, session, self.device_id) .await diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index bd863cbfb..9e4fb740a 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -147,6 +147,16 @@ impl SignalStore for InMemoryBackend { Ok(self.state.lock().await.sessions.contains_key(address)) } + async fn has_signal_state_for_user(&self, user: &str) -> Result { + fn matches(addr: &str, user: &str) -> bool { + addr.strip_prefix(user) + .is_some_and(|rest| rest.starts_with('@') || rest.starts_with(':')) + } + let state = self.state.lock().await; + Ok(state.sessions.keys().any(|k| matches(k, user)) + || state.identities.keys().any(|k| matches(k, user))) + } + async fn delete_session(&self, address: &str) -> Result<()> { self.state.lock().await.sessions.remove(address); Ok(()) @@ -683,6 +693,36 @@ mod tests { is_backend::(); } + #[tokio::test] + async fn has_signal_state_for_user_matches_by_user_prefix() { + let backend = InMemoryBackend::new(); + let user = "5511999990000"; + + assert!(!backend.has_signal_state_for_user(user).await.unwrap()); + + // Device 0 is keyed `user@server`. + backend + .put_session("5511999990000@s.whatsapp.net", b"sess") + .await + .unwrap(); + assert!(backend.has_signal_state_for_user(user).await.unwrap()); + + // A different user that this one is a prefix of must NOT match. + let other = InMemoryBackend::new(); + other + .put_session("55119999900001@s.whatsapp.net", b"sess") + .await + .unwrap(); + assert!(!other.has_signal_state_for_user(user).await.unwrap()); + + // Non-zero device is keyed `user:dev@server`; identity-only also counts. + let dev = InMemoryBackend::new(); + dev.put_identity("5511999990000:5@s.whatsapp.net", [7u8; 32]) + .await + .unwrap(); + assert!(dev.has_signal_state_for_user(user).await.unwrap()); + } + #[tokio::test] async fn store_sent_message_is_memory_bounded() { let backend = InMemoryBackend::new(); diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 2af5ce746..44e106024 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -275,6 +275,32 @@ impl SignalStoreCache { } } + /// Whether any session or identity is known for `user` (across device ids), + /// checking the in-memory cache first, then the durable backend. Lets a + /// caller skip a per-device migration scan for a user we've never had Signal + /// state with. Conservative on the cache side: any matching key counts + /// (even a stale/checked-out marker), so it never reports "none" when state + /// might exist. + pub async fn has_state_for_user(&self, user: &str, backend: &dyn SignalStore) -> Result { + fn matches(addr: &str, user: &str) -> bool { + addr.strip_prefix(user) + .is_some_and(|rest| rest.starts_with('@') || rest.starts_with(':')) + } + { + let state = self.sessions.lock().await; + if state.cache.keys().any(|k| matches(k, user)) { + return Ok(true); + } + } + { + let state = self.identities.lock().await; + if state.cache.keys().any(|k| matches(k, user)) { + return Ok(true); + } + } + Ok(backend.has_signal_state_for_user(user).await?) + } + // === Sessions (object cache — serialize only during flush) === /// Takes ownership of the cached session, leaving a `CheckedOut` marker. diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 12d91c962..f5520bfb5 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -130,6 +130,16 @@ pub trait SignalStore: Send + Sync { Ok(self.get_session(address).await?.is_some()) } + /// Whether any session or identity exists for `user` across all device ids. + /// Addresses are keyed `user@server` (device 0) or `user:dev@server`. Used + /// to skip the per-device PN->LID migration scan for users we've never had + /// Signal state with. Default is conservative (`true`) so a backend that + /// doesn't implement it keeps the caller's full per-device scan. + async fn has_signal_state_for_user(&self, user: &str) -> Result { + let _ = user; + Ok(true) + } + // --- PreKey Operations --- /// Store a pre-key.