perf(store): batch session/identity/sender-key flush into one transaction per category

storages/sqlite-storage/src/sqlite_store.rs

@@ -263,6 +263,15 @@ impl SqliteStore {
                     if is_retriable_sqlite_error(e) && attempt < MAX_RETRIES =>
                 {
                     let delay_ms = 10u64 * (1u64 << attempt.min(4));
+                    // Skip the first transient blip; warn from the second retry on so
+                    // sustained busy/locked contention doesn't go unobserved.
+                    if attempt >= 1 {
+                        warn!(
+                            "{op_name} busy/locked, retry {}/{} in {delay_ms}ms: {e}",
+                            attempt + 1,
+                            MAX_RETRIES + 1
+                        );
+                    }
                     tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
                 }
                 Ok(Err(e)) => return Err(e.into()),
@@ -1261,6 +1270,38 @@ impl SignalStore for SqliteStore {
             .await
     }
 
+    async fn put_identities_batch(&self, identities: &[(Arc<str>, [u8; 32])]) -> Result<()> {
+        if identities.is_empty() {
+            return Ok(());
+        }
+
+        let device_id = self.device_id;
+        // `Arc<Vec>` so each retry attempt bumps a refcount instead of re-cloning
+        // the whole batch.
+        let batch = Arc::new(identities.to_vec());
+        self.with_retry("put_identities_batch", || {
+            let batch = batch.clone();
+            Box::new(move |conn: &mut SqliteConnection| {
+                conn.transaction(|conn| {
+                    for (address, key) in batch.iter() {
+                        diesel::insert_into(identities::table)
+                            .values((
+                                identities::address.eq(address.as_ref()),
+                                identities::key.eq(&key[..]),
+                                identities::device_id.eq(device_id),
+                            ))
+                            .on_conflict((identities::address, identities::device_id))
+                            .do_update()
+                            .set(identities::key.eq(&key[..]))
+                            .execute(conn)?;
+                    }
+                    Ok(())
+                })
+            })
+        })
+        .await
+    }
+
     async fn load_identity(&self, address: &str) -> Result<Option<[u8; 32]>> {
         let blob = self
             .load_identity_for_device(address, self.device_id)
@@ -1355,6 +1396,36 @@ impl SignalStore for SqliteStore {
             .await
     }
 
+    async fn put_sessions_batch(&self, sessions: &[(Arc<str>, Bytes)]) -> Result<()> {
+        if sessions.is_empty() {
+            return Ok(());
+        }
+
+        let device_id = self.device_id;
+        let batch = Arc::new(sessions.to_vec());
+        self.with_retry("put_sessions_batch", || {
+            let batch = batch.clone();
+            Box::new(move |conn: &mut SqliteConnection| {
+                conn.transaction(|conn| {
+                    for (address, record) in batch.iter() {
+                        diesel::insert_into(sessions::table)
+                            .values((
+                                sessions::address.eq(address.as_ref()),
+                                sessions::record.eq(record.as_ref()),
+                                sessions::device_id.eq(device_id),
+                            ))
+                            .on_conflict((sessions::address, sessions::device_id))
+                            .do_update()
+                            .set(sessions::record.eq(record.as_ref()))
+                            .execute(conn)?;
+                    }
+                    Ok(())
+                })
+            })
+        })
+        .await
+    }
+
     async fn delete_session(&self, address: &str) -> Result<()> {
         self.delete_session_for_device(address, self.device_id)
             .await
@@ -1771,6 +1842,36 @@ impl SignalStore for SqliteStore {
             .await
     }
 
+    async fn put_sender_keys_batch(&self, sender_keys: &[(Arc<str>, Bytes)]) -> Result<()> {
+        if sender_keys.is_empty() {
+            return Ok(());
+        }
+
+        let device_id = self.device_id;
+        let batch = Arc::new(sender_keys.to_vec());
+        self.with_retry("put_sender_keys_batch", || {
+            let batch = batch.clone();
+            Box::new(move |conn: &mut SqliteConnection| {
+                conn.transaction(|conn| {
+                    for (address, record) in batch.iter() {
+                        diesel::insert_into(sender_keys::table)
+                            .values((
+                                sender_keys::address.eq(address.as_ref()),
+                                sender_keys::record.eq(record.as_ref()),
+                                sender_keys::device_id.eq(device_id),
+                            ))
+                            .on_conflict((sender_keys::address, sender_keys::device_id))
+                            .do_update()
+                            .set(sender_keys::record.eq(record.as_ref()))
+                            .execute(conn)?;
+                    }
+                    Ok(())
+                })
+            })
+        })
+        .await
+    }
+
     async fn get_sender_key(&self, address: &str) -> Result<Option<Vec<u8>>> {
         self.get_sender_key_for_device(address, self.device_id)
             .await
@@ -2967,6 +3068,90 @@ mod tests {
         assert!(empty.is_empty());
     }
 
+    #[tokio::test]
+    async fn put_signal_batches_persist_and_upsert() {
+        use std::sync::Arc;
+        let store = create_test_store().await;
+
+        let sessions: Vec<(Arc<str>, Bytes)> = (0..5u8)
+            .map(|i| {
+                (
+                    Arc::from(format!("user{i}@s.whatsapp.net").as_str()),
+                    Bytes::from(vec![i; 8]),
+                )
+            })
+            .collect();
+        store.put_sessions_batch(&sessions).await.unwrap();
+        for (addr, bytes) in &sessions {
+            assert_eq!(
+                store.get_session(addr).await.unwrap().as_deref(),
+                Some(bytes.as_ref())
+            );
+        }
+
+        let identities: Vec<(Arc<str>, [u8; 32])> = (0..5u8)
+            .map(|i| {
+                (
+                    Arc::from(format!("user{i}@s.whatsapp.net").as_str()),
+                    [i; 32],
+                )
+            })
+            .collect();
+        store.put_identities_batch(&identities).await.unwrap();
+        for (addr, key) in &identities {
+            assert_eq!(store.load_identity(addr).await.unwrap(), Some(*key));
+        }
+
+        let sender_keys: Vec<(Arc<str>, Bytes)> = (0..5u8)
+            .map(|i| {
+                (
+                    Arc::from(format!("g@g.us::user{i}").as_str()),
+                    Bytes::from(vec![i; 16]),
+                )
+            })
+            .collect();
+        store.put_sender_keys_batch(&sender_keys).await.unwrap();
+        for (addr, bytes) in &sender_keys {
+            assert_eq!(
+                store.get_sender_key(addr).await.unwrap().as_deref(),
+                Some(bytes.as_ref())
+            );
+        }
+
+        // Re-batching the same addresses upserts (on_conflict do_update).
+        let updated: Vec<(Arc<str>, Bytes)> = sessions
+            .iter()
+            .map(|(addr, _)| (addr.clone(), Bytes::from(vec![0xAA; 8])))
+            .collect();
+        store.put_sessions_batch(&updated).await.unwrap();
+        for (addr, _) in &sessions {
+            assert_eq!(
+                store.get_session(addr).await.unwrap().as_deref(),
+                Some([0xAA; 8].as_slice())
+            );
+        }
+
+        // Duplicate address within one batch: last value wins via on_conflict
+        // do_update inside the single transaction.
+        let dup: Arc<str> = Arc::from("dup@s.whatsapp.net");
+        store
+            .put_sessions_batch(&[
+                (dup.clone(), Bytes::from(vec![1u8; 4])),
+                (dup.clone(), Bytes::from(vec![2u8; 4])),
+            ])
+            .await
+            .unwrap();
+        assert_eq!(
+            store.get_session(&dup).await.unwrap().as_deref(),
+            Some([2u8; 4].as_slice())
+        );
+
+        // Empty batches short-circuit without error.
+        store.put_sessions_batch(&[]).await.unwrap();
+        store.put_identities_batch(&[]).await.unwrap();
+        store.put_sender_keys_batch(&[]).await.unwrap();
+    }
+
     #[test]
     fn test_parse_database_path_regular_path() {
         let path = "/var/lib/whatsapp/database.db";

wacore/src/store/signal_cache.rs

@@ -52,8 +52,6 @@ pub struct SignalStoreCache {
     sessions: Mutex<SessionStoreState>,
     identities: Mutex<ByteStoreState>,
     sender_keys: Mutex<SenderKeyStoreState>,
-    /// Avoids per-flush Vec allocation on the hot path (called after every message).
-    flush_encode_buf: Mutex<Vec<u8>>,
     /// Per-(group, sender) locks serializing each sender-key chain advance.
     /// Coordination only (like the client session locks): never time-evicted.
     sender_key_locks: Mutex<HashMap<Arc<str>, Arc<Mutex<()>>>>,
@@ -272,7 +270,6 @@ impl SignalStoreCache {
             sessions: Mutex::new(SessionStoreState::new()),
             identities: Mutex::new(ByteStoreState::new()),
             sender_keys: Mutex::new(SenderKeyStoreState::new()),
-            flush_encode_buf: Mutex::new(Vec::with_capacity(4096)),
             sender_key_locks: Mutex::new(HashMap::new()),
             max_entries,
         }
@@ -542,23 +539,28 @@ impl SignalStoreCache {
     ///   mutations to the same store are blocked until the flush completes.
     /// - Dirty sets are cleared only after successful writes.
     pub async fn flush(&self, backend: &dyn SignalStore) -> Result<()> {
-        // Flush sessions
+        // Flush sessions: one batched write for all dirty puts instead of one
+        // backend call (and one SQLite transaction) per session.
         {
             let mut state = self.sessions.lock().await;
             let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
             let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
 
-            let mut encode_buf = self.flush_encode_buf.lock().await;
+            let mut batch: Vec<(Arc<str>, bytes::Bytes)> = Vec::new();
             for address in &dirty_keys {
                 match state.cache.get(address.as_ref()) {
                     Some(SessionEntry::Present(record)) => {
-                        record.serialize_into(&mut encode_buf);
-                        backend.put_session(address, &encode_buf).await?;
+                        let mut buf = Vec::new();
+                        record.serialize_into(&mut buf);
+                        batch.push((address.clone(), bytes::Bytes::from(buf)));
                     }
                     Some(SessionEntry::CheckedOut) => continue,
                     _ => {}
                 }
             }
+            if !batch.is_empty() {
+                backend.put_sessions_batch(&batch).await?;
+            }
             for address in &deleted_keys {
                 backend.delete_session(address).await?;
             }
@@ -583,6 +585,7 @@ impl SignalStoreCache {
             let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
             let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
 
+            let mut batch: Vec<(Arc<str>, [u8; 32])> = Vec::new();
             for address in &dirty_keys {
                 if let Some(Some(data)) = state.cache.get(address.as_ref()) {
                     let key: [u8; 32] = data.as_ref().try_into().map_err(|_| {
@@ -591,9 +594,12 @@ impl SignalStoreCache {
                             data.len()
                         )
                     })?;
-                    backend.put_identity(address, key).await?;
+                    batch.push((address.clone(), key));
                 }
             }
+            if !batch.is_empty() {
+                backend.put_identities_batch(&batch).await?;
+            }
             for address in &deleted_keys {
                 backend.delete_identity(address).await?;
             }
@@ -612,20 +618,24 @@ impl SignalStoreCache {
             let mut state = self.sender_keys.lock().await;
             let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
 
+            let mut batch: Vec<(Arc<str>, bytes::Bytes)> = Vec::new();
             for name in &dirty_keys {
                 match state.cache.get(name.as_ref()) {
                     Some(Some(record)) => {
                         let bytes = record
                             .serialize()
                             .map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?;
-                        backend.put_sender_key(name, &bytes).await?;
+                        batch.push((name.clone(), bytes::Bytes::from(bytes)));
                     }
                     Some(None) => {
                         backend.delete_sender_key(name).await?;
                     }
                     None => {}
                 }
             }
+            if !batch.is_empty() {
+                backend.put_sender_keys_batch(&batch).await?;
+            }
 
             for key in &dirty_keys {
                 state.dirty.remove(key);

wacore/src/store/traits.rs

@@ -108,6 +108,20 @@ pub trait SignalStore: Send + Sync {
     /// Store an identity key for a remote address.
     async fn put_identity(&self, address: &str, key: [u8; 32]) -> Result<()>;
 
+    /// Store multiple identity keys in a single batch operation.
+    /// Default implementation falls back to individual `put_identity` calls.
+    /// Addresses are `Arc<str>` so callers (the flush path) pass shared keys
+    /// without allocating a `String` per entry.
+    async fn put_identities_batch(
+        &self,
+        identities: &[(std::sync::Arc<str>, [u8; 32])],
+    ) -> Result<()> {
+        for (address, key) in identities {
+            self.put_identity(address, *key).await?;
+        }
+        Ok(())
+    }
+
     /// Load an identity key for a remote address (always 32 bytes).
     async fn load_identity(&self, address: &str) -> Result<Option<[u8; 32]>>;
 
@@ -122,6 +136,15 @@ pub trait SignalStore: Send + Sync {
     /// Store an encrypted session.
     async fn put_session(&self, address: &str, session: &[u8]) -> Result<()>;
 
+    /// Store multiple encrypted sessions in a single batch operation.
+    /// Default implementation falls back to individual `put_session` calls.
+    async fn put_sessions_batch(&self, sessions: &[(std::sync::Arc<str>, Bytes)]) -> Result<()> {
+        for (address, session) in sessions {
+            self.put_session(address, session).await?;
+        }
+        Ok(())
+    }
+
     /// Delete a session.
     async fn delete_session(&self, address: &str) -> Result<()>;
 
@@ -195,6 +218,18 @@ pub trait SignalStore: Send + Sync {
     /// Store a sender key for group messaging.
     async fn put_sender_key(&self, address: &str, record: &[u8]) -> Result<()>;
 
+    /// Store multiple sender keys in a single batch operation.
+    /// Default implementation falls back to individual `put_sender_key` calls.
+    async fn put_sender_keys_batch(
+        &self,
+        sender_keys: &[(std::sync::Arc<str>, Bytes)],
+    ) -> Result<()> {
+        for (address, record) in sender_keys {
+            self.put_sender_key(address, record).await?;
+        }
+        Ok(())
+    }
+
     /// Get a sender key.
     async fn get_sender_key(&self, address: &str) -> Result<Option<Vec<u8>>>;