From 10e67a9a2e00dccf6a8883418e547ff6cefb6803 Mon Sep 17 00:00:00 2001 From: Jo <10510431+j178@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:29:38 +0800 Subject: [PATCH] Reduce blocking-pool overhead in file hooks --- .../prek/src/hooks/builtin_hooks/pattern.rs | 28 +++++++-- .../pre_commit_hooks/check_vcs_permalinks.rs | 26 ++++++--- .../pre_commit_hooks/detect_private_key.rs | 23 ++++++-- .../hooks/pre_commit_hooks/fix_end_of_file.rs | 57 ++++++++++--------- 4 files changed, 89 insertions(+), 45 deletions(-) diff --git a/crates/prek/src/hooks/builtin_hooks/pattern.rs b/crates/prek/src/hooks/builtin_hooks/pattern.rs index 55b70ca30..810c8dc3d 100644 --- a/crates/prek/src/hooks/builtin_hooks/pattern.rs +++ b/crates/prek/src/hooks/builtin_hooks/pattern.rs @@ -1,11 +1,11 @@ -use std::io::Write; +use std::io::{BufRead, BufReader, Write}; use std::path::Path; +use std::sync::Arc; use anyhow::{Context, Result}; use clap::Parser; use memchr::memchr_iter; use regex_automata::{MatchKind, meta::Regex, util::syntax}; -use tokio::io::{AsyncBufReadExt, BufReader}; use crate::hook::Hook; use crate::hooks::run_concurrent_file_checks; @@ -37,7 +37,7 @@ enum ScanMode { } struct Matcher { - regex: Regex, + regex: Arc, scan_mode: ScanMode, } @@ -70,7 +70,10 @@ impl Matcher { ScanMode::Lines }; - Ok(Self { regex, scan_mode }) + Ok(Self { + regex: Arc::new(regex), + scan_mode, + }) } } @@ -110,17 +113,30 @@ async fn check_file( async fn check_lines( file_base: &Path, filename: &Path, + patterns: &Arc, + policy: MatchPolicy, +) -> Result<(i32, Vec)> { + let file_path = file_base.join(filename); + let filename = filename.to_path_buf(); + let patterns = Arc::clone(patterns); + tokio::task::spawn_blocking(move || check_lines_sync(&file_path, &filename, &patterns, policy)) + .await? +} + +fn check_lines_sync( + file_path: &Path, + filename: &Path, patterns: &Regex, policy: MatchPolicy, ) -> Result<(i32, Vec)> { - let file = fs_err::tokio::File::open(file_base.join(filename)).await?; + let file = fs_err::File::open(file_path)?; let mut reader = BufReader::new(file); let mut matched = false; let mut output = Vec::new(); let mut line = Vec::new(); let mut line_number = 0; - while reader.read_until(b'\n', &mut line).await? != 0 { + while reader.read_until(b'\n', &mut line)? != 0 { line_number += 1; let contents = trim_line_ending(&line); if patterns.is_match(contents) { diff --git a/crates/prek/src/hooks/pre_commit_hooks/check_vcs_permalinks.rs b/crates/prek/src/hooks/pre_commit_hooks/check_vcs_permalinks.rs index 6ec336ac4..b74786387 100644 --- a/crates/prek/src/hooks/pre_commit_hooks/check_vcs_permalinks.rs +++ b/crates/prek/src/hooks/pre_commit_hooks/check_vcs_permalinks.rs @@ -1,12 +1,12 @@ use std::collections::BTreeSet; -use std::io::Write; +use std::io::{BufRead, BufReader, Write}; use std::path::Path; +use std::sync::Arc; use anyhow::Result; use clap::Parser; use memchr::memmem; use regex::bytes::{Match, Regex}; -use tokio::io::{AsyncBufReadExt, BufReader}; use crate::hook::Hook; use crate::hooks::run_concurrent_file_checks; @@ -83,7 +83,9 @@ pub(crate) async fn check_vcs_permalinks( filenames: &[&Path], ) -> Result<(i32, Vec)> { let args = Args::try_parse_from(hook.entry.expect_direct().split_with_args(&hook.args)?)?; - let matcher = GithubNonPermalinkMatcher::new(args.additional_github_domains); + let matcher = Arc::new(GithubNonPermalinkMatcher::new( + args.additional_github_domains, + )); let file_base = hook.project().relative_path(); run_concurrent_file_checks( @@ -97,10 +99,20 @@ pub(crate) async fn check_vcs_permalinks( async fn check_file( file_base: &Path, filename: &Path, - matcher: &GithubNonPermalinkMatcher, + matcher: &Arc, ) -> Result<(i32, Vec)> { let path = file_base.join(filename); - let file = fs_err::tokio::File::open(&path).await?; + let filename = filename.to_path_buf(); + let matcher = Arc::clone(matcher); + tokio::task::spawn_blocking(move || check_file_sync(&path, &filename, &matcher)).await? +} + +fn check_file_sync( + path: &Path, + filename: &Path, + matcher: &GithubNonPermalinkMatcher, +) -> Result<(i32, Vec)> { + let file = fs_err::File::open(path)?; let mut reader = BufReader::new(file); let mut retval = 0; @@ -108,7 +120,7 @@ async fn check_file( let mut line = Vec::new(); let mut line_number = 0; - while reader.read_until(b'\n', &mut line).await? != 0 { + while reader.read_until(b'\n', &mut line)? != 0 { line_number += 1; for m in matcher.find_non_permalink(&line) { retval = 1; @@ -226,7 +238,7 @@ mod tests { ) .await?; - let matcher = matcher(&["github.example.com"]); + let matcher = Arc::new(matcher(&["github.example.com"])); let relative = PathBuf::from("links.md"); let (code, output) = check_file(dir.path(), &relative, &matcher).await?; diff --git a/crates/prek/src/hooks/pre_commit_hooks/detect_private_key.rs b/crates/prek/src/hooks/pre_commit_hooks/detect_private_key.rs index 1e8eb42b1..9c3ff27fe 100644 --- a/crates/prek/src/hooks/pre_commit_hooks/detect_private_key.rs +++ b/crates/prek/src/hooks/pre_commit_hooks/detect_private_key.rs @@ -1,9 +1,9 @@ +use std::io::Read; use std::path::Path; use std::sync::LazyLock; use aho_corasick::AhoCorasick; use anyhow::Result; -use tokio::io::AsyncReadExt; use crate::hook::Hook; use crate::hooks::run_concurrent_file_checks; @@ -56,12 +56,24 @@ pub(crate) async fn detect_private_key(hook: &Hook, filenames: &[&Path]) -> Resu /// with `ATE KEY`, we keep the tail of the first read, prepend it to the second /// read, and search the combined window so `BEGIN RSA PRIVATE KEY` is still found. async fn check_file(file_base: &Path, filename: &Path) -> Result<(i32, Vec)> { - let mut file = fs_err::tokio::File::open(file_base.join(filename)).await?; + let file_path = file_base.join(filename); + // Keep a file's blocking I/O in one task instead of re-entering the blocking pool per read. + let found = tokio::task::spawn_blocking(move || check_file_sync(&file_path)).await??; + if found { + let error_message = format!("Private key found: {}\n", filename.display()); + Ok((1, error_message.into_bytes())) + } else { + Ok((0, Vec::new())) + } +} + +fn check_file_sync(file_path: &Path) -> Result { + let mut file = fs_err::File::open(file_path)?; let mut buf = [0u8; BUFFER_SIZE + CARRY_CAPACITY]; let mut carry_len = 0; loop { - let bytes_read = file.read(&mut buf[carry_len..]).await?; + let bytes_read = file.read(&mut buf[carry_len..])?; if bytes_read == 0 { break; } @@ -70,8 +82,7 @@ async fn check_file(file_base: &Path, filename: &Path) -> Result<(i32, Vec)> let search_buf = &buf[..search_len]; if PRIVATE_KEY_MATCHER.find(search_buf).is_some() { - let error_message = format!("Private key found: {}\n", filename.display()); - return Ok((1, error_message.into_bytes())); + return Ok(true); } // Move the tail of this chunk to the front of the buffer so a key marker @@ -82,7 +93,7 @@ async fn check_file(file_base: &Path, filename: &Path) -> Result<(i32, Vec)> } } - Ok((0, Vec::new())) + Ok(false) } #[cfg(test)] diff --git a/crates/prek/src/hooks/pre_commit_hooks/fix_end_of_file.rs b/crates/prek/src/hooks/pre_commit_hooks/fix_end_of_file.rs index 8025ae26f..8257a1edb 100644 --- a/crates/prek/src/hooks/pre_commit_hooks/fix_end_of_file.rs +++ b/crates/prek/src/hooks/pre_commit_hooks/fix_end_of_file.rs @@ -1,7 +1,7 @@ +use std::io::{Read, Seek, SeekFrom, Write}; use std::path::Path; use anyhow::Result; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, AsyncWriteExt, SeekFrom}; use crate::hook::Hook; use crate::hooks::run_concurrent_file_checks; @@ -18,43 +18,50 @@ pub(crate) async fn fix_end_of_file(hook: &Hook, filenames: &[&Path]) -> Result< async fn fix_file(file_base: &Path, filename: &Path) -> Result<(i32, Vec)> { let file_path = file_base.join(filename); + // Keep a file's blocking I/O in one task instead of re-entering the blocking pool per operation. + let modified = tokio::task::spawn_blocking(move || fix_file_sync(&file_path)).await??; + if modified { + Ok((1, format!("Fixing {}\n", filename.display()).into_bytes())) + } else { + Ok((0, Vec::new())) + } +} + +fn fix_file_sync(file_path: &Path) -> Result { // If the file is empty, do nothing and avoid opening a write handle. - let file_size = fs_err::tokio::metadata(&file_path).await?.len(); + let file_size = fs_err::metadata(file_path)?.len(); if file_size == 0 { - return Ok((0, Vec::new())); + return Ok(false); } - let mut file = fs_err::tokio::OpenOptions::new() + let mut file = fs_err::OpenOptions::new() .read(true) .write(true) - .open(file_path) - .await?; + .open(file_path)?; - match find_last_non_ending(&mut file).await? { + match find_last_non_ending(&mut file)? { (None, _) => { // File contains only line endings, so we can just set it to empty. - file.set_len(0).await?; - file.flush().await?; - file.shutdown().await?; - Ok((1, format!("Fixing {}\n", filename.display()).into_bytes())) + file.set_len(0)?; + file.flush()?; + Ok(true) } (Some(pos), None) => { // File has some content, but no line ending at the end. - file.seek(SeekFrom::Start(pos + 1)).await?; - file.write_all(b"\n").await?; - file.flush().await?; - file.shutdown().await?; - Ok((1, format!("Fixing {}\n", filename.display()).into_bytes())) + file.seek(SeekFrom::Start(pos + 1))?; + file.write_all(b"\n")?; + file.flush()?; + Ok(true) } (Some(pos), Some(line_ending)) => { // File has some content and at least one line ending. let new_size = pos + 1 + line_ending.len() as u64; if file_size == new_size { // File already has the correct line ending. - return Ok((0, Vec::new())); + return Ok(false); } - file.set_len(new_size).await?; - Ok((1, format!("Fixing {}\n", filename.display()).into_bytes())) + file.set_len(new_size)?; + Ok(true) } } } @@ -73,13 +80,13 @@ fn determine_line_ending(first: u8, second: u8) -> Option<&'static str> { /// Searches for the last non-line-ending character in the file. /// Returns the position of the last non-line-ending character and the line ending type. -async fn find_last_non_ending(reader: &mut T) -> Result<(Option, Option<&str>)> +fn find_last_non_ending(reader: &mut T) -> Result<(Option, Option<&str>)> where - T: AsyncRead + AsyncSeek + Unpin, + T: Read + Seek, { const MAX_SCAN_SIZE: usize = 4 * 1024; // 4KB - let data_len = reader.seek(SeekFrom::End(0)).await?; + let data_len = reader.seek(SeekFrom::End(0))?; if data_len == 0 { return Ok((None, None)); } @@ -92,10 +99,8 @@ where while read_len < data_len { let block_size = MAX_SCAN_SIZE.min(usize::try_from(data_len - read_len)?); // SAFETY: block_size is guaranteed to be less than or equal to MAX_SCAN_SIZE - reader - .seek(SeekFrom::Current(-i64::try_from(block_size).unwrap())) - .await?; - reader.read_exact(&mut buf[..block_size]).await?; + reader.seek(SeekFrom::Current(-i64::try_from(block_size).unwrap()))?; + reader.read_exact(&mut buf[..block_size])?; read_len += block_size as u64; let mut pos = block_size;