Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 30 additions & 21 deletions src/spiders/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,23 +24,30 @@ impl ResponseCache {
Ok(Self { cache_dir: path })
}

pub fn get(&self, url: &str) -> Option<CachedResponse> {
pub async fn get(&self, url: &str) -> Option<CachedResponse> {
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 {
Expand All @@ -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 = "<html>updated</html>".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,
"<html>updated</html>"
);

Expand All @@ -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());
}
}
60 changes: 33 additions & 27 deletions src/spiders/checkpoint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<CheckpointData> {
pub async fn restore(&self) -> Option<CheckpointData> {
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 {
Expand All @@ -65,29 +71,29 @@ 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,
vec!["https://example.com".to_string()]
);
}

#[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()
Expand All @@ -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());
}
}
10 changes: 5 additions & 5 deletions src/spiders/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ impl<S: Spider> CrawlerEngine<S> {
// 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);
Expand Down Expand Up @@ -302,12 +302,12 @@ impl<S: Spider> CrawlerEngine<S> {
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;
}
}

Expand Down Expand Up @@ -378,7 +378,7 @@ impl<S: Spider> CrawlerEngine<S> {

// 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(
Expand Down Expand Up @@ -434,7 +434,7 @@ impl<S: Spider> CrawlerEngine<S> {
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);
Expand Down
Loading