syncweb_core/daemon/

supervisor.rs

use std::time::{Duration, Instant};
use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
};

use iroh_docs::NamespaceId;
use n0_future::StreamExt;
use tokio::{
    sync::{broadcast, oneshot},
    time::sleep,
};

use crate::{
    daemon::current_timestamp,
    error::{Result, SyncwebError},
    filter::FilterEngine,
    net::NetworkLogger,
    sync::{ActiveSession, IntentControl, IntentHandle, SubscribeParams, SyncEngine, SyncEvent},
};

/// A synchronization intent together with the information needed to restart it.
#[non_exhaustive]
pub struct SupervisedIntent {
    pub namespace: NamespaceId,
    pub handle: Option<IntentHandle>,
    pub retry_count: u32,
    pub last_error: Option<String>,
    pub last_started_at: Option<u64>,
    /// Total bytes received by this intent since it was first started.
    pub bytes_transferred: u64,
    /// Total files received by this intent since it was first started.
    pub files_transferred: u64,
    session: Option<ActiveSession>,
}

/// Live bandwidth accounting for a supervised intent.
///
/// Records the delta between successive [`SyncEvent::Stats`] reports as
/// download events so `stats network` reflects transfers while the daemon is
/// still running (a continuous intent only returns on shutdown).
#[derive(Clone, Debug)]
pub struct BandwidthAccountant {
    logger: NetworkLogger,
    network_id: Option<String>,
    folder: String,
    recorded: Arc<AtomicU64>,
}

impl BandwidthAccountant {
    #[must_use]
    pub fn new(logger: NetworkLogger, network_id: Option<String>, folder: impl Into<String>) -> Self {
        Self {
            logger,
            network_id,
            folder: folder.into(),
            recorded: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Record any newly observed bytes as a download event.
    pub fn observe(&self, total_bytes: u64, files: u64) {
        let previous = self.recorded.load(Ordering::Relaxed);
        if total_bytes > previous
            && let Err(error) = self.logger.record_bandwidth_download(
                total_bytes.saturating_sub(previous),
                files,
                Some(&self.folder),
                None,
                self.network_id.as_deref(),
            )
        {
            tracing::warn!(%error, "failed to record bandwidth event");
        }
        self.recorded.store(total_bytes.max(previous), Ordering::Relaxed);
    }

    /// Total bytes recorded so far.
    #[must_use]
    pub fn total(&self) -> u64 {
        self.recorded.load(Ordering::Relaxed)
    }
}

impl std::fmt::Debug for SupervisedIntent {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("SupervisedIntent")
            .field("namespace", &self.namespace)
            .field("handle", &self.handle)
            .field("retry_count", &self.retry_count)
            .field("last_error", &self.last_error)
            .field("last_started_at", &self.last_started_at)
            .field("bytes_transferred", &self.bytes_transferred)
            .finish_non_exhaustive()
    }
}

/// Default retry window measured from an intent's first attempt.
///
/// Three months is the hard default; configure
/// [`DaemonConfig::retry_window`](crate::daemon::DaemonConfig::retry_window) to
/// override it.
pub const DEFAULT_RETRY_WINDOW: Duration = Duration::from_hours(2_160);

/// Starts synchronization intents and restarts failed intents with bounded
/// exponential backoff until a deadline measured from the first attempt.
#[derive(Clone, Copy, Debug)]
pub struct IntentSupervisor {
    retry_window: Duration,
    backoff_base: Duration,
    backoff_max: Duration,
}

/// Shared controls for intents currently owned by the daemon.
pub type IntentControls = Arc<Mutex<HashMap<NamespaceId, IntentControl>>>;

/// Internal startup options for a supervised daemon intent.
pub struct SupervisionOptions {
    controls: Option<IntentControls>,
    filter: Option<FilterEngine>,
    ready: Option<oneshot::Sender<()>>,
    bandwidth: Option<BandwidthAccountant>,
}

impl SupervisionOptions {
    pub const fn with_ready(
        controls: IntentControls,
        filter: Option<FilterEngine>,
        ready: oneshot::Sender<()>,
    ) -> Self {
        Self {
            controls: Some(controls),
            filter,
            ready: Some(ready),
            bandwidth: None,
        }
    }

