support prof.dump and prof.prefix

ci/run.sh

@@ -26,7 +26,7 @@ fi
 
 cargo build --target "${TARGET}"
 cargo test --target "${TARGET}"
-cargo test --target "${TARGET}" --features profiling
+_RJEM_MALLOC_CONF="prof:true" cargo test --target "${TARGET}" --features profiling
 cargo test --target "${TARGET}" --features debug
 cargo test --target "${TARGET}" --features stats
 cargo test --target "${TARGET}" --features 'debug profiling'

jemalloc-ctl/src/keys.rs

@@ -31,6 +31,8 @@ use crate::error::Result;
 use crate::std::str;
 use crate::{fmt, ops, raw};
 
+use super::ffi::CStr;
+
 /// A `Name` in the _MALLCTL NAMESPACE_.
 #[repr(transparent)]
 #[derive(PartialEq, Eq)]
@@ -105,7 +107,8 @@ impl Name {
             | b"opt.thp"
             | b"opt.prof_prefix"
             | b"thread.prof.name"
-            | b"prof.dump" => true,
+            | b"prof.dump"
+            | b"prof.prefix" => true,
             v if v.starts_with(b"arena.") && v.ends_with(b".dss") => true,
             v if v.starts_with(b"stats.arenas.") && v.ends_with(b".dss") => {
                 true
@@ -355,6 +358,58 @@ impl Access<&'static str> for Name {
     }
 }
 
