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

rust_rocksdb/
perf.rs
