1
use std::any::Any;
2
use std::collections::HashMap;
3
use std::fmt::Debug;
4
use std::sync::Arc;
5

            
6
use flume::{Receiver, Sender};
7

            
8
use crate::tasks::handle::{Handle, Id};
9
use crate::tasks::manager::{ManagedJob, Manager};
10
use crate::tasks::traits::Executable;
11
use crate::tasks::{Job, Keyed};
12

            
13
pub struct Jobs<Key> {
14
    last_task_id: u64,
15
    result_senders: HashMap<Id, Vec<Box<dyn AnySender>>>,
16
    keyed_jobs: HashMap<Key, Id>,
17
    queuer: Sender<Box<dyn Executable>>,
18
    queue: Receiver<Box<dyn Executable>>,
19
}
20

            
21
impl<Key> Debug for Jobs<Key>
22
where
23
    Key: Debug,
24
{
25
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26
        f.debug_struct("Jobs")
27
            .field("last_task_id", &self.last_task_id)
28
            .field("result_senders", &self.result_senders.len())
29
            .field("keyed_jobs", &self.keyed_jobs)
30
            .field("queuer", &self.queuer)
31
            .field("queue", &self.queue)
32
            .finish()
33
    }
34
}
35

            
36
impl<Key> Default for Jobs<Key> {
37
5197
    fn default() -> Self {
38
5197
        let (queuer, queue) = flume::unbounded();
39
5197

            
40
5197
        Self {
41
5197
            last_task_id: 0,
42
5197
            result_senders: HashMap::new(),
43
5197
            keyed_jobs: HashMap::new(),
44
5197
            queuer,
45
5197
            queue,
46
5197
        }
47
5197
    }
48
}
49

            
50
impl<Key> Jobs<Key>
51
where
52
    Key: Clone + std::hash::Hash + Eq + Send + Sync + Debug + 'static,
53
{
54
20782
    pub fn queue(&self) -> Receiver<Box<dyn Executable>> {
55
20782
        self.queue.clone()
56
20782
    }
57

            
58
257324
    pub fn enqueue<J: Job + 'static>(
59
257324
        &mut self,
60
257324
        job: J,
61
257324
        key: Option<Key>,
62
257324
        manager: Manager<Key>,
63
257324
    ) -> Handle<J::Output, J::Error> {
64
257324
        self.last_task_id = self.last_task_id.wrapping_add(1);
65
257324
        let id = Id(self.last_task_id);
66
257324
        self.queuer
67
257324
            .send(Box::new(ManagedJob {
68
257324
                id,
69
257324
                job,
70
257324
                manager,
71
257324
                key,
72
257324
            }))
73
257324
            .unwrap();
74
257324

            
75
257324
        self.create_new_task_handle(id)
76
257324
    }
77

            
78
270113
    pub fn create_new_task_handle<T: Send + Sync + 'static, E: Send + Sync + 'static>(
79
270113
        &mut self,
80
270113
        id: Id,
81
270113
    ) -> Handle<T, E> {
82
270113
        let (sender, receiver) = flume::bounded(1);
83
270113
        let senders = self.result_senders.entry(id).or_insert_with(Vec::default);
84
270113
        senders.push(Box::new(sender));
85
270113

            
86
270113
        Handle { id, receiver }
87
270113
    }
88

            
89
270112
    pub fn lookup_or_enqueue<J: Keyed<Key>>(
90
270112
        &mut self,
91
270112
        job: J,
92
270112
        manager: Manager<Key>,
93
270112
    ) -> Handle<<J as Job>::Output, <J as Job>::Error> {
94
270112
        let key = job.key();
95
270112
        if let Some(&id) = self.keyed_jobs.get(&key) {
96
12789
            self.create_new_task_handle(id)
97
        } else {
98
257323
            let handle = self.enqueue(job, Some(key.clone()), manager);
99
257323
            self.keyed_jobs.insert(key, handle.id);
100
257323
            handle
101
        }
102
270112
    }
103

            
104
    pub fn job_completed<T: Clone + Send + Sync + 'static, E: Send + Sync + 'static>(
105
        &mut self,
106
        id: Id,
107
        key: Option<&Key>,
108
        result: Result<T, E>,
109
    ) {
110
257324
        if let Some(key) = key {
111
257323
            self.keyed_jobs.remove(key);
112
257323
        }
113

            
114
257324
        if let Some(senders) = self.result_senders.remove(&id) {
115
257324
            let result = result.map_err(Arc::new);
116
527437
            for sender_handle in senders {
117
270113
                let sender = sender_handle
118
270113
                    .as_any()
119
270113
                    .downcast_ref::<flume::Sender<Result<T, Arc<E>>>>()
120
270113
                    .unwrap();
121
270113
                drop(sender.send(result.clone()));
122
270113
            }
123
        }
124
257324
    }
125
}
126

            
127
pub trait AnySender: Any + Send + Sync {
128
    fn as_any(&self) -> &'_ dyn Any;
129
}
130

            
131
impl<T> AnySender for flume::Sender<T>
132
where
133
    T: Send + Sync + 'static,
134
{
135
270113
    fn as_any(&self) -> &'_ dyn Any {
136
270113
        self
137
270113
    }
138
}