+impl<T: MibArg> Access<&'static CStr> for MibStr<T> {
+    fn read(&self) -> Result<&'static CStr> {
+        // this is safe because the only safe way to construct a `MibStr` is by
+        // validating that the key refers to a byte-string value
+        let s = unsafe { raw::read_str_mib(self.0.as_ref())? };
+        Ok(CStr::from_bytes_with_nul(s).unwrap())
+    }
+    fn write(&self, value: &'static CStr) -> Result<()> {
+        raw::write_str_mib(self.0.as_ref(), value.to_bytes_with_nul())
+    }
+    fn update(&self, value: &'static CStr) -> Result<&'static CStr> {
+        // this is safe because the only safe way to construct a `MibStr` is by
+        // validating that the key refers to a byte-string value
+        let s = unsafe {
+            raw::update_str_mib(self.0.as_ref(), value.to_bytes_with_nul())?
+        };
+        Ok(CStr::from_bytes_with_nul(s).unwrap())
+    }
+}
+
+impl Access<&'static CStr> for Name {
+    fn read(&self) -> Result<&'static CStr> {
+        assert!(
+            self.value_type_str(),
+            "the name \"{:?}\" does not refer to a byte string",
+            self
+        );
+        // this is safe because the key refers to a byte string:
+        let s = unsafe { raw::read_str(&self.0)? };
+        Ok(CStr::from_bytes_with_nul(s).unwrap())
+    }
+    fn write(&self, value: &'static CStr) -> Result<()> {
+        assert!(
+            self.value_type_str(),
+            "the name \"{:?}\" does not refer to a byte string",
+            self
+        );
+        raw::write_str(&self.0, value.to_bytes_with_nul())
+    }
+    fn update(&self, value: &'static CStr) -> Result<&'static CStr> {
+        assert!(
+            self.value_type_str(),
+            "the name \"{:?}\" does not refer to a byte string",
+            self
+        );
+        // this is safe because the key refers to a byte string:
+        let s =
+            unsafe { raw::update_str(&self.0, value.to_bytes_with_nul())? };
+        Ok(CStr::from_bytes_with_nul(s).unwrap())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{Access, AsName, Mib, MibStr};

jemalloc-ctl/src/lib.rs

@@ -74,7 +74,7 @@
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-use crate::std::{fmt, mem, num, ops, ptr, result, slice, str};
+use crate::std::{ffi, fmt, mem, num, ops, ptr, result, slice, str};
 #[cfg(not(feature = "use_std"))]
 use core as std;
 #[cfg(feature = "use_std")]
@@ -88,6 +88,7 @@ pub mod config;
 mod error;
 mod keys;
 pub mod opt;
+pub mod prof;
 pub mod raw;
 pub mod stats;
 #[cfg(feature = "use_std")]

jemalloc-ctl/src/macros.rs

@@ -43,7 +43,7 @@ macro_rules! types {
 
 /// Read
 macro_rules! r {
-    ($id:ident => $ret_ty:ty) => {
+    ($id:ident[ str: $byte_string:expr ] => $ret_ty:ty) => {
         paste::paste! {
             impl $id {
                 /// Reads value using string API.
@@ -72,6 +72,12 @@ macro_rules! r {
                     if cfg!(target_os = "macos") => return,
                     _ => (),
                 }
+                match $byte_string.as_slice() {
+                    b"opt.prof\0" |
+                    b"prof.active\0"
+                    if !cfg!(feature = "profiling") => return,
+                    _ => (),
+                }
 
                 let a = $id::read().unwrap();
 
@@ -92,7 +98,7 @@ macro_rules! r {
 
 /// Write
 macro_rules! w {
-    ($id:ident => $ret_ty:ty) => {
+    ($id:ident[ str: $byte_string:expr ] => $ret_ty:ty) => {
         paste::paste! {
             impl $id {
                 /// Writes `value` using string API.
@@ -114,24 +120,51 @@ macro_rules! w {
             #[test]
             #[cfg(not(target_arch = "mips64el"))]
             fn [<$id _write_test>]() {
+                /// Help test write
+                pub trait WriteTestDefault {
+                    fn default() -> Self;
+                }
+                macro_rules! impl_write_test_default {
+                    ($write_ty:ty, $val:expr) => {
+                        impl WriteTestDefault for $write_ty {
+                            fn default() -> $write_ty {
+                                $val
+                            }
+                        }
+                    };
+                }
+
+                use crate::ffi::CStr;
+                impl_write_test_default! {libc::size_t, 0}
+                impl_write_test_default! {u64, 0}
+                impl_write_test_default! {bool, false}
+                impl_write_test_default! {&'static CStr, CStr::from_bytes_with_nul(b"test\0").unwrap()}
+
                 match stringify!($id) {
                     "background_thread" |
                     "max_background_threads"
                         if cfg!(target_os = "macos") => return,
                     _ => (),
                 }
+                match $byte_string.as_slice() {
+                    b"prof.dump\0" |
+                    b"prof.active\0" |
+                    b"prof.prefix\0"
+                    if !cfg!(feature = "profiling") => return,
+                    _ => (),
+                }
 
-                let _ = $id::write($ret_ty::default()).unwrap();
+                let _ = $id::write(<$ret_ty as WriteTestDefault>::default()).unwrap();
 
                 let mib = $id::mib().unwrap();
-                let _ = mib.write($ret_ty::default()).unwrap();
+                let _ = mib.write(<$ret_ty as WriteTestDefault>::default()).unwrap();
 
                 #[cfg(feature = "use_std")]
                 println!(
                     concat!(
                         stringify!($id),
                         " (write): \"{}\""),
-                    $ret_ty::default()
+                        <$ret_ty as Default>::default()
                 );
 
             }
@@ -141,7 +174,7 @@ macro_rules! w {
 
 /// Update
 macro_rules! u {
-    ($id:ident  => $ret_ty:ty) => {
+    ($id:ident[ str: $byte_string:expr ] => $ret_ty:ty) => {
         paste::paste! {
             impl $id {
                 /// Updates key to `value` returning its old value using string API.
@@ -170,6 +203,11 @@ macro_rules! u {
                         if cfg!(target_os = "macos") => return,
                     _ => (),
                 }
+                match $byte_string.as_slice() {
+                    b"prof.active\0"
+                        if !cfg!(feature = "profiling") => return,
+                    _ => (),
+                }
 
                 let a = $id::update($ret_ty::default()).unwrap();
 
@@ -203,7 +241,7 @@ macro_rules! option {
             mib_docs: $(#[$doc_mib])*
         }
         $(
-            $ops!($id => $ret_ty);
+            $ops!($id[ str: $byte_string ] => $ret_ty);
         )*
     };
     // Non-string option:

jemalloc-ctl/src/opt.rs

@@ -215,3 +215,27 @@ option! {
     /// ```
     mib_docs: /// See [`background_thread`].
 }
+
+option! {
+    prof[ str: b"opt.prof\0", non_str: 2 ] => bool |
+    ops: r |
+    docs:
+    /// Memory profiling enabled/disabled. If enabled, profile memory allocation activity.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[global_allocator]
+    /// # static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+    /// #
+    /// # fn main() {
+    /// #[cfg(feature = "profiling")]
+    /// {
+    ///     use tikv_jemalloc_ctl::opt;
+    ///     let prof = opt::prof::read().unwrap();
+    ///     println!("Jemalloc profiling enabled: {}", prof);
+    /// }
+    /// # }
+    /// ```
+    mib_docs: /// See [`prof`].
+}

jemalloc-ctl/src/prof.rs

@@ -0,0 +1,86 @@
+//! `jemalloc`'s profiling utils.
+
+use crate::ffi::CStr;
+
+option! {
+    dump[ str: b"prof.dump\0", str: 2 ] => &'static CStr |
+    ops: w |
+    docs:
+    /// Dump a memory profile to the specified file, or if NULL is specified,
+    /// to a file according to the pattern <prefix>.<pid>.<seq>.m<mseq>.heap,
+    /// where <prefix> is controlled by the opt.prof_prefix and prof.prefix options.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[global_allocator]
+    /// # static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+    /// #
+    /// # fn main() {
+    /// #[cfg(feature = "profiling")]
+    /// {
+    ///     use tikv_jemalloc_ctl::prof;
+    ///     use std::ffi::CStr;
+    ///     let dump_file_name = CStr::from_bytes_with_nul(b"dump\0").unwrap();
+    ///     let dump = prof::dump::mib().unwrap();
+    ///     dump.write(dump_file_name).unwrap();
+    /// }
+    /// # }
+    /// ```
+    mib_docs: /// See [`dump`].
+}
+
+option! {
+    prefix[ str: b"prof.prefix\0", str: 2 ] => &'static CStr |
+    ops: w |
+    docs:
+    /// Set the filename prefix for profile dumps. See opt.prof_prefix for the default setting.
+    ///
+    /// This can be useful to differentiate profile dumps such as from forked processes.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[global_allocator]
+    /// # static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+    /// #
+    /// # fn main() {
+    /// #[cfg(feature = "profiling")]
+    /// {
+    ///     use tikv_jemalloc_ctl::prof;
+    ///     use std::ffi::CStr;
+    ///     let dump_file_name = CStr::from_bytes_with_nul(b"my_prefix\0").unwrap();
+    ///     let prefix = prof::prefix::mib().unwrap();
+    ///     prefix.write(dump_file_name).unwrap();
+    /// }
+    /// # }
+    /// ```
+    mib_docs: /// See [`prefix`].
+}
+
+option! {
+    active[ str: b"prof.active\0", non_str: 2 ] => bool |
+    ops: r, w, u |
+    docs:
+    /// Control whether sampling is currently active.
+    ///
+    /// See the `opt.prof_active` option for additional information,
+    /// as well as the interrelated `thread.prof.active` mallctl.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[global_allocator]
+    /// # static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+    /// #
+    /// # fn main() {
+    /// #[cfg(feature = "profiling")]
+    /// {
+    ///     use tikv_jemalloc_ctl::prof;
+    ///     let active = prof::active::mib().unwrap();
+    ///     active.write(true).unwrap();
+    /// }
+    /// # }
+    /// ```
+    mib_docs: /// See [`active`].
+}