265 lines
9.3 KiB
Rust
265 lines
9.3 KiB
Rust
use crate::network::{RoomState, gossip::IrohGossipState};
|
|
use iroh::{EndpointAddr, EndpointId};
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::Duration;
|
|
use tokio::sync::mpsc;
|
|
use tokio::task::JoinHandle;
|
|
use tokio::time::Instant;
|
|
|
|
const RECOVERY_COMMAND_CAPACITY: usize = 64;
|
|
const RECOVERY_DELAYS: [Duration; 7] = [
|
|
Duration::from_secs(1),
|
|
Duration::from_secs(2),
|
|
Duration::from_secs(4),
|
|
Duration::from_secs(8),
|
|
Duration::from_secs(15),
|
|
Duration::from_secs(30),
|
|
Duration::from_secs(60),
|
|
];
|
|
|
|
fn recovery_delay(attempt: usize) -> Duration {
|
|
RECOVERY_DELAYS[attempt.min(RECOVERY_DELAYS.len() - 1)]
|
|
}
|
|
|
|
enum RecoveryCommand {
|
|
Start {
|
|
peer_id: EndpointId,
|
|
addr: EndpointAddr,
|
|
},
|
|
Cancel(EndpointId),
|
|
}
|
|
|
|
struct RecoveryEntry {
|
|
addr: EndpointAddr,
|
|
attempt: usize,
|
|
next_attempt: Instant,
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
trait RecoveryRoom: Send + Sync {
|
|
async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String>;
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl RecoveryRoom for IrohGossipState {
|
|
async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String> {
|
|
RoomState::rebootstrap_peers(self, peers)
|
|
.await
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
}
|
|
|
|
/// Cloneable command side of the single per-session recovery coordinator.
|
|
/// `active` is shared with transport/event handlers so cancellation is visible
|
|
/// immediately even while the coordinator is awaiting an in-flight gossip call.
|
|
#[derive(Clone)]
|
|
pub(super) struct RecoveryCoordinator {
|
|
tx: mpsc::Sender<RecoveryCommand>,
|
|
active: Arc<Mutex<HashSet<EndpointId>>>,
|
|
}
|
|
|
|
impl RecoveryCoordinator {
|
|
pub(super) fn spawn(room_state: Arc<IrohGossipState>) -> (Self, JoinHandle<()>) {
|
|
Self::spawn_inner(room_state)
|
|
}
|
|
|
|
fn spawn_inner(room_state: Arc<dyn RecoveryRoom>) -> (Self, JoinHandle<()>) {
|
|
let (tx, rx) = mpsc::channel(RECOVERY_COMMAND_CAPACITY);
|
|
let active = Arc::new(Mutex::new(HashSet::new()));
|
|
let handle = Self {
|
|
tx,
|
|
active: active.clone(),
|
|
};
|
|
let task = tokio::spawn(run_coordinator(room_state, active, rx));
|
|
(handle, task)
|
|
}
|
|
|
|
/// Reserve one recovery slot before grace-expiry teardown begins. Returns
|
|
/// false when the peer is already recovering, preventing duplicate work.
|
|
pub(super) fn begin(&self, peer_id: EndpointId) -> bool {
|
|
self.active.lock().unwrap().insert(peer_id)
|
|
}
|
|
|
|
/// Activate the reserved slot with its retained authenticated address.
|
|
/// Uses a bounded non-blocking send while holding the active-set lock so a
|
|
/// concurrent cancellation is ordered before or after this command.
|
|
pub(super) fn activate(&self, peer_id: EndpointId, addr: EndpointAddr) -> Result<bool, ()> {
|
|
let mut active = self.active.lock().unwrap();
|
|
if !active.contains(&peer_id) {
|
|
return Ok(false);
|
|
}
|
|
if self
|
|
.tx
|
|
.try_send(RecoveryCommand::Start { peer_id, addr })
|
|
.is_err()
|
|
{
|
|
active.remove(&peer_id);
|
|
return Err(());
|
|
}
|
|
Ok(true)
|
|
}
|
|
|
|
pub(super) fn cancel(&self, peer_id: EndpointId) {
|
|
self.active.lock().unwrap().remove(&peer_id);
|
|
// Cancellation is governed by the shared active set, so it remains
|
|
// immediate even if the bounded command queue is temporarily full.
|
|
let _ = self.tx.try_send(RecoveryCommand::Cancel(peer_id));
|
|
}
|
|
|
|
pub(super) fn is_active(&self, peer_id: &EndpointId) -> bool {
|
|
self.active.lock().unwrap().contains(peer_id)
|
|
}
|
|
}
|
|
|
|
async fn run_coordinator(
|
|
room_state: Arc<dyn RecoveryRoom>,
|
|
active: Arc<Mutex<HashSet<EndpointId>>>,
|
|
mut rx: mpsc::Receiver<RecoveryCommand>,
|
|
) {
|
|
let mut entries: HashMap<EndpointId, RecoveryEntry> = HashMap::new();
|
|
|
|
loop {
|
|
// The shared active set is the authoritative cancellation gate. Prune
|
|
// here as well as on Cancel commands so a saturated command queue cannot
|
|
// leave an inactive, past-due entry spinning the timer loop.
|
|
let active_snapshot = active.lock().unwrap().clone();
|
|
entries.retain(|peer_id, _| active_snapshot.contains(peer_id));
|
|
let next_deadline = entries.values().map(|entry| entry.next_attempt).min();
|
|
let command = match next_deadline {
|
|
Some(deadline) => {
|
|
tokio::select! {
|
|
command = rx.recv() => command,
|
|
_ = tokio::time::sleep_until(deadline) => {
|
|
let now = Instant::now();
|
|
let active_snapshot = active.lock().unwrap().clone();
|
|
let due: Vec<(EndpointId, EndpointAddr)> = entries
|
|
.iter()
|
|
.filter(|(id, entry)| {
|
|
entry.next_attempt <= now && active_snapshot.contains(*id)
|
|
})
|
|
.map(|(id, entry)| (*id, entry.addr.clone()))
|
|
.collect();
|
|
|
|
if !due.is_empty() {
|
|
let addrs = due.iter().map(|(_, addr)| addr.clone()).collect();
|
|
if let Err(error) = room_state.rebootstrap_peers(addrs).await {
|
|
crate::log_msg(&format!(
|
|
"Background peer recovery attempt failed: {error}"
|
|
));
|
|
}
|
|
|
|
let scheduled_at = Instant::now();
|
|
for (peer_id, _) in due {
|
|
if !active.lock().unwrap().contains(&peer_id) {
|
|
entries.remove(&peer_id);
|
|
continue;
|
|
}
|
|
if let Some(entry) = entries.get_mut(&peer_id) {
|
|
entry.next_attempt = scheduled_at + recovery_delay(entry.attempt);
|
|
entry.attempt = entry.attempt.saturating_add(1);
|
|
}
|
|
}
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
None => rx.recv().await,
|
|
};
|
|
|
|
match command {
|
|
Some(RecoveryCommand::Start { peer_id, addr }) => {
|
|
if active.lock().unwrap().contains(&peer_id) {
|
|
entries.entry(peer_id).or_insert(RecoveryEntry {
|
|
addr,
|
|
attempt: 0,
|
|
next_attempt: Instant::now(),
|
|
});
|
|
}
|
|
}
|
|
Some(RecoveryCommand::Cancel(peer_id)) => {
|
|
entries.remove(&peer_id);
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use iroh::SecretKey;
|
|
|
|
struct RecordingRoom {
|
|
attempts: mpsc::UnboundedSender<Vec<EndpointAddr>>,
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl RecoveryRoom for RecordingRoom {
|
|
async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String> {
|
|
self.attempts.send(peers).map_err(|error| error.to_string())
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn retry_backoff_reaches_and_stays_at_sixty_seconds() {
|
|
let actual: Vec<u64> = (0..10)
|
|
.map(|attempt| recovery_delay(attempt).as_secs())
|
|
.collect();
|
|
assert_eq!(actual, vec![1, 2, 4, 8, 15, 30, 60, 60, 60, 60]);
|
|
}
|
|
|
|
#[test]
|
|
fn recovery_slots_are_deduplicated_and_cancel_immediately() {
|
|
let (tx, mut rx) = mpsc::channel(4);
|
|
let coordinator = RecoveryCoordinator {
|
|
tx,
|
|
active: Arc::new(Mutex::new(HashSet::new())),
|
|
};
|
|
let peer_id = SecretKey::generate().public();
|
|
|
|
assert!(coordinator.begin(peer_id));
|
|
assert!(
|
|
!coordinator.begin(peer_id),
|
|
"a peer gets only one recovery slot"
|
|
);
|
|
assert_eq!(
|
|
coordinator.activate(peer_id, EndpointAddr::from(peer_id)),
|
|
Ok(true)
|
|
);
|
|
assert!(matches!(
|
|
rx.try_recv(),
|
|
Ok(RecoveryCommand::Start { peer_id: id, .. }) if id == peer_id
|
|
));
|
|
|
|
coordinator.cancel(peer_id);
|
|
assert!(!coordinator.is_active(&peer_id));
|
|
assert!(matches!(
|
|
rx.try_recv(),
|
|
Ok(RecoveryCommand::Cancel(id)) if id == peer_id
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn coordinator_attempts_rebootstrap_immediately() {
|
|
let (attempts_tx, mut attempts_rx) = mpsc::unbounded_channel();
|
|
let (coordinator, task) = RecoveryCoordinator::spawn_inner(Arc::new(RecordingRoom {
|
|
attempts: attempts_tx,
|
|
}));
|
|
let peer_id = SecretKey::generate().public();
|
|
let addr = EndpointAddr::from(peer_id);
|
|
|
|
assert!(coordinator.begin(peer_id));
|
|
assert_eq!(coordinator.activate(peer_id, addr.clone()), Ok(true));
|
|
let attempted = tokio::time::timeout(Duration::from_secs(1), attempts_rx.recv())
|
|
.await
|
|
.expect("first recovery attempt should be immediate")
|
|
.expect("recording room remains subscribed");
|
|
assert_eq!(attempted, vec![addr]);
|
|
|
|
coordinator.cancel(peer_id);
|
|
task.abort();
|
|
}
|
|
}
|