From 3f15e9aebd574f7dd922e623e724eb7fde741b59 Mon Sep 17 00:00:00 2001 From: Ultradesu Date: Mon, 20 Jul 2026 02:00:40 +0300 Subject: [PATCH] added psql support --- Cargo.lock | 1 + Cargo.toml | 1 + crates/federation-net/src/engine.rs | 14 +++- crates/federation-net/src/lib.rs | 2 +- crates/music-dht/Cargo.toml | 1 + crates/music-dht/src/database.rs | 121 ++++++++++++++++++++++++++++ crates/music-dht/src/dht.rs | 7 +- crates/music-dht/src/lib.rs | 4 +- crates/music-dht/src/node.rs | 6 +- crates/music-dht/src/service.rs | 50 ++++++++++-- 10 files changed, 191 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6a66c15..a1535d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2151,6 +2151,7 @@ name = "music-dht" version = "0.1.0" dependencies = [ "anyhow", + "async-trait", "blake3", "data-encoding", "federation-net", diff --git a/Cargo.toml b/Cargo.toml index f315a00..4b27dca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ futures = "0.3" uuid = { version = "1", features = ["v7"] } unicode-normalization = "0.1" rusqlite = { version = "0.32", features = ["bundled"] } +async-trait = "0.1" thiserror = "2" tracing = "0.1" clap = { version = "4", features = ["derive"] } diff --git a/crates/federation-net/src/engine.rs b/crates/federation-net/src/engine.rs index 299806b..784e75a 100644 --- a/crates/federation-net/src/engine.rs +++ b/crates/federation-net/src/engine.rs @@ -8,7 +8,7 @@ use std::time::Duration; use iroh::endpoint::{Connection, RecvStream, SendStream, VarInt, presets}; use iroh::protocol::{AcceptError, ProtocolHandler, Router}; -use iroh::{Endpoint, EndpointAddr, EndpointId}; +use iroh::{Endpoint, EndpointAddr, EndpointId, SecretKey}; use serde::Serialize; use serde::de::DeserializeOwned; use tokio::sync::{Semaphore, mpsc, watch}; @@ -813,6 +813,18 @@ impl NetworkEngine { /// [`NetworkEvent`]s. pub async fn start(config: NetworkConfig) -> Result<(Self, NetworkEventReceiver)> { let secret_key = identity::load_or_create(&config.data_dir).await?; + Self::start_with_secret_key(config, secret_key).await + } + + /// Starts the engine with an application-provided identity key. + /// + /// Use this when the caller already stores the peer identity in durable + /// application storage. The regular [`NetworkEngine::start`] method + /// remains the file-backed default. + pub async fn start_with_secret_key( + config: NetworkConfig, + secret_key: SecretKey, + ) -> Result<(Self, NetworkEventReceiver)> { let endpoint = Endpoint::builder(presets::N0) .secret_key(secret_key) .bind() diff --git a/crates/federation-net/src/lib.rs b/crates/federation-net/src/lib.rs index d31b038..bb1a008 100644 --- a/crates/federation-net/src/lib.rs +++ b/crates/federation-net/src/lib.rs @@ -74,6 +74,6 @@ pub use ticket::{PeerTicket, TICKET_VERSION}; // Re-exported Iroh types that appear in the public API. pub use iroh::endpoint::{RecvStream, SendStream}; -pub use iroh::{EndpointAddr, EndpointId}; +pub use iroh::{EndpointAddr, EndpointId, SecretKey}; // Re-exported so applications can use the generic ticket helpers. pub use iroh_tickets::Ticket; diff --git a/crates/music-dht/Cargo.toml b/crates/music-dht/Cargo.toml index a347c1d..c9d8111 100644 --- a/crates/music-dht/Cargo.toml +++ b/crates/music-dht/Cargo.toml @@ -15,6 +15,7 @@ blake3 = { workspace = true } futures = { workspace = true } unicode-normalization = { workspace = true } rusqlite = { workspace = true } +async-trait = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } rand = { workspace = true } diff --git a/crates/music-dht/src/database.rs b/crates/music-dht/src/database.rs index ec803b3..f0246d7 100644 --- a/crates/music-dht/src/database.rs +++ b/crates/music-dht/src/database.rs @@ -8,6 +8,7 @@ use std::path::Path; use std::str::FromStr; use std::sync::{Arc, Mutex}; +use async_trait::async_trait; use federation_net::EndpointId; use rusqlite::{Connection, OptionalExtension, params}; @@ -54,6 +55,77 @@ CREATE TABLE IF NOT EXISTS known_peers ( ); "; +/// Persistence backend used by [`crate::MusicDhtService`] for local items, +/// replicated DHT records and learned peer contacts. +/// +/// The default service start path uses the bundled SQLite implementation. +/// Applications that already have durable storage can pass their own backend +/// to [`crate::MusicDhtService::start_with_storage`]. +#[async_trait] +pub trait MusicDhtStorage: std::fmt::Debug + Send + Sync { + /// Inserts or replaces a locally owned item record. + async fn upsert_local_item(&self, item: &LibraryItem) -> Result<()>; + + /// Lists locally owned items. Tombstones are excluded unless + /// `include_deleted` is set. + async fn list_local_items(&self, include_deleted: bool) -> Result>; + + /// Returns everything that must be republished: active records plus + /// tombstones that have not outlived [`TOMBSTONE_TTL`] yet. + async fn local_items_for_republish(&self, now_ms: u64) -> Result> { + let all = self.list_local_items(true).await?; + let tombstone_ttl = TOMBSTONE_TTL.as_millis() as u64; + Ok(all + .into_iter() + .filter(|item| { + !item.deleted || item.updated_at_ms.saturating_add(tombstone_ttl) > now_ms + }) + .collect()) + } + + /// Searches locally owned active items: exact normalized match, or all + /// query tokens present in the item's token set. + async fn search_local(&self, normalized_query: String) -> Result> { + let all = self.list_local_items(false).await?; + let query_tokens = tokenize(&normalized_query); + Ok(all + .into_iter() + .filter(|item| { + if item.normalized_name == normalized_query { + return true; + } + if query_tokens.is_empty() { + return false; + } + let artist_tokens = tokenize(&item.normalized_name); + query_tokens + .iter() + .all(|token| artist_tokens.iter().any(|t| t == token)) + }) + .collect()) + } + + /// Applies a validated incoming record to the replica store, following + /// the revision/tombstone rules. Returns `true` if the record was written + /// or refreshed. + async fn store_dht_record(&self, key: DhtKey, record: StoredRecord) -> Result; + + /// Returns non-expired replicas stored under `key`, including tombstones. + async fn dht_records_by_key(&self, key: DhtKey, now_ms: u64) -> Result>; + + /// Deletes expired replicas. Returns the number of removed rows. + async fn delete_expired_records(&self, now_ms: u64) -> Result; + + /// Inserts or refreshes a known peer contact. + async fn upsert_known_peer(&self, contact: &NodeContact) -> Result<()>; + + /// Deletes a persisted peer contact. + async fn delete_known_peer(&self, peer_id: EndpointId) -> Result<()>; + + /// Loads all persisted peer contacts. + async fn load_known_peers(&self) -> Result>; +} + /// Handle to the local SQLite database. /// /// Cheap to clone; all clones share one connection guarded by a mutex that is @@ -63,6 +135,12 @@ pub(crate) struct Database { conn: Arc>, } +impl std::fmt::Debug for Database { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Database").finish_non_exhaustive() + } +} + impl Database { /// Opens (creating if needed) the database at `path` and applies the /// schema. @@ -337,6 +415,49 @@ impl Database { } } +#[async_trait] +impl MusicDhtStorage for Database { + async fn upsert_local_item(&self, item: &LibraryItem) -> Result<()> { + Database::upsert_local_item(self, item).await + } + + async fn list_local_items(&self, include_deleted: bool) -> Result> { + Database::list_local_items(self, include_deleted).await + } + + async fn local_items_for_republish(&self, now_ms: u64) -> Result> { + Database::local_items_for_republish(self, now_ms).await + } + + async fn search_local(&self, normalized_query: String) -> Result> { + Database::search_local(self, normalized_query).await + } + + async fn store_dht_record(&self, key: DhtKey, record: StoredRecord) -> Result { + Database::store_dht_record(self, key, record).await + } + + async fn dht_records_by_key(&self, key: DhtKey, now_ms: u64) -> Result> { + Database::dht_records_by_key(self, key, now_ms).await + } + + async fn delete_expired_records(&self, now_ms: u64) -> Result { + Database::delete_expired_records(self, now_ms).await + } + + async fn upsert_known_peer(&self, contact: &NodeContact) -> Result<()> { + Database::upsert_known_peer(self, contact).await + } + + async fn delete_known_peer(&self, peer_id: EndpointId) -> Result<()> { + Database::delete_known_peer(self, peer_id).await + } + + async fn load_known_peers(&self) -> Result> { + Database::load_known_peers(self).await + } +} + fn item_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { let payload: Vec = row.get(0)?; postcard::from_bytes(&payload).map_err(|err| { diff --git a/crates/music-dht/src/dht.rs b/crates/music-dht/src/dht.rs index 2582da4..56f7bf1 100644 --- a/crates/music-dht/src/dht.rs +++ b/crates/music-dht/src/dht.rs @@ -11,7 +11,7 @@ use crate::record::{ /// Outcome of comparing an incoming record with the stored one. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum StoreDecision { +pub enum StoreDecision { /// No record stored yet, or the incoming one supersedes it: write it. Write, /// Same logical record: keep the stored row but extend its expiry. @@ -31,10 +31,7 @@ pub(crate) enum StoreDecision { /// only the expiry is refreshed (this is how republish extends TTL); /// * an older revision never replaces a newer one, in particular an old /// active record never resurrects a tombstone. -pub(crate) fn decide_store( - existing: Option<(u64, bool, u64)>, - incoming: &StoredRecord, -) -> StoreDecision { +pub fn decide_store(existing: Option<(u64, bool, u64)>, incoming: &StoredRecord) -> StoreDecision { let Some((revision, deleted, expires_at_ms)) = existing else { return StoreDecision::Write; }; diff --git a/crates/music-dht/src/lib.rs b/crates/music-dht/src/lib.rs index 2ba9052..00b17cf 100644 --- a/crates/music-dht/src/lib.rs +++ b/crates/music-dht/src/lib.rs @@ -82,6 +82,8 @@ pub use config::{ DEFAULT_EXPIRE_INTERVAL, DEFAULT_LOOKUP_TIMEOUT, DEFAULT_REPUBLISH_INTERVAL, DEFAULT_REQUEST_TIMEOUT, DEFAULT_TRANSPORT_TIMEOUT, MusicDhtConfig, MusicDhtConfigBuilder, }; +pub use database::MusicDhtStorage; +pub use dht::{StoreDecision, decide_store}; pub use error::{MusicDhtError, Result}; pub use message::{ DHT_PROTOCOL_VERSION, FindNodeRequest, FindNodeResponse, FindValueRequest, FindValueResponse, @@ -107,5 +109,5 @@ pub use service::{ // Re-exported types from the transport layer that appear in this API. pub use federation_net::{ ByteStream, EndpointAddr, EndpointId, NetworkId, PeerTicket, RecvStream, RendezvousConfig, - SendStream, StreamAcceptor, + SecretKey, SendStream, StreamAcceptor, }; diff --git a/crates/music-dht/src/node.rs b/crates/music-dht/src/node.rs index 944709c..405a881 100644 --- a/crates/music-dht/src/node.rs +++ b/crates/music-dht/src/node.rs @@ -14,7 +14,7 @@ use tokio::time::timeout; use tracing::{debug, info, warn}; use crate::config::MusicDhtConfig; -use crate::database::Database; +use crate::database::MusicDhtStorage; use crate::dht::validate_store; use crate::error::{MusicDhtError, Result}; use crate::message::{ @@ -77,7 +77,7 @@ pub(crate) struct LookupOutcome { /// Shared state of one DHT node. pub(crate) struct Node { pub engine: NetworkEngine, - pub db: Database, + pub db: Arc, pub config: MusicDhtConfig, pub node_id: NodeId, pub endpoint_id: EndpointId, @@ -98,7 +98,7 @@ pub(crate) struct Node { impl Node { pub fn new( engine: NetworkEngine, - db: Database, + db: Arc, config: MusicDhtConfig, events: mpsc::Sender, ) -> Self { diff --git a/crates/music-dht/src/service.rs b/crates/music-dht/src/service.rs index 4b6331e..ae683cc 100644 --- a/crates/music-dht/src/service.rs +++ b/crates/music-dht/src/service.rs @@ -6,14 +6,14 @@ use std::time::{Duration, Instant}; use federation_net::{ ByteStream, EndpointAddr, EndpointId, NetworkConfig, NetworkEngine, PeerTicket, SchemaId, - StreamAcceptor, + SecretKey, StreamAcceptor, }; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tracing::info; use crate::config::MusicDhtConfig; -use crate::database::Database; +use crate::database::{Database, MusicDhtStorage}; use crate::error::{MusicDhtError, Result}; use crate::node::{Node, record_supersedes}; use crate::normalization::{normalize_name, tokenize}; @@ -159,6 +159,42 @@ impl MusicDhtService { /// Starts the service: opens the database, starts the network engine, /// loads persisted contacts and spawns the maintenance tasks. pub async fn start(config: MusicDhtConfig) -> Result<(Self, MusicDhtEventReceiver)> { + let db = Arc::new(Database::open(&config.data_dir.join("state.sqlite3")).await?); + Self::start_with_storage(config, db).await + } + + /// Starts the service with an application-provided persistence backend. + /// + /// This is useful for servers that already have durable storage and do + /// not want the DHT state tied to a local SQLite file. The `data_dir` + /// in `config` is still used by the transport layer for the peer + /// identity. + pub async fn start_with_storage( + config: MusicDhtConfig, + storage: Arc, + ) -> Result<(Self, MusicDhtEventReceiver)> { + Self::start_with_storage_inner(config, storage, None).await + } + + /// Starts the service with application-provided persistence and identity. + /// + /// The identity controls the stable transport endpoint id. Supplying it + /// lets applications store both the DHT state and the peer identity in + /// their own durable database while keeping [`MusicDhtService::start`] + /// as the SQLite/file-backed default. + pub async fn start_with_storage_and_secret_key( + config: MusicDhtConfig, + storage: Arc, + secret_key: SecretKey, + ) -> Result<(Self, MusicDhtEventReceiver)> { + Self::start_with_storage_inner(config, storage, Some(secret_key)).await + } + + async fn start_with_storage_inner( + config: MusicDhtConfig, + storage: Arc, + secret_key: Option, + ) -> Result<(Self, MusicDhtEventReceiver)> { let mut engine_builder = NetworkConfig::builder() .data_dir(&config.data_dir) .network_id(config.network_id) @@ -176,11 +212,15 @@ impl MusicDhtService { let engine_config = engine_builder .build() .map_err(|err| MusicDhtError::Network(err.to_string()))?; - let (engine, net_events) = NetworkEngine::start(engine_config).await?; - let db = Database::open(&config.data_dir.join("state.sqlite3")).await?; + let (engine, net_events) = match secret_key { + Some(secret_key) => { + NetworkEngine::start_with_secret_key(engine_config, secret_key).await? + } + None => NetworkEngine::start(engine_config).await?, + }; let (event_tx, event_rx) = mpsc::channel(EVENT_CHANNEL_CAPACITY); - let node = Arc::new(Node::new(engine, db, config.clone(), event_tx)); + let node = Arc::new(Node::new(engine, storage, config.clone(), event_tx)); info!( endpoint_id = %node.endpoint_id, node_id = %node.node_id,