core_allocator/
lib.rs

1mod allocators;
2mod bind_core;
3mod reexport;
4mod resource;
5
6pub(crate) use resource::*;
7
8pub use allocators::*;
9pub use reexport::*;
10
11use std::fmt::{Debug, Formatter};
12
13use crate::bind_core::{bind_to_cpu_set, to_cpu_set};
14use anyhow::Result;
15use anyhow::{anyhow, Context};
16use log::*;
17
18use std::marker::PhantomData;
19use std::mem::forget;
20use std::rc::Rc;
21use std::sync::{Arc, TryLockError};
22
23pub trait CoreAllocator: Sync {
24    fn allocate_core(&self) -> Option<CoreGroup>;
25}
26
27pub struct CoreGroup {
28    _lock: Option<ResourceHandle>,
29    inner: CoreGroupInner,
30}
31
32impl CoreGroup {
33    pub fn any_core() -> Self {
34        Self {
35            _lock: None,
36            inner: CoreGroupInner::AnyCore,
37        }
38    }
39    pub fn cores(lock: ResourceHandle, cores: Vec<Arc<ManagedCore>>) -> Self {
40        Self {
41            _lock: Some(lock),
42            inner: CoreGroupInner::Cores(cores),
43        }
44    }
45    pub fn reserve(&self) -> Vec<ManagedCore> {
46        match &self.inner {
47            CoreGroupInner::AnyCore => {
48                vec![] // no dedicated cores
49            }
50            CoreGroupInner::Cores(cores) => cores.iter().map(|x| x.reserve().unwrap()).collect(),
51        }
52    }
53    pub fn bind_nth(&self, index: usize) -> Result<Cleanup> {
54        match &self.inner {
55            CoreGroupInner::AnyCore => Ok(Cleanup::new(None, None)),
56            CoreGroupInner::Cores(cores) => {
57                let core = cores.get(index).with_context(|| {
58                    format!("Could not find {}th core in the group {:?}", index, cores)
59                })?;
60                core.bind()
61            }
62        }
63    }
64
65    /// This may block due to Mutex
66    pub fn get_cores(&self) -> Option<Vec<usize>> {
67        match &self.inner {
68            CoreGroupInner::AnyCore => None,
69            CoreGroupInner::Cores(cores) => Some(cores.iter().map(|x| x.index).collect()),
70        }
71    }
72}
73
74enum CoreGroupInner {
75    AnyCore,
76    Cores(Vec<Arc<ManagedCore>>),
77}
78impl Debug for CoreGroup {
79    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
80        match &self.inner {
81            CoreGroupInner::AnyCore => f.write_str("AnyCore"),
82            CoreGroupInner::Cores(cores) => {
83                let mut numbers = vec![];
84                for c in cores {
85                    numbers.push(c.get_raw());
86                }
87                f.write_fmt(format_args!("Cores({:?})", numbers))
88            }
89        }
90    }
91}
92pub struct ManagedCore {
93    index: usize,
94    taken: Resource,
95}
96impl Debug for ManagedCore {
97    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
98        std::fmt::Debug::fmt(&self.index, f)
99    }
100}
101impl PartialEq for ManagedCore {
102    fn eq(&self, other: &Self) -> bool {
103        self.index.eq(&other.index)
104    }
105}
106impl ManagedCore {
107    pub fn new(index: usize) -> Self {
108        Self {
109            index,
110            taken: Resource::new(),
111        }
112    }
113    pub fn get_raw(&self) -> usize {
114        self.index
115    }
116    fn bind(&self) -> anyhow::Result<Cleanup> {
117        let guard = self.taken.try_lock().map_err(|x| match x {
118            TryLockError::Poisoned(_) => {
119                anyhow!("Poisoned")
120            }
121            TryLockError::WouldBlock => {
122                anyhow!("Core already bound: {}", self.index)
123            }
124        })?;
125        let prior = bind_to_cpu_set(to_cpu_set(Some(self.get_raw())))?;
126        Ok(Cleanup::new(Some(prior), Some(guard)))
127    }
128    pub fn reserve(&self) -> anyhow::Result<ManagedCore> {
129        let guard = self.taken.try_lock().map_err(|x| match x {
130            TryLockError::Poisoned(_) => {
131                anyhow!("Poisoned")
132            }
133            TryLockError::WouldBlock => {
134                anyhow!("Core already bound: {}", self.index)
135            }
136        })?;
137        forget(guard);
138        Ok(ManagedCore::new(self.index))
139    }
140}
141
142pub struct Cleanup<'a> {
143    prior_state: Option<CpuSet>,
144    alive: Option<ResourceHandle>,
145    _phantom: PhantomData<(Rc<()>, &'a ())>,
146}
147impl<'a> Cleanup<'a> {
148    pub fn new(prior_state: Option<CpuSet>, alive: Option<ResourceHandle>) -> Self {
149        Self {
150            prior_state,
151            alive,
152            _phantom: Default::default(),
153        }
154    }
155    pub fn detach(mut self) {
156        self.prior_state.take();
157        forget(self.alive.take());
158    }
159}
160
161impl<'a> Drop for Cleanup<'a> {
162    fn drop(&mut self) {
163        if let Some(prior) = self.prior_state.take() {
164            if let Err(err) = bind_to_cpu_set(prior) {
165                error!("Error restoring state: {:?}", err);
166            }
167        }
168    }
169}
core_allocator/lib.rs

core_allocator/
lib.rs
