diff --git a/src/api/main.py b/src/api/main.py index 4fb72e2..96c90c3 100644 --- a/src/api/main.py +++ b/src/api/main.py @@ -14,6 +14,7 @@ import logging import time +import threading from src.core.engine import NDLOCREngine from src.schemas.ocr import OCRResponse, OCRPage, OCRLine, OCRRequest, OCRJobResponse, OCRJobResult from src.schemas.openai import ChatCompletionRequest, ChatCompletionResponse, ChatCompletionChoice, ChatCompletionChoiceMessage @@ -27,21 +28,63 @@ class InMemoryJobStore: In-memory job store for OCR results. Used to track the status and final output of asynchronous background jobs. In a production environment, this should be replaced with a persistent store like Redis. + Includes time-to-live (TTL) eviction and LRU eviction to prevent unbounded memory growth. """ - def __init__(self): + def __init__(self, max_size: int = 1000, ttl_seconds: float = 3600.0): self._jobs: Dict[str, OCRJobResult] = {} + self._timestamps: Dict[str, float] = {} # job_id -> last access timestamp + self._lock = threading.Lock() + self._max_size = max_size + self._ttl_seconds = ttl_seconds + + def _evict_expired(self, now: float): + """Removes expired jobs from the store. Must be called under lock.""" + expired_ids = [ + job_id for job_id, ts in list(self._timestamps.items()) + if now - ts > self._ttl_seconds + ] + for job_id in expired_ids: + self._jobs.pop(job_id, None) + self._timestamps.pop(job_id, None) + + def _evict_lru(self): + """Removes the oldest (least recently accessed) job. Must be called under lock.""" + if not self._timestamps: + return + # Find key with the minimum timestamp + oldest_job_id = min(self._timestamps, key=self._timestamps.get) + self._jobs.pop(oldest_job_id, None) + self._timestamps.pop(oldest_job_id, None) def get(self, job_id: str) -> Optional[OCRJobResult]: - """Retrieves a job result by its ID.""" - return self._jobs.get(job_id) + """Retrieves a job result by its ID and updates its access time.""" + with self._lock: + now = time.time() + self._evict_expired(now) + if job_id in self._jobs: + self._timestamps[job_id] = now + return self._jobs[job_id] + return None def set(self, job_id: str, result: OCRJobResult): - """Stores or updates a job result.""" - self._jobs[job_id] = result + """Stores or updates a job result, evicting expired or LRU items if full.""" + with self._lock: + now = time.time() + self._evict_expired(now) + + # If still over limit and adding a new key would exceed max_size + if len(self._jobs) >= self._max_size and job_id not in self._jobs: + self._evict_lru() + + self._jobs[job_id] = result + self._timestamps[job_id] = now def exists(self, job_id: str) -> bool: """Checks if a job with the given ID exists.""" - return job_id in self._jobs + with self._lock: + now = time.time() + self._evict_expired(now) + return job_id in self._jobs def _engine_result_to_ocr_page(result: Dict[str, Any], index: int = 0) -> OCRPage: """ @@ -74,6 +117,8 @@ def _engine_result_to_ocr_page(result: Dict[str, Any], index: int = 0) -> OCRPag MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", 10 * 1024 * 1024)) # Default 10MB MAX_BODY_SIZE = int(os.getenv("MAX_BODY_SIZE", 15 * 1024 * 1024)) # Default 15MB MAX_PIXELS = int(os.getenv("MAX_PIXELS", 100_000_000)) # Default 100MP +JOB_STORE_MAX_SIZE = int(os.getenv("JOB_STORE_MAX_SIZE", "1000")) +JOB_STORE_TTL_SECONDS = float(os.getenv("JOB_STORE_TTL_SECONDS", "3600.0")) @asynccontextmanager async def lifespan(app: FastAPI): @@ -87,7 +132,7 @@ async def lifespan(app: FastAPI): # Models are loaded once and stored in app.state for sharing across requests. enable_tcy = os.getenv("ENABLE_TCY", "false").lower() == "true" app.state.engine = NDLOCREngine(device="cpu", enable_tcy=enable_tcy) - app.state.job_store = InMemoryJobStore() + app.state.job_store = InMemoryJobStore(max_size=JOB_STORE_MAX_SIZE, ttl_seconds=JOB_STORE_TTL_SECONDS) yield logger.info("Shutting down...") # Clean up engine resources (e.g., ThreadPoolExecutor) diff --git a/tests/test_job_store_limits.py b/tests/test_job_store_limits.py new file mode 100644 index 0000000..8ef276c --- /dev/null +++ b/tests/test_job_store_limits.py @@ -0,0 +1,81 @@ +import time +import threading +from src.api.main import InMemoryJobStore +from src.schemas.ocr import OCRJobResult + +def test_job_store_max_size_eviction(): + # Store with max_size=3 + store = InMemoryJobStore(max_size=3, ttl_seconds=3600.0) + + job1 = OCRJobResult(job_id="job1", status="pending") + job2 = OCRJobResult(job_id="job2", status="pending") + job3 = OCRJobResult(job_id="job3", status="pending") + job4 = OCRJobResult(job_id="job4", status="pending") + + store.set("job1", job1) + store.set("job2", job2) + store.set("job3", job3) + + # All should exist + assert store.exists("job1") + assert store.exists("job2") + assert store.exists("job3") + + # Access job1 and job2 to update their timestamps (making job3 the least recently used) + # Note: Because they are created very close in time, we want to make sure order is clear. + # To be extremely safe, we will simulate or just set them with slight artificial delays or let the system access them. + store.get("job1") + time.sleep(0.01) + store.get("job2") + time.sleep(0.01) + + # Now we insert job4, which should trigger LRU eviction on job3 + store.set("job4", job4) + + assert store.exists("job1") + assert store.exists("job2") + assert store.exists("job4") + assert not store.exists("job3") + +def test_job_store_ttl_eviction(): + # Store with max_size=10, ttl_seconds=0.05 + store = InMemoryJobStore(max_size=10, ttl_seconds=0.05) + + job1 = OCRJobResult(job_id="job1", status="pending") + store.set("job1", job1) + + assert store.exists("job1") + + # Sleep past TTL + time.sleep(0.06) + + # exists/get/set should trigger eviction and return None / False + assert not store.exists("job1") + assert store.get("job1") is None + +def test_job_store_thread_safety(): + # Ensure no race conditions under high concurrent set/get operations + store = InMemoryJobStore(max_size=10, ttl_seconds=10.0) + + num_threads = 20 + ops_per_thread = 100 + errors = [] + + def worker(thread_idx): + try: + for i in range(ops_per_thread): + job_id = f"thread-{thread_idx}-job-{i}" + job = OCRJobResult(job_id=job_id, status="pending") + store.set(job_id, job) + store.get(job_id) + store.exists(job_id) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"Encountered thread safety errors: {errors}"