linux_taskstats/
lib.rs

1#![allow(non_upper_case_globals)]
2#![allow(non_camel_case_types)]
3#![allow(non_snake_case)]
4
5#[allow(dead_code)]
6mod c_headers;
7#[cfg(feature = "format")]
8pub mod format;
9mod model;
10pub(crate) mod netlink;
11pub use model::*;
12
13pub use c_headers::taskstats;
14use c_headers::{
15    __u16, __u32, __u64, __u8, TASKSTATS_CMD_ATTR_DEREGISTER_CPUMASK, TASKSTATS_CMD_ATTR_PID,
16    TASKSTATS_CMD_ATTR_REGISTER_CPUMASK, TASKSTATS_CMD_ATTR_TGID, TASKSTATS_CMD_GET,
17    TASKSTATS_GENL_NAME, TASKSTATS_TYPE_AGGR_PID, TASKSTATS_TYPE_AGGR_TGID, TASKSTATS_TYPE_NULL,
18    TASKSTATS_TYPE_PID, TASKSTATS_TYPE_STATS, TASKSTATS_TYPE_TGID,
19};
20use log::{debug, warn};
21use netlink::Netlink;
22use netlink::NlPayload;
23use std::{mem, slice};
24use thiserror::Error;
25
26/// Errors possibly returned by `Client`
27#[derive(Debug, Error)]
28pub enum Error {
29    /// Error in netlink socket/protocol layer
30    #[error("error in netlink communication with kernel: {0}")]
31    Netlink(#[from] netlink::Error),
32    /// Failed to lookup family ID for taskstats
33    #[error("no family id corresponding to taskstats found")]
34    NoFamilyId,
35    /// Any unknown error
36    #[error("unknown error: {0}")]
37    Unknown(String),
38}
39
40pub type Result<T> = std::result::Result<T, Error>;
41
42/// Interface to access kernel taskstats API through the netlink socket.
43pub struct Client {
44    netlink: Netlink,
45    ts_family_id: u16,
46}
47
48impl Client {
49    /// Open netlink socket against kernel and create a new instance of `Client`
50    ///
51    /// # Errors
52    /// * when netlink socket initialization failed
53    /// * when kernel doesn't offer family id for taskstats
54    pub fn open() -> Result<Self> {
55        let netlink = Netlink::open()?;
56        let ts_family_id = Self::lookup_family_id(&netlink)?;
57        debug!("Found taskstats family id: {}", ts_family_id);
58        Ok(Self {
59            netlink,
60            ts_family_id,
61        })
62    }
63
64    fn lookup_family_id(netlink: &Netlink) -> Result<u16> {
65        netlink.send_cmd(
66            libc::GENL_ID_CTRL as u16,
67            libc::CTRL_CMD_GETFAMILY as u8,
68            libc::CTRL_ATTR_FAMILY_NAME as u16,
69            TASKSTATS_GENL_NAME,
70        )?;
71
72        let resp = netlink.recv_response()?;
73        for na in resp.payload_as_nlattrs() {
74            debug!("Family lookup: got nla_type: {}", na.header.nla_type);
75            if na.header.nla_type == libc::CTRL_ATTR_FAMILY_ID as u16 {
76                return Ok(*na.payload_as());
77            }
78        }
79        Err(Error::NoFamilyId)
80    }
81
82    /// Obtain taskstats for given task ID (e.g. single thread of a multithreaded process)
83    ///
84    /// # Arguments
85    /// * `tid` - Kernel task ID ("pid", "tid" and "task" are used interchangeably and refer to the
86    ///   standard Linux task defined by struct task_struct)
87    ///
88    /// # Return
89    /// * `TaskStats` storing the target task's stats
90    ///
91    /// # Errors
92    /// * when netlink socket failed
93    /// * when kernel responded error
94    /// * when the returned data couldn't be interpreted
95    pub fn pid_stats(&self, tid: u32) -> Result<TaskStats> {
96        self.send(TASKSTATS_CMD_ATTR_PID as u16, tid.as_buf())?;
97
98        let resp = self.netlink.recv_response()?;
99        for na in resp.payload_as_nlattrs() {
100            match na.header.nla_type as u32 {
101                TASKSTATS_TYPE_NULL => break,
102                TASKSTATS_TYPE_AGGR_PID => {
103                    for inner in na.payload_as_nlattrs() {
104                        match inner.header.nla_type as u32 {
105                            TASKSTATS_TYPE_PID => debug!("Received TASKSTATS_TYPE_PID"),
106                            TASKSTATS_TYPE_TGID => warn!("Received TASKSTATS_TYPE_TGID"),
107                            TASKSTATS_TYPE_STATS => {
108                                return Ok(TaskStats::from(inner.payload()));
109                            }
110                            unknown => warn!("Skipping unknown nla_type: {}", unknown),
111                        }
112                    }
113                }
114                unknown => warn!("Skipping unknown nla_type: {}", unknown),
115            }
116        }
117        Err(Error::Unknown(
118            "no TASKSTATS_TYPE_STATS found in response".to_string(),
119        ))
120    }
121
122    /// Obtain taskstats for given thread group ID (e.g. cumulated statistics of a multithreaded process)
123    ///
124    /// # Arguments
125    /// * `tgid` - Kernel thread group ID ("tgid", "process" and "thread group" are used
126    ///   interchangeably and refer to the traditional Unix process)
127    ///
128    /// # Return
129    /// * `TaskStats` storing the target thread group's aggregated stats
130    ///
131    /// # Errors
132    /// * when netlink socket failed
133    /// * when kernel responded error
134    /// * when the returned data couldn't be interpreted
135    pub fn tgid_stats(&self, tgid: u32) -> Result<TaskStats> {
136        self.send(TASKSTATS_CMD_ATTR_TGID as u16, tgid.as_buf())?;
137
138        let resp = self.netlink.recv_response()?;
139        for na in resp.payload_as_nlattrs() {
140            match na.header.nla_type as u32 {
141                TASKSTATS_TYPE_NULL => break,
142                TASKSTATS_TYPE_AGGR_TGID => {
143                    for inner in na.payload_as_nlattrs() {
144                        match inner.header.nla_type as u32 {
145                            TASKSTATS_TYPE_PID => warn!("Received TASKSTATS_TYPE_PID"),
146                            TASKSTATS_TYPE_TGID => debug!("Received TASKSTATS_TYPE_TGID"),
147                            TASKSTATS_TYPE_STATS => {
148                                return Ok(TaskStats::from(inner.payload()));
149                            }
150                            unknown => warn!("Skipping unknown nla_type: {}", unknown),
151                        }
152                    }
153                }
154                unknown => warn!("Skipping unknown nla_type: {}", unknown),
155            }
156        }
157        Err(Error::Unknown(
158            "no TASKSTATS_TYPE_STATS found in response".to_string(),
159        ))
160    }
161
162    /// Register listener with the specific cpumask
163    ///
164    /// # Arguments
165    /// * `cpu_mask` - cpumask is specified as an ascii string of comma-separated cpu ranges e.g.
166    ///   to listen to exit data from cpus 1,2,3,5,7,8 the cpumask would be "1-3,5,7-8".
167    pub fn register_cpumask(&self, cpu_mask: &str) -> Result<()> {
168        self.send(
169            TASKSTATS_CMD_ATTR_REGISTER_CPUMASK as u16,
170            cpu_mask.as_bytes(),
171        )?;
172        Ok(())
173    }
174
175    /// Deregister listener with the specific cpumask
176    /// If userspace forgets to deregister interest in cpus before closing the listening socket,
177    /// the kernel cleans up its interest set over time. However, for the sake of efficiency,
178    /// an explicit deregistration is advisable.
179    ///
180    /// # Arguments
181    /// * `cpu_mask` - cpumask is specified as an ascii string of comma-separated cpu ranges e.g.
182    ///   to listen to exit data from cpus 1,2,3,5,7,8 the cpumask would be "1-3,5,7-8".
183    pub fn deregister_cpumask(&self, cpu_mask: &str) -> Result<()> {
184        self.send(
185            TASKSTATS_CMD_ATTR_DEREGISTER_CPUMASK as u16,
186            cpu_mask.as_bytes(),
187        )?;
188        Ok(())
189    }
190
191    /// Listen registered cpumask's.
192    /// If no messages are available at the socket, the receive call
193    /// wait for a message to arrive, unless the socket is nonblocking.
194    ///
195    /// # Return
196    /// * `Ok(Vec<TaskStats>)`: vector with stats messages. If the current task is NOT the last
197    ///   one in its thread group, only one message is returned in the vector.
198    ///   However, if it is the last task, an additional element containing the per-thread
199    ///   group ID (tgid) statistics is also included. This additional element sums up
200    ///   the statistics for all threads within the thread group, both past and present
201    pub fn listen_registered(&self) -> Result<Vec<TaskStats>> {
202        let resp = self.netlink.recv_response()?;
203        let mut stats_vec = Vec::new();
204
205        for na in resp.payload_as_nlattrs() {
206            match na.header.nla_type as u32 {
207                TASKSTATS_TYPE_NULL => break,
208                TASKSTATS_TYPE_AGGR_PID | TASKSTATS_TYPE_AGGR_TGID => {
209                    for inner in na.payload_as_nlattrs() {
210                        match inner.header.nla_type as u32 {
211                            TASKSTATS_TYPE_PID => debug!("Received TASKSTATS_TYPE_PID"),
212                            TASKSTATS_TYPE_TGID => debug!("Received TASKSTATS_TYPE_TGID"),
213                            TASKSTATS_TYPE_STATS => {
214                                stats_vec.push(TaskStats::from(inner.payload()));
215                            }
216                            unknown => println!("Skipping unknown nla_type: {}", unknown),
217                        }
218                    }
219                }
220                unknown => println!("Skipping unknown nla_type: {}", unknown),
221            }
222        }
223        if !stats_vec.is_empty() {
224            return Ok(stats_vec);
225        }
226        Err(Error::Unknown(
227            "no TASKSTATS_TYPE_STATS found in response".to_string(),
228        ))
229    }
230
231    /// Set receiver buffer size in bytes (SO_RCVBUF socket option, see socket(7))
232    ///
233    /// # Arguments
234    /// * `payload` - buffer size in bytes. The kernel doubles this value
235    ///   (to allow space for bookkeeping overhead). The default value is set by the
236    ///   /proc/sys/net/core/rmem_default file, and the maximum allowed value is set by the
237    ///   /proc/sys/net/core/rmem_max file. The minimum (doubled) value for this option is 256.
238    pub fn set_rx_buf_sz<T>(&self, payload: T) -> Result<()> {
239        self.netlink
240            .set_rx_buf_sz(payload)
241            .map_err(|err| err.into())
242    }
243
244    /// Get receiver buffer size in bytes (SO_RCVBUF socket option, see socket(7))
245    ///
246    /// # Return
247    /// * `usize` buffer size in bytes.
248    ///   Kernel returns doubled value, that have been set using [set_rx_buf_sz]
249    pub fn get_rx_buf_sz(&self) -> Result<usize> {
250        self.netlink.get_rx_buf_sz().map_err(|err| err.into())
251    }
252
253    pub fn send(&self, taskstats_cmd: u16, data: &[u8]) -> Result<()> {
254        self.netlink.send_cmd(
255            self.ts_family_id,
256            TASKSTATS_CMD_GET as u8,
257            taskstats_cmd,
258            data,
259        )?;
260        Ok(())
261    }
262}
263
264trait AsBuf<T> {
265    fn as_buf(&self) -> &[u8];
266
267    fn as_buf_mut(&mut self) -> &mut [u8];
268}
269
270impl<T> AsBuf<T> for T {
271    #[inline]
272    fn as_buf(&self) -> &[u8] {
273        unsafe { slice::from_raw_parts(self as *const T as *const u8, mem::size_of::<T>()) }
274    }
275
276    #[inline]
277    fn as_buf_mut(&mut self) -> &mut [u8] {
278        unsafe { slice::from_raw_parts_mut(self as *mut T as *mut u8, mem::size_of::<T>()) }
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[cfg(test_priv)]
287    #[test]
288    fn test_pid_stats() {
289        let client = Client::open().unwrap();
290        let ts = client.pid_stats(std::process::id()).unwrap();
291
292        // Just asserts some fields which do likely have positive values
293        assert!(ts.delays.cpu.delay_total.as_nanos() > 0);
294        assert!(ts.cpu.virtual_time_total.as_nanos() > 0);
295    }
296
297    #[cfg(test_priv)]
298    #[test]
299    fn test_tgid_stats() {
300        let client = Client::open().unwrap();
301        let ts = client.tgid_stats(std::process::id()).unwrap();
302
303        // Just asserts some fields which do likely have positive values
304        assert!(ts.delays.cpu.delay_total.as_nanos() > 0);
305        assert!(ts.cpu.virtual_time_total.as_nanos() > 0);
306    }
307}
linux_taskstats/lib.rs

linux_taskstats/
lib.rs
