task-supervisor 0.5.2

Supervisor for long-lived tokio tasks: restarts with backoff, graceful shutdown, runtime control
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

mod common;

use std::sync::atomic::Ordering;
use std::time::Duration;

use task_supervisor::{SupervisorBuilder, SupervisorError, SupervisorHandleError};
use tokio::time::{pause, Instant};

use common::{sleep_ms, Cooperative, Stubborn};

#[tokio::test]
async fn shutdown_lets_cooperative_tasks_finish() {
    pause();
    let a = Cooperative::default();
    let b = Cooperative::default();
    let handle = SupervisorBuilder::new()
        .with_task("a", a.clone())
        .with_task("b", b.clone())
        .spawn();

    handle.shutdown().await.unwrap();
    assert!(a.stopped.load(Ordering::SeqCst));
    assert!(b.stopped.load(Ordering::SeqCst));
}

#[tokio::test]
async fn shutdown_is_bounded_by_shutdown_timeout() {
    pause();
    let handle = SupervisorBuilder::new()
        .with_stop_timeout(Duration::from_millis(200))
        .with_task("stubborn", Stubborn)
        .spawn();

    let started = Instant::now();
    handle.shutdown().await.unwrap();
    let elapsed = started.elapsed();
    assert!(elapsed >= Duration::from_millis(200), "{elapsed:?}");
    assert!(elapsed < Duration::from_millis(300), "{elapsed:?}");
}

#[tokio::test]
async fn dropping_last_handle_shuts_down_gracefully() {
    pause();
    let task = Cooperative::default();
    let handle = SupervisorBuilder::new()
        .with_task("t", task.clone())
        .spawn();
    let clone = handle.clone();

    drop(handle);
    sleep_ms(10).await;
    assert!(
        !task.stopped.load(Ordering::SeqCst),
        "one handle still alive"
    );

    drop(clone);
    sleep_ms(10).await;
    assert!(task.stopped.load(Ordering::SeqCst));
}

/// Regression: dropping a clone used to send Shutdown.
#[tokio::test]
async fn dropping_one_clone_keeps_supervisor_running() {
    pause();
    let handle = SupervisorBuilder::new()
        .with_task("t", Cooperative::default())
        .spawn();

    {
        let _short_lived = handle.clone();
    }
    sleep_ms(10).await;
    assert!(handle.task_status("t").await.is_ok());
}

#[tokio::test]
async fn wait_can_be_called_repeatedly_and_concurrently() {
    pause();
    let handle = SupervisorBuilder::new()
        .with_task("t", Cooperative::default())
        .spawn();

    let h = handle.clone();
    tokio::spawn(async move {
        sleep_ms(5).await;
        h.shutdown().await.unwrap();
    });

    let (r1, r2) = tokio::join!(handle.wait(), handle.wait());
    assert!(r1.is_ok() && r2.is_ok());
    assert!(handle.wait().await.is_ok());
}

/// Dropping the runtime under the supervisor is reported, not mistaken for a
/// clean exit.
#[test]
fn wait_reports_runtime_abort() {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();
    let handle =
        runtime.block_on(async { SupervisorBuilder::new().with_task("t", Stubborn).spawn() });
    drop(runtime);

    let other = tokio::runtime::Builder::new_current_thread()
        .build()
        .unwrap();
    assert_eq!(other.block_on(handle.wait()), Err(SupervisorError::Aborted));
}

/// Regression: a kill still waiting for its run to exit when shutdown starts
/// got `Closed`, although the task does end up dead.
#[tokio::test]
async fn kill_pending_at_shutdown_resolves_ok() {
    pause();
    let handle = SupervisorBuilder::new()
        .with_stop_timeout(Duration::from_millis(200))
        .with_task("t", Stubborn)
        .spawn();

    let h = handle.clone();
    let kill = tokio::spawn(async move { h.kill_task("t").await });
    sleep_ms(1).await;
    handle.shutdown().await.unwrap();

    assert_eq!(kill.await.unwrap(), Ok(()));
}

/// A restart interrupted by shutdown never produces a new run: it reports
/// `Closed`.
#[tokio::test]
async fn restart_pending_at_shutdown_reports_closed() {
    pause();
    let handle = SupervisorBuilder::new()
        .with_stop_timeout(Duration::from_millis(200))
        .with_task("t", Stubborn)
        .spawn();

    let h = handle.clone();
    let restart = tokio::spawn(async move { h.restart_task("t").await });
    sleep_ms(1).await;
    handle.shutdown().await.unwrap();

    assert_eq!(restart.await.unwrap(), Err(SupervisorHandleError::Closed));
}