task/manager/
spawner.rs

1// Spawner module - handles spawner registration and selection
2
3use super::*;
4
5use alloc::collections::BTreeMap;
6use embassy_executor::Spawner;
7
8impl Manager {
9    pub fn register_spawner(&'static self, spawner: Spawner) -> Result<usize> {
10        let mut inner = embassy_futures::block_on(self.0.write());
11
12        let identifier = Self::find_first_available_identifier(
13            &inner.spawners,
14            (usize::MIN..usize::MAX).step_by(1),
15        )
16        .ok_or(Error::TooManySpawners)?;
17
18        if inner.spawners.insert(identifier, spawner).is_some() {
19            unreachable!("Spawner identifier already exists");
20        }
21
22        Ok(identifier)
23    }
24
25    pub fn unregister_spawner(&'static self, identifier: usize) -> Result<()> {
26        let mut inner = embassy_futures::block_on(self.0.write());
27
28        inner
29            .spawners
30            .remove(&identifier)
31            .ok_or(Error::NoSpawnerAvailable)?;
32
33        Ok(())
34    }
35
36    /// Select the best spawner for a new task using load balancing algorithm
37    pub(crate) fn select_best_spawner(inner: &Inner) -> Result<usize> {
38        if inner.spawners.is_empty() {
39            return Err(Error::NoSpawnerAvailable);
40        }
41
42        let mut map = BTreeMap::new();
43
44        for identifier in inner.spawners.keys() {
45            map.insert(*identifier, 0); // Initialize all spawners with a load of 0
46        }
47
48        for metadata in inner.tasks.values() {
49            if let Some(load) = map.get_mut(&metadata.spawner_identifier) {
50                *load += 1; // Increment the load for the spawner
51            }
52        }
53
54        // Find the spawner with the lowest load score
55        let mut best_index = 0;
56        let mut best_score = usize::MAX;
57
58        for (identifier, spawner) in map.iter() {
59            if *spawner < best_score {
60                best_score = *spawner;
61                best_index = *identifier;
62            }
63        }
64
65        Ok(best_index)
66    }
67
68    pub async fn get_spawner(&self, task: TaskIdentifier) -> Result<usize> {
69        Self::get_task(&*self.0.read().await, task)
70            .map(|task| task.spawner_identifier)
71            .map_err(|_| Error::InvalidTaskIdentifier)
72    }
73}