Limit proxy's memory from TiFlash's side (#408)

Makefile

@@ -211,7 +211,7 @@ unset-override:
 
 pre-format: unset-override
 	@rustup component add rustfmt
-	@cargo install --force -q cargo-sort
+	@cargo install --locked --force -q cargo-sort
 
 ci_fmt_check:
 	M="fmt" ./proxy_scripts/ci_check.sh

components/raftstore/src/store/config.rs

@@ -396,7 +396,7 @@ impl Default for Config {
             apply_yield_duration: ReadableDuration::millis(500),
             apply_yield_write_size: ReadableSize::kb(32),
             perf_level: PerfLevel::Uninitialized,
-            evict_cache_on_memory_ratio: 0.0,
+            evict_cache_on_memory_ratio: 0.1,
             cmd_batch: true,
             cmd_batch_concurrent_ready_max_count: 1,
             raft_write_size_limit: ReadableSize::mb(1),

proxy_scripts/ci_check.sh

@@ -23,9 +23,9 @@ elif [[ $M == "testold" ]]; then
     cargo test --features "$ENABLE_FEATURES" --package tests --test failpoints cases::test_transaction
     cargo test --features "$ENABLE_FEATURES" --package tests --test failpoints cases::test_cmd_epoch_checker
     # cargo test --package tests --test failpoints cases::test_disk_full
-    cargo test --package tests --test failpoints cases::test_merge
+    # cargo test --package tests --test failpoints cases::test_merge
     # cargo test --package tests --test failpoints cases::test_snap
-    cargo test --package tests --test failpoints cases::test_import_service
+    # cargo test --package tests --test failpoints cases::test_import_service
 elif [[ $M == "testnew" ]]; then
     export ENGINE_LABEL_VALUE=tiflash
     export RUST_BACKTRACE=full

proxy_server/src/config.rs

@@ -273,6 +273,14 @@ impl Default for ProxyConfig {
             raftdb: RaftDbConfig::default(),
             storage: StorageConfig::default(),
             enable_io_snoop: false,
+            // Previously, we set `memory_usage_high_water` to 0.1, in order to make TiFlash to be
+            // always in a high-water situation. thus by setting
+            // `evict_cache_on_memory_ratio`, we can evict entry cache if there is a memory usage
+            // peak after restart. However there're some cases that the raftstore could
+            // take more than 5% of the total used memory, so TiFlash will reject
+            // msgAppend to every region. So, it actually not a good idea to make
+            // TiFlash Proxy always run in a high-water state, in order to reduce the
+            // memory usage peak after restart.
             readpool: ReadPoolConfig::default(),
             import: ImportConfig::default(),
         }

proxy_server/src/proxy.rs

@@ -267,12 +267,23 @@ pub unsafe fn run_proxy(
                 .long("only-decryption")
                 .help("Only do decryption in Proxy"),
         )
+        .arg(
+            Arg::with_name("memory-limit-size")
+                .long("memory-limit-size")
+                .help("Used as the maximum memory we can consume, in bytes")
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name("memory-limit-ratio")
+                .long("memory-limit-ratio")
+                .help("Used as the maximum memory we can consume, in percentage")
+                .takes_value(true),
+        )
         .get_matches_from(args);
 
     if matches.is_present("print-sample-config") {
         let config = TikvConfig::default();
         println!("{}", toml::to_string_pretty(&config).unwrap());
-        process::exit(0);
     }
 
     let mut unrecognized_keys = Vec::new();
@@ -308,6 +319,7 @@ pub unsafe fn run_proxy(
     if matches.is_present("only-decryption") {
         crate::run::run_tikv_only_decryption(config, proxy_config, engine_store_server_helper);
     } else {
+        // Log is enabled here.
         crate::run::run_tikv_proxy(config, proxy_config, engine_store_server_helper);
     }
 }

proxy_server/src/setup.rs

@@ -5,8 +5,8 @@ use std::borrow::ToOwned;
 use clap::ArgMatches;
 use collections::HashMap;
 pub use server::setup::initial_logger;
-use tikv::config::{MetricConfig, TikvConfig};
-use tikv_util::{self, logger};
+use tikv::config::{MetricConfig, TikvConfig, MEMORY_USAGE_LIMIT_RATE};
+use tikv_util::{self, config::ReadableSize, logger, sys::SysQuota};
 
 use crate::config::ProxyConfig;
 pub use crate::fatal;
@@ -142,4 +142,49 @@ pub fn overwrite_config_with_cmd_args(
                 .unwrap(),
         ),
     );
+
+    let mut memory_limit_set = config.memory_usage_limit.is_some();
+    if !memory_limit_set {
+        if let Some(s) = matches.value_of("memory-limit-size") {
+            let result: Result<u64, _> = s.parse();
+            if let Ok(memory_limit_size) = result {
+                info!(
+                    "overwrite memory_usage_limit by `memory-limit-size` to {}",
+                    memory_limit_size
+                );
+                config.memory_usage_limit = Some(ReadableSize(memory_limit_size));
+                memory_limit_set = true;
+            } else {
+                info!("overwrite memory_usage_limit by `memory-limit-size` failed"; "memory_limit_size" => s);
+            }
+        }
+    }
+
+    let total = SysQuota::memory_limit_in_bytes();
+    if !memory_limit_set {
+        if let Some(s) = matches.value_of("memory-limit-ratio") {
+            let result: Result<f64, _> = s.parse();
+            if let Ok(memory_limit_ratio) = result {
+                if memory_limit_ratio <= 0.0 || memory_limit_ratio > 1.0 {
+                    info!("overwrite memory_usage_limit meets error ratio"; "ratio" => memory_limit_ratio);
+                } else {
+                    let limit = (total as f64 * memory_limit_ratio) as u64;
+                    info!(
+                        "overwrite memory_usage_limit by `memory-limit-ratio`={} to {}",
+                        memory_limit_ratio, limit
+                    );
+                    config.memory_usage_limit = Some(ReadableSize(limit));
+                    memory_limit_set = true;
+                }
+            } else {
+                info!("overwrite memory_usage_limit meets error ratio"; "ratio" => s);
+            }
+        }
+    }
+
+    if !memory_limit_set && config.memory_usage_limit.is_none() {
+        let limit = (total as f64 * MEMORY_USAGE_LIMIT_RATE) as u64;
+        info!("overwrite memory_usage_limit failed, use TiKV's default"; "limit" => limit);
+        config.memory_usage_limit = Some(ReadableSize(limit));
+    }
 }

