diff --git a/.env.example b/.env.example index 17e0363..9e4a0a9 100644 --- a/.env.example +++ b/.env.example @@ -28,6 +28,7 @@ DATABASE_POOL_MODE=auto DATABASE_POOL_SIZE=5 DATABASE_POOL_MAX_OVERFLOW=10 DATABASE_POOL_TIMEOUT_SECONDS=30 +DATABASE_RESERVED_API_CONNECTIONS=3 # ------------------------------ # Frontend Dev Overrides diff --git a/app/api/routers/runs.py b/app/api/routers/runs.py index d074bb1..87dfa0a 100644 --- a/app/api/routers/runs.py +++ b/app/api/routers/runs.py @@ -221,11 +221,11 @@ def _spawn_background_execution( user_settings: Any, idempotency_key: str | None, ) -> None: - from app.db.session import get_session_factory + from app.db.background_session import background_session task = asyncio.create_task( ingest_service.execute_run( - session_factory=get_session_factory(), + session_factory=background_session, run_id=run.id, user_id=current_user.id, scholars=scholars, diff --git a/app/db/background_session.py b/app/db/background_session.py new file mode 100644 index 0000000..745f155 --- /dev/null +++ b/app/db/background_session.py @@ -0,0 +1,57 @@ +"""Semaphore-gated DB sessions for background tasks. + +Background tasks (ingestion, PDF resolution, scheduler) acquire a semaphore +permit before checking out a database connection. The semaphore cap is +``pool_size + max_overflow - reserved_for_api``, which guarantees that at +least *reserved_for_api* connections remain available for API request +handlers at all times. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db.session import get_session_factory +from app.logging_utils import structured_log +from app.settings import settings + +logger = logging.getLogger(__name__) + +_semaphore: asyncio.Semaphore | None = None + + +def _build_semaphore() -> asyncio.Semaphore: + pool_capacity = max(1, settings.database_pool_size) + max(0, settings.database_pool_max_overflow) + reserved = max(0, settings.database_reserved_api_connections) + limit = max(1, pool_capacity - reserved) + structured_log( + logger, + "info", + "db.background_semaphore_initialized", + pool_capacity=pool_capacity, + reserved_for_api=reserved, + background_limit=limit, + ) + return asyncio.Semaphore(limit) + + +def get_background_semaphore() -> asyncio.Semaphore: + global _semaphore + if _semaphore is None: + _semaphore = _build_semaphore() + return _semaphore + + +@asynccontextmanager +async def background_session() -> AsyncIterator[AsyncSession]: + """Yield a DB session after acquiring the background semaphore.""" + semaphore = get_background_semaphore() + async with semaphore: + session_factory = get_session_factory() + async with session_factory() as session: + yield session diff --git a/app/services/ingestion/queue_runner.py b/app/services/ingestion/queue_runner.py index 2220f58..9f7bd1b 100644 --- a/app/services/ingestion/queue_runner.py +++ b/app/services/ingestion/queue_runner.py @@ -5,8 +5,8 @@ from datetime import UTC, datetime from sqlalchemy import select +from app.db.background_session import background_session from app.db.models import QueueItemStatus, RunTriggerType, ScholarProfile, UserSetting -from app.db.session import get_session_factory from app.logging_utils import structured_log from app.services.ingestion import queue as queue_service from app.services.ingestion.application import ( @@ -57,8 +57,7 @@ class QueueJobRunner: async def drain_continuation_queue(self) -> None: now = datetime.now(UTC) - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: jobs = await queue_service.list_due_jobs( session, now=now, @@ -73,8 +72,7 @@ class QueueJobRunner: ) -> bool: if job.attempt_count < self._continuation_max_attempts: return False - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: dropped = await queue_service.mark_dropped( session, job_id=job.id, @@ -98,8 +96,7 @@ class QueueJobRunner: self, job: queue_service.ContinuationQueueJob, ) -> bool: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: queue_item = await queue_service.mark_retrying(session, job_id=job.id) await session.commit() if queue_item is None: @@ -110,8 +107,7 @@ class QueueJobRunner: self, job: queue_service.ContinuationQueueJob, ) -> bool: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: scholar_result = await session.execute( select(ScholarProfile.id).where( ScholarProfile.user_id == job.user_id, @@ -140,8 +136,7 @@ class QueueJobRunner: return False async def _reschedule_queue_job_lock_active(self, job: queue_service.ContinuationQueueJob) -> None: - session_factory = get_session_factory() - async with session_factory() as recovery_session: + async with background_session() as recovery_session: await queue_service.reschedule_job( recovery_session, job_id=job.id, @@ -167,8 +162,7 @@ class QueueJobRunner: self._tick_seconds, int(exc.safety_state.get("cooldown_remaining_seconds") or 0), ) - session_factory = get_session_factory() - async with session_factory() as recovery_session: + async with background_session() as recovery_session: await queue_service.reschedule_job( recovery_session, job_id=job.id, @@ -193,8 +187,7 @@ class QueueJobRunner: *, exc: Exception, ) -> None: - session_factory = get_session_factory() - async with session_factory() as recovery_session: + async with background_session() as recovery_session: queue_item = await queue_service.increment_attempt_count(recovery_session, job_id=job.id) if queue_item is None: await recovery_session.commit() @@ -238,8 +231,7 @@ class QueueJobRunner: ) async def _load_request_delay_for_user(self, user_id: int, *, floor: int) -> int: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: result = await session.execute( select(UserSetting.request_delay_seconds).where(UserSetting.user_id == user_id) ) @@ -253,8 +245,7 @@ class QueueJobRunner: request_delay_floor: int, ): request_delay_seconds = await self._load_request_delay_for_user(job.user_id, floor=request_delay_floor) - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: ingestion = ScholarIngestionService(source=self._source) try: return await ingestion.run_for_user( @@ -366,8 +357,7 @@ class QueueJobRunner: ) async def _finalize_queue_job_after_run(self, job: queue_service.ContinuationQueueJob, run_summary) -> None: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: if int(run_summary.failed_count) <= 0: await self._finalize_successful_queue_job(session, job, run_summary) else: diff --git a/app/services/ingestion/scheduler.py b/app/services/ingestion/scheduler.py index 6d2bec0..e44cb45 100644 --- a/app/services/ingestion/scheduler.py +++ b/app/services/ingestion/scheduler.py @@ -8,13 +8,13 @@ from typing import Any from sqlalchemy import select +from app.db.background_session import background_session from app.db.models import ( CrawlRun, RunTriggerType, User, UserSetting, ) -from app.db.session import get_session_factory from app.logging_utils import structured_log from app.services.ingestion.application import ( RunAlreadyInProgressError, @@ -148,8 +148,7 @@ class SchedulerService: await self._run_candidate(candidate) async def _load_candidate_rows(self) -> list[Any]: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: result = await session.execute( select( UserSetting.user_id, @@ -204,8 +203,7 @@ class SchedulerService: return candidates async def _is_due(self, candidate: _AutoRunCandidate, *, now: datetime) -> bool: - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: result = await session.execute( select(CrawlRun.start_dt) .where( @@ -227,8 +225,7 @@ class SchedulerService: *, candidate: _AutoRunCandidate, ): - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: ingestion = ScholarIngestionService(source=self._source) try: return await ingestion.run_for_user( @@ -294,8 +291,7 @@ class SchedulerService: if is_budget_cooldown_active(): return - session_factory = get_session_factory() - async with session_factory() as session: + async with background_session() as session: try: processed = await drain_ready_jobs( session, diff --git a/app/services/publications/pdf_queue_resolution.py b/app/services/publications/pdf_queue_resolution.py index 26e4b0b..ac71edd 100644 --- a/app/services/publications/pdf_queue_resolution.py +++ b/app/services/publications/pdf_queue_resolution.py @@ -4,8 +4,8 @@ import asyncio import logging from datetime import UTC, datetime, timedelta +from app.db.background_session import background_session from app.db.models import Publication, PublicationPdfJob -from app.db.session import get_session_factory from app.logging_utils import structured_log from app.services.publication_identifiers import application as identifier_service from app.services.publications.pdf_queue_common import ( @@ -51,8 +51,7 @@ async def _mark_attempt_started( publication_id: int, user_id: int, ) -> None: - session_factory = get_session_factory() - async with session_factory() as db_session: + async with background_session() as db_session: job = await db_session.get(PublicationPdfJob, publication_id) if job is None: job = queued_job(publication_id=publication_id, user_id=user_id) @@ -138,8 +137,7 @@ async def _persist_outcome( user_id: int, outcome: OaResolutionOutcome, ) -> None: - session_factory = get_session_factory() - async with session_factory() as db_session: + async with background_session() as db_session: publication = await db_session.get(Publication, publication_id) job = await db_session.get(PublicationPdfJob, publication_id) if publication is None or job is None: @@ -220,8 +218,7 @@ async def _run_resolution_task( openalex_api_key: str | None = None try: - session_factory = get_session_factory() - async with session_factory() as key_session: + async with background_session() as key_session: user_settings = await user_settings_service.get_or_create_settings(key_session, user_id=user_id) openalex_api_key = getattr(user_settings, "openalex_api_key", None) or settings.openalex_api_key except Exception: diff --git a/app/settings.py b/app/settings.py index c9098a1..40596d1 100644 --- a/app/settings.py +++ b/app/settings.py @@ -275,6 +275,8 @@ class Settings: crossref_max_lookups_per_request: int = _env_int("CROSSREF_MAX_LOOKUPS_PER_REQUEST", 8) openalex_api_key: str | None = os.getenv("OPENALEX_API_KEY") + database_reserved_api_connections: int = _env_int("DATABASE_RESERVED_API_CONNECTIONS", 3) + crossref_api_token: str | None = os.getenv("CROSSREF_API_TOKEN") crossref_api_mailto: str | None = os.getenv("CROSSREF_API_MAILTO") diff --git a/frontend/src/features/publications/composables/usePublicationData.ts b/frontend/src/features/publications/composables/usePublicationData.ts index 1ce8d11..ca4ddbb 100644 --- a/frontend/src/features/publications/composables/usePublicationData.ts +++ b/frontend/src/features/publications/composables/usePublicationData.ts @@ -277,12 +277,16 @@ export function usePublicationData() { // --- Run-triggered refresh watcher --- + let previousRunStatusKey: string | null = null; + watch( () => runStatus.latestRun, - async (nextRun, previousRun) => { - const nextRunId = nextRun && (nextRun.status === "running" || nextRun.status === "resolving") ? nextRun.id : null; - const previousRunId = previousRun && (previousRun.status === "running" || previousRun.status === "resolving") ? previousRun.id : null; - if (nextRunId === null || nextRunId === previousRunId) return; + async (nextRun) => { + const nextStatus = nextRun ? `${nextRun.id}:${nextRun.status}` : null; + if (nextStatus === previousRunStatusKey) return; + previousRunStatusKey = nextStatus; + const isActive = nextRun && (nextRun.status === "running" || nextRun.status === "resolving"); + if (!isActive) return; resetPageAndSnapshot(); await loadPublications(); }, diff --git a/frontend/src/pages/PublicationsPage.vue b/frontend/src/pages/PublicationsPage.vue index 9dcfec7..af05fa0 100644 --- a/frontend/src/pages/PublicationsPage.vue +++ b/frontend/src/pages/PublicationsPage.vue @@ -1,5 +1,5 @@