proc_macro_api/
pool.rs

1//! A pool of proc-macro server processes
2use std::sync::Arc;
3
4use rayon::iter::{IntoParallelIterator, ParallelIterator};
5
6use crate::{MacroDylib, ProcMacro, ServerError, process::ProcMacroServerProcess};
7
8#[derive(Debug, Clone)]
9pub(crate) struct ProcMacroServerPool {
10    workers: Arc<[ProcMacroServerProcess]>,
11    version: u32,
12}
13
14impl ProcMacroServerPool {
15    pub(crate) fn new(workers: Vec<ProcMacroServerProcess>) -> Self {
16        let version = workers[0].version();
17        Self { workers: workers.into(), version }
18    }
19}
20
21impl ProcMacroServerPool {
22    pub(crate) fn exited(&self) -> Option<&ServerError> {
23        for worker in &*self.workers {
24            worker.exited()?;
25        }
26        self.workers[0].exited()
27    }
28
29    pub(crate) fn pick_process(&self) -> Result<&ProcMacroServerProcess, ServerError> {
30        let mut best: Option<&ProcMacroServerProcess> = None;
31        let mut best_load = u32::MAX;
32
33        for w in self.workers.iter().filter(|w| w.exited().is_none()) {
34            let load = w.number_of_active_req();
35
36            if load == 0 {
37                return Ok(w);
38            }
39
40            if load < best_load {
41                best = Some(w);
42                best_load = load;
43            }
44        }
45
46        best.ok_or_else(|| ServerError {
47            message: "all proc-macro server workers have exited".into(),
48            io: None,
49        })
50    }
51
52    pub(crate) fn load_dylib(&self, dylib: &MacroDylib) -> Result<Vec<ProcMacro>, ServerError> {
53        let _span = tracing::info_span!("ProcMacroServer::load_dylib").entered();
54
55        let dylib_path = Arc::new(dylib.path.clone());
56        let dylib_last_modified =
57            std::fs::metadata(dylib_path.as_path()).ok().and_then(|m| m.modified().ok());
58
59        let (first, rest) = self.workers.split_first().expect("worker pool must not be empty");
60
61        let macros = first
62            .find_proc_macros(&dylib.path)?
63            .map_err(|e| ServerError { message: e, io: None })?;
64
65        rest.into_par_iter()
66            .map(|worker| {
67                worker
68                    .find_proc_macros(&dylib.path)?
69                    .map(|_| ())
70                    .map_err(|e| ServerError { message: e, io: None })
71            })
72            .collect::<Result<(), _>>()?;
73
74        Ok(macros
75            .into_iter()
76            .map(|(name, kind)| ProcMacro {
77                pool: self.clone(),
78                name: name.into(),
79                kind,
80                dylib_path: dylib_path.clone(),
81                dylib_last_modified,
82            })
83            .collect())
84    }
85
86    pub(crate) fn version(&self) -> u32 {
87        self.version
88    }
89}