perf(chat-store): avoid reprepare in contact and identity lookups

.github/workflows/avatar-restore.yml

@@ -56,13 +56,16 @@ jobs:
           text = mod.read_text()
           marker = 'mod native;'
           assert marker in text
-          mod.write_text(text.replace(marker, marker + '\n#[cfg(all(test, not(target_family = "wasm")))]\nmod probe;', 1))
+          if 'mod probe;' not in text:
+              mod.write_text(text.replace(marker, marker + '\n#[cfg(all(test, not(target_family = "wasm")))]\nmod probe;', 1))
           native = root / 'native.rs'
           text = native.read_text()
           marker = 'pub async fn read_persistent(key: &str) -> Option<Vec<u8>> {'
           assert marker in text
-          text = text.replace(marker, marker + '\n    #[cfg(test)]\n    PERSISTENT_READS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);', 1)
-          text += '\n#[cfg(test)]\npub(super) static PERSISTENT_READS: portable_atomic::AtomicU64 = portable_atomic::AtomicU64::new(0);\n'
+          if 'PERSISTENT_READS.fetch_add' not in text:
+              text = text.replace(marker, marker + '\n    #[cfg(test)]\n    PERSISTENT_READS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);', 1)
+          if 'static PERSISTENT_READS:' not in text:
+              text += '\n#[cfg(test)]\npub(super) static PERSISTENT_READS: portable_atomic::AtomicU64 = portable_atomic::AtomicU64::new(0);\n'
           native.write_text(text)
           PY
           cargo test -p oxidezap-gui --bin oxidezap --locked native_restore_probe -- --ignored --nocapture

crates/chat-store/src/lib.rs

@@ -38,6 +38,7 @@ mod error;
 mod fts;
 mod lid;
 mod materialize;
+mod point_lookup;
 mod queries;
 mod schema;
 mod spawn;

crates/chat-store/src/lid.rs

@@ -80,22 +80,26 @@ fn counterpart_of(
     use schema::lid_pn_mapping::dsl;
     let user = jid.user.as_str();
     if jid.is_lid() {
-        return Ok(dsl::lid_pn_mapping
-            .filter(dsl::device_id.eq(device_id).and(dsl::lid.eq(user)))
-            .select(dsl::phone_number)
-            .first::<String>(conn)
-            .optional()?
-            .map(|pn| Jid::new(pn, Server::Pn).to_string()));
-    }
-    Ok(dsl::lid_pn_mapping
-        .filter(dsl::device_id.eq(device_id).and(dsl::phone_number.eq(user)))
-        // The lid tiebreak keeps routing stable when updated_at ties —
-        // flapping between counterpart keys would re-split the thread.
-        .order((dsl::updated_at.desc(), dsl::lid.desc()))
-        .select(dsl::lid)
-        .first::<String>(conn)
+        return Ok(crate::point_lookup::First(
+            dsl::lid_pn_mapping
+                .filter(dsl::device_id.eq(device_id).and(dsl::lid.eq(user)))
+                .select(dsl::phone_number),
+        )
+        .get_result::<String>(conn)
         .optional()?
-        .map(|lid| Jid::new(lid, Server::Lid).to_string()))
+        .map(|pn| Jid::new(pn, Server::Pn).to_string()));
+    }
+    Ok(crate::point_lookup::First(
+        dsl::lid_pn_mapping
+            .filter(dsl::device_id.eq(device_id).and(dsl::phone_number.eq(user)))
+            // The lid tiebreak keeps routing stable when updated_at ties —
+            // flapping between counterpart keys would re-split the thread.
+            .order((dsl::updated_at.desc(), dsl::lid.desc()))
+            .select(dsl::lid),
+    )
+    .get_result::<String>(conn)
+    .optional()?
+    .map(|lid| Jid::new(lid, Server::Lid).to_string()))
 }
 
 /// Every key the peer's rows may live under: the given key, its PN, and all
@@ -710,3 +714,46 @@ pub(crate) fn merge_split_chat(
     cs.message_chats.insert(dest.to_string());
     Ok(dest.to_string())
 }
