Reduce blocking-pool overhead in file hooks

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<Regex>,
     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<Regex>,
+    policy: MatchPolicy,
+) -> Result<(i32, Vec<u8>)> {
+    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<u8>)> {
-    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) {

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<u8>)> {
     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,18 +99,28 @@ pub(crate) async fn check_vcs_permalinks(
 async fn check_file(
     file_base: &Path,
     filename: &Path,
-    matcher: &GithubNonPermalinkMatcher,
+    matcher: &Arc<GithubNonPermalinkMatcher>,
 ) -> Result<(i32, Vec<u8>)> {
     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<u8>)> {
+    let file = fs_err::File::open(path)?;
     let mut reader = BufReader::new(file);
 
     let mut retval = 0;
     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;
         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?;
 

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<u8>)> {
-    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<bool> {
+    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<u8>)>
         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<u8>)>
         }
     }
 
-    Ok((0, Vec::new()))
+    Ok(false)
 }
 
 #[cfg(test)]

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<u8>)> {
     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<bool> {
     // 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<T>(reader: &mut T) -> Result<(Option<u64>, Option<&str>)>
+fn find_last_non_ending<T>(reader: &mut T) -> Result<(Option<u64>, 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;