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

proxy_components/proxy_server/src/config.rs

@@ -267,10 +267,6 @@ pub struct ProxyConfig {
     #[online_config(skip)]
     pub enable_io_snoop: bool,
 
-    #[doc(hidden)]
-    #[online_config(skip)]
-    pub memory_usage_high_water: f64,
-
     #[online_config(submodule)]
     pub readpool: ReadPoolConfig,
 
@@ -295,7 +291,14 @@ impl Default for ProxyConfig {
             raftdb: RaftDbConfig::default(),
             storage: StorageConfig::default(),
             enable_io_snoop: false,
-            memory_usage_high_water: 0.1,
+            // 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(),
             engine_store: EngineStoreConfig::default(),
@@ -423,7 +426,6 @@ pub fn address_proxy_config(config: &mut TikvConfig, proxy_config: &ProxyConfig)
     config.storage.reserve_space = proxy_config.storage.reserve_space;
 
     config.enable_io_snoop = proxy_config.enable_io_snoop;
-    config.memory_usage_high_water = proxy_config.memory_usage_high_water;
 
     config.server.addr = proxy_config.server.addr.clone();
     config.server.advertise_addr = proxy_config.server.advertise_addr.clone();

proxy_components/proxy_server/src/proxy.rs

@@ -281,6 +281,18 @@ pub unsafe fn run_proxy(
                 .help("Set engine role label")
                 .takes_value(true),
         )
+        .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") {
@@ -322,6 +334,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_components/proxy_server/src/setup.rs

@@ -4,8 +4,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;
@@ -160,4 +160,49 @@ pub fn overwrite_config_with_cmd_args(
         });
         proxy_config.engine_store.enable_unips = enabled == 1;
     }
+
+    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/shared/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::utils::v1::*;
@@ -139,7 +140,7 @@ fn test_default_no_config_item() {
     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.1);
+    assert_eq!(config.memory_usage_high_water, 0.9);
     assert_eq!(config.server.reject_messages_on_memory_ratio, 0.05);
 
     assert_eq!(config.raft_store.enable_v2_compatible_learner, true);
@@ -331,3 +332,149 @@ enable-fast-add-peer = true
     info!("using proxy config"; "config" => ?proxy_config);
     assert_eq!(true, proxy_config.engine_store.enable_fast_add_peer);
 }
+
+#[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 = Some(ReadableSize(0));
+        proxy_config.rocksdb.defaultcf.block_cache_size = Some(ReadableSize(0));
+        proxy_config.rocksdb.lockcf.block_cache_size = Some(ReadableSize(0));
+        proxy_config.rocksdb.writecf.block_cache_size = Some(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 = Some(ReadableSize(0));
+        proxy_config.rocksdb.defaultcf.block_cache_size = Some(ReadableSize(0));
+        proxy_config.rocksdb.lockcf.block_cache_size = Some(ReadableSize(0));
+        proxy_config.rocksdb.writecf.block_cache_size = Some(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)));
+    }
+}

proxy_tests/proxy/shared/store.rs

@@ -64,7 +64,6 @@ mod config {
             ProxyConfig::from_file(path, Some(&mut proxy_unrecognized_keys)).unwrap();
         assert_eq!(proxy_config.raft_store.snap_handle_pool_size, 4);
         assert_eq!(proxy_config.server.engine_addr, "1.2.3.4:5");
-        assert_eq!(proxy_config.memory_usage_high_water, 0.65);
         assert!(proxy_unrecognized_keys.contains(&"nosense".to_string()));
         let v1 = vec!["a.b", "b"]
             .iter()