Automatic partition store memory manager

This removes the configuration option `worker.storage.num-partitions-to-share-memory-budget` and instead replaces with an automatic memory manager that checks and rebalances memory across partitions every 5 seconds.

- Removed rocksdb stall detection and hard-coded allow_stall to false in WBM since WBM almost always does the wrong thing and stalls forever under memory pressure. Removal of stall detection saves a timer on every write batch.
- Allowed more flexible configuration of rocksdb per database.
- Partition store manager now automatically rebalances its memory budget across open partition stores and runs a garbage collection pass on partitions that have been closed for some time to reclaim memory.
- Switched lz4 to zstd to make for consistency across all databases.
- Partition store manager reacts to live changes in memory budget
- Log server reacts to live memory budget changes by reconfiguring future write buffers to meet the new budget

Author: AhmedSoliman

Cargo.lock

@@ -25,7 +25,6 @@ use restate_types::config::{
     MetadataServerOptionsBuilder, WorkerOptionsBuilder,
 };
 use restate_types::config_loader::ConfigLoaderBuilder;
-use restate_types::live::Constant;
 use restate_types::logs::metadata::ProviderKind;
 use restate_types::retries::RetryPolicy;
 use std::time::Duration;
@@ -108,7 +107,7 @@ pub fn spawn_restate(config: Configuration) -> task_center::Handle {
     let live_config = Configuration::live();
 
     tc.block_on(async {
-        RocksDbManager::init(Constant::new(config.common));
+        RocksDbManager::init();
         prometheus.start_upkeep_task();
 
         TaskCenter::spawn(TaskKind::SystemBoot, "restate", async move {

benchmarks/src/lib.rs

@@ -776,8 +776,6 @@ mod tests {
     use restate_core::TestCoreEnvBuilder;
     use restate_core::{TaskCenter, TaskKind, TestCoreEnv};
     use restate_rocksdb::RocksDbManager;
-    use restate_types::config::CommonOptions;
-    use restate_types::live::Constant;
     use restate_types::logs::SequenceNumber;
     use restate_types::logs::metadata::{SegmentIndex, new_single_node_loglet_params};
     use restate_types::metadata::Precondition;
@@ -890,7 +888,7 @@ mod tests {
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
-        RocksDbManager::init(Constant::new(CommonOptions::default()));
+        RocksDbManager::init();
 
         let bifrost = Bifrost::init_local(node_env.metadata_writer).await;
 
@@ -1389,7 +1387,7 @@ mod tests {
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
-        RocksDbManager::init(Constant::new(CommonOptions::default()));
+        RocksDbManager::init();
         let bifrost = Bifrost::init_local(node_env.metadata_writer).await;
 
         // create an appender

crates/bifrost/src/bifrost.rs

@@ -10,15 +10,15 @@
 
 use std::sync::Arc;
 
-use restate_types::errors::MaybeRetryableError;
-use rocksdb::{BoundColumnFamily, DB, DBCompressionType, SliceTransform};
+use rocksdb::{BlockBasedOptions, BoundColumnFamily, Cache, DB, DBCompressionType, SliceTransform};
 use static_assertions::const_assert;
 
 use restate_rocksdb::{
     CfExactPattern, CfName, DbName, DbSpecBuilder, RocksDb, RocksDbManager, RocksError,
 };
-use restate_types::config::LocalLogletOptions;
-use restate_types::live::{LiveLoad, LiveLoadExt};
+use restate_types::config::{Configuration, LocalLogletOptions};
+use restate_types::errors::MaybeRetryableError;
+use restate_types::live::LiveLoad;
 use restate_types::storage::{StorageDecodeError, StorageEncodeError};
 
 use super::keys::{DATA_KEY_PREFIX_LENGTH, MetadataKey, MetadataKind};
@@ -74,24 +74,16 @@ impl RocksDbLogStore {
         let opts = options.live_load();
         let data_dir = opts.data_dir();
 
-        let db_spec = DbSpecBuilder::new(DbName::new(DB_NAME), data_dir, db_options())
-            .add_cf_pattern(
-                CfExactPattern::new(DATA_CF),
-                cf_data_options(opts.rocksdb_memory_budget()),
-            )
-            .add_cf_pattern(
-                CfExactPattern::new(METADATA_CF),
-                cf_metadata_options(opts.rocksdb_memory_budget()),
-            )
+        let db_spec = DbSpecBuilder::new(DbName::new(DB_NAME), data_dir, RocksConfigurator)
+            .add_cf_pattern(CfExactPattern::new(DATA_CF), RocksConfigurator)
+            .add_cf_pattern(CfExactPattern::new(METADATA_CF), RocksConfigurator)
             // not very important but it's to reduce the number of merges by flushing.
             // it's also a small cf so it should be quick.
             .add_to_flush_on_shutdown(CfExactPattern::new(METADATA_CF))
             .ensure_column_families(cfs)
             .build()
             .expect("valid spec");
-        let rocksdb = db_manager
-            .open_db(options.map(|options| &options.rocksdb).boxed(), db_spec)
-            .await?;
+        let rocksdb = db_manager.open_db(db_spec).await?;
         Ok(Self { rocksdb })
     }
 
@@ -132,62 +124,107 @@ impl RocksDbLogStore {
     }
 }
 
-fn db_options() -> rocksdb::Options {
-    let mut opts = rocksdb::Options::default();
-
-    // Enable atomic flushes.
-    // If WAL is disabled, this ensure we do not persist inconsistent data.
-    // If WAL is enabled, this ensures that flushing either cf flushes both.
-    // This is valuable because otherwise the metadata cf will flush rarely, and that would keep the WAL around
-    // until shutdown, full of data cf bytes that have already been flushed, wasting disk space.
-    opts.set_atomic_flush(true);
-
-    // This is Rocksdb's default, it's added here for clarity.
-    //
-    // Rationale: If WAL tail is corrupted, it's likely that it has failed during write, that said,
-    // we can use absolute consistency but on a single-node setup, we don't have a way to recover
-    // from it, so it's not useful for us.
-    opts.set_wal_recovery_mode(rocksdb::DBRecoveryMode::TolerateCorruptedTailRecords);
-    opts.set_wal_compression_type(DBCompressionType::Zstd);
-    // most reads are sequential
-    opts.set_advise_random_on_open(false);
-
-    opts
+struct RocksConfigurator;
+
+impl restate_rocksdb::configuration::DbConfigurator for RocksConfigurator {
+    fn get_db_options(
+        &self,
+        _db_name: &str,
+        env: &rocksdb::Env,
+        write_buffer_manager: &rocksdb::WriteBufferManager,
+    ) -> rocksdb::Options {
+        let mut db_options = restate_rocksdb::configuration::create_default_db_options(
+            env,
+            true, /* create_db_if_missing */
+            write_buffer_manager,
+        );
+        let local_loglet_config = &Configuration::pinned().bifrost.local;
+        // amend default options from rocksdb_manager
+        self.apply_db_opts_from_config(&mut db_options, &local_loglet_config.rocksdb);
+        // local loglet customizations
+
+        // Enable atomic flushes.
+        // If WAL is disabled, this ensure we do not persist inconsistent data.
+        // If WAL is enabled, this ensures that flushing either cf flushes both.
+        // This is valuable because otherwise the metadata cf will flush rarely, and that would keep the WAL around
+        // until shutdown, full of data cf bytes that have already been flushed, wasting disk space.
+        db_options.set_atomic_flush(true);
+
+        // This is Rocksdb's default, it's added here for clarity.
+        //
+        // Rationale: If WAL tail is corrupted, it's likely that it has failed during write, that said,
+        // we can use absolute consistency but on a single-node setup, we don't have a way to recover
+        // from it, so it's not useful for us.
+        db_options.set_wal_recovery_mode(rocksdb::DBRecoveryMode::TolerateCorruptedTailRecords);
+        db_options.set_wal_compression_type(DBCompressionType::Zstd);
+        // most reads are sequential
+        db_options.set_advise_random_on_open(false);
+
+        db_options
+    }
 }
 
-fn cf_data_options(
-    memory_budget: usize,
-) -> impl Fn(rocksdb::Options) -> rocksdb::Options + Send + Sync + 'static {
-    move |mut opts| {
-        // memory budget is in bytes. We divide the budget between the data cf and metadata cf.
-        // data 10% to metadata 90% to data.
-        let memtables_budget = (memory_budget as f64 * DATA_CF_BUDGET_RATIO).floor() as usize;
-        assert!(
-            memtables_budget > 0,
-            "memory budget should be greater than 0"
+impl restate_rocksdb::configuration::CfConfigurator for RocksConfigurator {
+    fn get_cf_options(
+        &self,
+        _db_name: &str,
+        cf_name: &str,
+        global_cache: &Cache,
+        write_buffer_manager: &rocksdb::WriteBufferManager,
+    ) -> rocksdb::Options {
+        let config = &Configuration::pinned().bifrost.local;
+        let mut cf_options =
+            restate_rocksdb::configuration::create_default_cf_options(Some(write_buffer_manager));
+        let block_options = restate_rocksdb::configuration::create_default_block_options(
+            &config.rocksdb,
+            Some(global_cache),
         );
 
-        set_memory_related_opts(&mut opts, memtables_budget);
-        opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
-        opts.set_num_levels(7);
-
-        opts.set_compression_per_level(&[
-            DBCompressionType::None,
-            DBCompressionType::None,
-            DBCompressionType::Lz4,
-            DBCompressionType::Lz4,
-            DBCompressionType::Lz4,
-            DBCompressionType::Lz4,
-            DBCompressionType::Zstd,
-        ]);
-
-        opts.set_prefix_extractor(SliceTransform::create_fixed_prefix(DATA_KEY_PREFIX_LENGTH));
-        opts.set_memtable_prefix_bloom_ratio(0.2);
-        //
-        opts
+        if cf_name == DATA_CF {
+            cf_data_options(&mut cf_options, &block_options, config);
+        } else if cf_name == METADATA_CF {
+            cf_metadata_options(&mut cf_options, &block_options, config);
+        }
+
+        cf_options
     }
 }
 
+fn cf_data_options(
+    opts: &mut rocksdb::Options,
+    block_options: &BlockBasedOptions,
+    local_loglet_config: &LocalLogletOptions,
+) {
+    opts.set_block_based_table_factory(block_options);
+
+    let memory_budget = local_loglet_config.rocksdb_memory_budget();
+
+    // memory budget is in bytes. We divide the budget between the data cf and metadata cf.
+    // data 10% to metadata 90% to data.
+    let memtables_budget = (memory_budget as f64 * DATA_CF_BUDGET_RATIO).floor() as usize;
+    assert!(
+        memtables_budget > 0,
+        "memory budget should be greater than 0"
+    );
+
+    set_memory_related_opts(opts, memtables_budget);
+    opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
+    opts.set_num_levels(7);
+
+    opts.set_compression_per_level(&[
+        DBCompressionType::None,
+        DBCompressionType::None,
+        DBCompressionType::Zstd,
+        DBCompressionType::Zstd,
+        DBCompressionType::Zstd,
+        DBCompressionType::Zstd,
+        DBCompressionType::Zstd,
+    ]);
+
+    opts.set_prefix_extractor(SliceTransform::create_fixed_prefix(DATA_KEY_PREFIX_LENGTH));
+    opts.set_memtable_prefix_bloom_ratio(0.2);
+}
+
 fn set_memory_related_opts(opts: &mut rocksdb::Options, memtables_budget: usize) {
     // We set the budget to allow 1 mutable + 3 immutable.
     opts.set_write_buffer_size(memtables_budget / 4);
@@ -206,34 +243,36 @@ fn set_memory_related_opts(opts: &mut rocksdb::Options, memtables_budget: usize)
 }
 
 fn cf_metadata_options(
-    memory_budget: usize,
-) -> impl Fn(rocksdb::Options) -> rocksdb::Options + Send + Sync + 'static {
-    move |mut opts| {
-        let memtables_budget =
-            (memory_budget as f64 * (1.0 - DATA_CF_BUDGET_RATIO)).floor() as usize;
-        assert!(
-            memtables_budget > 0,
-            "memory budget should be greater than 0"
-        );
-        set_memory_related_opts(&mut opts, memtables_budget);
-        //
-        // Set compactions per level
-        //
-        opts.set_num_levels(3);
-        opts.set_compression_per_level(&[
-            DBCompressionType::None,
-            DBCompressionType::None,
-            DBCompressionType::Lz4,
-        ]);
-        opts.set_memtable_whole_key_filtering(true);
-        opts.set_max_write_buffer_number(4);
-        opts.set_max_successive_merges(10);
-        // Merge operator for log state updates
-        opts.set_merge_operator(
-            "LogStateMerge",
-            log_state_full_merge,
-            log_state_partial_merge,
-        );
-        opts
-    }
+    opts: &mut rocksdb::Options,
+    block_options: &BlockBasedOptions,
+    local_loglet_config: &LocalLogletOptions,
+) {
+    opts.set_block_based_table_factory(block_options);
+
+    let memory_budget = local_loglet_config.rocksdb_memory_budget();
+
+    let memtables_budget = (memory_budget as f64 * (1.0 - DATA_CF_BUDGET_RATIO)).floor() as usize;
+    assert!(
+        memtables_budget > 0,
+        "memory budget should be greater than 0"
+    );
+    set_memory_related_opts(opts, memtables_budget);
+    //
+    // Set compactions per level
+    //
+    opts.set_num_levels(3);
+    opts.set_compression_per_level(&[
+        DBCompressionType::None,
+        DBCompressionType::None,
+        DBCompressionType::Zstd,
+    ]);
+    opts.set_memtable_whole_key_filtering(true);
+    opts.set_max_write_buffer_number(4);
+    opts.set_max_successive_merges(10);
+    // Merge operator for log state updates
+    opts.set_merge_operator(
+        "LogStateMerge",
+        log_state_full_merge,
+        log_state_partial_merge,
+    );
 }

crates/bifrost/src/providers/local_loglet/log_store.rs

@@ -323,7 +323,7 @@ mod tests {
             .await;
 
         let config = Live::from_value(Configuration::default());
-        RocksDbManager::init(config.clone().map(|c| &c.common));
+        RocksDbManager::init();
         let params = LogletParams::from("42".to_string());
 
         let local_loglet_config = config.map(|config| &config.bifrost.local);
@@ -356,7 +356,7 @@ mod tests {
             .await;
 
         let config = Live::from_value(Configuration::default());
-        RocksDbManager::init(config.clone().map(|c| &c.common));
+        RocksDbManager::init();
 
         let local_loglet_config = config.map(|config| &config.bifrost.local);
         let log_store = RocksDbLogStore::create(local_loglet_config.clone()).await?;

crates/bifrost/src/providers/local_loglet/mod.rs

@@ -442,7 +442,7 @@ mod tests {
     use restate_rocksdb::RocksDbManager;
     use restate_types::config::{Configuration, set_current_config};
     use restate_types::health::HealthStatus;
-    use restate_types::live::{Live, LiveLoadExt};
+    use restate_types::live::Live;
     use restate_types::logs::{Keys, LogletId};
     use restate_types::replication::{NodeSet, ReplicationProperty};
     use restate_types::{GenerationalNodeId, PlainNodeId};
@@ -467,7 +467,7 @@ mod tests {
         set_current_config(config.clone());
         let config = Live::from_value(config);
 
-        RocksDbManager::init(config.clone().map(|c| &c.common));
+        RocksDbManager::init();
 
         let mut node_env =
             TestCoreEnvBuilder::with_incoming_only_connector().add_mock_nodes_config();