rtkit_client_rs/
lib.rs

1mod low_level_zbus_api;
2mod types;
3
4pub use types::{NiceLevel, Priority};
5
6#[derive(thiserror::Error, Debug)]
7pub enum Error {
8    #[error("Permission denied: {0}")]
9    PermissionDenied(String),
10
11    #[error("Invalid PID or Thread ID")]
12    InvalidPidOrTid,
13
14    #[error("Got invalid max priority from RTKit: {0}")]
15    InvalidMaxPriority(i32),
16
17    #[error("Got invalid min nice level from RTKit: {0}")]
18    InvalidMinNiceLevel(i32),
19
20    #[error(transparent)]
21    ZbusError(#[from] zbus::Error),
22
23    #[error(transparent)]
24    ZbusMethodError(#[from] zbus::fdo::Error),
25}
26
27/// Helper for [make_thread_realtime], uses the current thread's tid.
28pub fn set_current_thread_priority(priority: Option<Priority>) -> Result<Priority, Error> {
29    let thread_id = nix::unistd::gettid().as_raw().try_into().unwrap();
30    set_thread_priority(thread_id, priority)
31}
32
33/// Elevates the realtime priority of a thread as close as possible to the requested priority.
34///
35/// If no priority is specified, the maximum priority allowed by the system will be used.
36///
37/// Returns the actual priority that was set.
38pub fn set_thread_priority(tid: u64, priority: Option<Priority>) -> Result<Priority, Error> {
39    let connection = zbus::blocking::Connection::system()?;
40    let proxy = low_level_zbus_api::RTKitProxyBlocking::builder(&connection)
41        .cache_properties(zbus::proxy::CacheProperties::No)
42        .build()?;
43
44    let max_priority = proxy.max_realtime_priority()?;
45    let max_priority = max_priority
46        .try_into()
47        .map_err(|_| Error::InvalidMaxPriority(max_priority))?;
48    let priority = priority.unwrap_or_default().min(max_priority);
49
50    proxy
51        .make_thread_realtime(tid, priority.value())
52        .map(|_| priority)
53        .map_err(|e| match e {
54            zbus::fdo::Error::AccessDenied(err) => Error::PermissionDenied(err),
55            // NOTE: RTKit returns "No such file or directory" for invalid tids.
56            zbus::fdo::Error::Failed(err) if err == "No such file or directory" => {
57                Error::InvalidPidOrTid
58            }
59            err => Error::ZbusMethodError(err),
60        })
61}
62
63/// Elevates the realtime priority of a specified processes thread as close as possible to the requested priority.
64///
65/// If no priority is specified, the maximum priority allowed by the system will be used.
66///
67/// Returns the actual priority that was set.
68pub fn set_process_thread_priority(
69    pid: u64,
70    tid: u64,
71    priority: Option<Priority>,
72) -> Result<Priority, Error> {
73    let connection = zbus::blocking::Connection::system()?;
74    let proxy = low_level_zbus_api::RTKitProxyBlocking::builder(&connection)
75        .cache_properties(zbus::proxy::CacheProperties::No)
76        .build()?;
77
78    let max_priority = proxy.max_realtime_priority()?;
79    let max_priority = max_priority
80        .try_into()
81        .map_err(|_| Error::InvalidMaxPriority(max_priority))?;
82    let priority = priority.unwrap_or_default().min(max_priority);
83
84    proxy
85        .make_thread_realtime_with_pid(pid, tid, priority.value())
86        .map(|_| priority)
87        .map_err(|e| match e {
88            zbus::fdo::Error::AccessDenied(err) => Error::PermissionDenied(err),
89            // NOTE: RTKit returns "No such file or directory" for invalid tids.
90            zbus::fdo::Error::Failed(err) if err == "No such file or directory" => {
91                Error::InvalidPidOrTid
92            }
93            err => Error::ZbusMethodError(err),
94        })
95}
96
97/// Helper for [set_thread_nice_level], uses the current thread's tid.
98pub fn set_current_thread_nice_level(nice_level: Option<NiceLevel>) -> Result<NiceLevel, Error> {
99    let thread_id = nix::unistd::gettid().as_raw().try_into().unwrap();
100    set_thread_nice_level(thread_id, nice_level)
101}
102
103/// Elevates the nice level of a thread as close as possible to the requested nice level.
104///
105/// If no nice level is specified, the maximum nice level allowed by the system will be used.
106///
107/// Returns the actual nice level that was set.
108pub fn set_thread_nice_level(tid: u64, nice_level: Option<NiceLevel>) -> Result<NiceLevel, Error> {
109    let connection = zbus::blocking::Connection::system()?;
110    let proxy = low_level_zbus_api::RTKitProxyBlocking::builder(&connection)
111        .cache_properties(zbus::proxy::CacheProperties::No)
112        .build()?;
113
114    let nice_level = nice_level.unwrap_or_default();
115
116    let max_nice_level = proxy.min_nice_level()?;
117    let max_nice_level = max_nice_level
118        .try_into()
119        .map_err(|_| Error::InvalidMinNiceLevel(max_nice_level))?;
120    let nice_level = nice_level.max(max_nice_level);
121
122    proxy
123        .make_thread_high_priority(tid, nice_level.value())
124        .map(|_| nice_level)
125        .map_err(|e| match e {
126            zbus::fdo::Error::AccessDenied(err) => Error::PermissionDenied(err),
127            // NOTE: RTKit returns "No such file or directory" for invalid tids.
128            zbus::fdo::Error::Failed(err) if err == "No such file or directory" => {
129                Error::InvalidPidOrTid
130            }
131            err => Error::ZbusMethodError(err),
132        })
133}
134
135/// Elevates the nice level of a specified processes thread as close as possible to the requested nice level.
136///
137/// If no nice level is specified, the maximum nice level allowed by the system will be used.
138/// Returns the actual nice level that was set.
139pub fn set_process_thread_nice_level(
140    pid: u64,
141    tid: u64,
142    nice_level: Option<NiceLevel>,
143) -> Result<NiceLevel, Error> {
144    let connection = zbus::blocking::Connection::system()?;
145    let proxy = low_level_zbus_api::RTKitProxyBlocking::builder(&connection)
146        .cache_properties(zbus::proxy::CacheProperties::No)
147        .build()?;
148
149    let nice_level = nice_level.unwrap_or_default();
150    let max_nice_level = proxy.min_nice_level()?;
151    let max_nice_level = max_nice_level
152        .try_into()
153        .map_err(|_| Error::InvalidMinNiceLevel(max_nice_level))?;
154    let nice_level = nice_level.max(max_nice_level);
155
156    proxy
157        .make_thread_high_priority_with_pid(pid, tid, nice_level.value())
158        .map(|_| nice_level)
159        .map_err(|e| match e {
160            zbus::fdo::Error::AccessDenied(err) => Error::PermissionDenied(err),
161            // NOTE: RTKit returns "No such file or directory" for invalid tids.
162            zbus::fdo::Error::Failed(err) if err == "No such file or directory" => {
163                Error::InvalidPidOrTid
164            }
165            err => Error::ZbusMethodError(err),
166        })
167}