+
+#[cfg(test)]
+mod point_lookup_tests {
+    use super::*;
+    use diesel::connection::SimpleConnection;
+
+    #[test]
+    fn counterpart_preserves_accounts_absence_and_latest_lid_ties() {
+        let mut conn = SqliteConnection::establish(":memory:").unwrap();
+        conn.batch_execute(
+            "CREATE TABLE lid_pn_mapping(device_id INTEGER NOT NULL,lid TEXT NOT NULL,phone_number TEXT NOT NULL,updated_at BIGINT NOT NULL,PRIMARY KEY(lid,device_id));\
+             INSERT INTO lid_pn_mapping VALUES(1,'100','500',1),(1,'200','500',2),(1,'300','500',2),(2,'100','900',9),(2,'400','500',9)",
+        ).unwrap();
+        // Repeated reads also exercise cached statements with different binds.
+        for _ in 0..3 {
+            for (account, input, expected) in [
+                (1, "100@lid", Some("500@s.whatsapp.net")),
+                (2, "100@lid", Some("900@s.whatsapp.net")),
+                (1, "500@s.whatsapp.net", Some("300@lid")),
+                (2, "500@s.whatsapp.net", Some("400@lid")),
+                (1, "400@lid", None),
+                (1, "999@s.whatsapp.net", None),
+                (3, "100@lid", None),
+            ] {
+                let jid = input.parse::<Jid>().unwrap();
+                assert_eq!(
+                    counterpart_of(&mut conn, account, &jid).unwrap().as_deref(),
+                    expected
+                );
+            }
+        }
+        assert_eq!(
+            counterpart_chat_key(&mut conn, 1, "100:5@lid")
+                .unwrap()
+                .as_deref(),
+            Some("500@s.whatsapp.net")
+        );
+        assert_eq!(
+            counterpart_chat_key(&mut conn, 1, "123@g.us").unwrap(),
+            None
+        );
+    }
+}

crates/chat-store/src/point_lookup.rs

@@ -0,0 +1,60 @@
+//! Fixed first-row limit for the measured contact and identity lookups.
+//!
+//! SQLite 3.53 specializes a bound LIMIT on every step after rebinding.
+//! These queries always need one row. A literal keeps limit-aware planning
+//! without a runtime limit to specialize. See docs/sqlite-point-lookup-performance.md.
+use diesel::prelude::*;
+use diesel::query_builder::{AstPass, Query, QueryFragment, QueryId};
+use diesel::sqlite::Sqlite;
+
+#[derive(QueryId)]
+pub(crate) struct First<Q>(pub(crate) Q);
+
+impl<Q: Query> Query for First<Q> {
+    type SqlType = Q::SqlType;
+}
+
+impl<Q: QueryFragment<Sqlite>> QueryFragment<Sqlite> for First<Q> {
+    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Sqlite>) -> QueryResult<()> {
+        self.0.walk_ast(out.reborrow())?;
+        out.push_sql(" LIMIT 1");
+        Ok(())
+    }
+}
+
+impl<Q> RunQueryDsl<SqliteConnection> for First<Q> {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::schema::contacts::dsl;
+    use diesel::connection::SimpleConnection;
+
+    #[test]
+    fn fixed_limit_keeps_filters_binds_and_optional_first_semantics() {
+        let mut conn = SqliteConnection::establish(":memory:").unwrap();
+        conn.batch_execute("CREATE TABLE contacts(device_id INTEGER NOT NULL,jid TEXT NOT NULL,push_name TEXT,full_name TEXT,first_name TEXT,business_name TEXT,PRIMARY KEY(device_id,jid)); INSERT INTO contacts(device_id,jid) VALUES(1,'one'),(1,'two'),(2,'one')").unwrap();
+        for (account, key) in [
+            (1, "one"),
+            (1, "two"),
+            (2, "one"),
+            (1, "absent"),
+            (3, "one"),
+        ] {
+            let query = || {
+                dsl::contacts
+                    .filter(dsl::device_id.eq(account).and(dsl::jid.eq(key)))
+                    .select(dsl::jid)
+            };
+            let expected = query().first::<String>(&mut conn).optional().unwrap();
+            let fixed = First(query());
+            let sql = diesel::debug_query::<Sqlite, _>(&fixed).to_string();
+            assert!(sql.contains(" LIMIT 1 -- binds:"), "{sql}");
+            assert!(!sql.contains("LIMIT ?"));
+            assert_eq!(
+                fixed.get_result::<String>(&mut conn).optional().unwrap(),
+                expected
+            );
+        }
+    }
+}

