fast_vk 0.3.1

Low-level VK API library designed for millions requests
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

use super::{Event, EventReceiver};

use std::panic;

use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;

pub struct TaskObserver {
    running_tasks: Arc<RwLock<usize>>,
    thread: Option<JoinHandle<()>>,
}

impl TaskObserver {
    pub fn new(mut receiver: EventReceiver) -> TaskObserver {
        let running_tasks = Arc::new(RwLock::new(0));

        let running_tasks_inner = Arc::clone(&running_tasks);
        let thread = Some(tokio::spawn(async move {
            loop {
                match receiver.recv().await {
                    Ok(event) => match event {
                        Event::GotWork => {
                            *running_tasks_inner.write().await += 1;
                        }
                        Event::DoneWork => {
                            *running_tasks_inner.write().await -= 1;
                        }
                    },
                    Err(_) => break,
                }
            }
        }));

        TaskObserver {
            thread,
            running_tasks,
        }
    }

    pub async fn running_task(&self) -> usize {
        // we can use unwrap safe because only drop function takes thread
        if self.thread.as_ref().unwrap().is_finished() {
            drop(&self);
        }

        *self.running_tasks.read().await
    }
}

impl Drop for TaskObserver {
    fn drop(&mut self) {
        let thread = self.thread.take().unwrap();

        if thread.is_finished() {
            let result = futures::executor::block_on(async { thread.await });

            if let Err(err) = result {
                if let Ok(reason) = err.try_into_panic() {
                    panic::resume_unwind(reason);
                }
            }
        } else {
            thread.abort();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::broadcast;

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn got_work_two_times() {
        let (event_sender, event_receiver) = broadcast::channel(2);

        let worker_watcher = TaskObserver::new(event_receiver);

        event_sender.send(Event::GotWork).unwrap();
        event_sender.send(Event::GotWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        assert_eq!(worker_watcher.running_task().await, 2);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn one_of_two_done_work() {
        let (event_sender, event_receiver) = broadcast::channel(2);

        let worker_watcher = TaskObserver::new(event_receiver);

        event_sender.send(Event::GotWork).unwrap();
        event_sender.send(Event::GotWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        event_sender.send(Event::DoneWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        assert_eq!(worker_watcher.running_task().await, 1);
    }

    #[tokio::test]
    async fn two_done_work() {
        let (event_sender, event_receiver) = broadcast::channel(2);

        let worker_watcher = TaskObserver::new(event_receiver);

        event_sender.send(Event::GotWork).unwrap();
        event_sender.send(Event::GotWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        event_sender.send(Event::DoneWork).unwrap();
        event_sender.send(Event::DoneWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        assert_eq!(worker_watcher.running_task().await, 0);
    }

    #[tokio::test]
    #[should_panic]
    async fn done_more_than_got() {
        let (event_sender, event_receiver) = broadcast::channel(3);

        let _worker_watcher = TaskObserver::new(event_receiver);

        event_sender.send(Event::GotWork).unwrap();
        event_sender.send(Event::GotWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;

        event_sender.send(Event::DoneWork).unwrap();
        event_sender.send(Event::DoneWork).unwrap();
        event_sender.send(Event::DoneWork).unwrap();
        tokio::time::sleep(std::time::Duration::from_nanos(1)).await;
    }
}