scx_stats/
server.rs

1use crate::StatsClient;
2use crate::{Meta, StatsData, StatsKind, StatsMeta};
3use anyhow::{Context, Result, anyhow, bail};
4use crossbeam::channel::{Receiver, RecvError, Select, Sender, unbounded};
5use log::{debug, error, warn};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::{BTreeMap, BTreeSet};
9use std::io::{BufRead, BufReader, Write};
10use std::os::unix::net::{UnixListener, UnixStream};
11use std::path::{Path, PathBuf};
12use std::sync::atomic::{AtomicBool, Ordering};
13use std::sync::{Arc, Mutex};
14use std::thread::spawn;
15
16pub trait StatsReader<Req, Res>:
17    FnMut(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value>
18{
19}
20impl<Req, Res, T: FnMut(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value>>
21    StatsReader<Req, Res> for T
22{
23}
24
25pub trait StatsReaderSend<Req, Res>:
26    FnMut(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value> + Send
27{
28}
29impl<
30    Req,
31    Res,
32    T: FnMut(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value> + Send,
33> StatsReaderSend<Req, Res> for T
34{
35}
36
37pub trait StatsReaderSync<Req, Res>:
38    Fn(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value> + Send + Sync
39{
40}
41impl<
42    Req,
43    Res,
44    T: Fn(&BTreeMap<String, String>, (&Sender<Req>, &Receiver<Res>)) -> Result<Value> + Send + Sync,
45> StatsReaderSync<Req, Res> for T
46{
47}
48
49pub trait StatsOpener<Req, Res>:
50    FnMut((&Sender<Req>, &Receiver<Res>)) -> Result<Box<dyn StatsReader<Req, Res>>> + Send
51{
52}
53impl<
54    Req,
55    Res,
56    T: FnMut((&Sender<Req>, &Receiver<Res>)) -> Result<Box<dyn StatsReader<Req, Res>>> + Send,
57> StatsOpener<Req, Res> for T
58{
59}
60
61pub trait StatsCloser<Req, Res>: FnOnce((&Sender<Req>, &Receiver<Res>)) + Send {}
62impl<Req, Res, T: FnOnce((&Sender<Req>, &Receiver<Res>)) + Send> StatsCloser<Req, Res> for T {}
63
64#[allow(clippy::type_complexity)]
65pub struct StatsOps<Req, Res> {
66    pub open: Box<dyn StatsOpener<Req, Res>>,
67    pub close: Option<Box<dyn StatsCloser<Req, Res>>>,
68}
69
70struct StatsOpenOps<Req, Res> {
71    map: BTreeMap<
72        String,
73        (
74            Arc<Mutex<StatsOps<Req, Res>>>,
75            Box<dyn StatsReader<Req, Res>>,
76            ChannelPair<Req, Res>,
77        ),
78    >,
79}
80
81impl<Req, Res> StatsOpenOps<Req, Res> {
82    fn new() -> Self {
83        Self {
84            map: BTreeMap::new(),
85        }
86    }
87}
88
89impl<Req, Res> std::ops::Drop for StatsOpenOps<Req, Res> {
90    fn drop(&mut self) {
91        for (ops, _, ch) in self.map.values_mut() {
92            if let Some(close) = ops.lock().unwrap().close.take() {
93                close((&ch.req, &ch.res));
94            }
95        }
96    }
97}
98
99#[derive(Clone, Debug, Serialize, Deserialize)]
100pub struct StatsRequest {
101    pub req: String,
102    #[serde(default)]
103    pub args: BTreeMap<String, String>,
104}
105
106impl StatsRequest {
107    pub fn new(req: &str, args: Vec<(String, String)>) -> Self {
108        Self {
109            req: req.to_string(),
110            args: args.into_iter().collect(),
111        }
112    }
113}
114
115#[derive(Clone, Debug, Serialize, Deserialize)]
116pub struct StatsResponse {
117    pub errno: i32,
118    pub args: BTreeMap<String, Value>,
119}
120
121pub struct StatsErrno(pub i32);
122
123impl std::fmt::Display for StatsErrno {
124    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
125        write!(f, "{}", std::io::Error::from_raw_os_error(self.0))
126    }
127}
128
129impl std::fmt::Debug for StatsErrno {
130    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131        write!(f, "{:?}", std::io::Error::from_raw_os_error(self.0))
132    }
133}
134
135struct ChannelPair<Req, Res> {
136    req: Sender<Req>,
137    res: Receiver<Res>,
138}
139
140impl<Req, Res> ChannelPair<Req, Res> {
141    fn bidi() -> (ChannelPair<Req, Res>, ChannelPair<Res, Req>) {
142        let (req, res) = (unbounded::<Req>(), unbounded::<Res>());
143        (
144            ChannelPair {
145                req: req.0,
146                res: res.1,
147            },
148            ChannelPair {
149                req: res.0,
150                res: req.1,
151            },
152        )
153    }
154}
155
156impl<Req, Res> Clone for ChannelPair<Req, Res> {
157    fn clone(&self) -> Self {
158        Self {
159            req: self.req.clone(),
160            res: self.res.clone(),
161        }
162    }
163}
164
165pub struct StatsServerData<Req, Res>
166where
167    Req: Send + 'static,
168    Res: Send + 'static,
169{
170    top: Option<String>,
171    meta: BTreeMap<String, StatsMeta>,
172    ops: BTreeMap<String, Arc<Mutex<StatsOps<Req, Res>>>>,
173}
174
175impl<Req, Res> Default for StatsServerData<Req, Res>
176where
177    Req: Send + 'static,
178    Res: Send + 'static,
179{
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185impl<Req, Res> StatsServerData<Req, Res>
186where
187    Req: Send + 'static,
188    Res: Send + 'static,
189{
190    pub fn new() -> Self {
191        Self {
192            top: None,
193            meta: BTreeMap::new(),
194            ops: BTreeMap::new(),
195        }
196    }
197
198    pub fn add_meta(mut self, meta: StatsMeta) -> Self {
199        if meta.attrs.top.is_some() && self.top.is_none() {
200            self.top = Some(meta.name.clone());
201        }
202        self.meta.insert(meta.name.clone(), meta);
203        self
204    }
205
206    pub fn add_ops(mut self, name: &str, ops: StatsOps<Req, Res>) -> Self {
207        self.ops.insert(name.to_string(), Arc::new(Mutex::new(ops)));
208        self
209    }
210
211    pub fn add_stats(self, name: &str, fetch: Box<dyn StatsReaderSend<Req, Res>>) -> Self {
212        let wrapped_fetch = Mutex::new(fetch);
213        let read: Box<dyn StatsReaderSync<Req, Res>> =
214            Box::new(move |args, chan| wrapped_fetch.lock().unwrap()(args, chan));
215        let wrapped_read = Arc::new(read);
216        let ops = StatsOps {
217            open: Box::new(move |_| {
218                let copy = wrapped_read.clone();
219                Ok(Box::new(move |args, chan| copy(args, chan)))
220            }),
221            close: None,
222        };
223
224        self.add_ops(name, ops)
225    }
226
227    fn visit_meta_inner(
228        &self,
229        name: &str,
230        visit: &mut impl FnMut(&StatsMeta) -> Result<()>,
231        nesting: &mut BTreeSet<String>,
232        visited: &mut BTreeSet<String>,
233    ) -> Result<()> {
234        let m = match self.meta.get(name) {
235            Some(v) => v,
236            None => bail!("unknown stats meta name {}", name),
237        };
238
239        if !nesting.insert(name.into()) {
240            bail!("loop in stats meta detected, {} already nested", name);
241        }
242        if !visited.insert(name.into()) {
243            return Ok(());
244        }
245
246        visit(m)?;
247
248        for (fname, field) in m.fields.iter() {
249            match &field.data {
250                StatsData::Array(StatsKind::Struct(inner)) => {
251                    self.visit_meta_inner(inner, visit, nesting, visited)?
252                }
253                StatsData::Dict {
254                    key: StatsKind::Struct(inner),
255                    datum: _,
256                } => bail!("{}.{} is a dict with struct {} as key", name, fname, inner),
257                StatsData::Dict {
258                    key: _,
259                    datum: StatsKind::Struct(inner),
260                } => self.visit_meta_inner(inner, visit, nesting, visited)?,
261                _ => {}
262            }
263        }
264
265        nesting.remove(name);
266        Ok(())
267    }
268
269    fn visit_meta(
270        &self,
271        from: &str,
272        visit: &mut impl FnMut(&StatsMeta) -> Result<()>,
273    ) -> Result<()> {
274        let mut nesting = BTreeSet::<String>::new();
275        let mut visited = BTreeSet::<String>::new();
276        self.visit_meta_inner(from, visit, &mut nesting, &mut visited)
277    }
278
279    fn verify_meta(&self) -> Result<()> {
280        if self.top.is_none() {
281            debug!("top-level stats metadata missing");
282            return Ok(());
283        }
284
285        // Null visit checks all nested stats are reachable without loops.
286        self.visit_meta(self.top.as_ref().unwrap(), &mut |_| Ok(()))
287    }
288
289    pub fn describe_meta<W: Write>(&self, w: &mut W, from: Option<&[&str]>) -> Result<()> {
290        let meta_names = match from {
291            Some(v) if !v.is_empty() => v,
292            Some(_) | None => {
293                let top = self
294                    .top
295                    .as_ref()
296                    .ok_or_else(|| anyhow!("don't know where to start"))?;
297                return self.describe_meta_inner(w, top);
298            }
299        };
300
301        for meta_name in meta_names {
302            self.describe_meta_inner(w, meta_name)?;
303        }
304
305        Ok(())
306    }
307
308    pub fn describe_meta_inner<W: Write>(&self, w: &mut W, from: &str) -> Result<()> {
309        let (mut nwidth, mut fwidth, mut dwidth) = (0usize, 0usize, 0usize);
310
311        self.visit_meta(from, &mut |m| {
312            nwidth = nwidth.max(m.name.len());
313            (fwidth, dwidth) = m.fields.iter().fold((fwidth, dwidth), |acc, (n, f)| {
314                (acc.0.max(n.len()), acc.1.max(f.data.to_string().len() + 2))
315            });
316            Ok(())
317        })?;
318
319        let mut first = true;
320        self.visit_meta(from, &mut |m| {
321            if !first {
322                writeln!(w)?;
323            }
324            first = false;
325
326            write!(w, "[{:nw$}]", m.name, nw = nwidth)?;
327            if let Some(desc) = &m.attrs.desc {
328                write!(w, " {desc}")?;
329            }
330            writeln!(w)?;
331
332            for (fname, f) in m.fields.iter() {
333                write!(
334                    w,
335                    "  {:fw$} {:dw$}",
336                    fname,
337                    format!("({})", f.data.to_string()),
338                    fw = fwidth,
339                    dw = dwidth
340                )?;
341                if let Some(desc) = &f.attrs.desc {
342                    write!(w, " : {desc}")?;
343                }
344                writeln!(w)?;
345            }
346            Ok(())
347        })
348    }
349}
350
351struct StatsServerInner<Req, Res>
352where
353    Req: Send + 'static,
354    Res: Send + 'static,
355{
356    listener: UnixListener,
357    data: Arc<Mutex<StatsServerData<Req, Res>>>,
358    inner_ch: ChannelPair<Req, Res>,
359    exit: Arc<AtomicBool>,
360}
361
362impl<Req, Res> StatsServerInner<Req, Res>
363where
364    Req: Send + 'static,
365    Res: Send + 'static,
366{
367    fn new(
368        listener: UnixListener,
369        data: Arc<Mutex<StatsServerData<Req, Res>>>,
370        inner_ch: ChannelPair<Req, Res>,
371        exit: Arc<AtomicBool>,
372    ) -> Self {
373        Self {
374            listener,
375            data,
376            inner_ch,
377            exit,
378        }
379    }
380
381    fn build_resp<T>(errno: i32, resp: &T) -> Result<StatsResponse>
382    where
383        T: Serialize,
384    {
385        Ok(StatsResponse {
386            errno,
387            args: [("resp".into(), serde_json::to_value(resp)?)]
388                .into_iter()
389                .collect(),
390        })
391    }
392
393    fn handle_request(
394        line: String,
395        data: &Arc<Mutex<StatsServerData<Req, Res>>>,
396        ch: &ChannelPair<Req, Res>,
397        open_ops: &mut StatsOpenOps<Req, Res>,
398    ) -> Result<StatsResponse> {
399        let req: StatsRequest = serde_json::from_str(&line)?;
400
401        match req.req.as_str() {
402            "stats" => {
403                let target = match req.args.get("target") {
404                    Some(v) => v,
405                    None => "top",
406                };
407
408                let ops =
409                    match data.lock().unwrap().ops.get(target) {
410                        Some(v) => v.clone(),
411                        None => Err(anyhow!("unknown stat target {:?}", req)
412                            .context(StatsErrno(libc::EINVAL)))?,
413                    };
414
415                if !open_ops.map.contains_key(target) {
416                    let read = (ops.lock().unwrap().open)((&ch.req, &ch.res))?;
417                    open_ops
418                        .map
419                        .insert(target.into(), (ops.clone(), read, ch.clone()));
420                }
421
422                let read = &mut open_ops.map.get_mut(target).unwrap().1;
423
424                let resp = read(&req.args, (&ch.req, &ch.res))?;
425
426                Self::build_resp(0, &resp)
427            }
428            "stats_meta" => Ok(Self::build_resp(0, &data.lock().unwrap().meta)?),
429            req => Err(anyhow!("unknown command {:?}", req).context(StatsErrno(libc::EINVAL)))?,
430        }
431    }
432
433    fn serve(
434        mut stream: UnixStream,
435        data: Arc<Mutex<StatsServerData<Req, Res>>>,
436        inner_ch: ChannelPair<Req, Res>,
437        exit: Arc<AtomicBool>,
438    ) -> Result<()> {
439        let mut stream_reader = BufReader::new(stream.try_clone()?);
440        let mut open_ops = StatsOpenOps::new();
441
442        loop {
443            let mut line = String::new();
444            stream_reader.read_line(&mut line)?;
445            if line.is_empty() {
446                return Ok(());
447            }
448            if exit.load(Ordering::Relaxed) {
449                debug!("server exiting due to exit");
450                return Ok(());
451            }
452
453            let resp = match Self::handle_request(line, &data, &inner_ch, &mut open_ops) {
454                Ok(v) => v,
455                Err(e) => {
456                    let errno = match e.downcast_ref::<StatsErrno>() {
457                        Some(e) if e.0 != 0 => e.0,
458                        _ => libc::EINVAL,
459                    };
460                    Self::build_resp(errno, &format!("{:?}", e))?
461                }
462            };
463
464            let output = serde_json::to_string(&resp)? + "\n";
465            stream.write_all(output.as_bytes())?;
466        }
467    }
468
469    fn proxy(inner_ch: ChannelPair<Req, Res>, add_res: Receiver<ChannelPair<Res, Req>>) {
470        let mut chs_cursor = 0;
471        let mut chs = BTreeMap::<u64, ChannelPair<Res, Req>>::new();
472        let mut ch_to_add: Option<ChannelPair<Res, Req>> = None;
473        let mut idx_to_drop: Option<u64> = None;
474
475        'outer: loop {
476            if let Some(new_ch) = ch_to_add.take() {
477                let idx = chs_cursor;
478                chs_cursor += 1;
479                chs.insert(idx, new_ch);
480                debug!("proxy: added new channel idx={}, total={}", idx, chs.len());
481            }
482
483            if let Some(idx) = idx_to_drop.take() {
484                debug!("proxy: dropping channel {}, total={}", idx, chs.len());
485                chs.remove(&idx).unwrap();
486            }
487
488            let mut sel = Select::new();
489            let inner_idx = sel.recv(&inner_ch.res);
490            let add_idx = sel.recv(&add_res);
491
492            let mut chs_sel_idx = BTreeMap::<usize, u64>::new();
493            for (idx, cp) in chs.iter() {
494                let sel_idx = sel.recv(&cp.res);
495                chs_sel_idx.insert(sel_idx, *idx);
496            }
497
498            'select: loop {
499                let oper = sel.select();
500                match oper.index() {
501                    sel_idx if sel_idx == add_idx => match oper.recv(&add_res) {
502                        Ok(ch) => {
503                            ch_to_add = Some(ch);
504                            debug!("proxy: received new channel from add_res");
505                            break 'select;
506                        }
507                        Err(RecvError) => {
508                            debug!("proxy: add_res disconnected, terminating");
509                            break 'outer;
510                        }
511                    },
512                    sel_idx if sel_idx == inner_idx => match oper.recv(&inner_ch.res) {
513                        Ok(_) => {
514                            error!("proxy: unexpected data in StatsServer.channels().0");
515                            panic!();
516                        }
517                        Err(RecvError) => break 'outer,
518                    },
519                    sel_idx => {
520                        let idx = chs_sel_idx.get(&sel_idx).unwrap();
521                        let pair = chs.get(idx).unwrap();
522
523                        let req = match oper.recv(&pair.res) {
524                            Ok(v) => v,
525                            Err(RecvError) => {
526                                idx_to_drop = Some(*idx);
527                                break 'select;
528                            }
529                        };
530
531                        if inner_ch.req.send(req).is_err() {
532                            break 'outer;
533                        }
534
535                        let resp = match inner_ch.res.recv() {
536                            Ok(v) => v,
537                            Err(RecvError) => break 'outer,
538                        };
539
540                        if pair.req.send(resp).is_err() {
541                            idx_to_drop = Some(*idx);
542                            break 'select;
543                        }
544                    }
545                }
546            }
547        }
548    }
549
550    fn listen(self) {
551        let inner_ch_copy = self.inner_ch.clone();
552        let (add_req, add_res) = unbounded::<ChannelPair<Res, Req>>();
553
554        spawn(move || Self::proxy(inner_ch_copy, add_res));
555
556        for stream in self.listener.incoming() {
557            if self.exit.load(Ordering::Relaxed) {
558                debug!("listener exiting");
559                break;
560            }
561            match stream {
562                Ok(stream) => {
563                    let data = self.data.clone();
564                    let exit = self.exit.clone();
565
566                    let (req_pair, res_pair) = ChannelPair::<Req, Res>::bidi();
567                    match add_req.send(res_pair) {
568                        Ok(()) => debug!("sent new channel to proxy"),
569                        Err(e) => warn!("StatsServer::proxy() failed ({e})"),
570                    }
571
572                    spawn(move || {
573                        if let Err(e) = Self::serve(stream, data, req_pair, exit) {
574                            warn!("stat communication errored ({e})");
575                        }
576                    });
577                }
578                Err(e) => warn!("failed to accept stat connection ({e})"),
579            }
580        }
581    }
582}
583
584pub struct StatsServer<Req, Res>
585where
586    Req: Send + 'static,
587    Res: Send + 'static,
588{
589    base_path: PathBuf,
590    sched_path: PathBuf,
591    stats_path: PathBuf,
592    path: Option<PathBuf>,
593
594    data: Arc<Mutex<StatsServerData<Req, Res>>>,
595
596    outer_ch: ChannelPair<Res, Req>,
597    inner_ch: Option<ChannelPair<Req, Res>>,
598    exit: Arc<AtomicBool>,
599}
600
601impl<Req, Res> StatsServer<Req, Res>
602where
603    Req: Send + 'static,
604    Res: Send + 'static,
605{
606    pub fn new(data: StatsServerData<Req, Res>) -> Self {
607        let (ich, och) = ChannelPair::<Req, Res>::bidi();
608
609        Self {
610            base_path: PathBuf::from("/var/run/scx"),
611            sched_path: PathBuf::from("root"),
612            stats_path: PathBuf::from("stats"),
613            path: None,
614            data: Arc::new(Mutex::new(data)),
615            outer_ch: och,
616            inner_ch: Some(ich),
617            exit: Arc::new(AtomicBool::new(false)),
618        }
619    }
620
621    pub fn set_base_path<P: AsRef<Path>>(mut self, path: P) -> Self {
622        self.base_path = PathBuf::from(path.as_ref());
623        self
624    }
625
626    pub fn set_sched_path<P: AsRef<Path>>(mut self, path: P) -> Self {
627        self.sched_path = PathBuf::from(path.as_ref());
628        self
629    }
630
631    pub fn set_stats_path<P: AsRef<Path>>(mut self, path: P) -> Self {
632        self.stats_path = PathBuf::from(path.as_ref());
633        self
634    }
635
636    pub fn set_path<P: AsRef<Path>>(mut self, path: P) -> Self {
637        self.path = Some(PathBuf::from(path.as_ref()));
638        self
639    }
640
641    pub fn launch(mut self) -> Result<Self> {
642        self.data.lock().unwrap().verify_meta()?;
643
644        if self.path.is_none() {
645            self.path = Some(self.base_path.join(&self.sched_path).join(&self.stats_path));
646        }
647        let path = &self.path.as_ref().unwrap();
648
649        if let Some(dir) = path.parent() {
650            std::fs::create_dir_all(dir).with_context(|| format!("creating {dir:?}"))?;
651        }
652
653        let res = std::fs::remove_file(path);
654        if let std::io::Result::Err(e) = &res
655            && e.kind() != std::io::ErrorKind::NotFound
656        {
657            res.with_context(|| format!("deleting {path:?}"))?;
658        }
659
660        let listener =
661            UnixListener::bind(path).with_context(|| format!("creating UNIX socket {path:?}"))?;
662
663        let inner = StatsServerInner::new(
664            listener,
665            self.data.clone(),
666            self.inner_ch.take().unwrap(),
667            self.exit.clone(),
668        );
669
670        spawn(move || inner.listen());
671        Ok(self)
672    }
673
674    pub fn channels(&self) -> (Sender<Res>, Receiver<Req>) {
675        (self.outer_ch.req.clone(), self.outer_ch.res.clone())
676    }
677}
678
679impl<Req, Res> std::ops::Drop for StatsServer<Req, Res>
680where
681    Req: Send + 'static,
682    Res: Send + 'static,
683{
684    fn drop(&mut self) {
685        self.exit.store(true, Ordering::Relaxed);
686        if let Some(path) = self.path.as_ref() {
687            let _ = StatsClient::new().set_path(path).connect(None);
688        }
689    }
690}
691
692pub trait ToJson {
693    fn to_json(&self) -> Result<Value>;
694}
695
696impl<T> ToJson for T
697where
698    T: Meta + Serialize,
699{
700    fn to_json(&self) -> Result<Value> {
701        Ok(serde_json::to_value(self)?)
702    }
703}
scx_stats/server.rs

scx_stats/
server.rs