proxy_tests/proxy/config.rs

@@ -8,6 +8,7 @@ use proxy_server::{
     proxy::{gen_proxy_config, gen_tikv_config},
     setup::overwrite_config_with_cmd_args,
 };
+use tikv::config::MEMORY_USAGE_LIMIT_RATE;
 use tikv_util::sys::SysQuota;
 
 use crate::proxy::*;
@@ -122,6 +123,11 @@ fn test_config_proxy_default_no_config_item() {
         std::cmp::min(256, (cpu_num * 8.0) as usize)
     );
     assert_eq!(config.server.status_thread_pool_size, 2);
+
+    assert_eq!(config.raft_store.evict_cache_on_memory_ratio, 0.1);
+    assert_eq!(config.memory_usage_high_water, 0.9);
+    // Seems #244 doesn't goes into this branch.
+    assert_eq!(config.server.reject_messages_on_memory_ratio, 0.2);
 }
 
 /// We test if the engine-label is set properly.
@@ -200,3 +206,149 @@ apply-low-priority-pool-size = 41
         config.raft_store.apply_batch_system.low_priority_pool_size
     );
 }
+
+#[test]
+fn test_memory_limit_overwrite() {
+    let app = App::new("RaftStore Proxy")
+        .arg(
+            Arg::with_name("memory-limit-size")
+                .long("memory-limit-size")
+                .help("Used as the maximum memory we can consume, in bytes")
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name("memory-limit-ratio")
+                .long("memory-limit-ratio")
+                .help("Used as the maximum memory we can consume, in percentage")
+                .takes_value(true),
+        );
+
+    let bootstrap = |args: Vec<&str>| {
+        let mut v: Vec<String> = vec![];
+        let matches = app.clone().get_matches_from(args);
+        let mut config = gen_tikv_config(&None, false, &mut v);
+        let mut proxy_config = gen_proxy_config(&None, false, &mut v);
+        proxy_config.raftdb.defaultcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.defaultcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.lockcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.writecf.block_cache_size = ReadableSize(0);
+        overwrite_config_with_cmd_args(&mut config, &mut proxy_config, &matches);
+        address_proxy_config(&mut config, &proxy_config);
+        config.compatible_adjust();
+        config
+    };
+
+    {
+        let args = vec![
+            "test_memory_limit_overwrite1",
+            "--memory-limit-size",
+            "12345",
+        ];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(12345)));
+    }
+
+    {
+        let args = vec![
+            "test_memory_limit_overwrite2",
+            "--memory-limit-size",
+            "12345",
+            "--memory-limit-ratio",
+            "0.9",
+        ];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(12345)));
+    }
+
+    let total = SysQuota::memory_limit_in_bytes();
+    {
+        let args = vec![
+            "test_memory_limit_overwrite3",
+            "--memory-limit-ratio",
+            "0.800000",
+        ];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        let limit = (total as f64 * 0.8) as u64;
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(limit)));
+    }
+
+    let default_limit = (total as f64 * MEMORY_USAGE_LIMIT_RATE) as u64;
+    {
+        let args = vec![
+            "test_memory_limit_overwrite4",
+            "--memory-limit-ratio",
+            "7.9",
+        ];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(default_limit)));
+    }
+
+    {
+        let args = vec![
+            "test_memory_limit_overwrite5",
+            "--memory-limit-ratio",
+            "'-0.9'",
+        ];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(default_limit)));
+    }
+
+    {
+        let args = vec!["test_memory_limit_overwrite6"];
+        let mut config = bootstrap(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(default_limit)));
+    }
+
+    let bootstrap2 = |args: Vec<&str>| {
+        let mut v: Vec<String> = vec![];
+        let matches = app.clone().get_matches_from(args);
+        let mut file = tempfile::NamedTempFile::new().unwrap();
+        write!(
+            file,
+            "
+memory-usage-limit = 42
+            "
+        )
+        .unwrap();
+        let path = file.path();
+        let cpath = Some(path.as_os_str());
+        let mut config = gen_tikv_config(&cpath, false, &mut v);
+        let mut proxy_config = gen_proxy_config(&cpath, false, &mut v);
+        proxy_config.raftdb.defaultcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.defaultcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.lockcf.block_cache_size = ReadableSize(0);
+        proxy_config.rocksdb.writecf.block_cache_size = ReadableSize(0);
+        overwrite_config_with_cmd_args(&mut config, &mut proxy_config, &matches);
+        address_proxy_config(&mut config, &proxy_config);
+        config.compatible_adjust();
+        config
+    };
+
+    {
+        let args = vec![
+            "test_memory_limit_nooverwrite3",
+            "--memory-limit-ratio",
+            "0.800000",
+        ];
+        let mut config = bootstrap2(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(42)));
+    }
+
+    {
+        let args = vec![
+            "test_memory_limit_nooverwrite1",
+            "--memory-limit-size",
+            "12345",
+        ];
+        let mut config = bootstrap2(args);
+        assert!(config.validate().is_ok());
+        assert_eq!(config.memory_usage_limit, Some(ReadableSize(42)));
+    }
+}