rust_rocksdb/
write_buffer_manager.rs

1use crate::cache::Cache;
2use crate::ffi;
3use libc::size_t;
4use std::ptr::NonNull;
5use std::sync::Arc;
6
7pub(crate) struct WriteBufferManagerWrapper {
8    pub(crate) inner: NonNull<ffi::rocksdb_write_buffer_manager_t>,
9}
10
11unsafe impl Send for WriteBufferManagerWrapper {}
12unsafe impl Sync for WriteBufferManagerWrapper {}
13
14impl Drop for WriteBufferManagerWrapper {
15    fn drop(&mut self) {
16        unsafe {
17            ffi::rocksdb_write_buffer_manager_destroy(self.inner.as_ptr());
18        }
19    }
20}
21
22#[derive(Clone)]
23pub struct WriteBufferManager(pub(crate) Arc<WriteBufferManagerWrapper>);
24
25impl WriteBufferManager {
26    /// <https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager>
27    /// Write buffer manager helps users control the total memory used by memtables across multiple column families and/or DB instances.
28    /// Users can enable this control by 2 ways:
29    ///
30    /// 1- Limit the total memtable usage across multiple column families and DBs under a threshold.
31    /// 2- Cost the memtable memory usage to block cache so that memory of RocksDB can be capped by the single limit.
32    /// The usage of a write buffer manager is similar to rate_limiter and sst_file_manager.
33    /// Users can create one write buffer manager object and pass it to all the options of column families or DBs whose memtable size they want to be controlled by this object.
34    ///
35    /// A memory limit is given when creating the write buffer manager object. RocksDB will try to limit the total memory to under this limit.
36    ///
37    /// a flush will be triggered on one column family of the DB you are inserting to,
38    ///
39    /// If mutable memtable size exceeds about 90% of the limit,
40    /// If the total memory is over the limit, more aggressive flush may also be triggered only if the mutable memtable size also exceeds 50% of the limit.
41    /// Both checks are needed because if already more than half memory is being flushed, triggering more flush may not help.
42    ///
43    /// The total memory is counted as total memory allocated in the arena, even if some of that may not yet be used by memtable.
44    ///
45    /// buffer_size: the memory limit in bytes.
46    /// allow_stall: If set true, it will enable stalling of all writers when memory usage exceeds buffer_size (soft limit).
47    ///             It will wait for flush to complete and memory usage to drop down
48    pub fn new_write_buffer_manager(buffer_size: size_t, allow_stall: bool) -> Self {
49        let inner = NonNull::new(unsafe {
50            ffi::rocksdb_write_buffer_manager_create(buffer_size, allow_stall)
51        })
52        .unwrap();
53        WriteBufferManager(Arc::new(WriteBufferManagerWrapper { inner }))
54    }
55
56    /// Users can set up RocksDB to cost memory used by memtables to block cache.
57    /// This can happen no matter whether you enable memtable memory limit or not.
58    /// This option is added to manage memory (memtables + block cache) under a single limit.
59    ///
60    /// buffer_size: the memory limit in bytes.
61    /// allow_stall: If set true, it will enable stalling of all writers when memory usage exceeds buffer_size (soft limit).
62    ///             It will wait for flush to complete and memory usage to drop down
63    /// cache: the block cache instance
64    pub fn new_write_buffer_manager_with_cache(
65        buffer_size: size_t,
66        allow_stall: bool,
67        cache: Cache,
68    ) -> Self {
69        let inner = NonNull::new(unsafe {
70            ffi::rocksdb_write_buffer_manager_create_with_cache(
71                buffer_size,
72                cache.0.inner.as_ptr(),
73                allow_stall,
74            )
75        })
76        .unwrap();
77        WriteBufferManager(Arc::new(WriteBufferManagerWrapper { inner }))
78    }
79
80    /// Returns the WriteBufferManager memory usage in bytes.
81    pub fn get_usage(&self) -> usize {
82        unsafe { ffi::rocksdb_write_buffer_manager_memory_usage(self.0.inner.as_ptr()) }
83    }
84
85    /// Returns the current buffer size in bytes.
86    pub fn get_buffer_size(&self) -> usize {
87        unsafe { ffi::rocksdb_write_buffer_manager_buffer_size(self.0.inner.as_ptr()) }
88    }
89
90    /// Set the buffer size in bytes.
91    pub fn set_buffer_size(&self, new_size: usize) {
92        unsafe {
93            ffi::rocksdb_write_buffer_manager_set_buffer_size(self.0.inner.as_ptr(), new_size);
94        }
95    }
96
97    /// Returns if WriteBufferManager is enabled.
98    pub fn enabled(&self) -> bool {
99        unsafe { ffi::rocksdb_write_buffer_manager_enabled(self.0.inner.as_ptr()) }
100    }
101
102    /// set the allow_stall flag.
103    pub fn set_allow_stall(&self, allow_stall: bool) {
104        unsafe {
105            ffi::rocksdb_write_buffer_manager_set_allow_stall(self.0.inner.as_ptr(), allow_stall);
106        }
107    }
108}