task-supervisor 0.4.1

Supervisor for long-lived tokio tasks
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

mod common;

use std::sync::Arc;
use std::{sync::atomic::Ordering, time::Duration};

use tokio::time::pause;

use task_supervisor::{SupervisorBuilder, SupervisorHandle, TaskStatus};

use common::{FailingTask, HealthyTask};

fn supervisor_handle() -> SupervisorHandle {
    SupervisorBuilder::new()
        .with_health_check_interval(Duration::from_millis(50))
        .with_max_restart_attempts(3)
        .with_base_restart_delay(Duration::from_millis(100))
        .build()
        .run()
}

#[tokio::test]
async fn test_healthy_task_remains_healthy() {
    pause();
    let handle = supervisor_handle();

    let run_flag = Arc::new(std::sync::atomic::AtomicBool::new(true));
    let task = HealthyTask {
        run_flag: run_flag.clone(),
    };
    handle.add_task("task", task).unwrap();

    tokio::time::sleep(Duration::from_millis(500)).await;
    let status = handle.get_task_status("task").await.unwrap().unwrap();
    assert_eq!(status, TaskStatus::Healthy);

    run_flag.store(false, Ordering::SeqCst);
    tokio::time::sleep(Duration::from_millis(100)).await;
    let status = handle.get_task_status("task").await.unwrap().unwrap();
    assert_eq!(status, TaskStatus::Completed);
}

#[tokio::test]
async fn test_no_heartbeat_task_gets_restarted() {
    pause();
    let handle = supervisor_handle();

    let run_flag = Arc::new(std::sync::atomic::AtomicBool::new(true));
    let task = HealthyTask {
        run_flag: run_flag.clone(),
    };

    handle.add_task("task", task).unwrap();
    tokio::time::sleep(Duration::from_millis(500)).await;

    let status = handle.get_task_status("task").await.unwrap().unwrap();
    assert_eq!(status, TaskStatus::Healthy);

    run_flag.store(false, Ordering::SeqCst);
    tokio::time::sleep(Duration::from_millis(150)).await;

    let status = handle.get_task_status("task").await.unwrap().unwrap();
    assert_eq!(status, TaskStatus::Completed);
}

#[tokio::test]
async fn test_multiple_tasks() {
    pause();
    let handle = supervisor_handle();

    let healthy_run_flag = Arc::new(std::sync::atomic::AtomicBool::new(true));
    let healthy_task = HealthyTask {
        run_flag: healthy_run_flag.clone(),
    };
    handle.add_task("healthy_task", healthy_task).unwrap();

    let failing_task = FailingTask {
        run_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
    };
    handle
        .add_task("failing_task", failing_task.clone())
        .unwrap();

    let no_heartbeat_run_flag = Arc::new(std::sync::atomic::AtomicBool::new(true));
    let no_heartbeat_task = HealthyTask {
        run_flag: no_heartbeat_run_flag.clone(),
    };
    handle
        .add_task("no_heartbeat_task", no_heartbeat_task)
        .unwrap();

    tokio::time::sleep(Duration::from_millis(1500)).await;

    let healthy_status = handle
        .get_task_status("healthy_task")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(healthy_status, TaskStatus::Healthy);

    let failing_status = handle
        .get_task_status("failing_task")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(failing_status, TaskStatus::Dead);

    let no_heartbeat_status = handle
        .get_task_status("no_heartbeat_task")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(no_heartbeat_status, TaskStatus::Healthy);

    healthy_run_flag.store(false, Ordering::SeqCst);
    no_heartbeat_run_flag.store(false, Ordering::SeqCst);
    tokio::time::sleep(Duration::from_millis(100)).await;

    let healthy_status = handle
        .get_task_status("healthy_task")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(healthy_status, TaskStatus::Completed);

    let no_heartbeat_status = handle
        .get_task_status("no_heartbeat_task")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(no_heartbeat_status, TaskStatus::Completed);
}