feat(cubestore): Enable jemalloc allocator for RocksDB on linux GNU (#…

…8104)

refs #7545

rust/cubestore/Cargo.toml

@@ -5,5 +5,6 @@ members = [
     "cubestore-sql-tests",
     "cubehll",
     "cubezetasketch",
-    "cuberpc"
+    "cuberpc",
+    "cuberockstore",
 ]

rust/cubestore/Dockerfile

@@ -23,6 +23,7 @@ WORKDIR /usr/src
 WORKDIR /build/cubestore
 COPY Cargo.toml .
 COPY Cargo.lock .
+COPY cuberockstore cuberockstore
 COPY cubehll cubehll
 COPY cubezetasketch cubezetasketch
 COPY cuberpc cuberpc

rust/cubestore/cuberockstore/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "cuberockstore"
+version = "0.1.0"
+authors = ["Cube Dev, Inc."]
+edition = "2021"
+license = "Apache-2.0"
+description = "Cube Rocks Store"
+
+[dependencies]
+
+[target.'cfg(all(target_os = "linux", target_env = "gnu"))'.dependencies]
+# jemalloc is supported only for gnu linux
+# TODO: Fix issue with cross building for linux musl & jemalloc
+rocksdb = { git = "https://github.com/cube-js/rust-rocksdb", branch = "cubestore", default-features = false, features = ["zstd", "snappy", "jemalloc"] }
+
+[target.'cfg(all(target_os = "linux", not(target_env = "gnu")))'.dependencies]
+rocksdb = { git = "https://github.com/cube-js/rust-rocksdb", branch = "cubestore", default-features = false, features = ["zstd", "snappy"] }
+
+[target.'cfg(not(target_os = "linux"))'.dependencies]
+rocksdb = { git = "https://github.com/cube-js/rust-rocksdb", branch = "cubestore", default-features = false, features = ["zstd", "snappy"] }

rust/cubestore/cuberockstore/src/lib.rs

@@ -0,0 +1,2 @@
+// re-export
+pub use rocksdb;

rust/cubestore/cubestore-sql-tests/Cargo.toml

@@ -53,7 +53,7 @@ tar = "0.4.38"
 
 [dev-dependencies]
 criterion = { version = "0.4.0", features = ["async_tokio", "html_reports"] }
-rocksdb = { git = "https://github.com/cube-js/rust-rocksdb", branch = "cubestore", default-features = false, features = ["zstd", "snappy"] }
+cuberockstore = { path = "../cuberockstore" }
 
 [[bench]]
 name = "in_process"

rust/cubestore/cubestore-sql-tests/benches/in_process.rs

@@ -1,6 +1,6 @@
 use criterion::{criterion_group, criterion_main, Criterion};
+use cuberockstore::rocksdb;
 use cubestore_sql_tests::cubestore_benches;
-use rocksdb::{Options, DB};
 use std::fs;
 use std::time::Duration;
 use tokio::runtime::Builder;
@@ -41,8 +41,8 @@ fn in_process_bench(criterion: &mut Criterion) {
             runtime.shutdown_timeout(Duration::from_secs(2));
         }
 
-        let _ = DB::destroy(&Options::default(), config.meta_store_path());
-        let _ = DB::destroy(&Options::default(), config.cache_store_path());
+        let _ = rocksdb::DB::destroy(&rocksdb::Options::default(), config.meta_store_path());
+        let _ = rocksdb::DB::destroy(&rocksdb::Options::default(), config.cache_store_path());
         let _ = fs::remove_dir_all(config.local_dir().clone());
     }
 }

rust/cubestore/cubestore/Cargo.toml

@@ -23,6 +23,7 @@ serde_derive = "1.0.115"
 serde = "1.0.115"
 serde_repr = "0.1"
 serde_bytes = "0.11.5"
+cuberockstore = { path = "../cuberockstore" }
 cubehll = { path = "../cubehll" }
 cubezetasketch = { path = "../cubezetasketch" }
 cuberpc = { path = "../cuberpc" }
@@ -41,7 +42,6 @@ simple_logger = "1.7.0"
 async-trait = "0.1.79"
 actix-rt = "2.7.0"
 regex = "1.3.9"
-rocksdb = { git = "https://github.com/cube-js/rust-rocksdb", branch = "cubestore", default-features = false, features = ["zstd", "snappy"] }
 uuid = { version = "0.8", features = ["serde", "v4"] }
 num = "0.3.0"
 enum_primitive = "0.1.1"

rust/cubestore/cubestore/src/cachestore/cache_item.rs

