sans_io_runtime/task/
group.rs

1use std::{collections::HashSet, fmt::Debug, marker::PhantomData, time::Instant};
2
3use crate::{collections::DynamicVec, Task, TaskSwitcher, TaskSwitcherChild};
4
5#[derive(Debug)]
6pub enum TaskGroupOutput<Out> {
7    TaskOutput(usize, Out),
8    OnResourceEmpty,
9}
10
11#[derive(Debug)]
12struct TaskContainer<T> {
13    task: T,
14    is_empty: bool,
15}
16
17impl<T> TaskContainer<T> {
18    fn new(task: T) -> Self {
19        Self {
20            task,
21            is_empty: false,
22        }
23    }
24}
25
26/// Represents a group of tasks.
27#[derive(Debug)]
28pub struct TaskGroup<In, Out, T: Task<In, Out>, const STACK_SIZE: usize> {
29    tasks: DynamicVec<Option<TaskContainer<T>>, STACK_SIZE>,
30    tasks_count: usize,
31    empty_tasks: HashSet<usize>,
32    switcher: TaskSwitcher,
33    _tmp: PhantomData<(In, Out)>,
34}
35
36impl<In, Out, T: Task<In, Out>, const STACK_SIZE: usize> Default
37    for TaskGroup<In, Out, T, STACK_SIZE>
38{
39    /// Creates a new task group with the specified worker ID.
40    fn default() -> Self {
41        Self {
42            tasks: DynamicVec::default(),
43            tasks_count: 0,
44            empty_tasks: HashSet::default(),
45            _tmp: Default::default(),
46            switcher: TaskSwitcher::new(0),
47        }
48    }
49}
50
51impl<In, Out, T: Task<In, Out>, const STACK_SIZE: usize> TaskGroup<In, Out, T, STACK_SIZE> {
52    /// Returns the number of tasks in the group.
53    pub fn tasks(&self) -> usize {
54        self.tasks_count
55    }
56
57    /// Check if we have task with index
58    pub fn has_task(&self, index: usize) -> bool {
59        matches!(self.tasks.get(index), Some(Some(_)))
60    }
61
62    /// Adds a task to the group.
63    pub fn add_task(&mut self, task: T) -> usize {
64        for (index, slot) in self.tasks.iter_mut().enumerate() {
65            if slot.is_none() {
66                *slot = Some(TaskContainer::new(task));
67                self.tasks_count += 1;
68                return index;
69            }
70        }
71
72        self.tasks.push(Some(TaskContainer::new(task)));
73        self.switcher.set_tasks(self.tasks.len());
74        self.tasks_count += 1;
75        self.tasks.len() - 1
76    }
77
78    /// Set task to a slot index
79    pub fn set_task(&mut self, index: usize, task: T) -> usize {
80        while self.tasks.len() <= index {
81            self.tasks.push(None);
82        }
83        self.tasks.set(index, Some(TaskContainer::new(task)));
84        self.switcher.set_tasks(self.tasks.len());
85        self.tasks_count += 1;
86        self.tasks.len() - 1
87    }
88
89    /// Remove a task from the group
90    pub fn remove_task(&mut self, index: usize) {
91        self.tasks.get_mut_or_panic(index).take();
92        self.empty_tasks.remove(&index);
93        self.tasks_count -= 1;
94        while let Some(None) = self.tasks.last() {
95            self.tasks.pop();
96        }
97        self.switcher.set_tasks(self.tasks.len());
98    }
99
100    /// Fire tick event to all tasks, after that we need to call pop_output util it return None
101    pub fn on_tick(&mut self, now: Instant) {
102        self.switcher.flag_all();
103        for index in 0..self.switcher.tasks() {
104            if let Some(Some(task)) = self.tasks.get_mut(index) {
105                task.task.on_tick(now);
106            }
107        }
108    }
109
110    /// Send event to correct task with index
111    pub fn on_event(&mut self, now: Instant, index: usize, input: In) -> Option<()> {
112        let task = self.tasks.get_mut(index)?.as_mut()?;
113        self.switcher.flag_task(index);
114        task.task.on_event(now, input);
115        Some(())
116    }
117
118    /// Gracefully destroys the task group.
119    pub fn on_shutdown(&mut self, now: Instant) {
120        self.switcher.flag_all();
121        for index in 0..self.switcher.tasks() {
122            log::info!("Group kill tasks {}/{}", index, self.switcher.tasks());
123            if let Some(Some(task)) = self.tasks.get_mut(index) {
124                task.task.on_shutdown(now);
125            }
126        }
127    }
128}
129
130impl<In, Out, T: Task<In, Out>, const STACK_SIZE: usize> TaskSwitcherChild<TaskGroupOutput<Out>>
131    for TaskGroup<In, Out, T, STACK_SIZE>
132{
133    type Time = T::Time;
134
135    fn empty_event(&self) -> TaskGroupOutput<Out> {
136        TaskGroupOutput::OnResourceEmpty
137    }
138
139    fn is_empty(&self) -> bool {
140        self.empty_tasks.len() == self.tasks()
141    }
142
143    /// Retrieves the output from the flagged processed task.
144    fn pop_output(&mut self, now: Self::Time) -> Option<TaskGroupOutput<Out>> {
145        while let Some(index) = self.switcher.current() {
146            let slot = self.tasks.get_mut(index);
147            if let Some(Some(slot)) = slot {
148                if let Some(out) = slot.task.pop_output(now) {
149                    return Some(TaskGroupOutput::TaskOutput(index, out));
150                } else {
151                    if !slot.is_empty {
152                        if slot.task.is_empty() {
153                            slot.is_empty = true;
154                            self.empty_tasks.insert(index);
155                            return Some(TaskGroupOutput::TaskOutput(
156                                index,
157                                slot.task.empty_event(),
158                            ));
159                        }
160                    } else {
161                        #[allow(clippy::collapsible_else_if)]
162                        if !slot.task.is_empty() {
163                            slot.is_empty = false;
164                            self.empty_tasks.remove(&index);
165                        }
166                    }
167                    self.switcher.finished(index);
168                }
169            } else {
170                self.switcher.finished(index);
171            }
172        }
173        None
174    }
175}
sans_io_runtime/task/group.rs

sans_io_runtime/task/
group.rs
