Skip to main content

pros_simulator/host/
multitasking.rs

1// use std::sync::Mutex;
2use std::{sync::Arc, time::Instant};
3
4use futures::{future::pending, FutureExt};
5use slab::Slab;
6use tokio::sync::{Mutex, OwnedMutexGuard};
7
8#[derive(Debug, Default)]
9pub struct HostMutex {
10    inner: Arc<Mutex<()>>,
11    lock: Option<OwnedMutexGuard<()>>,
12}
13
14#[derive(Debug, Default)]
15pub struct MutexPool {
16    mutexes: Slab<HostMutex>,
17}
18
19impl MutexPool {
20    /// Creates a mutex, returning its ID.
21    pub fn create_mutex(&mut self) -> usize {
22        self.mutexes.insert(HostMutex::default())
23    }
24    /// Creates a mutex, returning its ID.
25    pub fn delete_mutex(&mut self, mutex_id: usize) {
26        self.mutexes.remove(mutex_id);
27    }
28
29    /// Locks a mutex by ID, cancelling on timeout, and returning a boolean of whether the lock was
30    /// successful.
31    pub async fn lock(&mut self, mutex_id: usize, timeout: Option<Instant>) -> bool {
32        let sleep = timeout.map_or_else(
33            || pending().boxed(),
34            |i| tokio::time::sleep_until(i.into()).boxed(),
35        );
36
37        let mutex = self.mutexes.get_mut(mutex_id).unwrap();
38        tokio::select! {
39            biased;
40            lock = mutex.inner.clone().lock_owned() => {
41                mutex.lock = Some(lock);
42                true
43            }
44            _ = sleep => false,
45        }
46    }
47
48    pub fn unlock(&mut self, mutex_id: usize) {
49        let mutex = self.mutexes.get_mut(mutex_id).unwrap();
50        mutex.lock.take().unwrap();
51    }
52}