diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index c6e3d9461..4d4faff0d 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/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, [u8; 32])]) -> Result<()> { + if identities.is_empty() { + return Ok(()); + } + + let device_id = self.device_id; + // `Arc` 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> { 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, 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, 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>> { 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, 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, [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, 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, 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 = 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"; diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 145ddb866..19f7ffbdc 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -52,8 +52,6 @@ pub struct SignalStoreCache { sessions: Mutex, identities: Mutex, sender_keys: Mutex, - /// Avoids per-flush Vec allocation on the hot path (called after every message). - flush_encode_buf: Mutex>, /// Per-(group, sender) locks serializing each sender-key chain advance. /// Coordination only (like the client session locks): never time-evicted. sender_key_locks: Mutex, Arc>>>, @@ -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, 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, [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,13 +618,14 @@ impl SignalStoreCache { let mut state = self.sender_keys.lock().await; let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect(); + let mut batch: Vec<(Arc, 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?; @@ -626,6 +633,9 @@ impl SignalStoreCache { None => {} } } + if !batch.is_empty() { + backend.put_sender_keys_batch(&batch).await?; + } for key in &dirty_keys { state.dirty.remove(key); diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 51c3d68cc..d9b045321 100644 --- a/wacore/src/store/traits.rs +++ b/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` so callers (the flush path) pass shared keys + /// without allocating a `String` per entry. + async fn put_identities_batch( + &self, + identities: &[(std::sync::Arc, [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>; @@ -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, 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, 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>>;