diff --git a/src/spiders/cache.rs b/src/spiders/cache.rs index 199e137..d58202a 100644 --- a/src/spiders/cache.rs +++ b/src/spiders/cache.rs @@ -24,23 +24,30 @@ impl ResponseCache { Ok(Self { cache_dir: path }) } - pub fn get(&self, url: &str) -> Option { + pub async fn get(&self, url: &str) -> Option { let file_path = self.cache_path(url); - let data = std::fs::read_to_string(&file_path).ok()?; + // Async read so the Tokio worker thread is not blocked on disk I/O. + let data = tokio::fs::read_to_string(&file_path).await.ok()?; serde_json::from_str(&data).ok() } - pub fn put(&self, url: &str, response: &CachedResponse) -> Result<(), std::io::Error> { + pub async fn put(&self, url: &str, response: &CachedResponse) -> Result<(), std::io::Error> { let file_path = self.cache_path(url); + let cache_dir = self.cache_dir.clone(); let data = serde_json::to_string_pretty(response).map_err(std::io::Error::other)?; - // Write atomically: a temp file in the same directory keeps the rename - // same-filesystem, and `persist` performs an atomic replace on both - // POSIX (rename) and Windows (MoveFileExW with REPLACE_EXISTING). The - // temp file is auto-removed on drop if anything fails before persist. - let mut tmp = NamedTempFile::new_in(&self.cache_dir)?; - tmp.write_all(data.as_bytes())?; - tmp.persist(&file_path).map_err(|e| e.error)?; - Ok(()) + // The atomic temp-file + persist write is synchronous (NamedTempFile + // has no async API), so run it on the blocking pool to avoid stalling + // a Tokio worker thread. `persist` still does an atomic replace on both + // POSIX (rename) and Windows (MoveFileExW with REPLACE_EXISTING), and + // the temp file is auto-removed on drop if anything fails. + tokio::task::spawn_blocking(move || -> Result<(), std::io::Error> { + let mut tmp = NamedTempFile::new_in(&cache_dir)?; + tmp.write_all(data.as_bytes())?; + tmp.persist(&file_path).map_err(|e| e.error)?; + Ok(()) + }) + .await + .map_err(std::io::Error::other)? } fn cache_path(&self, url: &str) -> PathBuf { @@ -66,32 +73,34 @@ mod tests { } } - #[test] - fn put_then_get_roundtrips() { + #[tokio::test] + async fn put_then_get_roundtrips() { let dir = tempdir().unwrap(); let cache = ResponseCache::new(dir.path().to_str().unwrap()).unwrap(); cache .put("https://example.com", &sample("https://example.com")) + .await .unwrap(); - let got = cache.get("https://example.com").unwrap(); + let got = cache.get("https://example.com").await.unwrap(); assert_eq!(got.status, 200); assert_eq!(got.url, "https://example.com"); } - #[test] - fn put_overwrites_existing_entry_and_leaves_no_tmp_files() { + #[tokio::test] + async fn put_overwrites_existing_entry_and_leaves_no_tmp_files() { let dir = tempdir().unwrap(); let cache = ResponseCache::new(dir.path().to_str().unwrap()).unwrap(); // Two writes to the same URL — the second must atomically replace the // first (this is the path that fails on Windows with plain rename). cache .put("https://example.com", &sample("https://example.com")) + .await .unwrap(); let mut updated = sample("https://example.com"); updated.body = "updated".to_string(); - cache.put("https://example.com", &updated).unwrap(); + cache.put("https://example.com", &updated).await.unwrap(); assert_eq!( - cache.get("https://example.com").unwrap().body, + cache.get("https://example.com").await.unwrap().body, "updated" ); @@ -105,10 +114,10 @@ mod tests { assert!(entries[0].ends_with(".json")); } - #[test] - fn get_returns_none_for_missing() { + #[tokio::test] + async fn get_returns_none_for_missing() { let dir = tempdir().unwrap(); let cache = ResponseCache::new(dir.path().to_str().unwrap()).unwrap(); - assert!(cache.get("https://missing.example").is_none()); + assert!(cache.get("https://missing.example").await.is_none()); } } diff --git a/src/spiders/checkpoint.rs b/src/spiders/checkpoint.rs index 68213ba..d7e2e39 100644 --- a/src/spiders/checkpoint.rs +++ b/src/spiders/checkpoint.rs @@ -22,29 +22,35 @@ impl CheckpointManager { }) } - pub fn save(&self, data: &CheckpointData) -> Result<(), std::io::Error> { + pub async fn save(&self, data: &CheckpointData) -> Result<(), std::io::Error> { let file_path = self.checkpoint_dir.join("checkpoint.json"); + let dir = self.checkpoint_dir.clone(); let json = serde_json::to_string_pretty(data).map_err(std::io::Error::other)?; - // Write atomically via a temp file in the same directory: `persist` - // does an atomic replace on POSIX (rename) and Windows (MoveFileExW - // with REPLACE_EXISTING), and the temp file is auto-removed on drop if - // anything fails before persist — so a crash mid-write cannot corrupt - // or zero-out an existing checkpoint, and no orphan .tmp is left behind. - let mut tmp = NamedTempFile::new_in(&self.checkpoint_dir)?; - tmp.write_all(json.as_bytes())?; - tmp.persist(&file_path).map_err(|e| e.error)?; - Ok(()) + // The atomic temp-file + persist write is synchronous, so run it on the + // blocking pool to avoid stalling the async runtime (this runs in the + // main crawl loop). `persist` does an atomic replace on POSIX (rename) + // and Windows (MoveFileExW with REPLACE_EXISTING), and the temp file is + // auto-removed on drop if anything fails — so a crash mid-write cannot + // corrupt or zero-out an existing checkpoint, and no orphan .tmp leaks. + tokio::task::spawn_blocking(move || -> Result<(), std::io::Error> { + let mut tmp = NamedTempFile::new_in(&dir)?; + tmp.write_all(json.as_bytes())?; + tmp.persist(&file_path).map_err(|e| e.error)?; + Ok(()) + }) + .await + .map_err(std::io::Error::other)? } - pub fn restore(&self) -> Option { + pub async fn restore(&self) -> Option { let file_path = self.checkpoint_dir.join("checkpoint.json"); - let data = std::fs::read_to_string(&file_path).ok()?; + let data = tokio::fs::read_to_string(&file_path).await.ok()?; serde_json::from_str(&data).ok() } - pub fn cleanup(&self) { + pub async fn cleanup(&self) { let file_path = self.checkpoint_dir.join("checkpoint.json"); - let _ = std::fs::remove_file(file_path); + let _ = tokio::fs::remove_file(file_path).await; } pub fn exists(&self) -> bool { @@ -65,14 +71,14 @@ mod tests { } } - #[test] - fn save_then_restore_roundtrips() { + #[tokio::test] + async fn save_then_restore_roundtrips() { let dir = tempdir().unwrap(); let mgr = CheckpointManager::new(dir.path().to_str().unwrap()).unwrap(); assert!(!mgr.exists()); - mgr.save(&sample(42)).unwrap(); + mgr.save(&sample(42)).await.unwrap(); assert!(mgr.exists()); - let restored = mgr.restore().unwrap(); + let restored = mgr.restore().await.unwrap(); assert_eq!(restored.items_count, 42); assert_eq!( restored.pending_urls, @@ -80,14 +86,14 @@ mod tests { ); } - #[test] - fn save_overwrites_and_leaves_no_tmp_files() { + #[tokio::test] + async fn save_overwrites_and_leaves_no_tmp_files() { let dir = tempdir().unwrap(); let mgr = CheckpointManager::new(dir.path().to_str().unwrap()).unwrap(); - mgr.save(&sample(1)).unwrap(); + mgr.save(&sample(1)).await.unwrap(); // Overwrite — atomic replace must succeed on every platform. - mgr.save(&sample(2)).unwrap(); - assert_eq!(mgr.restore().unwrap().items_count, 2); + mgr.save(&sample(2)).await.unwrap(); + assert_eq!(mgr.restore().await.unwrap().items_count, 2); let entries: Vec<_> = std::fs::read_dir(dir.path()) .unwrap() @@ -97,13 +103,13 @@ mod tests { assert_eq!(entries, vec!["checkpoint.json".to_string()]); } - #[test] - fn cleanup_removes_checkpoint() { + #[tokio::test] + async fn cleanup_removes_checkpoint() { let dir = tempdir().unwrap(); let mgr = CheckpointManager::new(dir.path().to_str().unwrap()).unwrap(); - mgr.save(&sample(7)).unwrap(); + mgr.save(&sample(7)).await.unwrap(); assert!(mgr.exists()); - mgr.cleanup(); + mgr.cleanup().await; assert!(!mgr.exists()); } } diff --git a/src/spiders/engine.rs b/src/spiders/engine.rs index 9d6274d..8988f89 100644 --- a/src/spiders/engine.rs +++ b/src/spiders/engine.rs @@ -147,7 +147,7 @@ impl CrawlerEngine { // Check for checkpoint restore let resuming = if let Some(ref cp) = self.checkpoint { if cp.exists() { - if let Some(data) = cp.restore() { + if let Some(data) = cp.restore().await { let mut sched = self.scheduler.lock().await; for url in &data.pending_urls { let mut req = SpiderRequest::new(url); @@ -302,12 +302,12 @@ impl CrawlerEngine { seen_fingerprints: Vec::new(), // Fingerprints are internal to scheduler items_count, }; - let _ = cp.save(&data); + let _ = cp.save(&data).await; } } else { // Clean up checkpoint on successful completion if let Some(ref cp) = self.checkpoint { - cp.cleanup(); + cp.cleanup().await; } } @@ -378,7 +378,7 @@ impl CrawlerEngine { // Check dev cache if let Some(ref response_cache) = cache { - if let Some(cached) = response_cache.get(&url) { + if let Some(cached) = response_cache.get(&url).await { stats.lock().await.cache_hits += 1; let fetcher_resp = FetcherResponse::new( @@ -434,7 +434,7 @@ impl CrawlerEngine { url: response.url().to_string(), headers: response.headers().clone(), }; - let _ = response_cache.put(&url, &cached); + let _ = response_cache.put(&url, &cached).await; } let spider_resp = SpiderResponse::new(response);