1#![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
30pub struct RTKit {
32 connection: Connection,
33}
34
35impl RTKit {
36 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 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 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 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 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 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 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 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 pub fn current_thread_id() -> u64 {
184 unsafe { libc::syscall(libc::SYS_gettid) as u64 }
185 }
186
187 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 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}