From 39fa62ce051773c007ceeefc2cf3f7a2f3defaf7 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 | 25 +++++++++++++++++++++++++
 2 files changed, 27 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..74d88e0361a62 100644
--- a/src/binaries/query/main.rs
+++ b/src/binaries/query/main.rs
@@ -34,7 +34,9 @@ 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 tracing::info;
+use tracing::warn;
 
 #[databend_main]
 async fn main(_global_tracker: Arc<RuntimeTracker>) -> common_exception::Result<()> {
@@ -73,6 +75,10 @@ async fn main(_global_tracker: Arc<RuntimeTracker>) -> common_exception::Result<
     init_default_metrics_recorder();
     set_panic_hook();
 
+    #[cfg(not(target_os = "macos"))]
+    // 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,22 @@ fn run_cmd(conf: &Config) -> bool {
 
     true
 }
+
+#[cfg(not(target_os = "macos"))]
+fn try_set_max_open_files() {
+    let limits = get_own_limits().unwrap();
+    let max_open_files_limit = limits.max_open_files.soft;
+    if let Some(max_open_files) = max_open_files_limit {
+        if max_open_files < 65535 {
+            let set = sysinfo::set_open_files_limit(max_open_files.try_into().unwrap());
+            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);
+                }
+            }
+        }
+    }
+}