use std::time::Duration;
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
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,
sync::{ActiveSession, IntentControl, IntentHandle, SubscribeParams, SyncEngine, SyncEvent},
};
#[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>,
session: Option<ActiveSession>,
}
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)
.finish_non_exhaustive()
}
}
#[derive(Clone, Copy, Debug)]
pub struct IntentSupervisor {
max_retries: u32,
backoff_base: Duration,
backoff_max: Duration,
}
pub type IntentControls = Arc<Mutex<HashMap<NamespaceId, IntentControl>>>;
pub struct SupervisionOptions {
controls: Option<IntentControls>,
filter: Option<FilterEngine>,
ready: Option<oneshot::Sender<()>>,
}
impl SupervisionOptions {
pub const fn with_ready(
controls: IntentControls,
filter: Option<FilterEngine>,
ready: oneshot::Sender<()>,
) -> Self {
Self {
controls: Some(controls),
filter,
ready: Some(ready),
}
}
}
impl IntentSupervisor {
#[must_use]
pub const fn new(max_retries: u32, backoff_base: Duration, backoff_max: Duration) -> Self {
Self {
max_retries,
backoff_base,
backoff_max,
}
}
pub async fn run_intent(
&self,
sync: &SyncEngine,
namespace: NamespaceId,
params: SubscribeParams,
) -> SupervisedIntent {
self.run_intent_with_filter(sync, namespace, params, None).await
}
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()),
session: Some(session),
}
}
Err(error) => SupervisedIntent {
namespace,
handle: None,
retry_count: 0,
last_error: Some(error.to_string()),
last_started_at: None,
session: None,
},
}
}
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,
},
)
.await
}
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,
},
)
.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();
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 supervised.retry_count >= self.max_retries {
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::Stats(_) | SyncEvent::Paused | SyncEvent::Resumed) => {
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(())
}
#[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(3, 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(3, 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(3, 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(3, 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(3, 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(5, 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(1, Duration::from_millis(500), Duration::from_secs(5));
assert_eq!(supervisor.backoff_delay(1), Duration::from_millis(500));
}
}