Warn if cgroup memory limit is misaligned with rocksdb memory limit (#3925)

Author: jackkleeman

crates/rocksdb/src/db_manager.rs

@@ -16,7 +16,7 @@ use std::time::Duration;
 use parking_lot::RwLock;
 use rocksdb::{Cache, WriteBufferManager};
 use tokio_util::task::TaskTracker;
-use tracing::{debug, info, warn};
+use tracing::{debug, error, info, warn};
 
 use restate_core::{ShutdownError, TaskCenter, TaskKind, cancellation_watcher};
 use restate_serde_util::ByteCount;
@@ -67,6 +67,9 @@ impl RocksDbManager {
         }
         metric_definitions::describe_metrics();
         let opts = &Configuration::pinned().common;
+
+        check_memory_limit(opts);
+
         let cache = Cache::new_lru_cache(opts.rocksdb_total_memory_size.get());
         let write_buffer_manager = WriteBufferManager::new_write_buffer_manager_with_cache(
             opts.rocksdb_actual_total_memtables_size(),
@@ -421,6 +424,7 @@ impl DbWatchdog {
                 "[config update] Setting rocksdb total memory limit to {}",
                 ByteCount::from(new_common_opts.rocksdb_total_memory_size)
             );
+            check_memory_limit(new_common_opts);
             self.cache
                 .set_capacity(new_common_opts.rocksdb_total_memory_size.get());
             self.manager
@@ -460,4 +464,30 @@ impl DbWatchdog {
     }
 }
 
+fn check_memory_limit(opts: &CommonOptions) {
+    if let Some(process_memory_size) = opts.process_total_memory_size() {
+        let memory_ratio =
+            opts.rocksdb_total_memory_size.get() as f64 / process_memory_size.get() as f64;
+        if memory_ratio < 0.5 {
+            warn!(
+                "'rocksdb-total-memory-size' parameter is set to {}, less than half the process memory limit of {}. Roughly 75% of process memory should be given to RocksDB",
+                ByteCount::from(opts.rocksdb_total_memory_size),
+                ByteCount::from(process_memory_size),
+            )
+        } else if memory_ratio > 1.0 {
+            error!(
+                "'rocksdb-total-memory-size' parameter is set to {}, more than the process memory limit of {}. This guarantees an OOM under load; roughly 75% of process memory should be given to RocksDB",
+                ByteCount::from(opts.rocksdb_total_memory_size),
+                ByteCount::from(process_memory_size),
+            )
+        } else if memory_ratio > 0.9 {
+            error!(
+                "'rocksdb-total-memory-size' parameter is set to {}, more than 90% of the process memory limit of {}. This risks an OOM under load; roughly 75% of process memory should be given to RocksDB",
+                ByteCount::from(opts.rocksdb_total_memory_size),
+                ByteCount::from(process_memory_size),
+            )
+        }
+    }
+}
+
 static_assertions::assert_impl_all!(RocksDbManager: Send, Sync);

crates/types/src/config/common.rs

@@ -352,6 +352,15 @@ pub struct CommonOptions {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub storage_low_priority_bg_threads: Option<NonZeroUsize>,
 
+    /// # Total memory limit for this process
+    ///
+    /// This is intended to be determined automatically on Linux based on the cgroup limit,
+    /// and is used to emit warning logs if other memory limits are set too close to it.
+    #[serde_as(as = "Option<NonZeroByteCount>")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[cfg_attr(feature = "schemars", schemars(skip))]
+    pub process_total_memory_size: Option<NonZeroUsize>,
+
     /// # Total memory limit for rocksdb caches and memtables.
     ///
     /// This includes memory for uncompressed block cache and all memtables by all open databases.
@@ -497,6 +506,24 @@ impl CommonOptions {
         self.base_dir.as_ref()
     }
 
+    #[cfg(target_os = "linux")]
+    pub fn process_total_memory_size(&self) -> Option<NonZeroUsize> {
+        self.process_total_memory_size.or_else(|| {
+            [
+                "/sys/fs/cgroup/memory.max", // cgroup v2, takes precedence
+                "/sys/fs/cgroup/memory/memory.limit_in_bytes", // cgroup v1
+            ]
+            .iter()
+            .find_map(|path| std::fs::read_to_string(path).ok())
+            .and_then(|contents| contents.trim().parse().ok())
+        })
+    }
+
+    #[cfg(not(target_os = "linux"))]
+    pub fn process_total_memory_size(&self) -> Option<NonZeroUsize> {
+        self.process_total_memory_size
+    }
+
     pub fn rocksdb_actual_total_memtables_size(&self) -> usize {
         let sanitized = self.rocksdb_total_memtables_ratio.clamp(0.0, 1.0) as f64;
         let total_mem = self.rocksdb_total_memory_size.get() as f64;
@@ -597,8 +624,9 @@ impl Default for CommonOptions {
             default_thread_pool_size: None,
             storage_high_priority_bg_threads: None,
             storage_low_priority_bg_threads: None,
-            rocksdb_total_memtables_ratio: 0.5, // (50% of rocksdb-total-memory-size)
+            process_total_memory_size: None,
             rocksdb_total_memory_size: NonZeroUsize::new(6 * 1024 * 1024 * 1024).unwrap(), // 6GiB
+            rocksdb_total_memtables_ratio: 0.5, // (50% of rocksdb-total-memory-size)
             rocksdb_bg_threads: None,
             rocksdb_high_priority_bg_threads: NonZeroU32::new(2).unwrap(),
             rocksdb_perf_level: PerfStatsLevel::EnableCount,