@@ -4,7 +4,7 @@ use crate::metastore::{
 use crate::{base_rocks_secondary_index, rocks_table_new, CubeError};
 use chrono::serde::ts_seconds_option;
 use chrono::{DateTime, Duration, Utc};
-use rocksdb::WriteBatch;
+use cuberockstore::rocksdb::WriteBatch;
 use serde::{Deserialize, Deserializer, Serialize};
 
 #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]

rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs

@@ -22,15 +22,16 @@ use crate::util::WorkerLoop;
 use crate::{app_metrics, CubeError};
 use async_trait::async_trait;
 
+use cuberockstore::rocksdb::{BlockBasedOptions, Cache, Options, DB};
 use futures_timer::Delay;
-use rocksdb::{BlockBasedOptions, Cache, Options, DB};
 
 use crate::cachestore::cache_eviction_manager::{CacheEvictionManager, EvictionResult};
 use crate::cachestore::compaction::CompactionPreloadedState;
 use crate::cachestore::listener::RocksCacheStoreListener;
 use crate::cachestore::queue_item_payload::QueueItemPayloadRocksTable;
 use crate::table::{Row, TableValue};
 use chrono::{DateTime, Utc};
+use cuberockstore::rocksdb;
 use datafusion::cube_ext;
 use deepsize::DeepSizeOf;
 use itertools::Itertools;
@@ -87,7 +88,9 @@ impl RocksStoreDetails for RocksCacheStoreDetails {
 
         let mut opts = Options::default();
         opts.create_if_missing(true);
-        opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(13));
+        opts.set_prefix_extractor(cuberockstore::rocksdb::SliceTransform::create_fixed_prefix(
+            13,
+        ));
         opts.set_compaction_filter_factory(compaction::MetaStoreCacheCompactionFactory::new(
             compaction_state,
         ));
@@ -125,7 +128,9 @@ impl RocksStoreDetails for RocksCacheStoreDetails {
         let rocksdb_config = config.cachestore_rocksdb_config();
 
         let mut opts = Options::default();
-        opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(13));
+        opts.set_prefix_extractor(cuberockstore::rocksdb::SliceTransform::create_fixed_prefix(
+            13,
+        ));
 
         let block_opts = {
             let mut block_opts = BlockBasedOptions::default();

rust/cubestore/cubestore/src/cachestore/compaction.rs

@@ -4,10 +4,12 @@ use crate::metastore::{
 use crate::TableId;
 
 use chrono::{DateTime, Utc};
+use cuberockstore::rocksdb::compaction_filter::CompactionFilter;
+use cuberockstore::rocksdb::compaction_filter_factory::{
+    CompactionFilterContext, CompactionFilterFactory,
+};
+use cuberockstore::rocksdb::CompactionDecision;
 use log::{error, trace, warn};
-use rocksdb::compaction_filter::CompactionFilter;
-use rocksdb::compaction_filter_factory::{CompactionFilterContext, CompactionFilterFactory};
-use rocksdb::CompactionDecision;
 use std::collections::HashMap;
 use std::ffi::{CStr, CString};
 use std::sync::{Arc, Mutex};
@@ -248,7 +250,7 @@ mod tests {
     use crate::metastore::{BaseRocksSecondaryIndex, RocksTable};
     use crate::TableId;
     use chrono::Duration;
-    use rocksdb::compaction_filter::Decision;
+    use cuberockstore::rocksdb::compaction_filter::Decision;
     use serde::Serialize;
 
     fn get_test_filter_context() -> CompactionFilterContext {

rust/cubestore/cubestore/src/cachestore/queue_item.rs

@@ -6,7 +6,7 @@ use crate::table::{Row, TableValue};
 use crate::{base_rocks_secondary_index, rocks_table_new, CubeError};
 use chrono::serde::{ts_seconds, ts_seconds_option};
 use chrono::{DateTime, Duration, Utc};
-use rocksdb::WriteBatch;
+use cuberockstore::rocksdb::WriteBatch;
 use std::cmp::Ordering;
 use std::sync::Arc;
 

rust/cubestore/cubestore/src/cachestore/queue_item_payload.rs

@@ -4,7 +4,7 @@ use crate::metastore::{
 use crate::{base_rocks_secondary_index, rocks_table_new, CubeError};
 use chrono::serde::{ts_seconds, ts_seconds_option};
 use chrono::{DateTime, Duration, Utc};
-use rocksdb::WriteBatch;
+use cuberockstore::rocksdb::WriteBatch;
 use serde::{Deserialize, Deserializer, Serialize};
 
 #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]

rust/cubestore/cubestore/src/cachestore/queue_result.rs

@@ -6,7 +6,7 @@ use crate::metastore::{
 use crate::{base_rocks_secondary_index, rocks_table_new, CubeError};
 use chrono::serde::ts_seconds;
 use chrono::{DateTime, Duration, Utc};
-use rocksdb::WriteBatch;
+use cuberockstore::rocksdb::WriteBatch;
 use serde::{Deserialize, Deserializer, Serialize};
 
 #[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]

rust/cubestore/cubestore/src/config/mod.rs

@@ -43,13 +43,13 @@ use crate::telemetry::{
 };
 use crate::util::memory::{MemoryHandler, MemoryHandlerImpl};
 use crate::CubeError;
+use cuberockstore::rocksdb::{Options, DB};
 use datafusion::cube_ext;
 use datafusion::physical_plan::parquet::{LruParquetMetadataCache, NoopParquetMetadataCache};
 use futures::future::join_all;
 use log::Level;
 use log::{debug, error};
 use mockall::automock;
-use rocksdb::{Options, DB};
 use simple_logger::SimpleLogger;
 use std::fmt::Display;
 use std::future::Future;

rust/cubestore/cubestore/src/lib.rs

@@ -197,8 +197,8 @@ impl From<flexbuffers::DeserializationError> for CubeError {
     }
 }
 
-impl From<rocksdb::Error> for CubeError {
-    fn from(v: rocksdb::Error) -> Self {
+impl From<cuberockstore::rocksdb::Error> for CubeError {
+    fn from(v: cuberockstore::rocksdb::Error) -> Self {
         CubeError::from_error(v.into_string())
     }
 }

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -22,8 +22,8 @@ pub use rocks_table::*;
 use crate::cluster::node_name_by_partition;
 use crate::metastore::partition::partition_file_name;
 use async_trait::async_trait;
+use cuberockstore::rocksdb::{BlockBasedOptions, Cache, Env, MergeOperands, Options, DB};
 use log::info;
-use rocksdb::{BlockBasedOptions, Cache, Env, MergeOperands, Options, DB};
 use serde::{Deserialize, Serialize};
 use std::hash::Hash;
 use std::{env, io::Cursor, sync::Arc};
@@ -61,6 +61,7 @@ use chrono::{DateTime, Utc};
 use chunks::ChunkRocksTable;
 use core::fmt;
 use cubehll::HllSketch;
+use cuberockstore::rocksdb::backup::{BackupEngine, BackupEngineOptions};
 use cubezetasketch::HyperLogLogPlusPlus;
 use datafusion::cube_ext;
 use futures_timer::Delay;
@@ -72,7 +73,6 @@ use parquet::basic::{ConvertedType, Repetition};
 use parquet::{basic::Type, schema::types};
 use partition::{PartitionRocksIndex, PartitionRocksTable};
 use regex::Regex;
-use rocksdb::backup::{BackupEngine, BackupEngineOptions};
 
 use schema::{SchemaRocksIndex, SchemaRocksTable};
 use smallvec::alloc::fmt::Formatter;
@@ -1263,7 +1263,9 @@ impl RocksStoreDetails for RocksMetaStoreDetails {
         let rocksdb_config = config.metastore_rocksdb_config();
         let mut opts = Options::default();
         opts.create_if_missing(true);
-        opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(13));
+        opts.set_prefix_extractor(cuberockstore::rocksdb::SliceTransform::create_fixed_prefix(
+            13,
+        ));
         opts.set_merge_operator_associative("meta_store merge", meta_store_merge);
         // TODO(ovr): Decrease after additional fix for get_updates_since
         opts.set_wal_ttl_seconds(
@@ -1296,7 +1298,9 @@ impl RocksStoreDetails for RocksMetaStoreDetails {
     fn open_readonly_db(&self, path: &Path, config: &Arc<dyn ConfigObj>) -> Result<DB, CubeError> {
         let rocksdb_config = config.metastore_rocksdb_config();
         let mut opts = Options::default();
-        opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(13));
+        opts.set_prefix_extractor(cuberockstore::rocksdb::SliceTransform::create_fixed_prefix(
+            13,
+        ));
 
         let block_opts = {
             let mut block_opts = BlockBasedOptions::default();
@@ -4881,8 +4885,8 @@ mod tests {
     use super::*;
     use crate::config::{init_test_logger, Config};
     use crate::remotefs::{LocalDirRemoteFs, RemoteFs};
+    use cuberockstore::rocksdb::IteratorMode;
     use futures_timer::Delay;
-    use rocksdb::IteratorMode;
     use std::thread::sleep;
     use std::time::Duration;
     use std::{env, fs};

rust/cubestore/cubestore/src/metastore/rocks_store.rs

@@ -9,10 +9,12 @@ use async_trait::async_trait;
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use datafusion::cube_ext;
 
+use cuberockstore::rocksdb::backup::{BackupEngine, BackupEngineOptions, RestoreOptions};
+use cuberockstore::rocksdb::checkpoint::Checkpoint;
+use cuberockstore::rocksdb::{
+    DBCompressionType, Env, Snapshot, WriteBatch, WriteBatchIterator, DB,
+};
 use log::{info, trace};
-use rocksdb::backup::{BackupEngine, BackupEngineOptions, RestoreOptions};
-use rocksdb::checkpoint::Checkpoint;
-use rocksdb::{DBCompressionType, Env, Snapshot, WriteBatch, WriteBatchIterator, DB};
 use serde::{Deserialize, Serialize};
 use serde_repr::*;
 use std::collections::HashMap;
@@ -21,6 +23,7 @@ use std::io::{Cursor, Write};
 
 use crate::metastore::snapshot_info::SnapshotInfo;
 use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc};
+use cuberockstore::rocksdb;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, SystemTime};