diff --git a/crates/music-dht/src/database.rs b/crates/music-dht/src/database.rs index 4903da2..11be7c0 100644 --- a/crates/music-dht/src/database.rs +++ b/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; + /// 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> { + 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>; @@ -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 { + 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> { self.call(move |conn| { - 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) - } + let tx = conn.unchecked_transaction()?; + let mut stored = Vec::with_capacity(entries.len()); + for (key, record) in &entries { + stored.push(store_record_in_conn(&tx, key, record)?); } + 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> { + Database::store_dht_records(self, entries).await + } + async fn dht_records_by_key(&self, key: DhtKey, now_ms: u64) -> Result> { 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 { + 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; diff --git a/crates/music-dht/src/message.rs b/crates/music-dht/src/message.rs index 4814804..9206c67 100644 --- a/crates/music-dht/src/message.rs +++ b/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. -pub const DHT_PROTOCOL_VERSION: u16 = 4; +/// Version of the music-dht protocol. Batch limits are enforced by the +/// 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`]. -pub const MAX_RECORDS_PER_BATCH: usize = 64; +/// Maximum number of entries in a single [`StoreBatchRequest`]. Sized so a +/// 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)] diff --git a/crates/music-dht/src/node.rs b/crates/music-dht/src/node.rs index e325982..1fefba0 100644 --- a/crates/music-dht/src/node.rs +++ b/crates/music-dht/src/node.rs @@ -426,47 +426,17 @@ impl Node { } } - async fn answer_store(&self, request: StoreRecordRequest, sender: &EndpointId) -> bool { - if self.is_shutting_down() { - return false; - } - 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. + /// Validates every entry of a batch individually — one bad entry never + /// poisons the rest — and applies the valid ones in a single write + /// transaction. 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 { - stored += 1; + let key = entry.key; + 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> = 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> = 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 { - Ok(_) => local_replica = true, - Err(err) => warn!(error = %err, "failed to store own replica"), - } + own_replicas.push((key, record.clone())); } 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)| {