    /// Attach a live bandwidth accountant to the supervised intent.
    #[must_use]
    pub fn with_bandwidth(mut self, bandwidth: BandwidthAccountant) -> Self {
        self.bandwidth = Some(bandwidth);
        self
    }
}

impl IntentSupervisor {
    #[must_use]
    pub const fn new(retry_window: Duration, backoff_base: Duration, backoff_max: Duration) -> Self {
        Self {
            retry_window,
            backoff_base,
            backoff_max,
        }
    }

    /// Start one continuous intent and record any startup failure.
    pub async fn run_intent(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
    ) -> SupervisedIntent {
        self.run_intent_with_filter(sync, namespace, params, None).await
    }

    /// Start one filtered synchronization intent.
    pub async fn run_intent_with_filter(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
        filter_engine: Option<FilterEngine>,
    ) -> SupervisedIntent {
        let result = match filter_engine {
            Some(engine) => {
                sync.sync_with_filter(namespace, crate::sync::SessionMode::Continuous, params, engine)
                    .await
            }
            None => sync.subscribe(namespace, params).await,
        };
        match result {
            Ok(handle) => {
                let session = ActiveSession::register(namespace, handle.cancel_sender());
                SupervisedIntent {
                    namespace,
                    handle: Some(handle),
                    retry_count: 0,
                    last_error: None,
                    last_started_at: Some(current_timestamp()),
                    bytes_transferred: 0,
                    files_transferred: 0,
                    session: Some(session),
                }
            }
            Err(error) => SupervisedIntent {
                namespace,
                handle: None,
                retry_count: 0,
                last_error: Some(error.to_string()),
                last_started_at: None,
                bytes_transferred: 0,
                files_transferred: 0,
                session: None,
            },
        }
    }

    /// Supervise one intent until shutdown or until its retry budget is spent.
    ///
    /// A normal `Finished` or `Cancelled` event ends supervision. A `Failed`
    /// event, an unexpectedly closed stream, or a startup error consumes one
    /// retry and restarts the intent after backoff.
    ///
    /// # Errors
    ///
    /// Returns an error if supervision infrastructure fails while waiting for
    /// shutdown or retrying an intent.
    pub async fn supervise(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
        shutdown: broadcast::Receiver<()>,
    ) -> Result<SupervisedIntent> {
        self.supervise_inner(
            sync,
            namespace,
            params,
            shutdown,
            SupervisionOptions {
                controls: None,
                filter: None,
                ready: None,
                bandwidth: None,
            },
        )
        .await
    }

    /// Supervise an intent while exposing pause/resume controls to its owner.
    ///
    /// # Errors
    ///
    /// Returns an error if the shared control registry is poisoned or the
    /// intent supervision infrastructure cannot wait for shutdown.
    pub async fn supervise_with_controls(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
        shutdown: broadcast::Receiver<()>,
        controls: IntentControls,
        filter: Option<FilterEngine>,
    ) -> Result<SupervisedIntent> {
        self.supervise_inner(
            sync,
            namespace,
            params,
            shutdown,
            SupervisionOptions {
                controls: Some(controls),
                filter,
                ready: None,
                bandwidth: None,
            },
        )
        .await
    }

    pub(crate) async fn supervise_with_controls_and_ready(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
        shutdown: broadcast::Receiver<()>,
        options: SupervisionOptions,
    ) -> Result<SupervisedIntent> {
        self.supervise_inner(sync, namespace, params, shutdown, options).await
    }

