ab/frid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved DHT publication. Adjust batching, added validation per key

Browse Source
This commit is contained in:
Ultradesu
2026-07-23 16:43:10 +03:00
parent aaaad780f2
commit 512a818a6a
3 changed files with 150 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/music-dht/src/database.rs
+110 -36
View File
@@ -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
.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<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
+7 -4
View File
@@ -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)]
crates/music-dht/src/node.rs
+33 -43
View File
@@ -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<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 {
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)| {