Adding jemalloc-nodump-allocator option for lru-cache (#328)

Adding LRUCacheOptions::set_no_dump_allocator, so we can use JemallocNoDumpAllocator of rocksdb to exclude some LRU Cache, like block cache, from core dump.

librocksdb_sys/crocksdb/c.cc

@@ -173,6 +173,8 @@ using rocksdb::titandb::TitanOptions;
 using rocksdb::titandb::TitanReadOptions;
 using rocksdb::titandb::TitanBlobRunMode;
 
+using rocksdb::MemoryAllocator;
+
 using std::shared_ptr;
 
 extern "C" {
@@ -212,6 +214,7 @@ struct crocksdb_lru_cache_options_t {
   LRUCacheOptions rep;
 };
 struct crocksdb_cache_t           { shared_ptr<Cache>   rep; };
+struct crocksdb_memory_allocator_t { shared_ptr<MemoryAllocator> rep; };
 struct crocksdb_livefiles_t       { std::vector<LiveFileMetaData> rep; };
 struct crocksdb_column_family_handle_t  { ColumnFamilyHandle* rep; };
 struct crocksdb_envoptions_t      { EnvOptions        rep; };
@@ -3303,6 +3306,17 @@ void crocksdb_flushoptions_set_allow_write_stall(
   opt->rep.allow_write_stall = v;
 }
 
+crocksdb_memory_allocator_t* crocksdb_jemalloc_nodump_allocator_create(char** errptr) {
+  crocksdb_memory_allocator_t* allocator = new crocksdb_memory_allocator_t;
+  rocksdb::JemallocAllocatorOptions options;
+  SaveError(errptr, rocksdb::NewJemallocNodumpAllocator(options, &allocator->rep));
+  return allocator;
+}
+
+void crocksdb_memory_allocator_destroy(crocksdb_memory_allocator_t* allocator) {
+  delete allocator;
+}
+
 crocksdb_lru_cache_options_t* crocksdb_lru_cache_options_create() {
   return new crocksdb_lru_cache_options_t;
 }
@@ -3331,6 +3345,11 @@ void crocksdb_lru_cache_options_set_high_pri_pool_ratio(
   opt->rep.high_pri_pool_ratio = high_pri_pool_ratio;
 }
 
+void crocksdb_lru_cache_options_set_memory_allocator(
+    crocksdb_lru_cache_options_t* opt, crocksdb_memory_allocator_t* allocator) {
+    opt->rep.memory_allocator = allocator->rep;
+}
+
 crocksdb_cache_t* crocksdb_cache_create_lru(crocksdb_lru_cache_options_t* opt) {
   crocksdb_cache_t* c = new crocksdb_cache_t;
   c->rep = NewLRUCache(opt->rep);

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -76,6 +76,7 @@ typedef struct crocksdb_backup_engine_info_t   crocksdb_backup_engine_info_t;
 typedef struct crocksdb_restore_options_t crocksdb_restore_options_t;
 typedef struct crocksdb_lru_cache_options_t crocksdb_lru_cache_options_t;
 typedef struct crocksdb_cache_t           crocksdb_cache_t;
+typedef struct crocksdb_memory_allocator_t crocksdb_memory_allocator_t;
 typedef struct crocksdb_compactionfilter_t crocksdb_compactionfilter_t;
 typedef struct crocksdb_compactionfiltercontext_t
     crocksdb_compactionfiltercontext_t;
@@ -1322,6 +1323,12 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_flushoptions_set_wait(
 extern C_ROCKSDB_LIBRARY_API void crocksdb_flushoptions_set_allow_write_stall(
     crocksdb_flushoptions_t*, unsigned char);
 
+/* Memory allocator */
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_memory_allocator_t*
+crocksdb_jemalloc_nodump_allocator_create(char** errptr);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_memory_allocator_destroy(crocksdb_memory_allocator_t*);
+
 /* Cache */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_lru_cache_options_t*
@@ -1336,6 +1343,8 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_lru_cache_options_set_strict_capacity
     crocksdb_lru_cache_options_t*, bool);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_lru_cache_options_set_high_pri_pool_ratio(
     crocksdb_lru_cache_options_t*, double);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_lru_cache_options_set_memory_allocator(
+    crocksdb_lru_cache_options_t*, crocksdb_memory_allocator_t*);
 extern C_ROCKSDB_LIBRARY_API crocksdb_cache_t* crocksdb_cache_create_lru(
     crocksdb_lru_cache_options_t*);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_cache_destroy(crocksdb_cache_t* cache);

librocksdb_sys/src/lib.rs

@@ -29,6 +29,7 @@ pub enum DBWriteOptions {}
 pub enum DBReadOptions {}
 pub enum DBMergeOperator {}
 pub enum DBBlockBasedTableOptions {}
+pub enum DBMemoryAllocator {}
 pub enum DBLRUCacheOptions {}
 pub enum DBCache {}
 pub enum DBFilterPolicy {}
@@ -343,6 +344,10 @@ extern "C" {
         cf_descs: *const ColumnFamilyDescriptor,
     ) -> *mut Options;
 
+    // Memory Allocator
+    pub fn crocksdb_jemalloc_nodump_allocator_create() -> *mut DBMemoryAllocator;
+    pub fn crocksdb_memory_allocator_destroy(allocator: *mut DBMemoryAllocator);
+
     // Cache
     pub fn crocksdb_lru_cache_options_create() -> *mut DBLRUCacheOptions;
     pub fn crocksdb_lru_cache_options_destroy(opt: *mut DBLRUCacheOptions);
@@ -359,6 +364,10 @@ extern "C" {
         opt: *mut DBLRUCacheOptions,
         high_pri_pool_ratio: c_double,
     );
+    pub fn crocksdb_lru_cache_options_set_memory_allocator(
+        opt: *mut DBLRUCacheOptions,
+        allocator: *mut DBMemoryAllocator,
+    );
     pub fn crocksdb_cache_create_lru(opt: *mut DBLRUCacheOptions) -> *mut DBCache;
     pub fn crocksdb_cache_destroy(cache: *mut DBCache);
 

src/lib.rs

@@ -35,8 +35,8 @@ pub use metadata::{ColumnFamilyMetaData, LevelMetaData, SstFileMetaData};
 pub use perf_context::{get_perf_level, set_perf_level, IOStatsContext, PerfContext, PerfLevel};
 pub use rocksdb::{
     load_latest_options, run_ldb_tool, set_external_sst_file_global_seq_no, BackupEngine, CFHandle,
-    Cache, DBIterator, DBVector, Env, ExternalSstFileInfo, Kv, Range, SeekKey, SequentialFile,
-    SstFileWriter, Writable, WriteBatch, DB,
+    Cache, DBIterator, DBVector, Env, ExternalSstFileInfo, Kv, MemoryAllocator, Range, SeekKey,
+    SequentialFile, SstFileWriter, Writable, WriteBatch, DB,
 };
 pub use rocksdb_options::{
     BlockBasedOptions, CColumnFamilyDescriptor, ColumnFamilyOptions, CompactOptions,

src/rocksdb.rs

@@ -18,6 +18,7 @@ use crocksdb_ffi::{
     DBWriteBatch,
 };
 use libc::{self, c_char, c_int, c_void, size_t};
+use librocksdb_sys::DBMemoryAllocator;
 use metadata::ColumnFamilyMetaData;
 use rocksdb_options::{
     CColumnFamilyDescriptor, ColumnFamilyDescriptor, ColumnFamilyOptions, CompactOptions,
@@ -2376,6 +2377,30 @@ impl Drop for Cache {
     }
 }
 
+pub struct MemoryAllocator {
+    pub inner: *mut DBMemoryAllocator,
+}
+
+impl MemoryAllocator {
+    #[cfg(feature = "jemalloc")]
+    pub fn new_jemalloc_memory_allocator() -> Result<MemoryAllocator, String> {
+        unsafe {
+            let allocator = MemoryAllocator {
+                inner: ffi_try!(crocksdb_ffi::crocksdb_jemalloc_nodump_allocator_create()),
+            };
+            Ok(allocator)
+        }
+    }
+}
+
+impl Drop for MemoryAllocator {
+    fn drop(&mut self) {
+        unsafe {
+            crocksdb_ffi::crocksdb_memory_allocator_destroy(self.inner);
+        }
+    }
+}
+
 pub fn set_external_sst_file_global_seq_no(
     db: &DB,
     cf: &CFHandle,

src/rocksdb_options.rs

@@ -27,8 +27,8 @@ use event_listener::{new_event_listener, EventListener};
 use libc::{self, c_double, c_int, c_uchar, c_void, size_t};
 use merge_operator::MergeFn;
 use merge_operator::{self, full_merge_callback, partial_merge_callback, MergeOperatorCallback};
-use rocksdb::Cache;
 use rocksdb::Env;
+use rocksdb::{Cache, MemoryAllocator};
 use slice_transform::{new_slice_transform, SliceTransform};
 use std::ffi::{CStr, CString};
 use std::mem;
@@ -1877,6 +1877,15 @@ impl LRUCacheOptions {
             );
         }
     }
+
+    pub fn set_memory_allocator(&mut self, allocator: &MemoryAllocator) {
+        unsafe {
+            crocksdb_ffi::crocksdb_lru_cache_options_set_memory_allocator(
+                self.inner,
+                allocator.inner,
+            );
+        }
+    }
 }
 
 impl Drop for LRUCacheOptions {

tests/cases/test_rocksdb_options.rs

@@ -15,6 +15,7 @@ use rocksdb::crocksdb_ffi::{
     CompactionPriority, DBCompressionType, DBInfoLogLevel as InfoLogLevel, DBRateLimiterMode,
     DBStatisticsHistogramType as HistogramType, DBStatisticsTickerType as TickerType,
 };
+use rocksdb::rocksdb::MemoryAllocator;
 use rocksdb::{
     BlockBasedOptions, Cache, ColumnFamilyOptions, CompactOptions, DBOptions, Env,
     FifoCompactionOptions, IndexType, LRUCacheOptions, ReadOptions, SeekKey, SliceTransform,
@@ -318,6 +319,22 @@ fn test_set_lru_cache() {
     DB::open_cf(opts, path.path().to_str().unwrap(), vec!["default"]).unwrap();
 }
 
+#[cfg(feature = "jemalloc")]
+#[test]
+fn test_set_jemalloc_nodump_allocator_for_lru_cache() {
+    let path = TempDir::new("_rust_rocksdb_set_jemalloc_nodump_allocator").expect("");
+    let mut opts = DBOptions::new();
+    let mut cf_opts = ColumnFamilyOptions::new();
+    opts.create_if_missing(true);
+    let mut block_opts = BlockBasedOptions::new();
+    let mut cache_opts = LRUCacheOptions::new();
+    cache_opts.set_memory_allocator(MemoryAllocator::new_jemalloc_memory_allocator().unwrap());
+    cache_opts.set_capacity(8388608);
+    block_opts.set_block_cache(&Cache::new_lru_cache(cache_opts));
+    cf_opts.set_block_based_table_factory(&block_opts);
+    DB::open_cf(opts, path.path().to_str().unwrap(), vec!["default"]).unwrap();
+}
+
 #[test]
 fn test_set_cache_index_and_filter_blocks_with_high_priority() {
     let path = TempDir::new("_rust_rocksdb_set_cache_and_index_with_high_priority").expect("");