    async fn supervise_inner(
        &self,
        sync: &SyncEngine,
        namespace: NamespaceId,
        params: SubscribeParams,
        mut shutdown: broadcast::Receiver<()>,
        mut options: SupervisionOptions,
    ) -> Result<SupervisedIntent> {
        let mut supervised = self
            .run_intent_with_filter(sync, namespace, params.clone(), options.filter.clone())
            .await;
        let mut ready = options.ready.take();
        let first_attempt = Instant::now();

        loop {
            let Some(mut handle) = supervised.handle.take() else {
                if let Some(sender) = ready.take() {
                    let _ = sender.send(());
                }
                Self::remove_control(options.controls.as_ref(), namespace)?;
                if self.retry_window_exhausted(first_attempt) {
                    tracing::warn!(
                        %namespace,
                        retry_count = supervised.retry_count,
                        window_secs = self.retry_window.as_secs(),
                        "retry window elapsed; giving up on supervised intent"
                    );
                    return Ok(supervised);
                }
                let retry_number = supervised.retry_count.saturating_add(1);
                if !self
                    .wait_for_retry(self.backoff_delay(retry_number), &mut shutdown)
                    .await
                {
                    return Ok(supervised);
                }
                supervised = self
                    .run_intent_with_filter(sync, namespace, params.clone(), options.filter.clone())
                    .await;
                supervised.retry_count = retry_number;
                continue;
            };

            Self::insert_control(options.controls.as_ref(), namespace, handle.control())?;
            if let Some(sender) = ready.take() {
                let _ = sender.send(());
            }
            let session = supervised.session.take();
            let event = loop {
                tokio::select! {
                    shutdown_result = shutdown.recv() => {
                        match shutdown_result {
                            Ok(()) | Err(broadcast::error::RecvError::Closed) => {
                                if let Err(error) = handle.cancel() {
                                    supervised.last_error = Some(format!("failed to cancel intent: {error}"));
                                }
                                Self::remove_control(options.controls.as_ref(), namespace)?;
                                drop(session);
                                supervised.handle = None;
                                return Ok(supervised);
                            }
                            Err(broadcast::error::RecvError::Lagged(_)) => {}
                        }
                    }
                    next_event = handle.next() => break next_event,
                }
            };

            Self::remove_control(options.controls.as_ref(), namespace)?;
            match event {
                Some(SyncEvent::Started) => {
                    supervised.last_error = None;
                    supervised.last_started_at = Some(current_timestamp());
                    supervised.handle = Some(handle);
                    supervised.session = session;
                }
                Some(SyncEvent::Cancelled | SyncEvent::Finished) => {
                    supervised.handle = None;
                    drop(session);
                    return Ok(supervised);
                }
                Some(SyncEvent::Failed(error)) => {
                    supervised.last_error = Some(error);
                    supervised.handle = None;
                    drop(session);
                }
                Some(SyncEvent::Progress { .. } | SyncEvent::Paused | SyncEvent::Resumed) => {
                    supervised.handle = Some(handle);
                    supervised.session = session;
                }
                Some(SyncEvent::Stats(stats)) => {
                    supervised.bytes_transferred = stats.bytes_transferred;
                    if let Some(accountant) = options.bandwidth.as_ref() {
                        accountant.observe(stats.bytes_transferred, 0);
                    }
                    supervised.handle = Some(handle);
                    supervised.session = session;
                }
                None => {
                    supervised.last_error = Some("synchronization intent ended unexpectedly".to_owned());
                    supervised.handle = None;
                    drop(session);
                }
            }
        }
    }

    fn insert_control(controls: Option<&IntentControls>, namespace: NamespaceId, control: IntentControl) -> Result<()> {
        let Some(control_registry) = controls else {
            return Ok(());
        };
        control_registry
            .lock()
            .map_err(|error| SyncwebError::operation("daemon intent control mutex is poisoned", error))?
            .insert(namespace, control);
        Ok(())
    }

    fn remove_control(controls: Option<&IntentControls>, namespace: NamespaceId) -> Result<()> {
        let Some(control_registry) = controls else {
            return Ok(());
        };
        control_registry
            .lock()
            .map_err(|error| SyncwebError::operation("daemon intent control mutex is poisoned", error))?
            .remove(&namespace);
        Ok(())
    }

    /// Return whether the retry window measured from `first_attempt` has elapsed.
    #[must_use]
    pub fn retry_window_exhausted(&self, first_attempt: Instant) -> bool {
        first_attempt.elapsed() >= self.retry_window
    }

