Improved DHT publication. Adjust batching, added validation per key

crates/music-dht/src/database.rs

@@ -110,6 +110,17 @@ pub trait MusicDhtStorage: std::fmt::Debug + Send + Sync {
     /// or refreshed.
     async fn store_dht_record(&self, key: DhtKey, record: StoredRecord) -> Result<bool>;
 
+    /// Applies a batch of validated incoming records, one flag per entry in
+    /// order. Backends should override this with a single write transaction:
+    /// per-record commits dominate the cost of replicating a large library.
+    async fn store_dht_records(&self, entries: Vec<(DhtKey, StoredRecord)>) -> Result<Vec<bool>> {
+        let mut stored = Vec::with_capacity(entries.len());
+        for (key, record) in entries {
+            stored.push(self.store_dht_record(key, record).await?);
+        }
+        Ok(stored)
+    }
+
     /// Returns non-expired replicas stored under `key`, including tombstones.
     async fn dht_records_by_key(&self, key: DhtKey, now_ms: u64) -> Result<Vec<StoredRecord>>;
 
@@ -272,44 +283,24 @@ impl Database {
     /// the revision/tombstone rules. Returns `true` if the record was written
     /// or refreshed.
     pub async fn store_dht_record(&self, key: DhtKey, record: StoredRecord) -> Result<bool> {
+        self.call(move |conn| store_record_in_conn(conn, &key, &record))
+            .await
+    }
+
+    /// Applies a whole batch of records inside one transaction: one commit
+    /// instead of one per record.
+    pub async fn store_dht_records(
+        &self,
+        entries: Vec<(DhtKey, StoredRecord)>,
+    ) -> Result<Vec<bool>> {
         self.call(move |conn| {
-            let existing: Option<(i64, i64, i64)> = conn
+            let tx = conn.unchecked_transaction()?;
-                .query_row(
+            let mut stored = Vec::with_capacity(entries.len());
-                    "SELECT revision, deleted, expires_at_ms FROM dht_records
+            for (key, record) in &entries {
-                     WHERE dht_key = ?1 AND artist_id = ?2 AND owner_peer_id = ?3",
+                stored.push(store_record_in_conn(&tx, key, record)?);
-                    params![
-                        key.as_bytes().as_slice(),
-                        record.item.id.as_bytes().as_slice(),
-                        record.item.owner.to_string(),
-                    ],
-                    |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
-                )
-                .optional()?;
-            let existing =
-                existing.map(|(rev, del, exp)| (rev as u64, del != 0, exp as u64));
-            match decide_store(existing, &record) {
-                StoreDecision::Ignore => Ok(false),
-                StoreDecision::Write | StoreDecision::RefreshExpiry(_) => {
-                    let payload = postcard::to_stdvec(&record).map_err(|err| {
-                        rusqlite::Error::ToSqlConversionFailure(Box::new(err))
-                    })?;
-                    conn.execute(
-                        "INSERT OR REPLACE INTO dht_records
-                         (dht_key, artist_id, owner_peer_id, payload, revision, deleted, expires_at_ms)
-                         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
-                        params![
-                            key.as_bytes().as_slice(),
-                            record.item.id.as_bytes().as_slice(),
-                            record.item.owner.to_string(),
-                            payload,
-                            record.item.revision as i64,
-                            record.item.deleted as i64,
-                            record.expires_at_ms as i64,
-                        ],
-                    )?;
-                    Ok(true)
-                }
             }
+            tx.commit()?;
+            Ok(stored)
         })
         .await
     }
@@ -460,6 +451,10 @@ impl MusicDhtStorage for Database {
         Database::store_dht_record(self, key, record).await
     }
 
+    async fn store_dht_records(&self, entries: Vec<(DhtKey, StoredRecord)>) -> Result<Vec<bool>> {
+        Database::store_dht_records(self, entries).await
+    }
+
     async fn dht_records_by_key(&self, key: DhtKey, now_ms: u64) -> Result<Vec<StoredRecord>> {
         Database::dht_records_by_key(self, key, now_ms).await
     }
@@ -485,6 +480,51 @@ impl MusicDhtStorage for Database {
     }
 }
 
+/// Applies one validated record to the replica store, following the
+/// revision/tombstone rules. Returns `true` if the record was written or
+/// refreshed. Runs against a plain connection or an open transaction.
+fn store_record_in_conn(
+    conn: &Connection,
+    key: &DhtKey,
+    record: &StoredRecord,
+) -> rusqlite::Result<bool> {
+    let existing: Option<(i64, i64, i64)> = conn
+        .query_row(
+            "SELECT revision, deleted, expires_at_ms FROM dht_records
+             WHERE dht_key = ?1 AND artist_id = ?2 AND owner_peer_id = ?3",
+            params![
+                key.as_bytes().as_slice(),
+                record.item.id.as_bytes().as_slice(),
+                record.item.owner.to_string(),
+            ],
+            |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
+        )
+        .optional()?;
+    let existing = existing.map(|(rev, del, exp)| (rev as u64, del != 0, exp as u64));
+    match decide_store(existing, record) {
+        StoreDecision::Ignore => Ok(false),
+        StoreDecision::Write | StoreDecision::RefreshExpiry(_) => {
+            let payload = postcard::to_stdvec(record)
+                .map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
+            conn.execute(
+                "INSERT OR REPLACE INTO dht_records
+                 (dht_key, artist_id, owner_peer_id, payload, revision, deleted, expires_at_ms)
+                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
+                params![
+                    key.as_bytes().as_slice(),
+                    record.item.id.as_bytes().as_slice(),
+                    record.item.owner.to_string(),
+                    payload,
+                    record.item.revision as i64,
+                    record.item.deleted as i64,
+                    record.expires_at_ms as i64,
+                ],
+            )?;
+            Ok(true)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -590,6 +630,40 @@ mod tests {
         assert_eq!(found, vec![track]);
     }
 
+    #[tokio::test]
+    async fn record_batch_stores_in_one_call() {
+        let (_dir, db) = open_temp().await;
+        let owner = test_peer(1);
+        let net = federation_net::NetworkId::from_name("t");
+        let now = now_ms();
+        let entry = |name: &str, revision: u64| {
+            let item = item(owner, name, revision, false);
+            (
+                DhtKey::exact(&net, &item.normalized_name),
+                StoredRecord {
+                    item,
+                    publisher: owner,
+                    expires_at_ms: now + 60_000,
+                },
+            )
+        };
+
+        let stored = db
+            .store_dht_records(vec![entry("Massive Attack", 1), entry("Portishead", 1)])
+            .await
+            .expect("store batch");
+        assert_eq!(stored, vec![true, true]);
+        assert_eq!(db.dht_record_count(now).await.expect("count"), 2);
+
+        // A replay of the same batch changes nothing but reports it honestly.
+        let replayed = db
+            .store_dht_records(vec![entry("Massive Attack", 1), entry("Portishead", 2)])
+            .await
+            .expect("store batch");
+        assert_eq!(replayed, vec![false, true]);
+        assert_eq!(db.dht_record_count(now).await.expect("count"), 2);
+    }
+
     #[tokio::test]
     async fn expired_dht_record_is_not_returned() {
         let (_dir, db) = open_temp().await;

crates/music-dht/src/message.rs

@@ -8,14 +8,17 @@ use serde::{Deserialize, Serialize};
 use crate::record::{DhtKey, StoredRecord};
 use crate::routing::{NodeContact, NodeId};
 
-/// Version of the music-dht protocol.
+/// Version of the music-dht protocol. Batch limits are enforced by the
-pub const DHT_PROTOCOL_VERSION: u16 = 4;
+/// receiver, so raising them is a version bump too.
+pub const DHT_PROTOCOL_VERSION: u16 = 5;
 /// Maximum number of contacts in a single [`PeerExchange`].
 pub const MAX_PEER_EXCHANGE_CONTACTS: usize = 32;
 /// Maximum number of records in a single [`FindValueResponse`].
 pub const MAX_RECORDS_PER_RESPONSE: usize = 256;
-/// Maximum number of entries in a single [`StoreBatchRequest`].
+/// Maximum number of entries in a single [`StoreBatchRequest`]. Sized so a
-pub const MAX_RECORDS_PER_BATCH: usize = 64;
+/// batch of typical records nearly fills the byte budget: replicating a
+/// 20k-track library takes hundreds of requests, not tens of thousands.
+pub const MAX_RECORDS_PER_BATCH: usize = 512;
 
 /// Correlates a response with its request.
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]

crates/music-dht/src/node.rs

@@ -426,47 +426,17 @@ impl Node {
         }
     }
 
-    async fn answer_store(&self, request: StoreRecordRequest, sender: &EndpointId) -> bool {
+    /// Validates every entry of a batch individually — one bad entry never
-        if self.is_shutting_down() {
+    /// poisons the rest — and applies the valid ones in a single write
-            return false;
+    /// transaction.
-        }
-        let key = request.key;
-        match validate_store(request, &self.config.network_id, now_ms()) {
-            Ok(record) => {
-                let artist_id = record.item.id;
-                let deleted = record.item.deleted;
-                match self.db.store_dht_record(key, record).await {
-                    Ok(stored) => {
-                        if stored {
-                            debug!(
-                                item = %artist_id,
-                                tombstone = deleted,
-                                from = %sender,
-                                "stored DHT record"
-                            );
-                        }
-                        stored
-                    }
-                    Err(err) => {
-                        warn!(error = %err, "failed to store DHT record");
-                        false
-                    }
-                }
-            }
-            Err(reason) => {
-                warn!(from = %sender, reason = %reason, "rejected DHT store request");
-                false
-            }
-        }
-    }
-
-    /// Validates and stores every entry of a batch individually: one bad
-    /// entry never poisons the rest.
     async fn answer_store_batch(
         &self,
         request: StoreBatchRequest,
         sender: &EndpointId,
     ) -> StoreBatchResponse {
+        if self.is_shutting_down() {
+            return StoreBatchResponse { stored: 0 };
+        }
         if request.entries.len() > MAX_RECORDS_PER_BATCH {
             warn!(
                 from = %sender,
@@ -475,12 +445,26 @@ impl Node {
             );
             return StoreBatchResponse { stored: 0 };
         }
-        let mut stored = 0u32;
+        let received = request.entries.len();
+        let now = now_ms();
+        let mut valid = Vec::with_capacity(received);
         for entry in request.entries {
-            if self.answer_store(entry, sender).await {
+            let key = entry.key;
-                stored += 1;
+            match validate_store(entry, &self.config.network_id, now) {
+                Ok(record) => valid.push((key, record)),
+                Err(reason) => {
+                    warn!(from = %sender, reason = %reason, "rejected DHT store entry");
+                }
             }
         }
+        let stored = match self.db.store_dht_records(valid).await {
+            Ok(outcomes) => outcomes.into_iter().filter(|stored| *stored).count() as u32,
+            Err(err) => {
+                warn!(error = %err, "failed to store DHT record batch");
+                0
+            }
+        };
+        debug!(from = %sender, received, stored, "stored DHT record batch");
         StoreBatchResponse { stored }
     }
 
@@ -814,6 +798,8 @@ impl Node {
         let mut local_replica = false;
         // Batches under construction, keyed by the index into `contacts`.
         let mut per_peer: HashMap<usize, Vec<StoreRecordRequest>> = HashMap::new();
+        // Own replicas, applied in one write transaction at the end.
+        let mut own_replicas: Vec<(DhtKey, StoredRecord)> = Vec::new();
         // K-closest target sets are memoized per key: token keys repeat
         // heavily across the items of one artist or release.
         let mut targets_memo: HashMap<DhtKey, Vec<usize>> = HashMap::new();
@@ -845,10 +831,7 @@ impl Node {
                             <= distance(contacts[*farthest].node_id.as_bytes(), key.as_bytes())
                     });
                 if self_is_close {
-                    match self.db.store_dht_record(key, record.clone()).await {
+                    own_replicas.push((key, record.clone()));
-                        Ok(_) => local_replica = true,
-                        Err(err) => warn!(error = %err, "failed to store own replica"),
-                    }
                 }
 
                 for index in targets.iter() {
@@ -860,6 +843,13 @@ impl Node {
             }
         }
 
+        if !own_replicas.is_empty() {
+            match self.db.store_dht_records(own_replicas).await {
+                Ok(_) => local_replica = true,
+                Err(err) => warn!(error = %err, "failed to store own replicas"),
+            }
+        }
+
         // One concurrent pipeline per peer; batches within a pipeline are
         // sequential so a slow peer only throttles itself.
         let sends = per_peer.into_iter().map(|(index, entries)| {