use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use iroh::{Endpoint, EndpointAddr, PublicKey};
use rand::RngExt;
use tokio::sync::oneshot;
use crate::dht::chain_projection::with_lock;
use crate::dht::client;
use crate::dht::routing::{KEYSPACE_BITS, NodeId, RoutingTable, bucket_index};
pub const BUCKET_REFRESH_TICK: Duration = Duration::from_hours(1);
pub async fn run_bucket_refresh(
endpoint: Endpoint,
self_id: PublicKey,
routing: Arc<Mutex<RoutingTable>>,
mut stop_rx: oneshot::Receiver<()>,
interval: Duration,
refresh_clock: Arc<AtomicU64>,
metrics: Arc<crate::metrics::Metrics>,
) {
let self_node_id = NodeId::from_bytes(*self_id.as_bytes());
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
ticker.tick().await;
loop {
tokio::select! {
biased;
_ = &mut stop_rx => {
tracing::debug!("dht bucket-refresh: shutdown signal received");
return;
}
_ = ticker.tick() => {
let failures =
refresh_all_non_empty_buckets(&endpoint, self_node_id, &routing).await;
for _ in 0..failures {
metrics.dht_bucket_refresh_failure();
}
metrics.dht_routing_table_size(
with_lock(&routing, "dht routing table", |table| table.len()),
);
refresh_clock.store(now_us(), Ordering::Relaxed);
}
}
}
}
fn now_us() -> u64 {
match std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
Ok(d) => u64::try_from(d.as_micros()).unwrap_or(u64::MAX),
Err(err) => {
tracing::error!(
error = %err,
"wall-clock before UNIX_EPOCH; bucket-refresh timestamp falls back to 0, so \
admin_v1_status will report last_refresh=never despite refreshes running. \
Set the host clock."
);
0
}
}
}
async fn refresh_all_non_empty_buckets(
endpoint: &Endpoint,
self_node_id: NodeId,
routing: &Arc<Mutex<RoutingTable>>,
) -> usize {
let picks: Vec<BucketPick> = with_lock(routing, "dht routing table", |table| {
all_non_empty_picks(table)
});
if picks.is_empty() {
return 0;
}
let mut handles = Vec::with_capacity(picks.len());
for pick in picks {
let endpoint_cloned = endpoint.clone();
let routing_cloned = Arc::clone(routing);
handles.push(tokio::spawn(async move {
refresh_one_bucket(&endpoint_cloned, self_node_id, &routing_cloned, pick).await
}));
}
let mut failures = 0;
for h in handles {
if !h.await.unwrap_or(false) {
failures += 1;
}
}
failures
}
fn all_non_empty_picks(table: &RoutingTable) -> Vec<BucketPick> {
(0..KEYSPACE_BITS)
.filter_map(|idx| {
let peer = bucket_iter(table, idx).last().copied()?;
let target = random_target_in_bucket(table.self_id(), idx);
Some(BucketPick {
bucket_index: idx,
peer,
target,
})
})
.collect()
}
async fn refresh_one_bucket(
endpoint: &Endpoint,
self_node_id: NodeId,
routing: &Arc<Mutex<RoutingTable>>,
pick: BucketPick,
) -> bool {
let BucketPick {
bucket_index,
peer,
target,
} = pick;
let target_pk = match PublicKey::from_bytes(peer.as_bytes()) {
Ok(k) => k,
Err(e) => {
tracing::warn!(
peer = %peer,
error = %e,
"dht bucket-refresh: routing-table peer not a valid public key"
);
return false;
}
};
let addr = EndpointAddr::new(target_pk);
match client::find_node(endpoint, addr, target, self_node_id).await {
Ok(resp) => {
with_lock(routing, "dht routing table", |table| {
for nid in resp.closer_nodes.into_inner() {
if nid != self_node_id {
table.insert(nid);
}
}
});
true
}
Err(e) => {
tracing::debug!(
bucket = bucket_index,
error = %e,
"dht bucket-refresh: FindNode failed; will retry on next tick"
);
false
}
}
}
struct BucketPick {
bucket_index: usize,
peer: NodeId,
target: NodeId,
}
#[cfg(test)]
fn pick_bucket(table: &RoutingTable, bucket: usize) -> Option<BucketPick> {
let peer = bucket_iter(table, bucket).last().copied()?;
let target = random_target_in_bucket(table.self_id(), bucket);
Some(BucketPick {
bucket_index: bucket,
peer,
target,
})
}
fn bucket_iter(table: &RoutingTable, target_bucket: usize) -> impl Iterator<Item = &NodeId> {
let self_id = *table.self_id();
table.iter_peers().filter(move |peer| {
let d = crate::dht::routing::xor_distance(&self_id, peer);
bucket_index(&d) == Some(target_bucket)
})
}
fn random_target_in_bucket(self_id: &NodeId, bucket_index: usize) -> NodeId {
let bit_from_msb = KEYSPACE_BITS.saturating_sub(1).saturating_sub(bucket_index);
let byte_idx = bit_from_msb / 8;
let bit_within = bit_from_msb % 8;
let mut target = *self_id.as_bytes();
let flip_mask = 1u8 << (7 - bit_within);
if let Some(b) = target.get_mut(byte_idx) {
*b ^= flip_mask;
}
let mut rng = rand::rng();
for slot in target.iter_mut().skip(byte_idx + 1) {
*slot = rng.random();
}
NodeId::from_bytes(target)
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::indexing_slicing
)]
mod tests;