scx_stats/
client.rs

1use crate::StatsErrno;
2use crate::StatsRequest;
3use crate::StatsResponse;
4use anyhow::Result;
5use anyhow::anyhow;
6use anyhow::bail;
7use log::trace;
8use serde::Deserialize;
9use std::io::BufRead;
10use std::io::BufReader;
11use std::io::Write;
12use std::io::{self};
13use std::os::unix::net::UnixStream;
14use std::path::Path;
15use std::path::PathBuf;
16use std::time::Duration;
17
18pub struct StatsClient {
19    base_path: PathBuf,
20    sched_path: PathBuf,
21    stats_path: PathBuf,
22    path: Option<PathBuf>,
23
24    stream: Option<UnixStream>,
25    reader: Option<BufReader<UnixStream>>,
26}
27
28impl Default for StatsClient {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl StatsClient {
35    pub fn new() -> Self {
36        Self {
37            base_path: PathBuf::from("/var/run/scx"),
38            sched_path: PathBuf::from("root"),
39            stats_path: PathBuf::from("stats"),
40            path: None,
41
42            stream: None,
43            reader: None,
44        }
45    }
46
47    pub fn set_base_path<P: AsRef<Path>>(mut self, path: P) -> Self {
48        self.base_path = PathBuf::from(path.as_ref());
49        self
50    }
51
52    pub fn set_sched_path<P: AsRef<Path>>(mut self, path: P) -> Self {
53        self.sched_path = PathBuf::from(path.as_ref());
54        self
55    }
56
57    pub fn set_stats_path<P: AsRef<Path>>(mut self, path: P) -> Self {
58        self.stats_path = PathBuf::from(path.as_ref());
59        self
60    }
61
62    pub fn set_path<P: AsRef<Path>>(mut self, path: P) -> Self {
63        self.path = Some(PathBuf::from(path.as_ref()));
64        self
65    }
66
67    pub fn connect(mut self, timeout_ms: Option<u64>) -> Result<Self> {
68        if self.path.is_none() {
69            self.path = Some(self.base_path.join(&self.sched_path).join(&self.stats_path));
70        }
71        let path = self.path.as_ref().unwrap();
72
73        let stream = UnixStream::connect(path)?;
74        // Apply the same timeout to both writer and reader sides if provided
75        if let Some(ms) = timeout_ms {
76            let dur = Duration::from_millis(ms);
77            stream.set_write_timeout(Some(dur))?;
78            stream.set_read_timeout(Some(dur))?;
79        }
80
81        self.stream = Some(stream.try_clone()?);
82        self.reader = Some(BufReader::new(stream));
83        Ok(self)
84    }
85
86    pub fn send_request<T>(&mut self, req: &StatsRequest) -> Result<T>
87    where
88        T: for<'a> Deserialize<'a>,
89    {
90        if self.stream.is_none() {
91            bail!("not connected");
92        }
93
94        let req = serde_json::to_string(&req)? + "\n";
95        trace!("Sending: {}", req.trim());
96        // Attempt write with timeout
97        if let Err(e) = self.stream.as_ref().unwrap().write_all(req.as_bytes()) {
98            if e.kind() == io::ErrorKind::TimedOut || e.kind() == io::ErrorKind::WouldBlock {
99                return Err(anyhow!("write timed out"));
100            } else {
101                return Err(e.into());
102            }
103        }
104
105        let mut line = String::new();
106        match self.reader.as_mut().unwrap().read_line(&mut line) {
107            Ok(0) => return Err(anyhow!("connection closed")),
108            Ok(_) => { /* proceed */ }
109            Err(e) => {
110                if e.kind() == io::ErrorKind::TimedOut || e.kind() == io::ErrorKind::WouldBlock {
111                    return Err(anyhow!("read timed out"));
112                } else {
113                    return Err(e.into());
114                }
115            }
116        }
117
118        trace!("Received: {}", line.trim());
119        let mut resp: StatsResponse = serde_json::from_str(&line)?;
120
121        let (errno, resp) = (
122            resp.errno,
123            resp.args.remove("resp").unwrap_or(serde_json::Value::Null),
124        );
125
126        if errno != 0 {
127            Err(anyhow!("{}", resp).context(StatsErrno(errno)))?;
128        }
129
130        Ok(serde_json::from_value(resp)?)
131    }
132
133    pub fn request<T>(&mut self, req: &str, args: Vec<(String, String)>) -> Result<T>
134    where
135        T: for<'a> Deserialize<'a>,
136    {
137        self.send_request(&StatsRequest::new(req, args))
138    }
139}
scx_stats/client.rs

scx_stats/
client.rs
