rtkit_rs/
lib.rs

1// SPDX-License-Identifier: MIT
2// SPDX-FileCopyrightText: Copyright (c) 2025 Asymptotic Inc.
3// SPDX-FileCopyrightText: Copyright (c) 2025 Sanchayan Maity
4
5#![warn(missing_docs)]
6#![doc = include_str!("../README.md")]
7
8use zbus::blocking::Connection;
9use zbus::zvariant::Value;
10use zbus::Result;
11
12const RTKIT_OBJECT_PATH: &str = "/org/freedesktop/RealtimeKit1";
13const RTKIT_SERVICE_NAME: &str = "org.freedesktop.RealtimeKit1";
14const RTKIT_INTERFACE: &str = "org.freedesktop.RealtimeKit1";
15
16fn is_rtkit_available(connection: &Connection) -> Result<bool> {
17    let message = connection.call_method(
18        Some("org.freedesktop.DBus"),
19        "/org/freedesktop/DBus",
20        Some("org.freedesktop.DBus"),
21        "ListNames",
22        &(),
23    )?;
24
25    let names: Vec<String> = message.body().deserialize()?;
26
27    Ok(names.contains(&"org.freedesktop.RealtimeKit1".to_string()))
28}
29
30/// The top-level structure providing access to the crate's functionality.
31pub struct RTKit {
32    connection: Connection,
33}
34
35impl RTKit {
36    /// Create an instance of the `RTKit` structure. This makes a connection to the system D-Bus
37    /// daemon, and ensures that the `rtkit` daemon is available.
38    ///
39    /// Returns an `RTKit` structure if the connection succeeds and the daemon is available, or an
40    /// error otherwise.
41    pub fn new() -> anyhow::Result<RTKit> {
42        let connection = Connection::system()?;
43
44        is_rtkit_available(&connection)?;
45
46        Ok(RTKit { connection })
47    }
48
49    /// Returns the maximum permitted real-time priority value.
50    pub fn max_realtime_priority(&self) -> anyhow::Result<i32> {
51        match self.connection.call_method(
52            Some(RTKIT_SERVICE_NAME),
53            RTKIT_OBJECT_PATH,
54            Some("org.freedesktop.DBus.Properties"),
55            "Get",
56            &("org.freedesktop.RealtimeKit1", "MaxRealtimePriority"),
57        ) {
58            Ok(message) => {
59                let body = message.body().clone().to_owned();
60                let variant: Result<Value> = body.deserialize();
61                match variant {
62                    Ok(value) => Ok(i32::try_from(&value).unwrap()),
63                    Err(e) => Err(e.into()),
64                }
65            }
66            Err(e) => Err(e.into()),
67        }
68    }
69
70    /// Returns the minimum permitted nice level value.
71    pub fn min_nice_level(&self) -> anyhow::Result<i32> {
72        match self.connection.call_method(
73            Some(RTKIT_SERVICE_NAME),
74            RTKIT_OBJECT_PATH,
75            Some("org.freedesktop.DBus.Properties"),
76            "Get",
77            &("org.freedesktop.RealtimeKit1", "MinNiceLevel"),
78        ) {
79            Ok(message) => {
80                let body = message.body().clone().to_owned();
81                let variant: Result<Value> = body.deserialize();
82                match variant {
83                    Ok(value) => Ok(i32::try_from(&value).unwrap()),
84                    Err(e) => Err(e.into()),
85                }
86            }
87            Err(e) => Err(e.into()),
88        }
89    }
90
91    /// Returns the maximum time (in microseconds) that may be set for `RLIMIT_RTTIME`. This is the
92    /// maximum time a real-time thread may continuously occupy the CPU before being blocked by a
93    /// system call.
94    ///
95    /// Applications _must_ set an `RTLIMIT_RTTIME` before attempting to request real-time
96    /// priority.
97    pub fn rttime_usec_max(&self) -> anyhow::Result<i64> {
98        match self.connection.call_method(
99            Some(RTKIT_SERVICE_NAME),
100            RTKIT_OBJECT_PATH,
101            Some("org.freedesktop.DBus.Properties"),
102            "Get",
103            &("org.freedesktop.RealtimeKit1", "RTTimeUSecMax"),
104        ) {
105            Ok(message) => {
106                let body = message.body().clone().to_owned();
107                let variant: Result<Value> = body.deserialize();
108                match variant {
109                    Ok(value) => Ok(i64::try_from(&value).unwrap()),
110                    Err(e) => Err(e.into()),
111                }
112            }
113            Err(e) => Err(e.into()),
114        }
115    }
116
117    /// Requests a nice level of `priority` for the specified thread id (this is a non-real-time
118    /// scheduling level).
119    pub fn make_thread_high_priority(&self, thread_id: u64, priority: i32) -> anyhow::Result<()> {
120        self.connection.call_method(
121            Some(RTKIT_SERVICE_NAME),
122            RTKIT_OBJECT_PATH,
123            Some(RTKIT_INTERFACE),
124            "MakeThreadHighPriority",
125            &(thread_id, priority),
126        )?;
127
128        Ok(())
129    }
130
131    /// Requests a nice level of `priority` for the specified thread id of a specified process id
132    /// (this is a non-real-time scheduling level).
133    pub fn make_thread_high_priority_with_pid(
134        &self,
135        process_id: u64,
136        thread_id: u64,
137        priority: i32,
138    ) -> anyhow::Result<()> {
139        self.connection.call_method(
140            Some(RTKIT_SERVICE_NAME),
141            RTKIT_OBJECT_PATH,
142            Some(RTKIT_INTERFACE),
143            "MakeThreadHighPriorityWithPID",
144            &(process_id, thread_id, priority),
145        )?;
146
147        Ok(())
148    }
149
150    /// Requests a real-time priority of `priority` for the specified thread id.
151    pub fn make_thread_realtime(&self, thread_id: u64, priority: u32) -> anyhow::Result<()> {
152        self.connection.call_method(
153            Some(RTKIT_SERVICE_NAME),
154            RTKIT_OBJECT_PATH,
155            Some(RTKIT_INTERFACE),
156            "MakeThreadRealtime",
157            &(thread_id, priority),
158        )?;
159
160        Ok(())
161    }
162
163    /// Requests a real-time priority of `priority` for the specified thread id of a specified
164    /// process id.
165    pub fn make_thread_realtime_with_pid(
166        &self,
167        process_id: u64,
168        thread_id: u64,
169        priority: u32,
170    ) -> anyhow::Result<()> {
171        self.connection.call_method(
172            Some(RTKIT_SERVICE_NAME),
173            RTKIT_OBJECT_PATH,
174            Some(RTKIT_INTERFACE),
175            "MakeThreadRealtimeWithPID",
176            &(process_id, thread_id, priority),
177        )?;
178
179        Ok(())
180    }
181
182    /// A convenience method to return the calling thread's thread id.
183    pub fn current_thread_id() -> u64 {
184        unsafe { libc::syscall(libc::SYS_gettid) as u64 }
185    }
186
187    /// A convenience method to return the current process id.
188    pub fn current_process_id() -> u64 {
189        std::process::id() as u64
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    fn get_sched_attr() -> anyhow::Result<libc::sched_attr> {
198        unsafe {
199            let mut attr: libc::sched_attr = std::mem::MaybeUninit::zeroed().assume_init();
200
201            let ret = libc::syscall(
202                libc::SYS_sched_getattr,
203                0,
204                &mut attr as *mut libc::sched_attr,
205                std::mem::size_of::<libc::sched_attr>(),
206                0,
207            );
208
209            if ret < 0 {
210                Err(std::io::Error::last_os_error().into())
211            } else {
212                Ok(attr)
213            }
214        }
215    }
216
217    #[test]
218    fn test_property() {
219        let rtkit = RTKit::new().unwrap();
220
221        // Test for default values
222        assert_eq!(rtkit.max_realtime_priority().unwrap(), 20);
223        assert_eq!(rtkit.min_nice_level().unwrap(), -15);
224        assert_eq!(rtkit.rttime_usec_max().unwrap(), 200000);
225    }
226
227    #[test]
228    fn test_thread_id_retrieval() {
229        assert!(RTKit::current_thread_id() > 0);
230    }
231
232    #[test]
233    fn test_process_id_retrieval() {
234        assert!(RTKit::current_process_id() > 0);
235    }
236
237    #[test]
238    fn test_make_thread_high_priority() {
239        let rtkit = RTKit::new().unwrap();
240
241        let thread_id = RTKit::current_thread_id();
242        assert!(rtkit.make_thread_high_priority(thread_id, -10).is_ok());
243
244        let attr = get_sched_attr().unwrap();
245        assert_eq!(attr.sched_nice, -10);
246    }
247
248    #[test]
249    fn test_make_thread_high_priority_with_pid() {
250        let rtkit = RTKit::new().unwrap();
251
252        let pid = RTKit::current_process_id();
253        let thread_id = RTKit::current_thread_id();
254        rtkit
255            .make_thread_high_priority_with_pid(pid, thread_id, -10)
256            .unwrap();
257
258        let attr = get_sched_attr().unwrap();
259        assert_eq!(attr.sched_nice, -10);
260    }
261
262    #[test]
263    fn test_make_thread_realtime() {
264        let rtkit = RTKit::new().unwrap();
265        let rttime_max = rtkit.rttime_usec_max().unwrap() as u64;
266
267        let rlim = libc::rlimit {
268            rlim_cur: rttime_max,
269            rlim_max: rttime_max,
270        };
271
272        let ret = unsafe { libc::setrlimit(libc::RLIMIT_RTTIME, &rlim) };
273        assert_eq!(ret, 0);
274
275        let thread_id = RTKit::current_thread_id();
276        assert!(rtkit.make_thread_realtime(thread_id, 10).is_ok());
277
278        let attr = get_sched_attr().unwrap();
279        assert!(attr.sched_policy > libc::SCHED_OTHER as u32);
280        assert_eq!(attr.sched_priority, 10);
281    }
282
283    #[test]
284    fn test_make_thread_realtime_with_pid() {
285        let rtkit = RTKit::new().unwrap();
286        let rttime_max = rtkit.rttime_usec_max().unwrap() as u64;
287
288        let rlim = libc::rlimit {
289            rlim_cur: rttime_max,
290            rlim_max: rttime_max,
291        };
292
293        let ret = unsafe { libc::setrlimit(libc::RLIMIT_RTTIME, &rlim) };
294        assert_eq!(ret, 0);
295
296        let process_id = RTKit::current_process_id();
297        let thread_id = RTKit::current_thread_id();
298        assert!(rtkit
299            .make_thread_realtime_with_pid(process_id, thread_id, 10)
300            .is_ok());
301
302        let attr = get_sched_attr().unwrap();
303        assert!(attr.sched_policy > libc::SCHED_OTHER as u32);
304        assert_eq!(attr.sched_priority, 10);
305    }
306}