fix(tc-token): atomic newer-wins store; drop tc_token_lock and close the cross-source race

src/client.rs

@@ -654,11 +654,6 @@ pub struct Client {
     /// Single-flights signed pre-key rotation so overlapping post-login tasks
     /// (from reconnect churn) can't run the rotate/upload/prune flow concurrently.
     pub(crate) signed_pre_key_rotation_lock: Arc<async_lock::Mutex<()>>,
-    /// Serializes the history-sync tc-token get-then-store: chunks ingest
-    /// concurrently, but the newer-wins guard is a non-atomic read-then-write,
-    /// so without this two same-contact candidates could interleave and let an
-    /// older privacy token overwrite a fresher one.
-    pub(crate) tc_token_lock: Arc<async_lock::Mutex<()>>,
     /// Notifier for when offline sync (ib offline stanza) is received.
     /// WhatsApp Web waits for this before sending passive tasks (prekey upload, active IQ, presence).
     pub(crate) offline_sync_notifier: Arc<event_listener::Event>,

src/client/lifecycle.rs

@@ -218,7 +218,6 @@ impl Client {
             initial_app_state_keys_received: Arc::new(AtomicBool::new(false)),
             prekey_upload_lock: Arc::new(async_lock::Mutex::new(())),
             signed_pre_key_rotation_lock: Arc::new(async_lock::Mutex::new(())),
-            tc_token_lock: Arc::new(async_lock::Mutex::new(())),
             offline_sync_notifier: Arc::new(event_listener::Event::new()),
             offline_sync_completed: Arc::new(AtomicBool::new(false)),
             offline_sync_finish_started: Arc::new(AtomicBool::new(false)),

src/history_sync.rs

@@ -407,24 +407,9 @@ impl Client {
 
         let backend = self.persistence_manager.backend();
 
-        // Serialize the get-then-store below: history-sync chunks ingest
-        // concurrently, so without this two same-contact candidates could
-        // read the same baseline and the older one's unconditional write could
-        // land last, clobbering a fresher privacy token.
-        let _guard = self.tc_token_lock.lock().await;
-
-        // Skip only when an existing *real* token is newer than this candidate; a
-        // byte-less placeholder stamps token_timestamp with a sender epoch and
-        // must never block the first real token from history sync.
-        if let Ok(Some(existing)) = backend.get_tc_token(token_key).await
-            && !existing.token.is_empty()
-            && (existing.token_timestamp as u64) > candidate.tc_token_timestamp
-        {
-            return;
-        }
-
-        // Two atomic upserts — token fields, then the sender bucket (advance-only)
-        // — so a concurrent post-send issuance is never clobbered by this write.
+        // Newer-wins lives in the store now (atomic, lock-free), so concurrent
+        // history-sync chunks and the privacy path converge without a
+        // get-then-store or a lock.
         if let Err(e) = backend
             .store_received_tc_token(
                 token_key,

storages/sqlite-storage/src/sqlite_store.rs

@@ -3004,40 +3004,55 @@ impl ProtocolStore for SqliteStore {
         token: &[u8],
         token_timestamp: i64,
     ) -> Result<()> {
-        let pool = self.pool.clone();
         let device_id = self.device_id;
         let jid = jid.to_string();
         let token = token.to_vec();
         let now = wacore::time::now_secs();
-        tokio::task::spawn_blocking(move || -> Result<()> {
-            let mut conn = pool
-                .get()
-                .map_err(|e| StoreError::Connection(Box::new(e)))?;
-            // On conflict update only the token fields so a sender_timestamp
-            // written concurrently by the issuance path survives.
-            diesel::insert_into(tc_tokens::table)
-                .values((
-                    tc_tokens::jid.eq(&jid),
-                    tc_tokens::token.eq(&token),
-                    tc_tokens::token_timestamp.eq(token_timestamp),
-                    tc_tokens::sender_timestamp.eq(None::<i64>),
-                    tc_tokens::device_id.eq(device_id),
-                    tc_tokens::updated_at.eq(now),
-                ))
-                .on_conflict((tc_tokens::jid, tc_tokens::device_id))
-                .do_update()
-                .set((
-                    tc_tokens::token.eq(&token),
-                    tc_tokens::token_timestamp.eq(token_timestamp),
-                    tc_tokens::updated_at.eq(now),
-                ))
-                .execute(&mut conn)
-                .map_err(|e| StoreError::Database(Box::new(e)))?;
-            Ok(())
+        // IMMEDIATE so the read + conditional write is atomic against concurrent
+        // writers (WAL + busy_timeout serialize them): this is the lock-free
+        // newer-wins that lets history-sync and the privacy path converge without
+        // clobbering a fresher token. with_retry rides out transient SQLITE_BUSY.
+        self.with_retry("store_received_tc_token", || {
+            let jid = jid.clone();
+            let token = token.clone();
+            Box::new(move |conn: &mut SqliteConnection| {
+                conn.immediate_transaction(|conn| -> QueryResult<()> {
+                    let existing: Option<(Vec<u8>, i64)> = tc_tokens::table
+                        .filter(tc_tokens::jid.eq(&jid))
+                        .filter(tc_tokens::device_id.eq(device_id))
+                        .select((tc_tokens::token, tc_tokens::token_timestamp))
+                        .first(conn)
+                        .optional()?;
+                    let write = match &existing {
+                        Some((existing_token, existing_ts)) => {
+                            existing_token.is_empty() || token_timestamp >= *existing_ts
+                        }
+                        None => true,
+                    };
+                    if write {
+                        diesel::insert_into(tc_tokens::table)
+                            .values((
+                                tc_tokens::jid.eq(&jid),
+                                tc_tokens::token.eq(&token),
+                                tc_tokens::token_timestamp.eq(token_timestamp),
+                                tc_tokens::sender_timestamp.eq(None::<i64>),
+                                tc_tokens::device_id.eq(device_id),
+                                tc_tokens::updated_at.eq(now),
+                            ))
+                            .on_conflict((tc_tokens::jid, tc_tokens::device_id))
+                            .do_update()
+                            .set((
+                                tc_tokens::token.eq(&token),
+                                tc_tokens::token_timestamp.eq(token_timestamp),
+                                tc_tokens::updated_at.eq(now),
+                            ))
+                            .execute(conn)?;
+                    }
+                    Ok(())
+                })
+            })
         })
         .await
-        .map_err(|e| StoreError::Database(Box::new(e)))??;
-        Ok(())
     }
 
     async fn touch_tc_token_sender_timestamp(
@@ -4272,6 +4287,51 @@ mod tests {
         assert_eq!(c.sender_timestamp, Some(6000), "touch is advance-only");
     }
 
+    #[tokio::test]
+    async fn store_received_tc_token_is_newer_wins() {
+        let store = create_test_store().await;
+
+        // First real token at t=5000.
+        store
+            .store_received_tc_token("c@lid", &[1, 1, 1], 5000)
+            .await
+            .unwrap();
+
+        // A stale write (older timestamp) must not clobber the fresher token —
+        // this is the atomic newer-wins that replaces the tc_token_lock.
+        store
+            .store_received_tc_token("c@lid", &[2, 2, 2], 3000)
+            .await
+            .unwrap();
+        let e = store.get_tc_token("c@lid").await.unwrap().unwrap();
+        assert_eq!(e.token, vec![1, 1, 1], "older write must not overwrite");
+        assert_eq!(e.token_timestamp, 5000);
+
+        // A newer write wins.
+        store
+            .store_received_tc_token("c@lid", &[3, 3, 3], 7000)
+            .await
+            .unwrap();
+        let e = store.get_tc_token("c@lid").await.unwrap().unwrap();
+        assert_eq!(e.token, vec![3, 3, 3]);
+        assert_eq!(e.token_timestamp, 7000);
+
+        // A byte-less placeholder never blocks the first real token, even when
+        // that token's timestamp is older than the placeholder's sender epoch.
+        store
+            .touch_tc_token_sender_timestamp("p@lid", 9000)
+            .await
+            .unwrap();
+        store
+            .store_received_tc_token("p@lid", &[4, 4, 4], 6000)
+            .await
+            .unwrap();
+        let e = store.get_tc_token("p@lid").await.unwrap().unwrap();
+        assert_eq!(e.token, vec![4, 4, 4], "placeholder must accept real token");
+        assert_eq!(e.token_timestamp, 6000);
+        assert_eq!(e.sender_timestamp, Some(9000), "sender bucket preserved");
+    }
+
     #[tokio::test]
     async fn test_delete_expired_two_window_pruning() {
         let store = create_test_store().await;

wacore/src/store/in_memory.rs

@@ -614,9 +614,13 @@ impl ProtocolStore for InMemoryBackend {
         let mut s = self.state.lock().await;
         match s.tc_tokens.get_mut(jid) {
             Some(entry) => {
-                entry.token = token.to_vec();
-                entry.token_timestamp = token_timestamp;
-                // sender_timestamp left untouched
+                // Newer-wins (see the trait doc): don't let a stale write
+                // clobber a fresher token.
+                if entry.token.is_empty() || token_timestamp >= entry.token_timestamp {
+                    entry.token = token.to_vec();
+                    entry.token_timestamp = token_timestamp;
+                    // sender_timestamp left untouched
+                }
             }
             None => {
                 s.tc_tokens.insert(
@@ -1247,6 +1251,55 @@ mod tests {
         assert_eq!(fresh.sender_timestamp, None);
     }
 
+    #[tokio::test]
+    async fn store_received_tc_token_is_newer_wins() {
+        let backend = InMemoryBackend::new();
+
+        // First real token at t=5000.
+        backend
+            .store_received_tc_token("c", &[1, 1, 1], 5000)
+            .await
+            .unwrap();
+
+        // A stale write (older timestamp) must not clobber the fresher token —
+        // this is what lets concurrent history-sync chunks converge lock-free.
+        backend
+            .store_received_tc_token("c", &[2, 2, 2], 3000)
+            .await
+            .unwrap();
+        let e = backend.get_tc_token("c").await.unwrap().unwrap();
+        assert_eq!(e.token, vec![1, 1, 1], "older write must not overwrite");
+        assert_eq!(e.token_timestamp, 5000);
+
+        // A newer write wins.
+        backend
+            .store_received_tc_token("c", &[3, 3, 3], 7000)
+            .await
+            .unwrap();
+        let e = backend.get_tc_token("c").await.unwrap().unwrap();
+        assert_eq!(e.token, vec![3, 3, 3]);
+        assert_eq!(e.token_timestamp, 7000);
+
+        // A byte-less placeholder (sender epoch t=9000) never blocks a real token,
+        // even when the real token's timestamp is older than the placeholder's.
+        backend
+            .touch_tc_token_sender_timestamp("p", 9000)
+            .await
+            .unwrap();
+        backend
+            .store_received_tc_token("p", &[4, 4, 4], 6000)
+            .await
+            .unwrap();
+        let e = backend.get_tc_token("p").await.unwrap().unwrap();
+        assert_eq!(
+            e.token,
+            vec![4, 4, 4],
+            "placeholder must accept first real token"
+        );
+        assert_eq!(e.token_timestamp, 6000);
+        assert_eq!(e.sender_timestamp, Some(9000), "sender bucket preserved");
+    }
+
     #[tokio::test]
     async fn prune_respects_sender_and_token_windows() {
         let backend = InMemoryBackend::new();

wacore/src/store/traits.rs

@@ -496,18 +496,30 @@ pub trait ProtocolStore: Send + Sync {
     /// `sender_timestamp`. The symmetric counterpart of
     /// [`touch_tc_token_sender_timestamp`](Self::touch_tc_token_sender_timestamp):
     /// each writer owns its own field, so the notification path never drops a
-    /// sender bucket that the issuance path wrote concurrently. Same atomicity
-    /// requirement — the default read-modify-write is for third-party backends.
+    /// sender bucket that the issuance path wrote concurrently.
+    ///
+    /// **Newer-wins**: the token pair is overwritten only when the stored token
+    /// is a byte-less placeholder or the incoming `token_timestamp` is at least
+    /// as new — a stale write must not clobber a fresher real token. Doing this
+    /// in the store (atomically for the built-in backends) is what lets the
+    /// concurrent history-sync and privacy-notification writers converge without
+    /// a lock. Same atomicity requirement as the sender bucket — the default
+    /// read-modify-write here is a best-effort for third-party backends.
     async fn store_received_tc_token(
         &self,
         jid: &str,
         token: &[u8],
         token_timestamp: i64,
     ) -> Result<()> {
-        let sender_timestamp = self
-            .get_tc_token(jid)
-            .await?
-            .and_then(|existing| existing.sender_timestamp);
+        let existing = self.get_tc_token(jid).await?;
+        // Keep a fresher real token; a placeholder never blocks the first real one.
+        if let Some(existing) = &existing
+            && !existing.token.is_empty()
+            && token_timestamp < existing.token_timestamp
+        {
+            return Ok(());
+        }
+        let sender_timestamp = existing.and_then(|existing| existing.sender_timestamp);
         self.put_tc_token(
             jid,
             &TcTokenEntry {