From 745242ea3ed59e6ee8328f94a72cc7261a02dffd Mon Sep 17 00:00:00 2001
From: Zijie Lu <wslzj40@gmail.com>
Date: Fri, 21 Oct 2022 14:07:06 +0800
Subject: [PATCH] feat(query): get the system open file limits and set

---
 src/binaries/Cargo.toml    |  2 ++
 src/binaries/query/main.rs | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/binaries/Cargo.toml b/src/binaries/Cargo.toml
index f21975cb2cc5b..f9e1422d9c85f 100644
--- a/src/binaries/Cargo.toml
+++ b/src/binaries/Cargo.toml
@@ -49,10 +49,12 @@ databend-query = { path = "../query/service" }
 # Crates.io dependencies
 anyhow = { workspace = true }
 clap = { workspace = true }
+limits-rs = "0.1.0"
 openraft = { workspace = true }
 sentry = "0.27.0"
 serde = { workspace = true }
 serde_json = { workspace = true }
+sysinfo = "0.26.5"
 tokio-stream = "0.1.10"
 tonic = "0.8.1"
 tracing = "0.1.36"
diff --git a/src/binaries/query/main.rs b/src/binaries/query/main.rs
index c1aa756833690..ff08603417a86 100644
--- a/src/binaries/query/main.rs
+++ b/src/binaries/query/main.rs
@@ -34,7 +34,10 @@ use databend_query::servers::MySQLHandler;
 use databend_query::servers::Server;
 use databend_query::servers::ShutdownHandle;
 use databend_query::GlobalServices;
+use limits_rs::get_own_limits;
+use sysinfo::set_open_files_limit;
 use tracing::info;
+use tracing::warn;
 
 #[databend_main]
 async fn main(_global_tracker: Arc<RuntimeTracker>) -> common_exception::Result<()> {
@@ -73,6 +76,9 @@ async fn main(_global_tracker: Arc<RuntimeTracker>) -> common_exception::Result<
     init_default_metrics_recorder();
     set_panic_hook();
 
+    // Try to set max open files.
+    try_set_max_open_files();
+
     GlobalServices::init(conf.clone()).await?;
     let mut shutdown_handle = ShutdownHandle::create()?;
 
@@ -251,3 +257,24 @@ fn run_cmd(conf: &Config) -> bool {
 
     true
 }
+
+fn try_set_max_open_files() {
+    let limits = get_own_limits().unwrap();
+    let max_open_files_limit = limits.max_open_files.soft;
+    match max_open_files_limit {
+        Some(max_open_files) => {
+            if max_open_files < 65535 {
+                let set = sysinfo::set_open_files_limit(max_open_files);
+                match set {
+                    true => {
+                        warn!("Open files limit has been set to {}", max_open_files);
+                    }
+                    false => {
+                        warn!("Open files limit set to {} failed", max_open_files);
+                    }
+                }
+            }
+        }
+        None => {}
+    }
+}