fix(chat-store): correlate amendment operation acknowledgements

crates/chat-store/src/store/event.rs

@@ -9,7 +9,7 @@ use waproto::whatsapp as wa;
 
 use crate::materialize::{KIND_UNDECRYPTABLE, unavailable_kind};
 use crate::schema;
-use crate::store::ack::{AckApplied, DeferredAcks, apply_server_ack};
+use crate::store::ack::{AckApplied, DeferredAcks, apply_server_ack_with_operations};
 use crate::store::chat_rows::{
     ChatBump, bump_chat, chat_row, delete_chat_rows, ensure_chat, recompute_chat_preview,
     remaining_messages,
@@ -45,7 +45,9 @@ pub(super) fn apply_event(
         }
         Event::Receipt(receipt) => apply_receipt(conn, device_id, receipt, cs),
         Event::ServerAck(ack) => {
-            if let AckApplied::Deferrable(chat) = apply_server_ack(conn, device_id, ack, cs)? {
+            if let AckApplied::Deferrable(chat) =
+                apply_server_ack_with_operations(conn, device_id, ack, cs, deferred)?
+            {
                 deferred.defer(ack, chat, wacore::time::now_utc().timestamp_millis());
             }
             Ok(())

crates/chat-store/src/store/mod.rs

@@ -121,6 +121,10 @@ impl WriterSender {
 
 pub(crate) enum WriterMsg {
     Event(Arc<Event>),
+    Operation {
+        chat: Jid,
+        operation_id: String,
+    },
     InboundDurability {
         event: Arc<Event>,
         done: oneshot::Sender<std::result::Result<(), String>>,
@@ -219,24 +223,27 @@ pub(crate) struct BarrierOutcome {
 struct BarrierError {
     first: Arc<str>,
     write_loss: bool,
+    ack_loss: bool,
 }
 
 impl BarrierOutcome {
     fn publish(
         error: Option<String>,
         write_loss: bool,
+        ack_loss: bool,
         unreceived: &Arc<Mutex<Option<Arc<BarrierError>>>>,
     ) -> Self {
         let mut current = unreceived.lock().unwrap_or_else(|e| e.into_inner());
         if let Some(error) = error {
-            // Keep the first unseen error and permanent write-admission loss.
+            // Keep the first unseen error and both kinds of permanent loss.
             // This bounded summary survives later transient failures. A fresh
             // generation prevents an older response from clearing newer loss.
             *current = Some(Arc::new(BarrierError {
                 first: current
                     .as_ref()
                     .map_or_else(|| Arc::from(error), |old| Arc::clone(&old.first)),
                 write_loss: write_loss || current.as_ref().is_some_and(|old| old.write_loss),
+                ack_loss: ack_loss || current.as_ref().is_some_and(|old| old.ack_loss),
             }));
         }
         Self {
@@ -262,6 +269,9 @@ impl BarrierOutcome {
         if error.write_loss && !message.contains("writer admission overflow") {
             message.push_str("; writer admission overflow: one or more writes were not accepted");
         }
+        if error.ack_loss && !message.contains("unmatched ACK capacity exceeded") {
+            message.push_str("; unmatched ACK capacity exceeded: one or more incoming acknowledgements could not be retained");
+        }
         Err(message)
     }
 }
@@ -636,7 +646,12 @@ impl ChatStore {
         self.changes.subscribe()
     }
 
-    /// Record a message this client just sent. Goes through the writer queue so
+    /// Record an outgoing message before sending with the same pre-generated id.
+    /// Do not send if admission fails; call `mark_send_failed` if the send fails.
+    /// Recording after the network send is supported only within the bounded
+    /// early-ACK window (64 unknown ACKs, 60 seconds). At capacity, admitted
+    /// waiters remain and excess ACKs are rejected; flush/close reports the loss.
+    /// Goes through the writer queue so
     /// it cannot race the server ack / receipts that follow it in event order.
     /// Status starts at [`MessageStatus::Pending`](crate::types::MessageStatus::Pending)
     /// and is lifted by acks/receipts. `timestamp` is the optimistic display
@@ -668,6 +683,32 @@ impl ChatStore {
             .map_err(|error| ChatStoreError::Store(StoreError::Validation(error.to_string())))
     }
 
+    /// Register the stanza id returned by a successful amendment send.
+    /// The id belongs to the operation, not the message it changes. No message
+    /// row is created. Call this before the target recorder, using the fresh
+    /// `SendResult.message_id`, never the target id. Existing recorders remain
+    /// available for callers without an operation id.
+    ///
+    /// Correlation is bounded to 256 operations for 60 seconds in this writer;
+    /// reconnect preserves it, closing or restarting the store does not.
+    /// Queue admission is not a commit acknowledgement or server acceptance.
+    /// Unknown acks still share the 64-entry early-message queue until their
+    /// operation is recorded. Chatless acks cannot be safely correlated and
+    /// keep the ordinary early-row deferral. Evicting an operation at capacity
+    /// warns; its late acks become unknown again. More than 64 unknown acks
+    /// exceed the admission bound: existing waiters stay, excess ACKs are lost,
+    /// and flush/close reports the loss. Reusing operation
+    /// ids across chats or for message rows is unsupported: known collisions
+    /// are refused, and future rows cannot retroactively disambiguate an ack.
+    pub fn record_operation(&self, chat: &Jid, operation_id: &str) -> Result<()> {
+        self.tx
+            .send(WriterMsg::Operation {
+                chat: chat.clone(),
+                operation_id: operation_id.to_owned(),
+            })
+            .map_err(|error| ChatStoreError::Store(StoreError::Validation(error.to_string())))
+    }
+
     /// Record an edit this client just sent for one of its own messages.
     ///
     /// This is the local counterpart of an inbound `MESSAGE_EDIT`: it updates
@@ -1025,7 +1066,7 @@ impl ChatStore {
     /// Commit accepted writes preceding this barrier, retrying retained writes
     /// in their original order. Rollback retains the failed batch and everything
     /// behind it. A later flush can recover after the underlying fault is fixed.
-    /// Reports queue admission failures and post-commit durability failures too.
+    /// Reports write/unknown-ACK admission failures and post-commit durability failures too.
     /// The error is temporal across callers, not an individual write receipt.
     /// Cancellation retains it until a public flush/close future returns it.
     /// Handling or discarding the returned result belongs to the caller.
@@ -1421,20 +1462,25 @@ mod barrier_tests {
     #[test]
     fn a_new_failure_preserves_an_unreceived_admission_error() {
         let unreceived = Arc::new(Mutex::new(None));
-        let cancelled =
-            BarrierOutcome::publish(Some("writer admission overflow".into()), false, &unreceived);
+        let cancelled = BarrierOutcome::publish(
+            Some("writer admission overflow".into()),
+            false,
+            false,
+            &unreceived,
+        );
         drop(cancelled);
         let later = BarrierOutcome::publish(
             Some("backend durability barrier".into()),
             false,
+            false,
             &unreceived,
         );
         assert_eq!(
             later.into_result().unwrap_err(),
             "writer admission overflow"
         );
         assert!(
-            BarrierOutcome::publish(None, false, &unreceived)
+            BarrierOutcome::publish(None, false, false, &unreceived)
                 .into_result()
                 .is_ok()
         );
@@ -1443,11 +1489,16 @@ mod barrier_tests {
     #[test]
     fn an_older_response_cannot_clear_a_new_failure_with_different_text() {
         let unreceived = Arc::new(Mutex::new(None));
-        let older =
-            BarrierOutcome::publish(Some("writer admission overflow".into()), false, &unreceived);
+        let older = BarrierOutcome::publish(
+            Some("writer admission overflow".into()),
+            false,
+            false,
+            &unreceived,
+        );
         let newer = BarrierOutcome::publish(
             Some("backend durability barrier".into()),
             false,
+            false,
             &unreceived,
         );
         assert_eq!(
@@ -1459,7 +1510,7 @@ mod barrier_tests {
             "writer admission overflow"
         );
         assert!(
-            BarrierOutcome::publish(None, false, &unreceived)
+            BarrierOutcome::publish(None, false, false, &unreceived)
                 .into_result()
                 .is_ok()
         );
@@ -1468,17 +1519,26 @@ mod barrier_tests {
     #[test]
     fn later_admission_loss_survives_cancellation_and_transient_failures() {
         let unreceived = Arc::new(Mutex::new(None));
-        let older = BarrierOutcome::publish(Some("backend failure".into()), false, &unreceived);
-        let cancelled = BarrierOutcome::publish(Some("lost writes".into()), true, &unreceived);
+        let older =
+            BarrierOutcome::publish(Some("backend failure".into()), false, false, &unreceived);
+        let cancelled =
+            BarrierOutcome::publish(Some("lost writes".into()), true, false, &unreceived);
         drop(cancelled);
-        let latest =
-            BarrierOutcome::publish(Some("another backend failure".into()), false, &unreceived);
+        let cancelled = BarrierOutcome::publish(Some("lost ACKs".into()), false, true, &unreceived);
+        drop(cancelled);
+        let latest = BarrierOutcome::publish(
+            Some("another backend failure".into()),
+            false,
+            false,
+            &unreceived,
+        );
         assert_eq!(older.into_result().unwrap_err(), "backend failure");
         let message = latest.into_result().unwrap_err();
         assert!(message.contains("backend failure"));
         assert!(message.contains("writer admission overflow"));
+        assert!(message.contains("unmatched ACK capacity exceeded"));
         assert!(
-            BarrierOutcome::publish(None, false, &unreceived)
+            BarrierOutcome::publish(None, false, false, &unreceived)
                 .into_result()
                 .is_ok()
         );
@@ -1487,17 +1547,17 @@ mod barrier_tests {
     #[test]
     fn receiving_an_older_error_keeps_a_newer_identical_error() {
         let unreceived = Arc::new(Mutex::new(None));
-        let older = BarrierOutcome::publish(Some("failed".into()), false, &unreceived);
-        let newer = BarrierOutcome::publish(Some("failed".into()), false, &unreceived);
+        let older = BarrierOutcome::publish(Some("failed".into()), false, false, &unreceived);
+        let newer = BarrierOutcome::publish(Some("failed".into()), false, false, &unreceived);
         assert!(older.into_result().is_err());
         assert!(
-            BarrierOutcome::publish(None, false, &unreceived)
+            BarrierOutcome::publish(None, false, false, &unreceived)
                 .into_result()
                 .is_err()
         );
         assert!(newer.into_result().is_err());
         assert!(
-            BarrierOutcome::publish(None, false, &unreceived)
+            BarrierOutcome::publish(None, false, false, &unreceived)
                 .into_result()
                 .is_ok()
         );

crates/chat-store/src/store/writer.rs

@@ -19,7 +19,9 @@ use whatsapp_rust_sqlite_storage::{CommitBarrierError, SharedSqlite};
 
 use crate::error::db_err;
 use crate::schema;
-use crate::store::ack::{AckApplied, DeferredAcks, apply_server_ack, lock_deferred_acks};
+use crate::store::ack::{
+    AckApplied, DeferredAcks, apply_server_ack_with_operations, lock_deferred_acks,
+};
 use crate::store::chat_rows::{ChatBump, bump_chat};
 use crate::store::edit::apply_edit;
 use crate::store::event::apply_event;
@@ -82,8 +84,20 @@ pub(super) async fn writer_loop(
                         "writer admission overflow: one or more writes were not accepted".into()
                     });
                 }
-                let outcome =
-                    BarrierOutcome::publish(pending_error.take(), write_loss, &unreceived);
+                let ack_loss = lock_deferred_acks(&deferred).take_admission_failure();
+                if ack_loss {
+                    let loss = "unmatched ACK capacity exceeded: one or more incoming acknowledgements could not be retained";
+                    pending_error = Some(match pending_error.take() {
+                        Some(error) => format!("{error}; {loss}"),
+                        None => loss.to_owned(),
+                    });
+                }
+                let outcome = BarrierOutcome::publish(
+                    pending_error.take(),
+                    write_loss,
+                    ack_loss,
+                    &unreceived,
+                );
                 if stopping {
                     rx.close();
                     drop(pending);
@@ -352,6 +366,13 @@ fn apply_writer_msg(
 ) -> QueryResult<()> {
     match msg {
         WriterMsg::Event(event) => apply_event(conn, device_id, event, cs, deferred),
+        WriterMsg::Operation { chat, operation_id } => deferred.record_operation(
+            conn,
+            device_id,
+            &chat.to_string(),
+            operation_id,
+            wacore::time::now_utc().timestamp_millis(),
+        ),
         WriterMsg::InboundDurability { event, .. } => {
             apply_event(conn, device_id, event, cs, deferred)
         }
@@ -429,18 +450,21 @@ fn apply_writer_msg(
                 // The row this send's ack was waiting for now exists. Applying
                 // it here also corrects the optimistic timestamp we just wrote
                 // to the server's, before anything renders the row.
-                if let Some(ack) = deferred.take_matching(
+                while let Some(ack) = deferred.take_matching(
                     msg_id,
                     &chat_str,
                     wacore::time::now_utc().timestamp_millis(),
-                ) && let AckApplied::Deferrable(_) = apply_server_ack(conn, device_id, &ack, cs)?
-                {
-                    // The row exists, so this should not happen; say so rather
-                    // than let the ack vanish the way it used to.
-                    warn!(
-                        target: "ChatStore/Ack",
-                        "Held ack for {msg_id} matched no row even after its insert"
-                    );
+                ) {
+                    if let AckApplied::Deferrable(_) =
+                        apply_server_ack_with_operations(conn, device_id, &ack, cs, deferred)?
+                    {
+                        // The row exists, so this should not happen; say so rather
+                        // than let the ack vanish the way it used to.
+                        warn!(
+                            target: "ChatStore/Ack",
+                            "Held ack for {msg_id} matched no row even after its insert"
+                        );
+                    }
                 }
             }
             cs.message_chats.insert(chat_str);