tikv_jemalloc_ctl/
keys.rs

1//! Key types to index the _MALLCTL NAMESPACE_.
2//!
3//! The [`Name`] and [`Mib`]/[`MibStr`] types are provided as safe indices into
4//! the _MALLCTL NAMESPACE_. These are constructed from null-terminated strings
5//! via the [`AsName`] trait. The [`Access`] trait provides provides safe access
6//! into the `_MALLCTL NAMESPACE_`.
7//!
8//! # Example
9//!
10//! ```
11//! #[global_allocator]
12//! static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
13//!
14//! fn main() {
15//!     use tikv_jemalloc_ctl::{Access, AsName, Name, Mib};
16//!     use libc::{c_uint, c_char};
17//!     let name = b"arenas.nbins\0".name();
18//!     let nbins: c_uint = name.read().unwrap();
19//!     let mut mib: Mib<[usize; 4]> = b"arenas.bin.0.size\0".name().mib().unwrap();
20//!     for i in 0..4 {
21//!         mib[2] = i;
22//!         let bin_size: usize = mib.read().unwrap();
23//!         println!("arena bin {} has size {}", i, bin_size);
24//!     }
25//! }
26//! ```
27
28#![allow(clippy::uninlined_format_args)]
29
30use crate::error::Result;
31use crate::std::str;
32use crate::{fmt, ops, raw};
33
34/// A `Name` in the _MALLCTL NAMESPACE_.
35#[repr(transparent)]
36#[derive(PartialEq, Eq)]
37pub struct Name([u8]);
38
39/// Converts a null-terminated byte-string into a [`Name`].
40pub trait AsName {
41    /// Converts a null-terminated byte-string into a [`Name`].
42    fn name(&self) -> &Name;
43}
44
45impl AsName for [u8] {
46    fn name(&self) -> &Name {
47        assert!(
48            !self.is_empty(),
49            "cannot create Name from empty byte-string"
50        );
51        assert_eq!(
52            *self.last().unwrap(),
53            b'\0',
54            "cannot create Name from non-null-terminated byte-string \"{}\"",
55            str::from_utf8(self).unwrap()
56        );
57        unsafe { &*(self as *const Self as *const Name) }
58    }
59}
60
61impl AsName for str {
62    fn name(&self) -> &Name {
63        self.as_bytes().name()
64    }
65}
66
67impl Name {
68    /// Returns the [`Mib`] of `self`.
69    pub fn mib<T: MibArg>(&self) -> Result<Mib<T>> {
70        let mut mib: Mib<T> = Mib::default();
71        raw::name_to_mib(&self.0, mib.0.as_mut())?;
72        Ok(mib)
73    }
74
75    /// Returns the [`MibStr`] of `self` which is a key whose value is a string.
76    pub fn mib_str<T: MibArg>(&self) -> Result<MibStr<T>> {
77        assert!(
78            self.value_type_str(),
79            "key \"{}\" does not refer to a string",
80            self
81        );
82        let mut mib: MibStr<T> = MibStr::default();
83        raw::name_to_mib(&self.0, mib.0.as_mut())?;
84        Ok(mib)
85    }
86
87    /// Returns `true` if `self` is a key in the _MALLCTL NAMESPCE_ referring to
88    /// a null-terminated string.
89    pub fn value_type_str(&self) -> bool {
90        // remove the null-terminator:
91        let name = self.0.split_at(self.0.len() - 1).0;
92        if name.is_empty() {
93            return false;
94        }
95        debug_assert_ne!(*name.last().unwrap(), b'\0');
96
97        match name {
98            b"version"
99            | b"config.malloc_conf"
100            | b"opt.metadata_thp"
101            | b"opt.dss"
102            | b"opt.percpu_arena"
103            | b"opt.stats_print_opts"
104            | b"opt.junk"
105            | b"opt.thp"
106            | b"opt.prof_prefix"
107            | b"thread.prof.name"
108            | b"prof.dump" => true,
109            v if v.starts_with(b"arena.") && v.ends_with(b".dss") => true,
110            v if v.starts_with(b"stats.arenas.") && v.ends_with(b".dss") => {
111                true
112            }
113            _ => false,
114        }
115    }
116
117    /// Returns the name as null-terminated byte-string.
118    pub fn as_bytes(&self) -> &'static [u8] {
119        unsafe { &*(self as *const Self as *const [u8]) }
120    }
121}
122
123impl fmt::Debug for Name {
124    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
125        write!(f, "{}", str::from_utf8(&self.0).unwrap())
126    }
127}
128
129impl fmt::Display for Name {
130    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
131        write!(f, "{}", str::from_utf8(&self.0).unwrap())
132    }
133}
134
135/// Management Information Base of a non-string value.
136#[repr(transparent)]
137#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)]
138pub struct Mib<T: MibArg>(T);
139
140/// Management Information Base of a string value.
141#[repr(transparent)]
142#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)]
143pub struct MibStr<T: MibArg>(T);
144
145impl<T: MibArg> AsRef<[usize]> for Mib<T> {
146    fn as_ref(&self) -> &[usize] {
147        self.0.as_ref()
148    }
149}
150
151impl<T: MibArg> AsMut<[usize]> for Mib<T> {
152    fn as_mut(&mut self) -> &mut [usize] {
153        self.0.as_mut()
154    }
155}
156
157impl<T: MibArg> ops::Index<usize> for Mib<T> {
158    type Output = usize;
159    fn index(&self, idx: usize) -> &Self::Output {
160        &self.0.as_ref()[idx]
161    }
162}
163
164impl<T: MibArg> ops::IndexMut<usize> for Mib<T> {
165    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
166        &mut self.0.as_mut()[idx]
167    }
168}
169
170impl<T: MibArg> ops::Index<usize> for MibStr<T> {
171    type Output = usize;
172    fn index(&self, idx: usize) -> &Self::Output {
173        &self.0.as_ref()[idx]
174    }
175}
176
177impl<T: MibArg> ops::IndexMut<usize> for MibStr<T> {
178    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
179        &mut self.0.as_mut()[idx]
180    }
181}
182
183/// Safe read access to the _MALLCTL NAMESPACE_.
184pub trait Access<T> {
185    /// Read the key at `self`.
186    fn read(&self) -> Result<T>;
187    /// Write `value` at the key `self`.
188    fn write(&self, value: T) -> Result<()>;
189    /// Write `value` at the key `self` returning its previous value.
190    fn update(&self, value: T) -> Result<T>;
191}
192
193macro_rules! impl_access {
194    ($id:ty) => {
195        impl<T: MibArg> Access<$id> for Mib<T> {
196            fn read(&self) -> Result<$id> {
197                unsafe { raw::read_mib(self.0.as_ref()) }
198            }
199            fn write(&self, value: $id) -> Result<()> {
200                unsafe { raw::write_mib(self.0.as_ref(), value) }
201            }
202            fn update(&self, value: $id) -> Result<$id> {
203                unsafe { raw::update_mib(self.0.as_ref(), value) }
204            }
205        }
206        impl Access<$id> for Name {
207            fn read(&self) -> Result<$id> {
208                unsafe { raw::read(&self.0) }
209            }
210            fn write(&self, value: $id) -> Result<()> {
211                unsafe { raw::write(&self.0, value) }
212            }
213            fn update(&self, value: $id) -> Result<$id> {
214                unsafe { raw::update(&self.0, value) }
215            }
216        }
217    };
218}
219
220impl_access!(u32);
221impl_access!(u64);
222impl_access!(isize);
223impl_access!(usize);
224
225impl<T: MibArg> Access<bool> for Mib<T> {
226    fn read(&self) -> Result<bool> {
227        unsafe {
228            let v: u8 = raw::read_mib(self.0.as_ref())?;
229            assert!(v == 0 || v == 1);
230            Ok(v == 1)
231        }
232    }
233    fn write(&self, value: bool) -> Result<()> {
234        unsafe { raw::write_mib(self.0.as_ref(), value) }
235    }
236    fn update(&self, value: bool) -> Result<bool> {
237        unsafe {
238            let v: u8 = raw::update_mib(self.0.as_ref(), value as u8)?;
239            Ok(v == 1)
240        }
241    }
242}
243
244impl Access<bool> for Name {
245    fn read(&self) -> Result<bool> {
246        unsafe {
247            let v: u8 = raw::read(&self.0)?;
248            assert!(v == 0 || v == 1);
249            Ok(v == 1)
250        }
251    }
252    fn write(&self, value: bool) -> Result<()> {
253        unsafe { raw::write(&self.0, value) }
254    }
255    fn update(&self, value: bool) -> Result<bool> {
256        unsafe {
257            let v: u8 = raw::update(&self.0, value as u8)?;
258            Ok(v == 1)
259        }
260    }
261}
262
263impl<T: MibArg> Access<&'static [u8]> for MibStr<T> {
264    fn read(&self) -> Result<&'static [u8]> {
265        // this is safe because the only safe way to construct a `MibStr` is by
266        // validating that the key refers to a byte-string value
267        unsafe { raw::read_str_mib(self.0.as_ref()) }
268    }
269    fn write(&self, value: &'static [u8]) -> Result<()> {
270        raw::write_str_mib(self.0.as_ref(), value)
271    }
272    fn update(&self, value: &'static [u8]) -> Result<&'static [u8]> {
273        // this is safe because the only safe way to construct a `MibStr` is by
274        // validating that the key refers to a byte-string value
275        unsafe { raw::update_str_mib(self.0.as_ref(), value) }
276    }
277}
278
279impl Access<&'static [u8]> for Name {
280    fn read(&self) -> Result<&'static [u8]> {
281        assert!(
282            self.value_type_str(),
283            "the name \"{:?}\" does not refer to a byte string",
284            self
285        );
286        // this is safe because the key refers to a byte string:
287        unsafe { raw::read_str(&self.0) }
288    }
289    fn write(&self, value: &'static [u8]) -> Result<()> {
290        assert!(
291            self.value_type_str(),
292            "the name \"{:?}\" does not refer to a byte string",
293            self
294        );
295        raw::write_str(&self.0, value)
296    }
297    fn update(&self, value: &'static [u8]) -> Result<&'static [u8]> {
298        assert!(
299            self.value_type_str(),
300            "the name \"{:?}\" does not refer to a byte string",
301            self
302        );
303        // this is safe because the key refers to a byte string:
304        unsafe { raw::update_str(&self.0, value) }
305    }
306}
307
308impl<T: MibArg> Access<&'static str> for MibStr<T> {
309    fn read(&self) -> Result<&'static str> {
310        // this is safe because the only safe way to construct a `MibStr` is by
311        // validating that the key refers to a byte-string value
312        let s = unsafe { raw::read_str_mib(self.0.as_ref())? };
313        Ok(str::from_utf8(s).unwrap())
314    }
315    fn write(&self, value: &'static str) -> Result<()> {
316        raw::write_str_mib(self.0.as_ref(), value.as_bytes())
317    }
318    fn update(&self, value: &'static str) -> Result<&'static str> {
319        // this is safe because the only safe way to construct a `MibStr` is by
320        // validating that the key refers to a byte-string value
321        let s =
322            unsafe { raw::update_str_mib(self.0.as_ref(), value.as_bytes())? };
323        Ok(str::from_utf8(s).unwrap())
324    }
325}
326
327impl Access<&'static str> for Name {
328    fn read(&self) -> Result<&'static str> {
329        assert!(
330            self.value_type_str(),
331            "the name \"{:?}\" does not refer to a byte string",
332            self
333        );
334        // this is safe because the key refers to a byte string:
335        let s = unsafe { raw::read_str(&self.0)? };
336        Ok(str::from_utf8(s).unwrap())
337    }
338    fn write(&self, value: &'static str) -> Result<()> {
339        assert!(
340            self.value_type_str(),
341            "the name \"{:?}\" does not refer to a byte string",
342            self
343        );
344        raw::write_str(&self.0, value.as_bytes())
345    }
346    fn update(&self, value: &'static str) -> Result<&'static str> {
347        assert!(
348            self.value_type_str(),
349            "the name \"{:?}\" does not refer to a byte string",
350            self
351        );
352        // this is safe because the key refers to a byte string:
353        let s = unsafe { raw::update_str(&self.0, value.as_bytes())? };
354        Ok(str::from_utf8(s).unwrap())
355    }
356}
357
358#[cfg(test)]
359mod tests {
360    use super::{Access, AsName, Mib, MibStr};
361    #[test]
362    fn bool_rw() {
363        let name = b"thread.tcache.enabled\0".name();
364        let tcache: bool = name.read().unwrap();
365
366        let new_tcache = !tcache;
367
368        name.write(new_tcache).unwrap();
369
370        let mib: Mib<[usize; 3]> = name.mib().unwrap();
371        let r: bool = mib.read().unwrap();
372        assert_eq!(r, new_tcache);
373    }
374
375    #[test]
376    fn u32_r() {
377        let name = b"arenas.bin.0.nregs\0".name();
378        let v: u32 = name.read().unwrap();
379
380        let mib: Mib<[usize; 4]> = name.mib().unwrap();
381        let r: u32 = mib.read().unwrap();
382        assert_eq!(r, v);
383    }
384
385    #[test]
386    fn size_t_r() {
387        let name = b"arenas.lextent.0.size\0".name();
388        let v: libc::size_t = name.read().unwrap();
389
390        let mib: Mib<[usize; 4]> = name.mib().unwrap();
391        let r: libc::size_t = mib.read().unwrap();
392        assert_eq!(r, v);
393    }
394
395    #[test]
396    fn ssize_t_rw() {
397        let name = b"arenas.dirty_decay_ms\0".name();
398        let v: libc::ssize_t = name.read().unwrap();
399        name.write(v).unwrap();
400
401        let mib: Mib<[usize; 2]> = name.mib().unwrap();
402        let r: libc::ssize_t = mib.read().unwrap();
403        assert_eq!(r, v);
404    }
405
406    #[test]
407    fn u64_rw() {
408        let name = b"epoch\0".name();
409        let epoch: u64 = name.read().unwrap();
410        name.write(epoch).unwrap();
411
412        let mib: Mib<[usize; 1]> = name.mib().unwrap();
413        let epoch: u64 = mib.read().unwrap();
414        mib.write(epoch).unwrap();
415    }
416
417    #[test]
418    fn str_rw() {
419        let name = b"arena.0.dss\0".name();
420        let dss: &'static [u8] = name.read().unwrap();
421        name.write(dss).unwrap();
422
423        let mib: MibStr<[usize; 3]> = name.mib_str().unwrap();
424        let dss2: &'static [u8] = mib.read().unwrap();
425        mib.write(dss2).unwrap();
426
427        assert_eq!(dss, dss2);
428    }
429}
430
431pub trait MibArg:
432    Copy
433    + Clone
434    + PartialEq
435    + Default
436    + fmt::Debug
437    + AsRef<[usize]>
438    + AsMut<[usize]>
439{
440}
441impl<T> MibArg for T where
442    T: Copy
443        + Clone
444        + PartialEq
445        + Default
446        + fmt::Debug
447        + AsRef<[usize]>
448        + AsMut<[usize]>
449{
450}