rust_rocksdb/
perf.rs

1// Copyright 2020 Tran Tuan Linh
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use libc::{c_int, c_uchar};
16use std::{
17    cell::{Cell, RefCell},
18    marker::PhantomData,
19};
20
21use crate::cache::Cache;
22use crate::ffi_util::from_cstr_and_free;
23use crate::{DB, DBCommon, ThreadMode, TransactionDB};
24use crate::{Error, db::DBInner, ffi};
25
26#[derive(Debug, Copy, Clone, PartialEq, Eq)]
27#[repr(i32)]
28pub enum PerfStatsLevel {
29    /// Unknown settings
30    Uninitialized = 0,
31    /// Disable perf stats
32    Disable = 1,
33    /// Enables only count stats
34    EnableCount = 2,
35    /// Enables count stats and time spent waiting for RocksDB work
36    EnableWait = 3,
37    /// Count stats and enable time stats except for mutexes
38    EnableTimeExceptForMutex = 4,
39    /// Other than time, also measure CPU time counters. Still don't measure
40    /// time (neither wall time nor CPU time) for mutexes
41    EnableTimeAndCPUTimeExceptForMutex = 5,
42    /// Enables count and time stats
43    EnableTime = 6,
44    /// N.B must always be the last value!
45    OutOfBound = 7,
46}
47
48// Include the generated PerfMetric enum from perf_enum.rs
49include!("perf_enum.rs");
50
51/// Sets the perf stats level for current thread.
52pub fn set_perf_stats(lvl: PerfStatsLevel) {
53    unsafe {
54        ffi::rocksdb_set_perf_level(lvl as c_int);
55    }
56}
57
58/// Thread local context for gathering performance counter efficiently
59/// and transparently.
60pub struct PerfContext {
61    pub(crate) inner: *mut ffi::rocksdb_perfcontext_t,
62    reusable: bool,
63}
64
65thread_local! {
66    static ACTIVE_MANUAL_PERF_CONTEXTS: Cell<usize> = const { Cell::new(0) };
67    static REUSABLE_PERF_CONTEXT: RefCell<PerfContext> = RefCell::new({
68        let inner = unsafe { ffi::rocksdb_perfcontext_create() };
69        assert!(!inner.is_null(), "Could not create Perf Context");
70        PerfContext {
71            inner,
72            reusable: true,
73        }
74    });
75}
76
77impl Default for PerfContext {
78    fn default() -> Self {
79        let ctx = unsafe { ffi::rocksdb_perfcontext_create() };
80        assert!(!ctx.is_null(), "Could not create Perf Context");
81        ACTIVE_MANUAL_PERF_CONTEXTS.with(|count| count.set(count.get() + 1));
82
83        Self {
84            inner: ctx,
85            reusable: false,
86        }
87    }
88}
89
90impl Drop for PerfContext {
91    fn drop(&mut self) {
92        if !self.reusable {
93            ACTIVE_MANUAL_PERF_CONTEXTS.with(|count| count.set(count.get() - 1));
94        }
95        unsafe {
96            ffi::rocksdb_perfcontext_destroy(self.inner);
97        }
98    }
99}
100
101impl PerfContext {
102    /// Reset context
103    #[inline]
104    pub fn reset(&mut self) {
105        unsafe {
106            ffi::rocksdb_perfcontext_reset(self.inner);
107        }
108    }
109
110    /// Get the report on perf
111    pub fn report(&self, exclude_zero_counters: bool) -> String {
112        unsafe {
113            let ptr =
114                ffi::rocksdb_perfcontext_report(self.inner, c_uchar::from(exclude_zero_counters));
115            from_cstr_and_free(ptr)
116        }
117    }
118
119    /// Returns value of a metric
120    #[inline]
121    pub fn metric(&self, id: PerfMetric) -> u64 {
122        unsafe { ffi::rocksdb_perfcontext_metric(self.inner, id as c_int) }
123    }
124}
125
126/// Runs `f` with a reusable thread-local [`PerfContext`].
127///
128/// The context is reset before each call. This avoids allocating and destroying
129/// the C wrapper for every measurement.
130///
131/// # Panics
132///
133/// Panics if called reentrantly or while a manually created [`PerfContext`] is
134/// alive on the same thread. RocksDB returns the same underlying context to all
135/// wrappers on a thread, so resetting it would discard the manual measurement.
136pub fn with_thread_local<F, R>(f: F) -> R
137where
138    F: FnOnce(&mut PerfContext) -> R,
139{
140    ACTIVE_MANUAL_PERF_CONTEXTS.with(|count| {
141        assert_eq!(
142            count.get(),
143            0,
144            "with_thread_local cannot run while a manual PerfContext is alive on the same thread"
145        );
146    });
147    REUSABLE_PERF_CONTEXT.with(|ctx| {
148        let mut ctx = ctx.try_borrow_mut().unwrap_or_else(|_| {
149            panic!("with_thread_local cannot be called reentrantly on the same thread")
150        });
151        ctx.reset();
152        f(&mut ctx)
153    })
154}
155
156/// Memory usage stats
157pub struct MemoryUsageStats {
158    /// Approximate memory usage of all the mem-tables
159    pub mem_table_total: u64,
160    /// Approximate memory usage of un-flushed mem-tables
161    pub mem_table_unflushed: u64,
162    /// Approximate memory usage of all the table readers
163    pub mem_table_readers_total: u64,
164    /// Approximate memory usage by cache
165    pub cache_total: u64,
166}
167
168/// Wrap over memory_usage_t. Hold current memory usage of the specified DB instances and caches
169pub struct MemoryUsage {
170    inner: *mut ffi::rocksdb_memory_usage_t,
171}
172
173impl Drop for MemoryUsage {
174    fn drop(&mut self) {
175        unsafe {
176            ffi::rocksdb_approximate_memory_usage_destroy(self.inner);
177        }
178    }
179}
180
181impl MemoryUsage {
182    /// Approximate memory usage of all the mem-tables
183    pub fn approximate_mem_table_total(&self) -> u64 {
184        unsafe { ffi::rocksdb_approximate_memory_usage_get_mem_table_total(self.inner) }
185    }
186
187    /// Approximate memory usage of un-flushed mem-tables
188    pub fn approximate_mem_table_unflushed(&self) -> u64 {
189        unsafe { ffi::rocksdb_approximate_memory_usage_get_mem_table_unflushed(self.inner) }
190    }
191
192    /// Approximate memory usage of all the table readers
193    pub fn approximate_mem_table_readers_total(&self) -> u64 {
194        unsafe { ffi::rocksdb_approximate_memory_usage_get_mem_table_readers_total(self.inner) }
195    }
196
197    /// Approximate memory usage by cache
198    pub fn approximate_cache_total(&self) -> u64 {
199        unsafe { ffi::rocksdb_approximate_memory_usage_get_cache_total(self.inner) }
200    }
201}
202
203/// Creates [`MemoryUsage`] from DBs and caches.
204///
205/// Most users should call [`get_memory_usage_stats`] instead.
206///
207/// A `MemoryUsageBuilder` must not outlive the `DB`s added to it:
208///
209/// ```compile_fail,E0597
210/// use rust_rocksdb::{perf::MemoryUsageBuilder, DB};
211///
212/// let mut builder = MemoryUsageBuilder::new().unwrap();
213/// {
214///     let db = DB::open_default("foo").unwrap();
215///     builder.add_db(&db);
216/// }
217/// let _memory_usage = builder.build().unwrap();
218/// ```
219pub struct MemoryUsageBuilder<'a> {
220    inner: *mut ffi::rocksdb_memory_consumers_t,
221    base_dbs: Vec<*mut ffi::rocksdb_t>,
222    // must not outlive the DBs/caches that are added
223    _marker: PhantomData<&'a ()>,
224}
225
226impl Drop for MemoryUsageBuilder<'_> {
227    fn drop(&mut self) {
228        unsafe {
229            ffi::rocksdb_memory_consumers_destroy(self.inner);
230        }
231        for base_db in &self.base_dbs {
232            unsafe {
233                ffi::rocksdb_transactiondb_close_base_db(*base_db);
234            }
235        }
236    }
237}
238
239impl<'a> MemoryUsageBuilder<'a> {
240    /// Create new instance
241    pub fn new() -> Result<Self, Error> {
242        let mc = unsafe { ffi::rocksdb_memory_consumers_create() };
243        if mc.is_null() {
244            Err(Error::new(
245                "Could not create MemoryUsage builder".to_owned(),
246            ))
247        } else {
248            Ok(Self {
249                inner: mc,
250                base_dbs: Vec::new(),
251                _marker: PhantomData,
252            })
253        }
254    }
255
256    /// Add a DB instance to collect memory usage from it and add up in total stats
257    pub fn add_tx_db<T: ThreadMode>(&mut self, db: &'a TransactionDB<T>) {
258        unsafe {
259            let base_db = ffi::rocksdb_transactiondb_get_base_db(db.inner);
260            ffi::rocksdb_memory_consumers_add_db(self.inner, base_db);
261            // rocksdb_transactiondb_get_base_db allocates a struct that must be freed
262            self.base_dbs.push(base_db);
263        }
264    }
265
266    /// Add a DB instance to collect memory usage from it and add up in total stats
267    pub fn add_db<T: ThreadMode, D: DBInner>(&mut self, db: &'a DBCommon<T, D>) {
268        unsafe {
269            ffi::rocksdb_memory_consumers_add_db(self.inner, db.inner.inner());
270        }
271    }
272
273    /// Add a cache to collect memory usage from it and add up in total stats
274    pub fn add_cache(&mut self, cache: &'a Cache) {
275        unsafe {
276            ffi::rocksdb_memory_consumers_add_cache(self.inner, cache.0.inner.as_ptr());
277        }
278    }
279
280    /// Build up MemoryUsage
281    pub fn build(&self) -> Result<MemoryUsage, Error> {
282        unsafe {
283            let mu = ffi_try!(ffi::rocksdb_approximate_memory_usage_create(self.inner));
284            Ok(MemoryUsage { inner: mu })
285        }
286    }
287}
288
289/// Get memory usage stats from DB instances and Cache instances
290pub fn get_memory_usage_stats(
291    dbs: Option<&[&DB]>,
292    caches: Option<&[&Cache]>,
293) -> Result<MemoryUsageStats, Error> {
294    let mut builder = MemoryUsageBuilder::new()?;
295    if let Some(dbs_) = dbs {
296        for db in dbs_ {
297            builder.add_db(db);
298        }
299    }
300    if let Some(caches_) = caches {
301        for cache in caches_ {
302            builder.add_cache(cache);
303        }
304    }
305
306    let mu = builder.build()?;
307    Ok(MemoryUsageStats {
308        mem_table_total: mu.approximate_mem_table_total(),
309        mem_table_unflushed: mu.approximate_mem_table_unflushed(),
310        mem_table_readers_total: mu.approximate_mem_table_readers_total(),
311        cache_total: mu.approximate_cache_total(),
312    })
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318    use crate::{DB, Options};
319    use std::panic;
320    use tempfile::TempDir;
321
322    #[test]
323    fn perf_stats_level_matches_rocksdb_11_1_2() {
324        assert_eq!(PerfStatsLevel::Uninitialized as i32, 0);
325        assert_eq!(PerfStatsLevel::Disable as i32, 1);
326        assert_eq!(PerfStatsLevel::EnableCount as i32, 2);
327        assert_eq!(PerfStatsLevel::EnableWait as i32, 3);
328        assert_eq!(PerfStatsLevel::EnableTimeExceptForMutex as i32, 4);
329        assert_eq!(PerfStatsLevel::EnableTimeAndCPUTimeExceptForMutex as i32, 5);
330        assert_eq!(PerfStatsLevel::EnableTime as i32, 6);
331        assert_eq!(PerfStatsLevel::OutOfBound as i32, 7);
332    }
333
334    #[test]
335    fn with_thread_local_reuses_context_after_unwind() {
336        let first = with_thread_local(|ctx| ctx.inner);
337        let panic_result = panic::catch_unwind(|| {
338            with_thread_local(|_| panic!("test panic"));
339        });
340        assert!(panic_result.is_err());
341        let second = with_thread_local(|ctx| ctx.inner);
342
343        assert_eq!(first, second);
344    }
345
346    #[test]
347    fn with_thread_local_resets_context_before_use() {
348        let temp_dir = TempDir::new().unwrap();
349        let mut opts = Options::default();
350        opts.create_if_missing(true);
351        let db = DB::open(&opts, temp_dir.path()).unwrap();
352        db.put(b"key", b"value").unwrap();
353
354        set_perf_stats(PerfStatsLevel::EnableCount);
355        let comparison_count = with_thread_local(|ctx| {
356            db.get(b"key").unwrap();
357            ctx.metric(PerfMetric::UserKeyComparisonCount)
358        });
359        assert!(comparison_count > 0);
360        assert_eq!(
361            with_thread_local(|ctx| ctx.metric(PerfMetric::UserKeyComparisonCount)),
362            0
363        );
364        set_perf_stats(PerfStatsLevel::Disable);
365    }
366
367    #[test]
368    #[should_panic(expected = "with_thread_local cannot be called reentrantly on the same thread")]
369    fn with_thread_local_rejects_reentrant_use() {
370        with_thread_local(|_| with_thread_local(|_| ()));
371    }
372
373    #[test]
374    #[should_panic(
375        expected = "with_thread_local cannot run while a manual PerfContext is alive on the same thread"
376    )]
377    fn with_thread_local_rejects_active_manual_context() {
378        let _manual = PerfContext::default();
379        with_thread_local(|_| ());
380    }
381
382    #[test]
383    fn test_perf_context_with_db_operations() {
384        let temp_dir = TempDir::new().unwrap();
385        let mut opts = Options::default();
386        opts.create_if_missing(true);
387        let db = DB::open(&opts, temp_dir.path()).unwrap();
388
389        // Insert data with deletions to test internal key/delete skipping
390        let n = 10;
391        for i in 0..n {
392            let k = vec![i as u8];
393            db.put(&k, &k).unwrap();
394            if i % 2 == 0 {
395                db.delete(&k).unwrap();
396            }
397        }
398
399        set_perf_stats(PerfStatsLevel::EnableCount);
400        let mut ctx = PerfContext::default();
401
402        // Use iterator with explicit seek to trigger metrics
403        let mut iter = db.raw_iterator();
404        iter.seek_to_first();
405        let mut valid_count = 0;
406        while iter.valid() {
407            valid_count += 1;
408            iter.next();
409        }
410
411        // Check counts - should have 5 valid entries (odd numbers: 1,3,5,7,9)
412        assert_eq!(
413            valid_count, 5,
414            "Iterator should find 5 valid entries (odd numbers)"
415        );
416
417        // Check internal skip metrics
418        let internal_key_skipped = ctx.metric(PerfMetric::InternalKeySkippedCount);
419        let internal_delete_skipped = ctx.metric(PerfMetric::InternalDeleteSkippedCount);
420
421        // In RocksDB, when iterating over deleted keys in SST files:
422        // - We should skip the deletion markers (n/2 = 5 deletes)
423        // - Total internal keys skipped should be >= number of deletions
424        assert!(
425            internal_key_skipped >= (n / 2) as u64,
426            "internal_key_skipped ({}) should be >= {} (deletions)",
427            internal_key_skipped,
428            n / 2
429        );
430        assert_eq!(
431            internal_delete_skipped,
432            (n / 2) as u64,
433            "internal_delete_skipped ({internal_delete_skipped}) should equal {} (deleted entries)",
434            n / 2
435        );
436        assert_eq!(
437            ctx.metric(PerfMetric::SeekInternalSeekTime),
438            0,
439            "Time metrics should be 0 with EnableCount"
440        );
441
442        // Test reset
443        ctx.reset();
444        assert_eq!(ctx.metric(PerfMetric::InternalKeySkippedCount), 0);
445        assert_eq!(ctx.metric(PerfMetric::InternalDeleteSkippedCount), 0);
446
447        // Change perf level to EnableTime
448        set_perf_stats(PerfStatsLevel::EnableTime);
449
450        // Iterate backwards
451        let mut iter = db.raw_iterator();
452        iter.seek_to_last();
453        let mut backward_count = 0;
454        while iter.valid() {
455            backward_count += 1;
456            iter.prev();
457        }
458        assert_eq!(
459            backward_count, 5,
460            "Backward iteration should also find 5 valid entries"
461        );
462
463        // Check accumulated metrics after second iteration
464        let key_skipped_after = ctx.metric(PerfMetric::InternalKeySkippedCount);
465        let delete_skipped_after = ctx.metric(PerfMetric::InternalDeleteSkippedCount);
466
467        // After both iterations, we should have accumulated more skipped keys
468        assert!(
469            key_skipped_after >= internal_key_skipped,
470            "After second iteration, internal_key_skipped ({key_skipped_after}) should be >= first iteration ({internal_key_skipped})",
471        );
472        assert_eq!(
473            delete_skipped_after,
474            (n / 2) as u64,
475            "internal_delete_skipped should still be {} after second iteration",
476            n / 2
477        );
478
479        // Disable perf stats
480        set_perf_stats(PerfStatsLevel::Disable);
481    }
482}
rust_rocksdb/perf.rs

rust_rocksdb/
perf.rs
