feat: batch write store - #281
Conversation
📝 WalkthroughWalkthroughAdds client connection liveness and weak self-references, replaces SignalProtocol adapter with a batched/cached variant and blocking crypto tasks, introduces terminal transport-send error handling, refactors group encryption to SenderKeyName flow, downgrades receipt logging, and extends benchmark stores with a backend-serialized mode. Changes
Sequence Diagram(s)(omitted — changes are across many areas but do not introduce a single new multi-component sequential flow suitable for a concise diagram) Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
aa67874 to
a0d4ec2
Compare
Bencher Report![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/send.rs (1)
349-365:⚠️ Potential issue | 🟠 MajorGroup stanza encryption still runs on the async runtime.
prepare_group_stanzaperforms Signal encryption; this is CPU-heavy and should be offloaded viatokio::task::spawn_blocking(including the retry path) to avoid stalling the runtime.As per coding guidelines: All blocking I/O (like
ureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime.Also applies to: 419-436
🤖 Fix all issues with AI agents
In `@src/client.rs`:
- Around line 239-245: The current Client::shared() method uses expect on
upgrading a Weak self reference; change shared() to return Option<Arc<Self>> (or
Result<Arc<Self>, SomeError>) instead of panicking so callers can handle missing
self; update all call sites of shared() accordingly. Also replace any places
that call self.self_weak.set(...) and rely on expect/panic (the other occurrence
around the self initialization logic) to handle the Err case gracefully
(log/return an error or propagate a Result) rather than unwrapping—ensure all
references to shared() and self_weak.set(...) are updated to use the new
non-panicking return types and proper error handling.
In `@src/message.rs`:
- Around line 973-992: The code currently awaits
device_guard.backend.delete_session(&address_str).await while holding
device_guard, which can block other tasks; route this through the device command
path instead. Create a DeviceCommand variant (e.g.,
DeleteSession(address_str.clone())) and call
PersistenceManager::process_command(...) or the existing persistence command API
to perform the delete after cloning the needed address/signal_address, drop
device_guard before awaiting the command, and then call
adapter.invalidate_session(&signal_address).await only after the command
completes; ensure you remove the direct backend.delete_session call and use the
DeviceCommand/PersistenceManager::process_command flow to modify device state.
In `@src/store/signal_adapter.rs`:
- Around line 79-94: The flush method currently clears cache.dirty_sessions
up-front which loses pending updates if any SessionStore::store_session call
fails; change flush so dirty_sessions is not removed before writes succeed:
acquire the cache via self.cache.lock().await, clone or collect the list of
dirty addresses (or their (address, record.clone()) pairs) without mutating
cache.dirty_sessions, release the lock, perform SessionStore::store_session(&mut
self.inner, &address, &record).await for each, and only after a successful store
remove that address from cache.dirty_sessions (or, on failure, re-add it) so
failed writes remain marked dirty for retry; reference the flush function,
cache.dirty_sessions, cache.sessions and SessionStore::store_session when making
this change.
In `@src/store/signal.rs`:
- Around line 207-236: The code treats stored empty identity bytes as present;
change handling so empty byte arrays are treated as None: after loading
existing_identity_bytes from self.backend.load_identity(&address_str), normalize
it to e.g. let existing_nonempty = existing_identity_bytes.as_deref().filter(|b|
!b.is_empty()); use existing_nonempty for the equality check (replace the
is_some_and(...) test) and when computing IdentityChange::from_changed pass
existing_nonempty.is_some() so that empty payloads are not considered "existing"
for both equality and change detection before calling put_identity and returning
the IdentityChange.
| let deleted = if let Err(err) = | ||
| device_guard.backend.delete_session(&address_str).await | ||
| { | ||
| log::warn!( | ||
| "Failed to delete stale session for {}: {:?}", | ||
| signal_address, | ||
| err | ||
| ); | ||
| false | ||
| } else { | ||
| log::info!( | ||
| "Deleted stale session for {} to allow re-establishment", | ||
| signal_address | ||
| ); | ||
| } | ||
| true | ||
| }; | ||
| drop(device_guard); | ||
| if deleted { | ||
| adapter.invalidate_session(&signal_address).await; | ||
| } |
There was a problem hiding this comment.
Avoid holding the device write-lock across an async delete; route through the command path.
delete_session is awaited while holding device_guard, which can stall other tasks and conflicts with the device-command guideline. Clone the backend (or use PersistenceManager::process_command) and drop the lock before the await.
As per coding guidelines: src/**/*.rs: Use DeviceCommand + PersistenceManager::process_command() for all device state modifications; never modify Device state directly.
Suggested fix (lock-free delete)
- let device_guard = device_arc.write().await;
let address_str = signal_address.to_string();
- let deleted = if let Err(err) =
- device_guard.backend.delete_session(&address_str).await
- {
+ let backend = {
+ let device = device_arc.read().await;
+ Arc::clone(&device.backend)
+ };
+ let deleted = if let Err(err) = backend.delete_session(&address_str).await {
log::warn!(
"Failed to delete stale session for {}: {:?}",
signal_address,
err
);
false
} else {
log::info!(
"Deleted stale session for {} to allow re-establishment",
signal_address
);
true
};
- drop(device_guard);🤖 Prompt for AI Agents
In `@src/message.rs` around lines 973 - 992, The code currently awaits
device_guard.backend.delete_session(&address_str).await while holding
device_guard, which can block other tasks; route this through the device command
path instead. Create a DeviceCommand variant (e.g.,
DeleteSession(address_str.clone())) and call
PersistenceManager::process_command(...) or the existing persistence command API
to perform the delete after cloning the needed address/signal_address, drop
device_guard before awaiting the command, and then call
adapter.invalidate_session(&signal_address).await only after the command
completes; ensure you remove the direct backend.delete_session call and use the
DeviceCommand/PersistenceManager::process_command flow to modify device state.
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Fix all issues with AI agents
In `@src/client.rs`:
- Around line 1917-1922: The is_connected check currently treats a failed
try_lock on noise_socket as a false (disconnected) which yields false negatives;
update the is_connected implementation to avoid treating lock contention as a
disconnect—either rely solely on the atomic connection_alive flag or consider a
failed try_lock as "still connected" (e.g., treat Err from
noise_socket.try_lock() as true) so the result is monotonic; modify the logic in
is_connected (referencing connection_alive and noise_socket) and ensure
wait_for_socket / wait_for_connected behavior still relies on the updated
monotonic is_connected and notify_waiters semantics.
In `@src/message.rs`:
- Around line 1033-1036: After adapter.flush().await fails the current code logs
a warning but still returns any_success=true which misreports that session state
was persisted; update the error path so that if adapter.flush().await returns
Err(flush_err) you mutate any_success = false (or otherwise clear/overwrite the
success flag returned from this function) before returning, and include the
flush_err in the process log; locate the flush call and the variables
any_success and dispatched_undecryptable in the surrounding function and ensure
the function returns the corrected any_success value (alternatively, implement a
retry/receipt-trigger path in that error branch if your design prefers proactive
retry instead of flipping the flag).
- Around line 764-783: When decrypt_session_message_with_blocking returns Err (a
JoinError), add a warning log that explicitly notes that the spawn_blocking task
failed and any unflushed adapter state from previous successful iterations may
be lost; include identifying context like info.id and info.source.sender and the
JoinError (e) in the log and mention that adapter was moved into the task so its
buffered writes couldn't be flushed, so operators are aware of potential
data-loss risk. Ensure this log is emitted in the Err(e) branch where adapter
was previously moved and before returning the tuple, referencing
decrypt_session_message_with_blocking, adapter, and JoinError/spawn_blocking in
the message.
In `@src/send.rs`:
- Around line 176-203: The peer/DM path currently calls stanza_result? before
flushing the BatchedSignalProtocolStoreAdapter, so any store updates are dropped
on error; move the store_adapter.flush().await so it always runs regardless of
stanza_result (e.g., match on stanza_result, call store_adapter.flush().await in
both Ok and Err branches, then propagate the error or return the stanza),
referencing BatchedSignalProtocolStoreAdapter, prepare_peer_stanza,
store_adapter.flush, and the spawn_blocking block; apply the same change to the
other occurrence mentioned (around the 492-534 region).
- Line 2: prepare_group_stanza is being executed on the async runtime even
though it performs CPU-bound per-recipient encryption via group_encrypt and
message_encrypt; wrap the call to prepare_group_stanza in
tokio::task::spawn_blocking and await the JoinHandle to offload encryption to a
blocking thread, and ensure any errors/results are propagated back to the async
context (update both call sites that currently call prepare_group_stanza
directly so they mirror the existing peer/DM spawn_blocking pattern).
In `@src/store/signal_adapter.rs`:
- Around line 79-101: The flush method snapshots dirty_sessions then unlocks, so
concurrent store_* calls can replace a session after snapshotting and be lost
when flush blindly removes the dirty flag; to fix, when removing the dirty flag
in flush (for SessionStore::store_session / self.cache.dirty_sessions), check
that the cached session still equals the record you just flushed (or use a
stored version/token) and only remove the dirty marker if they match; apply the
same guarded-remove pattern to identities and sender keys (the other flush
blocks referenced) so newer in-memory updates remain dirty and will be persisted
later.
🧹 Nitpick comments (2)
src/message.rs (2)
829-845:delete_identitybypasses the batched adapter and goes directly to the backend.
delete_session(line 976) correctly goes throughadapter.delete_session(...), butdelete_identityhere callsbackend.delete_identity(...)directly, then separately invalidates the adapter cache. This is inconsistent — if the batched adapter has a cached identity that differs from what the backend now holds, there's a window for stale reads from other concurrent paths.Consider routing this through the adapter (if it supports
delete_identity) for consistency with howdelete_sessionis handled.As per coding guidelines:
src/**/*.rs: UseDeviceCommand+PersistenceManager::process_command()for all device state modifications; never modify Device state directly.
1075-1078: Group decrypt still acquires a write lock ondevice_arcper message.The session path now uses the batched adapter for reduced lock contention, but
process_group_enc_batchstill holdsdevice_arc.write().awaitfor eachgroup_decryptcall. If the batched adapter also wraps the sender key store (which the relevant snippet insignal_adapter.rsshowsCachedSenderKeyAdapter), consider using it here too for consistency and reduced lock pressure.
| self.connection_alive.load(Ordering::Acquire) | ||
| && self | ||
| .noise_socket | ||
| .try_lock() | ||
| .is_ok_and(|guard| guard.is_some()) | ||
| } |
There was a problem hiding this comment.
Avoid false negatives from try_lock in is_connected.
try_lock failure signals contention, not a disconnect. Returning false can cause wait_for_socket/wait_for_connected to time out after notify_waiters. Consider treating lock contention as connected (or rely solely on connection_alive) so the check is monotonic.
Suggested fix
- self.connection_alive.load(Ordering::Acquire)
- && self
- .noise_socket
- .try_lock()
- .is_ok_and(|guard| guard.is_some())
+ if !self.connection_alive.load(Ordering::Acquire) {
+ return false;
+ }
+ match self.noise_socket.try_lock() {
+ Ok(guard) => guard.is_some(),
+ Err(_) => true, // lock contention shouldn't be treated as disconnected
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.connection_alive.load(Ordering::Acquire) | |
| && self | |
| .noise_socket | |
| .try_lock() | |
| .is_ok_and(|guard| guard.is_some()) | |
| } | |
| if !self.connection_alive.load(Ordering::Acquire) { | |
| return false; | |
| } | |
| match self.noise_socket.try_lock() { | |
| Ok(guard) => guard.is_some(), | |
| Err(_) => true, // lock contention shouldn't be treated as disconnected | |
| } | |
| } |
🤖 Prompt for AI Agents
In `@src/client.rs` around lines 1917 - 1922, The is_connected check currently
treats a failed try_lock on noise_socket as a false (disconnected) which yields
false negatives; update the is_connected implementation to avoid treating lock
contention as a disconnect—either rely solely on the atomic connection_alive
flag or consider a failed try_lock as "still connected" (e.g., treat Err from
noise_socket.try_lock() as true) so the result is monotonic; modify the logic in
is_connected (referencing connection_alive and noise_socket) and ensure
wait_for_socket / wait_for_connected behavior still relies on the updated
monotonic is_connected and notify_waiters semantics.
| let (next_adapter, parsed_message, decrypt_res) = | ||
| match decrypt_session_message_with_blocking( | ||
| parsed_message, | ||
| signal_address.clone(), | ||
| adapter, | ||
| ) | ||
| .await | ||
| { | ||
| Ok(result) => result, | ||
| Err(e) => { | ||
| log::error!( | ||
| "spawn_blocking failed while decrypting message {} from {}: {}", | ||
| info.id, | ||
| info.source.sender, | ||
| e | ||
| ); | ||
| return (any_success, any_duplicate, dispatched_undecryptable); | ||
| } | ||
| }; | ||
| adapter = next_adapter; |
There was a problem hiding this comment.
On JoinError, accumulated adapter writes from prior iterations are silently lost.
If the spawn_blocking task panics or is cancelled, the adapter (moved into the closure) is irrecoverable, so you can't flush. However, any session state mutations from previous successful decryptions in this batch are also lost because they haven't been flushed yet.
This is low-probability (JoinError is exceptional) and self-healing (retry will re-establish state), but worth a log line noting the data-loss risk so operators aren't surprised.
Suggested: warn about lost writes
Err(e) => {
log::error!(
- "spawn_blocking failed while decrypting message {} from {}: {}",
+ "spawn_blocking failed while decrypting message {} from {}: {}. \
+ Any batched session writes from prior messages in this batch are lost.",
info.id,
info.source.sender,
e
);
return (any_success, any_duplicate, dispatched_undecryptable);
}🤖 Prompt for AI Agents
In `@src/message.rs` around lines 764 - 783, When
decrypt_session_message_with_blocking returns Err (a JoinError), add a warning
log that explicitly notes that the spawn_blocking task failed and any unflushed
adapter state from previous successful iterations may be lost; include
identifying context like info.id and info.source.sender and the JoinError (e) in
the log and mention that adapter was moved into the task so its buffered writes
couldn't be flushed, so operators are aware of potential data-loss risk. Ensure
this log is emitted in the Err(e) branch where adapter was previously moved and
before returning the tuple, referencing decrypt_session_message_with_blocking,
adapter, and JoinError/spawn_blocking in the message.
| if let Err(flush_err) = adapter.flush().await { | ||
| log::warn!("Failed to flush cached Signal stores after decrypt batch: {flush_err}"); | ||
| } | ||
| (any_success, any_duplicate, dispatched_undecryptable) |
There was a problem hiding this comment.
Flush failure silently loses session state updates while still reporting success.
If flush() fails, the function returns any_success = true (messages were dispatched to the event bus), but session ratchet state was never persisted. The next message from this sender will fail decryption with a ratchet mismatch or duplicate-message error.
Consider whether a flush failure should flip any_success to false or at least trigger a retry receipt proactively, since the session state is now inconsistent.
🤖 Prompt for AI Agents
In `@src/message.rs` around lines 1033 - 1036, After adapter.flush().await fails
the current code logs a warning but still returns any_success=true which
misreports that session state was persisted; update the error path so that if
adapter.flush().await returns Err(flush_err) you mutate any_success = false (or
otherwise clear/overwrite the success flag returned from this function) before
returning, and include the flush_err in the process log; locate the flush call
and the variables any_success and dispatched_undecryptable in the surrounding
function and ensure the function returns the corrected any_success value
(alternatively, implement a retry/receipt-trigger path in that error branch if
your design prefers proactive retry instead of flipping the flag).
| @@ -1,5 +1,5 @@ | |||
| use crate::client::Client; | |||
| use crate::store::signal_adapter::SignalProtocolStoreAdapter; | |||
| use crate::store::signal_adapter::BatchedSignalProtocolStoreAdapter; | |||
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd /tmp && find . -type f -name "send.rs" 2>/dev/null | head -5Repository: jlucaso1/whatsapp-rust
Length of output: 48
🏁 Script executed:
fd -a "send.rs"Repository: jlucaso1/whatsapp-rust
Length of output: 134
🏁 Script executed:
rg -l "prepare_group_stanza"Repository: jlucaso1/whatsapp-rust
Length of output: 149
🏁 Script executed:
rg -n "prepare_group_stanza|spawn_blocking" src/send.rs | head -30Repository: jlucaso1/whatsapp-rust
Length of output: 685
🏁 Script executed:
sed -n '175,215p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 2114
🏁 Script executed:
sed -n '340,365p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 911
🏁 Script executed:
sed -n '495,535p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 2112
🏁 Script executed:
rg -A 20 "pub fn prepare_group_stanza" wacore/src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 48
🏁 Script executed:
rg -B 5 -A 30 "fn prepare_group_stanza" wacore/src/Repository: jlucaso1/whatsapp-rust
Length of output: 2155
🏁 Script executed:
rg -l "prepare_group_stanza" wacore/src/Repository: jlucaso1/whatsapp-rust
Length of output: 112
🏁 Script executed:
cat wacore/src/send.rs | head -100Repository: jlucaso1/whatsapp-rust
Length of output: 3922
🏁 Script executed:
wc -l wacore/src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 89
🏁 Script executed:
rg -n "group_encrypt|message_encrypt|sender_key" wacore/src/send.rs | head -20Repository: jlucaso1/whatsapp-rust
Length of output: 1352
🏁 Script executed:
sed -n '206,380p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 7562
🏁 Script executed:
sed -n '350,365p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 550
🏁 Script executed:
sed -n '415,435p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 1128
🏁 Script executed:
sed -n '365,430p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 3243
🏁 Script executed:
sed -n '396,412p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 909
🏁 Script executed:
sed -n '195,205p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 543
🏁 Script executed:
sed -n '257,270p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 835
🏁 Script executed:
sed -n '440,450p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 505
🏁 Script executed:
sed -n '436,450p' src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 658
🏁 Script executed:
sed -n '395,450p' src/send.rs | grep -A 10 "Err(e)"Repository: jlucaso1/whatsapp-rust
Length of output: 1143
🏁 Script executed:
rg -B 2 "pub async fn prepare_group_stanza" wacore/src/send.rs | head -10Repository: jlucaso1/whatsapp-rust
Length of output: 138
🏁 Script executed:
sed -n '754,760p' wacore/src/send.rsRepository: jlucaso1/whatsapp-rust
Length of output: 320
Wrap group encryption in spawn_blocking to prevent async runtime stalling.
The peer and DM paths both use tokio::task::spawn_blocking to offload encryption work (lines 182 and 504), but the group path calls prepare_group_stanza directly on the async runtime at lines 351 and 421. Since prepare_group_stanza performs per-recipient encryption with group_encrypt and message_encrypt calls, these CPU-bound operations should be wrapped in spawn_blocking to avoid blocking reactor threads under load.
Per coding guidelines: "All blocking I/O (like ureq calls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped in tokio::task::spawn_blocking to avoid stalling the async runtime."
🤖 Prompt for AI Agents
In `@src/send.rs` at line 2, prepare_group_stanza is being executed on the async
runtime even though it performs CPU-bound per-recipient encryption via
group_encrypt and message_encrypt; wrap the call to prepare_group_stanza in
tokio::task::spawn_blocking and await the JoinHandle to offload encryption to a
blocking thread, and ensure any errors/results are propagated back to the async
context (update both call sites that currently call prepare_group_stanza
directly so they mirror the existing peer/DM spawn_blocking pattern).
| let device_store_arc = self.persistence_manager.get_device_arc().await; | ||
| let mut store_adapter = SignalProtocolStoreAdapter::new(device_store_arc); | ||
|
|
||
| wacore::send::prepare_peer_stanza( | ||
| &mut store_adapter.session_store, | ||
| &mut store_adapter.identity_store, | ||
| to, | ||
| encryption_jid, | ||
| message, | ||
| request_id, | ||
| ) | ||
| .await? | ||
| let message_for_encrypt = message.clone(); | ||
| let to_for_encrypt = to.clone(); | ||
| let encryption_jid_for_encrypt = encryption_jid.clone(); | ||
| // Peer encryption is a single-session operation. With the batched | ||
| // cache, store ops resolve from memory so the work is CPU-bound. | ||
| let (mut store_adapter, stanza_result) = tokio::task::spawn_blocking(move || { | ||
| let runtime = tokio::runtime::Handle::current(); | ||
| let mut store_adapter = BatchedSignalProtocolStoreAdapter::new(device_store_arc); | ||
| let stanza_result = runtime.block_on(async { | ||
| wacore::send::prepare_peer_stanza( | ||
| &mut store_adapter.session_store, | ||
| &mut store_adapter.identity_store, | ||
| to_for_encrypt, | ||
| encryption_jid_for_encrypt, | ||
| &message_for_encrypt, | ||
| request_id, | ||
| ) | ||
| .await | ||
| }); | ||
| (store_adapter, stanza_result) | ||
| }) | ||
| .await | ||
| .map_err(|e| anyhow!("spawn_blocking failed during peer encryption: {e}"))?; | ||
| let stanza = stanza_result?; | ||
| store_adapter.flush().await?; | ||
| stanza | ||
| } else if to.is_group() { |
There was a problem hiding this comment.
Flush cached stores even when prepare_* returns error.
With the batched cache, any state changes made before an error are dropped if we exit without flushing. The group path already flushes on error; peer/DM should do the same (or add a guard) to keep session/identity updates durable.
Example pattern (apply similarly to DM)
- let stanza = stanza_result?;
- store_adapter.flush().await?;
- stanza
+ let stanza = match stanza_result {
+ Ok(s) => s,
+ Err(e) => {
+ if let Err(flush_err) = store_adapter.flush().await {
+ log::warn!("Failed to flush cached Signal stores after error: {flush_err}");
+ }
+ return Err(e.into());
+ }
+ };
+ store_adapter.flush().await?;
+ stanzaAlso applies to: 492-534
🤖 Prompt for AI Agents
In `@src/send.rs` around lines 176 - 203, The peer/DM path currently calls
stanza_result? before flushing the BatchedSignalProtocolStoreAdapter, so any
store updates are dropped on error; move the store_adapter.flush().await so it
always runs regardless of stanza_result (e.g., match on stanza_result, call
store_adapter.flush().await in both Ok and Err branches, then propagate the
error or return the stanza), referencing BatchedSignalProtocolStoreAdapter,
prepare_peer_stanza, store_adapter.flush, and the spawn_blocking block; apply
the same change to the other occurrence mentioned (around the 492-534 region).
| async fn flush(&mut self) -> Result<(), SignalProtocolError> { | ||
| let pending_writes: Vec<_> = { | ||
| let cache = self.cache.lock().await; | ||
| cache | ||
| .dirty_sessions | ||
| .iter() | ||
| .filter_map(|address| { | ||
| cache | ||
| .sessions | ||
| .get(address) | ||
| .and_then(|opt| opt.as_ref()) | ||
| .map(|record| (address.clone(), record.clone())) | ||
| }) | ||
| .collect() | ||
| }; | ||
|
|
||
| for (address, record) in pending_writes { | ||
| SessionStore::store_session(&mut self.inner, &address, &record).await?; | ||
| self.cache.lock().await.dirty_sessions.remove(&address); | ||
| } | ||
|
|
||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
Flush can drop newer updates during concurrent writes.
Each flush snapshots dirty_* then unlocks. If a new store_* happens for the same key during the flush, the subsequent dirty_* removal can clear the flag for the newer value, so it never gets persisted. Please guard dirty-flag removal with a version/token check (or compare the cached value to the flushed value) so newer writes remain dirty.
One way to guard dirty-flag removal (apply similarly to identities/sender keys)
- for (address, record) in pending_writes {
- SessionStore::store_session(&mut self.inner, &address, &record).await?;
- self.cache.lock().await.dirty_sessions.remove(&address);
- }
+ for (address, record) in pending_writes {
+ SessionStore::store_session(&mut self.inner, &address, &record).await?;
+ let mut cache = self.cache.lock().await;
+ let still_same = cache
+ .sessions
+ .get(&address)
+ .and_then(|opt| opt.as_ref())
+ .is_some_and(|current| current == &record);
+ if still_same {
+ cache.dirty_sessions.remove(&address);
+ }
+ }Also applies to: 137-159, 179-205
🤖 Prompt for AI Agents
In `@src/store/signal_adapter.rs` around lines 79 - 101, The flush method
snapshots dirty_sessions then unlocks, so concurrent store_* calls can replace a
session after snapshotting and be lost when flush blindly removes the dirty
flag; to fix, when removing the dirty flag in flush (for
SessionStore::store_session / self.cache.dirty_sessions), check that the cached
session still equals the record you just flushed (or use a stored version/token)
and only remove the dirty marker if they match; apply the same guarded-remove
pattern to identities and sender keys (the other flush blocks referenced) so
newer in-memory updates remain dirty and will be persisted later.
1 participant
Summary by CodeRabbit
New Features
Bug Fixes
Performance
Changes