Skip to main content

syncweb_core/daemon/
supervisor.rs

1use std::time::Duration;
2use std::{
3    collections::HashMap,
4    sync::{Arc, Mutex},
5};
6
7use iroh_docs::NamespaceId;
8use n0_future::StreamExt;
9use tokio::{
10    sync::{broadcast, oneshot},
11    time::sleep,
12};
13
14use crate::{
15    daemon::current_timestamp,
16    error::{Result, SyncwebError},
17    filter::FilterEngine,
18    sync::{ActiveSession, IntentControl, IntentHandle, SubscribeParams, SyncEngine, SyncEvent},
19};
20
21/// A synchronization intent together with the information needed to restart it.
22#[non_exhaustive]
23pub struct SupervisedIntent {
24    pub namespace: NamespaceId,
25    pub handle: Option<IntentHandle>,
26    pub retry_count: u32,
27    pub last_error: Option<String>,
28    pub last_started_at: Option<u64>,
29    session: Option<ActiveSession>,
30}
31
32impl std::fmt::Debug for SupervisedIntent {
33    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        formatter
35            .debug_struct("SupervisedIntent")
36            .field("namespace", &self.namespace)
37            .field("handle", &self.handle)
38            .field("retry_count", &self.retry_count)
39            .field("last_error", &self.last_error)
40            .field("last_started_at", &self.last_started_at)
41            .finish_non_exhaustive()
42    }
43}
44
45/// Starts synchronization intents and restarts failed intents with bounded
46/// exponential backoff.
47#[derive(Clone, Copy, Debug)]
48pub struct IntentSupervisor {
49    max_retries: u32,
50    backoff_base: Duration,
51    backoff_max: Duration,
52}
53
54/// Shared controls for intents currently owned by the daemon.
55pub type IntentControls = Arc<Mutex<HashMap<NamespaceId, IntentControl>>>;
56
57/// Internal startup options for a supervised daemon intent.
58pub struct SupervisionOptions {
59    controls: Option<IntentControls>,
60    filter: Option<FilterEngine>,
61    ready: Option<oneshot::Sender<()>>,
62}
63
64impl SupervisionOptions {
65    pub const fn with_ready(
66        controls: IntentControls,
67        filter: Option<FilterEngine>,
68        ready: oneshot::Sender<()>,
69    ) -> Self {
70        Self {
71            controls: Some(controls),
72            filter,
73            ready: Some(ready),
74        }
75    }
76}
77
78impl IntentSupervisor {
79    #[must_use]
80    pub const fn new(max_retries: u32, backoff_base: Duration, backoff_max: Duration) -> Self {
81        Self {
82            max_retries,
83            backoff_base,
84            backoff_max,
85        }
86    }
87
88    /// Start one continuous intent and record any startup failure.
89    pub async fn run_intent(
90        &self,
91        sync: &SyncEngine,
92        namespace: NamespaceId,
93        params: SubscribeParams,
94    ) -> SupervisedIntent {
95        self.run_intent_with_filter(sync, namespace, params, None).await
96    }
97
98    /// Start one filtered synchronization intent.
99    pub async fn run_intent_with_filter(
100        &self,
101        sync: &SyncEngine,
102        namespace: NamespaceId,
103        params: SubscribeParams,
104        filter_engine: Option<FilterEngine>,
105    ) -> SupervisedIntent {
106        let result = match filter_engine {
107            Some(engine) => {
108                sync.sync_with_filter(namespace, crate::sync::SessionMode::Continuous, params, engine)
109                    .await
110            }
111            None => sync.subscribe(namespace, params).await,
112        };
113        match result {
114            Ok(handle) => {
115                let session = ActiveSession::register(namespace, handle.cancel_sender());
116                SupervisedIntent {
117                    namespace,
118                    handle: Some(handle),
119                    retry_count: 0,
120                    last_error: None,
121                    last_started_at: Some(current_timestamp()),
122                    session: Some(session),
123                }
124            }
125            Err(error) => SupervisedIntent {
126                namespace,
127                handle: None,
128                retry_count: 0,
129                last_error: Some(error.to_string()),
130                last_started_at: None,
131                session: None,
132            },
133        }
134    }
135
136    /// Supervise one intent until shutdown or until its retry budget is spent.
137    ///
138    /// A normal `Finished` or `Cancelled` event ends supervision. A `Failed`
139    /// event, an unexpectedly closed stream, or a startup error consumes one
140    /// retry and restarts the intent after backoff.
141    ///
142    /// # Errors
143    ///
144    /// Returns an error if supervision infrastructure fails while waiting for
145    /// shutdown or retrying an intent.
146    pub async fn supervise(
147        &self,
148        sync: &SyncEngine,
149        namespace: NamespaceId,
150        params: SubscribeParams,
151        shutdown: broadcast::Receiver<()>,
152    ) -> Result<SupervisedIntent> {
153        self.supervise_inner(
154            sync,
155            namespace,
156            params,
157            shutdown,
158            SupervisionOptions {
159                controls: None,
160                filter: None,
161                ready: None,
162            },
163        )
164        .await
165    }
166
167    /// Supervise an intent while exposing pause/resume controls to its owner.
168    ///
169    /// # Errors
170    ///
171    /// Returns an error if the shared control registry is poisoned or the
172    /// intent supervision infrastructure cannot wait for shutdown.
173    pub async fn supervise_with_controls(
174        &self,
175        sync: &SyncEngine,
176        namespace: NamespaceId,
177        params: SubscribeParams,
178        shutdown: broadcast::Receiver<()>,
179        controls: IntentControls,
180        filter: Option<FilterEngine>,
181    ) -> Result<SupervisedIntent> {
182        self.supervise_inner(
183            sync,
184            namespace,
185            params,
186            shutdown,
187            SupervisionOptions {
188                controls: Some(controls),
189                filter,
190                ready: None,
191            },
192        )
193        .await
194    }
195
196    pub(crate) async fn supervise_with_controls_and_ready(
197        &self,
198        sync: &SyncEngine,
199        namespace: NamespaceId,
200        params: SubscribeParams,
201        shutdown: broadcast::Receiver<()>,
202        options: SupervisionOptions,
203    ) -> Result<SupervisedIntent> {
204        self.supervise_inner(sync, namespace, params, shutdown, options).await
205    }
206
207    async fn supervise_inner(
208        &self,
209        sync: &SyncEngine,
210        namespace: NamespaceId,
211        params: SubscribeParams,
212        mut shutdown: broadcast::Receiver<()>,
213        mut options: SupervisionOptions,
214    ) -> Result<SupervisedIntent> {
215        let mut supervised = self
216            .run_intent_with_filter(sync, namespace, params.clone(), options.filter.clone())
217            .await;
218        let mut ready = options.ready.take();
219
220        loop {
221            let Some(mut handle) = supervised.handle.take() else {
222                if let Some(sender) = ready.take() {
223                    let _ = sender.send(());
224                }
225                Self::remove_control(options.controls.as_ref(), namespace)?;
226                if supervised.retry_count >= self.max_retries {
227                    return Ok(supervised);
228                }
229                let retry_number = supervised.retry_count.saturating_add(1);
230                if !self
231                    .wait_for_retry(self.backoff_delay(retry_number), &mut shutdown)
232                    .await
233                {
234                    return Ok(supervised);
235                }
236                supervised.retry_count = retry_number;
237                supervised = self
238                    .run_intent_with_filter(sync, namespace, params.clone(), options.filter.clone())
239                    .await;
240                supervised.retry_count = retry_number;
241                continue;
242            };
243
244            Self::insert_control(options.controls.as_ref(), namespace, handle.control())?;
245            if let Some(sender) = ready.take() {
246                let _ = sender.send(());
247            }
248            let session = supervised.session.take();
249            let event = loop {
250                tokio::select! {
251                    shutdown_result = shutdown.recv() => {
252                        match shutdown_result {
253                            Ok(()) | Err(broadcast::error::RecvError::Closed) => {
254                                if let Err(error) = handle.cancel() {
255                                    supervised.last_error = Some(format!("failed to cancel intent: {error}"));
256                                }
257                                Self::remove_control(options.controls.as_ref(), namespace)?;
258                                drop(session);
259                                supervised.handle = None;
260                                return Ok(supervised);
261                            }
262                            Err(broadcast::error::RecvError::Lagged(_)) => {}
263                        }
264                    }
265                    next_event = handle.next() => break next_event,
266                }
267            };
268
269            Self::remove_control(options.controls.as_ref(), namespace)?;
270            match event {
271                Some(SyncEvent::Started) => {
272                    supervised.retry_count = 0;
273                    supervised.last_error = None;
274                    supervised.last_started_at = Some(current_timestamp());
275                    supervised.handle = Some(handle);
276                    supervised.session = session;
277                }
278                Some(SyncEvent::Cancelled | SyncEvent::Finished) => {
279                    supervised.handle = None;
280                    drop(session);
281                    return Ok(supervised);
282                }
283                Some(SyncEvent::Failed(error)) => {
284                    supervised.last_error = Some(error);
285                    supervised.handle = None;
286                    drop(session);
287                }
288                Some(SyncEvent::Progress { .. } | SyncEvent::Stats(_) | SyncEvent::Paused | SyncEvent::Resumed) => {
289                    supervised.handle = Some(handle);
290                    supervised.session = session;
291                }
292                None => {
293                    supervised.last_error = Some("synchronization intent ended unexpectedly".to_owned());
294                    supervised.handle = None;
295                    drop(session);
296                }
297            }
298        }
299    }
300
301    fn insert_control(controls: Option<&IntentControls>, namespace: NamespaceId, control: IntentControl) -> Result<()> {
302        let Some(control_registry) = controls else {
303            return Ok(());
304        };
305        control_registry
306            .lock()
307            .map_err(|error| SyncwebError::operation("daemon intent control mutex is poisoned", error))?
308            .insert(namespace, control);
309        Ok(())
310    }
311
312    fn remove_control(controls: Option<&IntentControls>, namespace: NamespaceId) -> Result<()> {
313        let Some(control_registry) = controls else {
314            return Ok(());
315        };
316        control_registry
317            .lock()
318            .map_err(|error| SyncwebError::operation("daemon intent control mutex is poisoned", error))?
319            .remove(&namespace);
320        Ok(())
321    }
322
323    /// Return the delay for a one-based retry number.
324    #[must_use]
325    pub fn backoff_delay(&self, retry_number: u32) -> Duration {
326        if retry_number == 0 {
327            return Duration::ZERO;
328        }
329        let mut delay = self.backoff_base;
330        for _ in 1..retry_number {
331            delay = match delay.checked_mul(2) {
332                Some(value) if value <= self.backoff_max => value,
333                _ => self.backoff_max,
334            };
335            if delay == self.backoff_max {
336                break;
337            }
338        }
339        delay.min(self.backoff_max)
340    }
341
342    async fn wait_for_retry(&self, delay: Duration, shutdown: &mut broadcast::Receiver<()>) -> bool {
343        tokio::select! {
344            () = sleep(delay) => true,
345            shutdown_result = shutdown.recv() => {
346                !matches!(shutdown_result, Ok(()) | Err(broadcast::error::RecvError::Closed))
347            }
348        }
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    #[test]
357    fn backoff_delay_starts_at_base() {
358        let supervisor = IntentSupervisor::new(3, Duration::from_secs(1), Duration::from_mins(1));
359        assert_eq!(supervisor.backoff_delay(1), Duration::from_secs(1));
360    }
361
362    #[test]
363    fn backoff_delay_doubles() {
364        let supervisor = IntentSupervisor::new(3, Duration::from_secs(1), Duration::from_mins(1));
365        assert_eq!(supervisor.backoff_delay(2), Duration::from_secs(2));
366        assert_eq!(supervisor.backoff_delay(3), Duration::from_secs(4));
367    }
368
369    #[test]
370    fn backoff_delay_capped() {
371        let supervisor = IntentSupervisor::new(3, Duration::from_secs(1), Duration::from_secs(3));
372        assert_eq!(supervisor.backoff_delay(2), Duration::from_secs(2));
373        assert_eq!(supervisor.backoff_delay(3), Duration::from_secs(3));
374        assert_eq!(supervisor.backoff_delay(10), Duration::from_secs(3));
375    }
376
377    #[test]
378    fn backoff_delay_zero_retry() {
379        let supervisor = IntentSupervisor::new(3, Duration::from_secs(1), Duration::from_mins(1));
380        assert_eq!(supervisor.backoff_delay(0), Duration::ZERO);
381    }
382
383    #[test]
384    fn backoff_delay_zero_base() {
385        let supervisor = IntentSupervisor::new(3, Duration::ZERO, Duration::from_mins(1));
386        assert_eq!(supervisor.backoff_delay(5), Duration::ZERO);
387    }
388
389    #[test]
390    fn auto_backoff_delay_respects_max_retries() {
391        let supervisor = IntentSupervisor::new(5, Duration::from_secs(2), Duration::from_secs(20));
392        let delays: Vec<_> = (1..=8).map(|n| supervisor.backoff_delay(n)).collect();
393        let expected = [
394            Duration::from_secs(2),
395            Duration::from_secs(4),
396            Duration::from_secs(8),
397            Duration::from_secs(16),
398            Duration::from_secs(20),
399            Duration::from_secs(20),
400            Duration::from_secs(20),
401            Duration::from_secs(20),
402        ];
403        for (delay, want) in delays.iter().zip(expected.iter()) {
404            assert_eq!(delay, want);
405        }
406    }
407
408    #[test]
409    fn backoff_delay_starting_value_equals_base() {
410        let supervisor = IntentSupervisor::new(1, Duration::from_millis(500), Duration::from_secs(5));
411        assert_eq!(supervisor.backoff_delay(1), Duration::from_millis(500));
412    }
413}