ckb-stop-handler 1.1.1

Stop handler utilities for graceful shutdown
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

use crate::{
    broadcast_exit_signals, new_crossbeam_exit_rx, new_tokio_exit_rx, register_thread,
    wait_all_ckb_services_exit,
};
use ckb_async_runtime::{Handle, new_global_runtime};
use ckb_channel::select;
use rand::Rng;
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};
use std::time::Duration;
use tokio_util::sync::CancellationToken;

fn send_ctrlc_later(duration: Duration) {
    std::thread::spawn(move || {
        std::thread::sleep(duration);

        // send SIGINT to myself on Linux and MacOS platform
        unsafe {
            libc::raise(libc::SIGINT);
            println!("[ $$ sent SIGINT to myself $$ ]");
        }
    });
}

#[derive(Default)]
struct TestStopMemo {
    spawned_threads_count: Arc<AtomicI64>,
    stopped_threads_count: Arc<AtomicI64>,

    spawned_tokio_task_count: Arc<AtomicI64>,
    stopped_tokio_task_count: Arc<AtomicI64>,
}

impl TestStopMemo {
    fn start_many_threads(&self) {
        for i in 0..rand::thread_rng().gen_range(3..7) {
            let join = std::thread::spawn({
                let stopped_threads_count = Arc::clone(&self.stopped_threads_count);
                move || {
                    let ticker = ckb_channel::tick(Duration::from_millis(500));
                    let deadline = ckb_channel::after(Duration::from_millis(
                        (rand::thread_rng().gen_range(1.0..5.0) * 1000.0) as u64,
                    ));

                    let stop = new_crossbeam_exit_rx();

                    loop {
                        select! {
                            recv(ticker) -> _ => {
                                println!("thread {} received tick signal", i);
                            },
                            recv(stop) -> _ => {
                                println!("thread {} received crossbeam exit signal", i);
                                stopped_threads_count.fetch_add(1, Ordering::SeqCst);
                                return;
                            },
                            recv(deadline) -> _ =>{
                                println!("thread {} finish its job", i);
                                stopped_threads_count.fetch_add(1, Ordering::SeqCst);
                                return
                            }
                        }
                    }
                }
            });

            self.spawned_threads_count.fetch_add(1, Ordering::SeqCst);
            register_thread(&format!("test thread {}", i), join);
        }
    }

    fn start_many_tokio_tasks(&self, handle: &Handle) {
        for i in 0..rand::thread_rng().gen_range(3..7) {
            let stop: CancellationToken = new_tokio_exit_rx();

            handle.spawn({
                let spawned_tokio_task_count = Arc::clone(&self.spawned_tokio_task_count);
                let stopped_tokio_task_count = Arc::clone(&self.stopped_tokio_task_count);
                async move {
                    spawned_tokio_task_count.fetch_add(1, Ordering::SeqCst);

                    let mut interval = tokio::time::interval(Duration::from_millis(500));

                    let duration = Duration::from_millis(
                        (rand::thread_rng().gen_range(1.0..5.0) * 1000.0) as u64,
                    );
                    let deadline = tokio::time::sleep(duration);
                    tokio::pin!(deadline);

                    loop {
                        tokio::select! {
                            _ = &mut deadline =>{
                                println!("tokio task {} finish its job", i);
                                stopped_tokio_task_count.fetch_add(1, Ordering::SeqCst);
                                break;
                            }
                            _ = interval.tick()=> {
                                println!("tokio task {} received tick signal", i);
                            },
                            _ = stop.cancelled() => {
                                println!("tokio task {} receive exit signal", i);
                                stopped_tokio_task_count.fetch_add(1, Ordering::SeqCst);
                                break
                            },
                            else => break,
                        }
                    }
                }
            });
        }
    }
}

#[test]
fn basic() {
    let (mut handle, mut stop_recv, _runtime) = new_global_runtime(None);

    ctrlc::set_handler(move || {
        broadcast_exit_signals();
    })
    .expect("Error setting Ctrl-C handler");

    send_ctrlc_later(Duration::from_secs(3));

    let test_memo = TestStopMemo::default();

    test_memo.start_many_threads();
    test_memo.start_many_tokio_tasks(&handle);

    handle.drop_guard();
    wait_all_ckb_services_exit();
    handle.block_on(async move {
        stop_recv.recv().await;
    });

    assert_eq!(
        test_memo.spawned_threads_count.load(Ordering::SeqCst),
        test_memo.stopped_threads_count.load(Ordering::SeqCst),
    );
    assert_eq!(
        test_memo.spawned_tokio_task_count.load(Ordering::SeqCst),
        test_memo.stopped_tokio_task_count.load(Ordering::SeqCst),
    );
}