1
use std::{any::Any, collections::HashMap, fmt::Debug, sync::Arc};
2

            
3
use flume::{Receiver, Sender};
4

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

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

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

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

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

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

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

            
74
188303
        self.create_new_task_handle(id)
75
188303
    }
76

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

            
85
196940
        Handle { id, receiver }
86
196940
    }
87

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

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

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

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

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