Furumi/furumi-tui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

Route similarity search through the DHT

Browse Source
This commit is contained in:
Aleksandr Bogomiakov
2026-08-10 19:40:37 +01:00
parent e95d2e7fe1
commit 87cb7fe74c
11 changed files with 374 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/federation/capabilities.rs
+1
View File
@@ -171,6 +171,7 @@ mod tests {
"catalog",
"audio",
"similarity",
"similarity_dht",
"device_sync",
"jam",
] {
src/federation/mod.rs
+125 -4
View File
@@ -25,6 +25,8 @@ use std::sync::atomic::{AtomicI64, Ordering};
use std::time::Duration;
use anyhow::{Context, Result};
use music_dht::similarity_dht::SimilarityDht;
use music_dht::similarity_lsh::SIMILARITY_DHT_ALPN;
use music_dht::{
ByteStream, ByteStreamConnectionStats, EndpointId, ItemKind, ItemSpec, LibraryItem,
MusicDhtConfig, MusicDhtService, NetworkId, PeerTicket, PublishStats, RendezvousConfig,
@@ -413,6 +415,7 @@ pub struct NetworkLibrarySource {
struct Running {
service: Arc<MusicDhtService>,
similarity_dht: Arc<SimilarityDht>,
network_name: String,
network_id: NetworkId,
tasks: Vec<tokio::task::JoinHandle<()>>,
@@ -697,12 +700,15 @@ impl Federation {
.stream_protocol(AUDIO_ALPN)
// ...and browse each other's per-artist catalogs over this one.
.stream_protocol(CATALOG_ALPN)
// Anonymous, bounded direct embedding queries.
.stream_protocol(SIMILARITY_ALPN)
// Anonymous, bounded direct embedding queries have their own
// versioned contract and survive catalog-schema upgrades.
.schema_independent_stream_protocol(SIMILARITY_ALPN)
// Personal-device sync (likes, playlists, trusted devices).
.stream_protocol(crate::devices::SYNC_ALPN)
// Capability-scoped shared playback control.
.stream_protocol(crate::jam::JAM_ALPN)
// Signed LSH summaries form their own upgrade-safe DHT overlay.
.schema_independent_stream_protocol(SIMILARITY_DHT_ALPN)
// Informational application/protocol versions.
.schema_independent_stream_protocol(capabilities::CAPABILITIES_ALPN)
.build()
@@ -717,6 +723,25 @@ impl Federation {
"federation started"
);
let similarity_dht = SimilarityDht::open(
Arc::clone(&service),
self.data_dir.join("similarity-routing.sqlite3"),
)
.await
.map_err(|err| anyhow::anyhow!("failed to start the similarity DHT: {err}"))?;
let similarity_dht_acceptor = service
.stream_acceptor(SIMILARITY_DHT_ALPN)
.map_err(|err| anyhow::anyhow!("failed to take similarity DHT acceptor: {err}"))?;
let similarity_dht_serve_task =
tokio::spawn(Arc::clone(&similarity_dht).serve(similarity_dht_acceptor));
let similarity_dht_maintenance_task =
tokio::spawn(Arc::clone(&similarity_dht).maintenance());
let similarity_dht_sync_task = tokio::spawn(similarity_route_sync_loop(
Arc::clone(&similarity_dht),
Arc::clone(&self.similarity),
Arc::clone(&self.library),
));
// Drain DHT events into the log; the channel is bounded.
let event_task = tokio::spawn(async move {
while let Some(event) = events.recv().await {
@@ -797,6 +822,7 @@ impl Federation {
*guard = Some(Running {
service,
similarity_dht,
network_name,
network_id,
tasks: vec![
@@ -811,6 +837,9 @@ impl Federation {
jam_poll_task,
capabilities_serve_task,
capabilities_probe_task,
similarity_dht_serve_task,
similarity_dht_maintenance_task,
similarity_dht_sync_task,
],
});
self.set_error(None);
@@ -835,6 +864,20 @@ impl Federation {
.context("federation is not running")
}
async fn similarity_services(&self) -> Result<(Arc<MusicDhtService>, Arc<SimilarityDht>)> {
self.running
.lock()
.await
.as_ref()
.map(|running| {
(
Arc::clone(&running.service),
Arc::clone(&running.similarity_dht),
)
})
.context("federation is not running")
}
fn ensure_connected_devices_enabled(&self) -> Result<()> {
let settings = self.settings();
anyhow::ensure!(
@@ -1175,8 +1218,15 @@ impl Federation {
self.similarity.network_allowed(),
"similarity federation has no consent"
);
let service = self.service().await?;
similarity::search(service, query, limit, Arc::clone(&self.transport_stats)).await
let (service, similarity_dht) = self.similarity_services().await?;
similarity::search(
service,
similarity_dht,
query,
limit,
Arc::clone(&self.transport_stats),
)
.await
}
/// Resolves a share-link content id to one playable federated track.
@@ -2565,6 +2615,77 @@ fn sort_fed_appearances(appearances: &mut [FedAppearsOn]) {
});
}
async fn similarity_route_sync_loop(
routing: Arc<SimilarityDht>,
similarity: Arc<crate::similarity::Manager>,
library: Arc<Library>,
) {
let mut interval = tokio::time::interval(SYNC_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut published_marker: Option<(String, blake3::Hash)> = None;
loop {
interval.tick().await;
if !similarity.network_allowed() {
if published_marker.take().is_some() {
routing.clear_local_signatures();
tracing::info!("local similarity DHT publication disabled");
}
continue;
}
let status = similarity.status();
let Some(profile_id) = status.active_profile else {
continue;
};
if status.phase != crate::similarity::Phase::Ready {
continue;
}
let library = Arc::clone(&library);
let profile_for_task = profile_id.clone();
let loaded = tokio::task::spawn_blocking(move || {
let signatures = library.similarity_routing_signatures(&profile_for_task)?;
let mut hasher = blake3::Hasher::new();
for signature in &signatures {
hasher.update(signature);
}
Ok::<_, anyhow::Error>((signatures, hasher.finalize()))
})
.await;
let (signatures, fingerprint) = match loaded {
Ok(Ok(loaded)) => loaded,
Ok(Err(error)) => {
tracing::warn!(%error, %profile_id, "similarity routing signatures unavailable");
continue;
}
Err(error) => {
tracing::warn!(%error, "similarity routing signature task failed");
continue;
}
};
let marker = (profile_id.clone(), fingerprint);
if published_marker.as_ref() == Some(&marker) {
continue;
}
match routing
.sync_local_signatures(profile_id.clone(), signatures)
.await
{
Ok(stats) => {
tracing::info!(
profile = %profile_id,
records = stats.records,
keys = stats.keys,
remote_nodes = stats.remote_nodes,
"local similarity DHT index synchronized"
);
published_marker = Some(marker);
}
Err(error) => {
tracing::warn!(%error, %profile_id, "similarity DHT synchronization failed");
}
}
}
}
async fn stop_running(running: Option<Running>) {
let Some(running) = running else { return };
for task in &running.tasks {
src/federation/similarity.rs
+111 -35
View File
@@ -11,16 +11,21 @@ use std::time::Duration;
use anyhow::{Context as _, Result};
use futures_util::stream::{self, StreamExt as _};
use music_dht::similarity::{self as wire, SimilarityHit, SimilarityRequest, SimilarityResponse};
use music_dht::{ByteStream, EndpointId, ItemId, ItemKind, MusicDhtService, StreamAcceptor};
use music_dht::similarity_dht::SimilarityDht;
use music_dht::{
ByteStream, EndpointId, ItemId, ItemKind, MusicDhtService, PeerTicket, StreamAcceptor,
};
use crate::federation::{FedSearchResults, FedTrack, TransportStats};
use crate::similarity::{Manager, QueryVector};
pub use music_dht::similarity::SIMILARITY_ALPN;
const MAX_QUERY_PEERS: usize = 16;
const QUERY_CONCURRENCY: usize = 6;
const INITIAL_QUERY_PEERS: usize = 16;
const MAX_QUERY_PEERS: usize = 48;
const QUERY_CONCURRENCY: usize = 8;
const QUERY_TIMEOUT: Duration = Duration::from_secs(5);
const ROUTING_TIMEOUT: Duration = Duration::from_secs(5);
const MAX_PER_ARTIST: usize = 3;
const MAX_NEAR_DUPLICATE_SIGNATURE_DISTANCE: u32 = 8;
@@ -122,20 +127,49 @@ async fn serve_one(
pub async fn search(
service: Arc<MusicDhtService>,
routing: Arc<SimilarityDht>,
query: QueryVector,
limit: usize,
transport: Arc<TransportStats>,
) -> Result<FedSearchResults> {
let own = service.endpoint_id();
let mut peers = Vec::new();
let routed = match tokio::time::timeout(
ROUTING_TIMEOUT,
routing.find_peers(&query.profile_id, &query.vector, MAX_QUERY_PEERS),
)
.await
{
Ok(Ok(peers)) => peers,
Err(_) => {
tracing::debug!("similarity DHT lookup timed out; using known peers");
Vec::new()
}
Ok(Err(error)) => {
tracing::debug!(%error, "similarity DHT lookup unavailable; using known peers");
Vec::new()
}
};
let mut seen = HashSet::new();
let mut peers: Vec<QueryPeer> = routed
.into_iter()
.filter_map(|ticket| {
let owner = ticket.endpoint_id();
(owner != own && seen.insert(owner)).then_some(QueryPeer {
owner,
ticket: Some(ticket),
})
})
.collect();
for peer in service
.connected_peers()
.into_iter()
.chain(service.known_peers().into_iter().map(|peer| peer.peer_id))
{
if peer != own && seen.insert(peer) {
peers.push(peer);
peers.push(QueryPeer {
owner: peer,
ticket: None,
});
}
if peers.len() >= MAX_QUERY_PEERS {
break;
@@ -148,30 +182,40 @@ pub async fn search(
limit.clamp(1, wire::MAX_SIMILARITY_RESULTS),
)?);
let responses = stream::iter(peers.into_iter().map(|peer| {
let service = Arc::clone(&service);
let request = Arc::clone(&request);
let transport = Arc::clone(&transport);
async move {
tokio::time::timeout(
QUERY_TIMEOUT,
query_peer(service, peer, &request, transport),
)
.await
.map_err(|_| anyhow::anyhow!("similarity peer timed out"))?
}
}))
.buffer_unordered(QUERY_CONCURRENCY)
.collect::<Vec<_>>()
.await;
let mut hits = Vec::new();
let initial = peers.len().min(INITIAL_QUERY_PEERS);
let responses = query_peers(
Arc::clone(&service),
&peers[..initial],
Arc::clone(&request),
Arc::clone(&transport),
)
.await;
let mut successful = 0usize;
for response in responses {
match response {
Ok(peer_hits) => hits.extend(peer_hits),
Ok(peer_hits) => {
successful += 1;
hits.extend(peer_hits);
}
Err(err) => tracing::debug!(%err, "similarity peer query skipped"),
}
}
if initial < peers.len() && (hits.len() < limit || successful < initial.min(4)) {
for response in query_peers(
Arc::clone(&service),
&peers[initial..],
Arc::clone(&request),
Arc::clone(&transport),
)
.await
{
match response {
Ok(peer_hits) => hits.extend(peer_hits),
Err(err) => tracing::debug!(%err, "fallback similarity peer query skipped"),
}
}
}
hits.sort_by(|left, right| right.1.total_cmp(&left.1));
let mut dedup = HashSet::new();
let mut embedding_signatures = vec![query_signature];
@@ -224,22 +268,54 @@ pub async fn search(
})
}
type PeerHits = Vec<(
FedTrack,
f32,
Option<[u8; wire::SIMILARITY_SIGNATURE_BYTES]>,
)>;
#[derive(Clone)]
struct QueryPeer {
owner: EndpointId,
ticket: Option<PeerTicket>,
}
async fn query_peers(
service: Arc<MusicDhtService>,
peers: &[QueryPeer],
request: Arc<SimilarityRequest>,
transport: Arc<TransportStats>,
) -> Vec<Result<PeerHits>> {
stream::iter(peers.iter().cloned().map(|peer| {
let service = Arc::clone(&service);
let request = Arc::clone(&request);
let transport = Arc::clone(&transport);
async move {
tokio::time::timeout(
QUERY_TIMEOUT,
query_peer(service, peer, &request, transport),
)
.await
.map_err(|_| anyhow::anyhow!("similarity peer timed out"))?
}
}))
.buffer_unordered(QUERY_CONCURRENCY)
.collect()
.await
}
async fn query_peer(
service: Arc<MusicDhtService>,
owner: EndpointId,
peer: QueryPeer,
request: &SimilarityRequest,
transport: Arc<TransportStats>,
) -> Result<
Vec<(
FedTrack,
f32,
Option<[u8; wire::SIMILARITY_SIGNATURE_BYTES]>,
)>,
> {
let mut stream = service
.open_stream(owner, SIMILARITY_ALPN)
.await
.map_err(|err| anyhow::anyhow!("cannot reach similarity peer: {err}"))?;
) -> Result<PeerHits> {
let owner = peer.owner;
let mut stream = match peer.ticket {
Some(ticket) => service.open_stream_to(&ticket, SIMILARITY_ALPN).await,
None => service.open_stream(owner, SIMILARITY_ALPN).await,
}
.map_err(|err| anyhow::anyhow!("cannot reach similarity peer: {err}"))?;
super::record_stream_transport(&transport, "similarity", "outbound", "open", &stream);
let response = wire::exchange(&mut stream, request).await?;
super::record_stream_transport(&transport, "similarity", "outbound", "done", &stream);
src/library/mod.rs
+75 -2
View File
@@ -207,6 +207,7 @@ CREATE TABLE IF NOT EXISTS track_embeddings (
profile_id TEXT NOT NULL REFERENCES similarity_profiles(profile_id) ON DELETE CASCADE,
dimensions INTEGER NOT NULL,
vector BLOB NOT NULL,
routing_signature BLOB,
source_content_id TEXT,
computed_at_ms INTEGER NOT NULL,
PRIMARY KEY (track_id, profile_id)
@@ -1551,15 +1552,17 @@ impl Library {
"embedding contains a non-finite value"
);
let bytes = embedding_to_bytes(vector);
let routing_signature = music_dht::similarity_lsh::routing_signature(vector)?;
let conn = self.lock();
conn.execute(
"INSERT INTO track_embeddings (
track_id, profile_id, dimensions, vector,
track_id, profile_id, dimensions, vector, routing_signature,
source_content_id, computed_at_ms
) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
ON CONFLICT(track_id, profile_id) DO UPDATE SET
dimensions = excluded.dimensions,
vector = excluded.vector,
routing_signature = excluded.routing_signature,
source_content_id = excluded.source_content_id,
computed_at_ms = excluded.computed_at_ms",
params![
@@ -1567,6 +1570,7 @@ impl Library {
profile_id,
vector.len() as i64,
bytes,
routing_signature.as_slice(),
track.content_id.as_deref(),
now_ms_i64(),
],
@@ -1633,6 +1637,65 @@ impl Library {
Ok(embeddings)
}
/// Loads the compact DHT-routing signatures for every current local
/// embedding. Rows created before similarity routing existed are
/// backfilled in place from their durable vectors.
pub fn similarity_routing_signatures(&self, profile_id: &str) -> Result<Vec<[u8; 32]>> {
let mut conn = self.lock();
let transaction = conn.transaction()?;
let missing = {
let mut statement = transaction.prepare(
"SELECT e.track_id, e.dimensions, e.vector
FROM track_embeddings e
JOIN tracks t ON t.id = e.track_id
WHERE e.profile_id = ?1
AND e.source_content_id IS t.content_id
AND (e.routing_signature IS NULL OR length(e.routing_signature) != 32)
ORDER BY e.track_id",
)?;
statement
.query_map([profile_id], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, Vec<u8>>(2)?,
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?
};
for (track_id, dimensions, vector_bytes) in missing {
let vector = embedding_from_bytes(dimensions, &vector_bytes)?;
let signature = music_dht::similarity_lsh::routing_signature(&vector)?;
transaction.execute(
"UPDATE track_embeddings
SET routing_signature = ?3
WHERE track_id = ?1 AND profile_id = ?2",
params![track_id, profile_id, signature.as_slice()],
)?;
}
transaction.commit()?;
let mut statement = conn.prepare(
"SELECT e.routing_signature
FROM track_embeddings e
JOIN tracks t ON t.id = e.track_id
WHERE e.profile_id = ?1
AND e.source_content_id IS t.content_id
ORDER BY e.track_id",
)?;
let stored = statement
.query_map([profile_id], |row| row.get::<_, Vec<u8>>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
let signatures = stored
.into_iter()
.map(|signature| {
<[u8; 32]>::try_from(signature)
.map_err(|_| anyhow::anyhow!("invalid similarity routing signature length"))
})
.collect::<Result<Vec<_>>>()?;
Ok(signatures)
}
pub fn similarity_storage_stats(&self, profile_id: &str) -> Result<SimilarityStorageStats> {
let conn = self.lock();
let total_tracks = conn.query_row("SELECT COUNT(*) FROM tracks", [], |row| {
@@ -3444,6 +3507,16 @@ fn ensure_schema_migrations(conn: &Connection) -> Result<()> {
[],
)?;
}
let embedding_columns = table_columns(conn, "track_embeddings")?;
if !embedding_columns
.iter()
.any(|column| column == "routing_signature")
{
conn.execute(
"ALTER TABLE track_embeddings ADD COLUMN routing_signature BLOB",
[],
)?;
}
let mut rows = conn.prepare("SELECT id, title FROM playlists WHERE sync_id IS NULL")?;
let missing = rows
.query_map([], |row| {
src/library/tests.rs
+27 -4
View File
@@ -701,24 +701,47 @@ fn similarity_embeddings_round_trip_and_keep_profiles_separate() {
.into_iter()
.find(|track| track.id == track_id)
.unwrap();
lib.store_similarity_embedding(&track, "profile-a", &[0.1, 0.2, 0.3])
let first = [0.26726124, 0.5345225, 0.8017837];
let second = [0.8017837, 0.5345225, 0.26726124];
lib.store_similarity_embedding(&track, "profile-a", &first)
.unwrap();
lib.store_similarity_embedding(&track, "profile-b", &[0.3, 0.2, 0.1])
lib.store_similarity_embedding(&track, "profile-b", &second)
.unwrap();
assert_eq!(
lib.similarity_embedding(track_id, "profile-a").unwrap(),
Some(vec![0.1, 0.2, 0.3])
Some(first.to_vec())
);
assert_eq!(
lib.similarity_embedding(track_id, "profile-b").unwrap(),
Some(vec![0.3, 0.2, 0.1])
Some(second.to_vec())
);
let stats = lib.similarity_storage_stats("profile-a").unwrap();
assert_eq!(stats.total_tracks, 1);
assert_eq!(stats.embedded_tracks, 1);
assert_eq!(stats.stored_vectors, 2);
assert_eq!(stats.stored_bytes, 24);
lib.lock()
.execute(
"UPDATE track_embeddings SET routing_signature = NULL WHERE profile_id = 'profile-a'",
[],
)
.unwrap();
assert_eq!(
lib.similarity_routing_signatures("profile-a")
.unwrap()
.len(),
1
);
let stored_signature_bytes: i64 = lib
.lock()
.query_row(
"SELECT length(routing_signature) FROM track_embeddings WHERE profile_id = 'profile-a'",
[],
|row| row.get(0),
)
.unwrap();
assert_eq!(stored_signature_bytes, 32);
}
#[test]
src/ui/federation.rs
+1
View File
@@ -400,6 +400,7 @@ fn protocol_label(id: &str) -> &str {
"catalog" => "Catalog",
"audio" => "Audio transfer",
"similarity" => "Similarity search",
"similarity_dht" => "Similarity DHT",
"device_sync" => "Device sync",
"jam" => "Jam",
other => other,