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

mod common;

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

use tokio::time::{advance, pause};

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

use common::{CompletingTask, HealthyTask, ImmediateCompleteTask};

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

#[tokio::test]
async fn test_add_task_dynamically() {
    pause();

    let handle = supervisor_handle();

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

    tokio::time::sleep(Duration::from_millis(150)).await;
    let status = handle.get_task_status("task").await.unwrap().unwrap();
    assert_eq!(status, TaskStatus::Completed);
    assert_eq!(task.run_count.load(std::sync::atomic::Ordering::SeqCst), 1);
}

#[tokio::test]
async fn test_restart_task() {
    pause();

    let handle = supervisor_handle();

    let task = CompletingTask {
        run_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
    };

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

    handle.restart("task").unwrap();
    tokio::time::sleep(Duration::from_secs(3)).await;

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

    // Count should be 2: Initial + restart
    assert_eq!(task.run_count.load(std::sync::atomic::Ordering::SeqCst), 2);
}

#[tokio::test]
async fn test_kill_task() {
    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();

    advance(Duration::from_millis(50)).await;
    handle.kill_task("task").unwrap();
    advance(Duration::from_millis(100)).await;

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

#[tokio::test]
async fn test_supervisor_shutdown() {
    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();
    advance(Duration::from_millis(50)).await;

    handle.shutdown().unwrap();
    advance(Duration::from_millis(200)).await;

    let result = handle.get_task_status("task").await;
    assert!(matches!(result, Err(SupervisorHandleError::SendError)));
}

#[tokio::test]
async fn test_error_handling() {
    pause();

    let handle = supervisor_handle();

    handle.shutdown().unwrap();
    advance(Duration::from_millis(500)).await;

    let result = handle.add_task("task", ImmediateCompleteTask);
    assert!(matches!(result, Err(SupervisorHandleError::SendError)));

    let result = handle.get_task_status("non_existent").await;
    assert!(matches!(result, Err(SupervisorHandleError::SendError)));
}