fix: use async walkdir (#3193)

2026-11-06 00:09:22 +00:00 · 2024-04-08 17:06:54 +05:30
parent 47db7660bb
commit 7472a0fcf0
16 changed files with 263 additions and 72 deletions
@@ -339,7 +339,7 @@ checksum = "ad5228fd1a6b5d0f60d636776c2a70acc9fc667034bb4ac02ec4259f0eeeab6c"
 dependencies = [
 "actix-service",
 "actix-web",
- "futures-lite",
+ "futures-lite 1.13.0",
 "pin-project",
 "prometheus",
 "quanta",
@@ -844,6 +844,19 @@ dependencies = [
 "nom",
 ]

+[[package]]
+name = "async-channel"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f28243a43d821d11341ab73c80bed182dc015c514b951616cf79bd4af39af0c3"
+dependencies = [
+ "concurrent-queue",
+ "event-listener 5.3.0",
+ "event-listener-strategy 0.5.1",
+ "futures-core",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "async-compression"
 version = "0.4.7"
@@ -862,6 +875,28 @@ dependencies = [
 "zstd-safe",
 ]

+[[package]]
+name = "async-fs"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc19683171f287921f2405677dd2ed2549c3b3bda697a563ebc3a121ace2aba1"
+dependencies = [
+ "async-lock",
+ "blocking",
+ "futures-lite 2.3.0",
+]
+
+[[package]]
+name = "async-lock"
+version = "3.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d034b430882f8381900d3fe6f0aaa3ad94f2cb4ac519b429692a1bc2dda4ae7b"
+dependencies = [
+ "event-listener 4.0.3",
+ "event-listener-strategy 0.4.0",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "async-nats"
 version = "0.34.0"
@@ -928,6 +963,12 @@ dependencies = [
 "syn 2.0.58",
 ]

+[[package]]
+name = "async-task"
+version = "4.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbb36e985947064623dbd357f727af08ffd077f93d696782f3c56365fa2e2799"
+
 [[package]]
 name = "async-trait"
 version = "0.1.79"
@@ -953,6 +994,16 @@ dependencies = [
 "tungstenite",
 ]

+[[package]]
+name = "async-walkdir"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f6338023cbfc0555eccb8e83d3d4dcf1183b51ca9140a03b1dbb8a559193db"
+dependencies = [
+ "async-fs",
+ "futures-lite 2.3.0",
+]
+
 [[package]]
 name = "atoi"
 version = "2.0.0"
@@ -962,6 +1013,12 @@ dependencies = [
 "num-traits",
 ]

+[[package]]
+name = "atomic-waker"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
+
 [[package]]
 name = "atty"
 version = "0.2.14"
@@ -1549,6 +1606,22 @@ dependencies = [
 "generic-array",
 ]

+[[package]]
+name = "blocking"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a37913e8dc4ddcc604f0c6d3bf2887c995153af3611de9e23c352b44c1b9118"
+dependencies = [
+ "async-channel",
+ "async-lock",
+ "async-task",
+ "fastrand 2.0.2",
+ "futures-io",
+ "futures-lite 2.3.0",
+ "piper",
+ "tracing",
+]
+
 [[package]]
 name = "borsh"
 version = "1.4.0"
@@ -2152,6 +2225,15 @@ dependencies = [
 "sha1",
 ]

+[[package]]
+name = "concurrent-queue"
+version = "2.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d16048cd947b08fa32c24458a22f5dc5e835264f689f4f5653210c69fd107363"
+dependencies = [
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "config"
 version = "0.1.0"
@@ -2162,6 +2244,7 @@ dependencies = [
 "arrow",
 "arrow-json",
 "arrow-schema",
+ "async-walkdir",
 "base64 0.21.7",
 "byteorder",
 "bytes",
@@ -2202,7 +2285,6 @@ dependencies = [
 "tracing-log",
 "tracing-subscriber",
 "utoipa",
- "walkdir",
 ]

 [[package]]
@@ -3176,6 +3258,48 @@ version = "2.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0"

+[[package]]
+name = "event-listener"
+version = "4.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67b215c49b2b248c855fb73579eb1f4f26c38ffdc12973e20e07b91d78d5646e"
+dependencies = [
+ "concurrent-queue",
+ "parking",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "event-listener"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d9944b8ca13534cdfb2800775f8dd4902ff3fc75a50101466decadfdf322a24"
+dependencies = [
+ "concurrent-queue",
+ "parking",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "event-listener-strategy"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "958e4d70b6d5e81971bebec42271ec641e7ff4e170a6fa605f2b8a8b65cb97d3"
+dependencies = [
+ "event-listener 4.0.3",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "event-listener-strategy"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "332f51cb23d20b0de8458b86580878211da09bcd4503cb579c225b3d124cabb3"
+dependencies = [
+ "event-listener 5.3.0",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "exitcode"
 version = "1.1.2"
@@ -3396,6 +3520,19 @@ dependencies = [
 "waker-fn",
 ]

+[[package]]
+name = "futures-lite"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52527eb5074e35e9339c6b4e8d12600c7128b68fb25dcb9fa9dec18f7c25f3a5"
+dependencies = [
+ "fastrand 2.0.2",
+ "futures-core",
+ "futures-io",
+ "parking",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "futures-macro"
 version = "0.3.30"
@@ -3935,6 +4072,7 @@ version = "0.1.0"
 dependencies = [
 "arrow",
 "arrow-schema",
+ "async-walkdir",
 "byteorder",
 "bytes",
 "chrono",
@@ -3951,7 +4089,6 @@ dependencies = [
 "snafu 0.7.5",
 "tokio",
 "wal",
- "walkdir",
 ]

 [[package]]
@@ -4959,6 +5096,7 @@ dependencies = [
 "arrow-schema",
 "async-recursion",
 "async-trait",
+ "async-walkdir",
 "awc",
 "base64 0.21.7",
 "blake3",
@@ -5041,7 +5179,6 @@ dependencies = [
 "utoipa",
 "utoipa-swagger-ui",
 "vrl",
- "walkdir",
 "zstd",
 ]

@@ -5470,6 +5607,17 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"

+[[package]]
+name = "piper"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "668d31b1c4eba19242f2088b2bf3316b82ca31082a8335764db4e083db7485d4"
+dependencies = [
+ "atomic-waker",
+ "fastrand 2.0.2",
+ "futures-io",
+]
+
 [[package]]
 name = "pkcs1"
 version = "0.7.5"
@@ -6969,7 +7117,7 @@ dependencies = [
 "crc",
 "crossbeam-queue",
 "either",
- "event-listener",
+ "event-listener 2.5.3",
 "futures-channel",
 "futures-core",
 "futures-intrusive",
@@ -167,7 +167,7 @@ chrono = { version = "0.4", default-features = false, features = ["clock"] }
 expect-test = "1.4"
 base64 = "0.21"
 float-cmp = "0.9"
-walkdir = "2"
+async-walkdir = "1.0.0"

 [workspace]
 members = ["src/config", "src/infra", "src/ingester", "src/wal", "src/proto"]
@@ -280,5 +280,5 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"] }
 url = "2.2"
 utoipa = { version = "4", features = ["actix_extras", "openapi_extensions"] }
 utoipa-swagger-ui = { version = "4", features = ["actix-web"] }
-walkdir = "2"
+async-walkdir = "1.0.0"
 zstd = "0.13"
@@ -40,6 +40,26 @@ pub async fn run() -> Result<(), anyhow::Error> {
    };

    let locker = infra::dist_lock::lock(SCHEMA_MIGRATION_KEY, 0).await?;
+
+    // check again
+
+    match upgrade_schema_row_per_version().await {
+        std::result::Result::Ok(true) => {
+            log::info!("[Schema:Migration]: Starting schema migration");
+        }
+        std::result::Result::Ok(false) => {
+            log::info!("[Schema:Migration]: Schema migration already done");
+            return Ok(());
+        }
+        Err(err) => {
+            log::error!(
+                "[Schema:Migration]: Error checking schema migration status: {}",
+                err
+            );
+            return Err(err);
+        }
+    };
+
    let default_end_dt = "0".to_string();
    let cc = infra_db::get_coordinator().await;
    cc.add_start_dt_column().await?;
@@ -50,5 +50,5 @@ tracing.workspace = true
 tracing-log.workspace = true
 tracing-subscriber.workspace = true
 utoipa.workspace = true
-walkdir.workspace = true
+async-walkdir.workspace = true
 lettre.workspace = true
@@ -19,6 +19,9 @@ use std::{
    path::Path,
 };

+use async_walkdir::WalkDir;
+use futures::StreamExt;
+
 #[inline(always)]
 pub fn get_file_meta(file: &str) -> Result<Metadata, std::io::Error> {
    let file = File::open(file)?;
@@ -40,37 +43,51 @@ pub fn put_file_contents(file: &str, contents: &[u8]) -> Result<(), std::io::Err
 }

 #[inline(always)]
-pub fn scan_files<P: AsRef<Path>>(root: P, ext: &str) -> Vec<String> {
-    walkdir::WalkDir::new(root)
-        .into_iter()
-        .filter_map(|entry| {
-            let entry = entry.ok()?;
-            let path = entry.path();
-            if path.is_file() {
-                let path_ext = path.extension()?.to_str()?;
-                if path_ext == ext {
-                    Some(path.to_str().unwrap().to_string())
+pub async fn scan_files<P: AsRef<Path>>(root: P, ext: &str) -> Vec<String> {
+    let mut wd = WalkDir::new(root);
+    let mut resp = Vec::new();
+    loop {
+        match wd.next().await {
+            Some(Ok(entry)) => {
+                let path = entry.path();
+                if path.is_file() {
+                    if let Some(e) = path.extension() {
+                        if e == ext {
+                            resp.push(path.to_str().unwrap().to_string())
+                        }
+                    }
                } else {
-                    None
+                    continue;
                }
-            } else {
-                None
            }
-        })
-        .collect()
-}
-
-pub fn clean_empty_dirs(dir: &str) -> Result<(), std::io::Error> {
-    let mut dirs = Vec::new();
-    for entry in walkdir::WalkDir::new(dir) {
-        let entry = entry?;
-        if entry.path().display().to_string() == dir {
-            continue;
-        }
-        if entry.file_type().is_dir() {
-            dirs.push(entry.path().to_str().unwrap().to_string());
+            Some(Err(_)) => {}
+            None => break,
        }
    }
+    resp
+}
+
+pub async fn clean_empty_dirs(dir: &str) -> Result<(), std::io::Error> {
+    let mut dirs = Vec::new();
+    let mut wd = async_walkdir::WalkDir::new(dir);
+
+    loop {
+        match wd.next().await {
+            Some(Ok(entry)) => {
+                if entry.path().display().to_string() == dir {
+                    continue;
+                }
+                if let Ok(ft) = entry.file_type().await {
+                    if ft.is_dir() {
+                        dirs.push(entry.path().to_str().unwrap().to_string())
+                    };
+                }
+            }
+            Some(Err(_)) => {}
+            None => break,
+        }
+    }
+
    dirs.sort_by_key(|b| std::cmp::Reverse(b.len()));
    for dir in dirs {
        if let Ok(entries) = std::fs::read_dir(&dir) {
@@ -166,7 +166,7 @@ impl Metrics for Querier {
        };

        let pattern = format!("{wal_dir}/files/{org_id}/metrics/{stream_name}/");
-        let files = scan_files(&pattern, "parquet");
+        let files = scan_files(&pattern, "parquet").await;

        if arrow_files.is_empty() && files.is_empty() {
            return Ok(Response::new(resp));
@@ -71,7 +71,7 @@ impl FileData {
        log::info!("Loading disk cache start");
        std::fs::create_dir_all(&self.root_dir).expect("create cache dir success");
        let wal_dir = Path::new(&self.root_dir).canonicalize().unwrap();
-        let files = scan_files(&self.root_dir, "parquet");
+        let files = scan_files(&self.root_dir, "parquet").await;
        for file in files {
            let local_path = Path::new(&file).canonicalize().unwrap();
            let file_key = local_path
@@ -23,7 +23,7 @@ serde.workspace = true
 serde_json.workspace = true
 snafu.workspace = true
 tokio.workspace = true
-walkdir.workspace = true
+async-walkdir.workspace = true

 [[bin]]
 name = "wal-reader"
@@ -16,14 +16,16 @@
 use std::{
    fs::{create_dir_all, File},
    io::{BufRead, BufReader},
-    path::PathBuf,
+    path::{Path, PathBuf},
    sync::Arc,
 };

+use async_walkdir::WalkDir;
 use config::{
    utils::{schema::infer_json_schema_from_values, schema_ext::SchemaExt},
    CONFIG,
 };
+use futures::StreamExt;
 use snafu::ResultExt;

 use crate::{errors::*, immutable, memtable, writer::WriterKey};
@@ -47,12 +49,12 @@ use crate::{errors::*, immutable, memtable, writer::WriterKey};
 //    files actually wrote to disk completely, need to continue step 5
 pub(crate) async fn check_uncompleted_parquet_files() -> Result<()> {
    // 1. get all .lock files
-    let wal_dir = PathBuf::from(&CONFIG.common.data_wal_dir).join("logs");
+    let wal_dir_path_buf = PathBuf::from(&CONFIG.common.data_wal_dir).join("logs");
+
+    let wal_dir = wal_dir_path_buf.as_path();
    // create wal dir if not exists
-    create_dir_all(&wal_dir).context(OpenDirSnafu {
-        path: wal_dir.clone(),
-    })?;
-    let lock_files = scan_files(wal_dir, "lock");
+    create_dir_all(wal_dir).context(OpenDirSnafu { path: wal_dir })?;
+    let lock_files = scan_files(wal_dir, "lock").await;

    // 2. check if there is a .wal file with same name, delete it and rename the .par to .parquet
    for lock_file in lock_files.iter() {
@@ -88,12 +90,12 @@ pub(crate) async fn check_uncompleted_parquet_files() -> Result<()> {
    }

    // 4. delete all the .par files
-    let parquet_dir = PathBuf::from(&CONFIG.common.data_wal_dir).join("files");
+    let parquet_dir_path_buf = PathBuf::from(&CONFIG.common.data_wal_dir).join("files");
+
+    let parquet_dir = parquet_dir_path_buf.as_path();
    // create wal dir if not exists
-    create_dir_all(&parquet_dir).context(OpenDirSnafu {
-        path: parquet_dir.clone(),
-    })?;
-    let par_files = scan_files(parquet_dir, "par");
+    create_dir_all(parquet_dir).context(OpenDirSnafu { path: parquet_dir })?;
+    let par_files = scan_files(parquet_dir, "par").await;
    for par_file in par_files.iter() {
        log::warn!("delete uncompleted par file: {:?}", par_file);
        std::fs::remove_file(par_file).context(DeleteFileSnafu { path: par_file })?;
@@ -107,7 +109,7 @@ pub(crate) async fn replay_wal_files() -> Result<()> {
    create_dir_all(&wal_dir).context(OpenDirSnafu {
        path: wal_dir.clone(),
    })?;
-    let wal_files = scan_files(&wal_dir, "wal");
+    let wal_files = scan_files(&wal_dir, "wal").await;
    if wal_files.is_empty() {
        return Ok(());
    }
@@ -214,22 +216,26 @@ pub(crate) async fn replay_wal_files() -> Result<()> {
    Ok(())
 }

-fn scan_files(root_dir: impl Into<PathBuf>, ext: &str) -> Vec<PathBuf> {
-    walkdir::WalkDir::new(root_dir.into())
-        .into_iter()
-        .filter_map(|entry| {
-            let entry = entry.ok()?;
-            let path = entry.path();
-            if path.is_file() {
-                let path_ext = path.extension()?.to_str()?;
-                if path_ext == ext {
-                    Some(PathBuf::from(path))
+async fn scan_files(root_dir: &Path, ext: &str) -> Vec<PathBuf> {
+    let mut wd = WalkDir::new(root_dir);
+    let mut resp = Vec::new();
+    loop {
+        match wd.next().await {
+            Some(Ok(entry)) => {
+                let path = entry.path();
+                if path.is_file() {
+                    if let Some(e) = path.extension() {
+                        if e == ext {
+                            resp.push(PathBuf::from(path.to_str().unwrap()))
+                        }
+                    }
                } else {
-                    None
+                    continue;
                }
-            } else {
-                None
            }
-        })
-        .collect()
+            Some(Err(_)) => {}
+            None => break,
+        }
+    }
+    resp
 }
@@ -62,7 +62,7 @@ pub async fn move_file_list_to_storage(check_in_use: bool) -> Result<(), anyhow:
        .unwrap();

    let pattern = format!("{}file_list/", &CONFIG.common.data_wal_dir);
-    let files = scan_files(&pattern, "json");
+    let files = scan_files(&pattern, "json").await;

    for file in files {
        let local_file = file.to_owned();
@@ -70,7 +70,7 @@ pub async fn move_files_to_storage() -> Result<(), anyhow::Error> {
        .unwrap();

    let pattern = format!("{}files/", &CONFIG.common.data_wal_dir);
-    let files = scan_files(&pattern, "json");
+    let files = scan_files(&pattern, "json").await;

    // use multiple threads to upload files
    let mut tasks = Vec::new();
@@ -77,7 +77,7 @@ pub async fn move_files_to_storage() -> Result<(), anyhow::Error> {
        .unwrap();

    let pattern = wal_dir.join("files/");
-    let mut files = scan_files(&pattern, "parquet");
+    let mut files = scan_files(&pattern, "parquet").await;
    if files.is_empty() {
        return Ok(());
    }
@@ -69,8 +69,8 @@ async fn load_ingest_wal_used_bytes() -> Result<(), anyhow::Error> {
        Err(_) => return Ok(()),
    };
    let pattern = format!("{}files/", &CONFIG.common.data_wal_dir);
-    let mut files = scan_files(&pattern, "parquet");
-    files.extend(scan_files(&pattern, "json"));
+    let mut files = scan_files(&pattern, "parquet").await;
+    files.extend(scan_files(&pattern, "json").await);
    let mut sizes = HashMap::new();
    for file in files {
        let local_file = file.to_owned();
@@ -188,7 +188,7 @@ pub async fn init() -> Result<(), anyhow::Error> {
            log::error!("Failed to create wal dir: {}", e);
        }
        // clean empty sub dirs
-        _ = clean_empty_dirs(&CONFIG.common.data_wal_dir);
+        let _ = clean_empty_dirs(&CONFIG.common.data_wal_dir).await;
    }

    tokio::task::spawn(async move { files::run().await });
@@ -75,7 +75,7 @@ pub async fn set(key: &str, meta: Option<FileMeta>, deleted: bool) -> Result<(),
 async fn get_in_wal() -> Result<Vec<FileKey>, anyhow::Error> {
    let mut result = Vec::with_capacity(1024);
    let pattern = format!("{}file_list/", &CONFIG.common.data_wal_dir);
-    let files = scan_files(&pattern, "json");
+    let files = scan_files(&pattern, "json").await;
    let mut line_no = 0;
    for file in files {
        line_no += 1;
@@ -536,7 +536,7 @@ async fn get_file_list_inner(
        "{}/files/{}/{stream_type}/{}/",
        wal_dir, &sql.org_id, &sql.stream_name
    );
-    let files = scan_files(&pattern, file_ext);
+    let files = scan_files(&pattern, file_ext).await;
    if files.is_empty() {
        return Ok(vec![]);
    }