    /// Return the delay for a one-based retry number.
    #[must_use]
    pub fn backoff_delay(&self, retry_number: u32) -> Duration {
        const MIN_BACKOFF: Duration = Duration::from_millis(100);
        if retry_number == 0 {
            return Duration::ZERO;
        }
        let mut delay = self.backoff_base.max(MIN_BACKOFF);
        for _ in 1..retry_number {
            delay = match delay.checked_mul(2) {
                Some(value) if value <= self.backoff_max => value,
                _ => self.backoff_max,
            };
            if delay == self.backoff_max {
                break;
            }
        }
        delay.min(self.backoff_max)
    }

    async fn wait_for_retry(&self, delay: Duration, shutdown: &mut broadcast::Receiver<()>) -> bool {
        tokio::select! {
            () = sleep(delay) => true,
            shutdown_result = shutdown.recv() => {
                !matches!(shutdown_result, Ok(()) | Err(broadcast::error::RecvError::Closed))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backoff_delay_starts_at_base() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_secs(1),
            Duration::from_mins(1),
        );
        assert_eq!(supervisor.backoff_delay(1), Duration::from_secs(1));
    }

    #[test]
    fn backoff_delay_doubles() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_secs(1),
            Duration::from_mins(1),
        );
        assert_eq!(supervisor.backoff_delay(2), Duration::from_secs(2));
        assert_eq!(supervisor.backoff_delay(3), Duration::from_secs(4));
    }

    #[test]
    fn backoff_delay_capped() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_secs(1),
            Duration::from_secs(3),
        );
        assert_eq!(supervisor.backoff_delay(2), Duration::from_secs(2));
        assert_eq!(supervisor.backoff_delay(3), Duration::from_secs(3));
        assert_eq!(supervisor.backoff_delay(10), Duration::from_secs(3));
    }

    #[test]
    fn backoff_delay_zero_retry() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_secs(1),
            Duration::from_mins(1),
        );
        assert_eq!(supervisor.backoff_delay(0), Duration::ZERO);
    }

    #[test]
    fn backoff_delay_zero_base() {
        let supervisor = IntentSupervisor::new(Duration::from_hours(2_160), Duration::ZERO, Duration::from_mins(1));
        assert_eq!(supervisor.backoff_delay(1), Duration::from_millis(100));
        assert_eq!(supervisor.backoff_delay(5), Duration::from_millis(1600));
    }

    #[test]
    fn auto_backoff_delay_respects_max_retries() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_secs(2),
            Duration::from_secs(20),
        );
        let delays: Vec<_> = (1..=8).map(|n| supervisor.backoff_delay(n)).collect();
        let expected = [
            Duration::from_secs(2),
            Duration::from_secs(4),
            Duration::from_secs(8),
            Duration::from_secs(16),
            Duration::from_secs(20),
            Duration::from_secs(20),
            Duration::from_secs(20),
            Duration::from_secs(20),
        ];
        for (delay, want) in delays.iter().zip(expected.iter()) {
            assert_eq!(delay, want);
        }
    }

    #[test]
    fn backoff_delay_starting_value_equals_base() {
        let supervisor = IntentSupervisor::new(
            Duration::from_hours(2_160),
            Duration::from_millis(500),
            Duration::from_secs(5),
        );
        assert_eq!(supervisor.backoff_delay(1), Duration::from_millis(500));
    }

    #[test]
    fn retry_window_is_measured_from_first_attempt() {
        let supervisor = IntentSupervisor::new(Duration::from_mins(1), Duration::from_secs(1), Duration::from_mins(1));
        let now = Instant::now();
        assert!(!supervisor.retry_window_exhausted(now));
        let started_ago = now.checked_sub(Duration::from_secs(61)).expect("instant underflow");
        assert!(supervisor.retry_window_exhausted(started_ago));
        assert_eq!(DEFAULT_RETRY_WINDOW, Duration::from_hours(2_160));
    }
}