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