Ahmed Elshahat
f0db3d7d14
CheckpointManager.save() and ResponseCacheManager.put() write to a temp file and then move it into place with Path.rename(). On Windows, os.rename cannot overwrite an existing destination and raises FileExistsError (WinError 183), so every write after the first one fails: checkpoint saving raises and breaks resume, while the development response cache swallows the error and keeps returning the stale entry. Path.replace() (os.replace) overwrites the destination atomically on every platform and behaves identically to rename() on POSIX, so this is a no-op on Linux and macOS and only fixes the broken overwrite on Windows. Add a regression test for the cache overwrite path; the checkpoint overwrite is already covered by test_multiple_saves_overwrite.
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
import pickle
|
|
from pathlib import Path
|
|
from dataclasses import dataclass, field
|
|
|
|
import anyio
|
|
from anyio import Path as AsyncPath
|
|
|
|
from scrapling.core.utils import log
|
|
from scrapling.core._types import Set, List, Optional, TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from scrapling.spiders.request import Request
|
|
|
|
|
|
@dataclass
|
|
class CheckpointData:
|
|
"""Container for checkpoint state."""
|
|
|
|
requests: List["Request"] = field(default_factory=list)
|
|
seen: Set[bytes] = field(default_factory=set)
|
|
|
|
|
|
class CheckpointManager:
|
|
"""Manages saving and loading checkpoint state to/from disk."""
|
|
|
|
CHECKPOINT_FILE = "checkpoint.pkl"
|
|
|
|
def __init__(self, crawldir: str | Path | AsyncPath, interval: float = 300.0):
|
|
self.crawldir = AsyncPath(crawldir)
|
|
self._checkpoint_path = self.crawldir / self.CHECKPOINT_FILE
|
|
self.interval = interval
|
|
if not isinstance(interval, (int, float)):
|
|
raise TypeError("Checkpoints interval must be integer or float.")
|
|
else:
|
|
if interval < 0:
|
|
raise ValueError("Checkpoints interval must be equal or greater than 0.")
|
|
|
|
async def has_checkpoint(self) -> bool:
|
|
"""Check if a checkpoint exists."""
|
|
return await self._checkpoint_path.exists()
|
|
|
|
async def save(self, data: CheckpointData) -> None:
|
|
"""Save checkpoint data to disk atomically."""
|
|
await self.crawldir.mkdir(parents=True, exist_ok=True)
|
|
|
|
temp_path = self._checkpoint_path.with_suffix(".tmp")
|
|
|
|
try:
|
|
serialized = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL)
|
|
async with await anyio.open_file(temp_path, "wb") as f:
|
|
await f.write(serialized)
|
|
|
|
await temp_path.replace(self._checkpoint_path)
|
|
|
|
log.info(f"Checkpoint saved: {len(data.requests)} requests, {len(data.seen)} seen URLs")
|
|
except Exception as e:
|
|
# Clean up temp file if it exists
|
|
if await temp_path.exists():
|
|
await temp_path.unlink()
|
|
log.error(f"Failed to save checkpoint: {e}")
|
|
raise
|
|
|
|
async def load(self) -> Optional[CheckpointData]:
|
|
"""Load checkpoint data from disk.
|
|
|
|
Returns None if no checkpoint exists or if loading fails.
|
|
"""
|
|
if not await self.has_checkpoint():
|
|
return None
|
|
|
|
try:
|
|
async with await anyio.open_file(self._checkpoint_path, "rb") as f:
|
|
content = await f.read()
|
|
data: CheckpointData = pickle.loads(content)
|
|
|
|
log.info(f"Checkpoint loaded: {len(data.requests)} requests, {len(data.seen)} seen URLs")
|
|
return data
|
|
|
|
except Exception as e:
|
|
log.error(f"Failed to load checkpoint (starting fresh): {e}")
|
|
return None
|
|
|
|
async def cleanup(self) -> None:
|
|
"""Delete checkpoint file after successful completion."""
|
|
try:
|
|
if await self._checkpoint_path.exists():
|
|
await self._checkpoint_path.unlink()
|
|
log.debug("Checkpoint file cleaned up")
|
|
except Exception as e:
|
|
log.warning(f"Failed to cleanup checkpoint file: {e}")
|