crates/chat-store/src/queries.rs

@@ -2166,18 +2166,20 @@ impl ChatStore {
         let row: Option<(ContactRow, Option<ContactLabels>)> = self
             .db()
             .read(move |conn| {
-                let row: Option<ContactRow> = dsl::contacts
-                    .filter(dsl::device_id.eq(device_id).and(dsl::jid.eq(&jid_str)))
-                    .select((
-                        dsl::jid,
-                        dsl::push_name,
-                        dsl::full_name,
-                        dsl::first_name,
-                        dsl::business_name,
-                    ))
-                    .first(conn)
-                    .optional()
-                    .map_err(db_err)?;
+                let row: Option<ContactRow> = crate::point_lookup::First(
+                    dsl::contacts
+                        .filter(dsl::device_id.eq(device_id).and(dsl::jid.eq(&jid_str)))
+                        .select((
+                            dsl::jid,
+                            dsl::push_name,
+                            dsl::full_name,
+                            dsl::first_name,
+                            dsl::business_name,
+                        )),
+                )
+                .get_result(conn)
+                .optional()
+                .map_err(db_err)?;
                 let labels = match &row {
                     Some((jid, _, _, _, _)) => {
                         labels_for(conn, device_id, std::slice::from_ref(jid))?.remove(jid)

crates/chat-store/tests/point_lookup.rs

@@ -0,0 +1,49 @@
+mod common;
+use common::*;
+
+#[tokio::test]
+async fn contact_lookup_keeps_account_labels_optional_fields_and_missing_rows() {
+    let (db, store) = test_store().await;
+    db.database().provision_device(2).await.unwrap();
+    db.shared().run(|conn| {
+        diesel::sql_query("INSERT INTO contacts(device_id,jid,push_name,full_name,first_name,business_name) VALUES(1,'559900000001@s.whatsapp.net','push','full','first','business'),(2,'559900000001@s.whatsapp.net','other','other','other','other'),(1,'559900000002@s.whatsapp.net',NULL,NULL,NULL,NULL)")
+            .execute(conn).map_err(db_err)?;
+        diesel::sql_query("INSERT INTO contact_labels(device_id,jid,alias,tags) VALUES(1,'559900000001@s.whatsapp.net','alias','[\"tag\"]'),(2,'559900000001@s.whatsapp.net','other','[]')")
+            .execute(conn).map_err(db_err)?;
+        Ok(())
+    }).await.unwrap();
+    for _ in 0..3 {
+        let contact = store
+            .contact(&jid("559900000001:7@s.whatsapp.net"))
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(contact.jid, jid(PEER));
+        assert_eq!(contact.push_name.as_deref(), Some("push"));
+        assert_eq!(contact.full_name.as_deref(), Some("full"));
+        assert_eq!(contact.first_name.as_deref(), Some("first"));
+        assert_eq!(contact.business_name.as_deref(), Some("business"));
+        assert_eq!(contact.alias.as_deref(), Some("alias"));
+        assert_eq!(contact.tags, vec!["tag"]);
+        let blank = store
+            .contact(&jid("559900000002@s.whatsapp.net"))
+            .await
+            .unwrap()
+            .unwrap();
+        assert!(
+            blank.full_name.is_none()
+                && blank.push_name.is_none()
+                && blank.first_name.is_none()
+                && blank.business_name.is_none()
+                && blank.alias.is_none()
+        );
+        assert!(blank.tags.is_empty());
+        assert!(
+            store
+                .contact(&jid("559900000003@s.whatsapp.net"))
+                .await
+                .unwrap()
+                .is_none()
+        );
+    }
+}