1use 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}