"""Service-level concurrency: racing promotions, promote-vs-rollback, timeouts.""" from __future__ import annotations import asyncio import os import time from collections.abc import AsyncIterator, Iterable from datetime import UTC, datetime, timedelta from uuid import uuid4 import pytest from qdrant_client import AsyncQdrantClient from sqlalchemy.ext.asyncio import AsyncEngine from ragcore.application.indexes import ( IndexBuildError, IndexBuildService, IndexLockTimeoutError, IndexOperationStatus, ) from ragcore.config import Settings from ragcore.infrastructure.database import create_database_engine, create_session_factory from ragcore.infrastructure.database.indexes import PostgresIndexOperationStore from ragcore.infrastructure.qdrant import ( CollectionKind, CollectionProfile, IndexPayload, IndexPoint, ) from ragcore.infrastructure.qdrant.rebuild import QdrantRebuildStore pytestmark = pytest.mark.integration # Real build latency in this environment (collection create + 8 payload indexes + # per-point upsert/retrieve/set_payload round trips) can run into the tens of # seconds, so timeouts here are generous. What is asserted is boundedness and # correctness of the outcome, not a tight wall-clock bound. _GENEROUS_TIMEOUT = timedelta(seconds=60) class FixtureSource: def __init__(self, points: Iterable[IndexPoint]) -> None: self.points = tuple(points) async def iter_points(self, profile: CollectionProfile) -> AsyncIterator[IndexPoint]: for point in self.points: yield point class PassingSmoke: async def run(self, profile: CollectionProfile) -> None: return None def make_profile() -> CollectionProfile: # A fresh profile name (not just version) per call keeps each test's alias # namespace isolated, so leftover/slow state from one test can never block # another test's lock acquisition. return CollectionProfile( kind=CollectionKind.DENSE, profile=f"concurrency-{uuid4().hex[:8]}", version=uuid4().hex[:12], dimensions=4, model_digest="a" * 64, ) def make_point(profile: CollectionProfile, index: int) -> IndexPoint: point_id = uuid4() return IndexPoint( point_id=point_id, dense_vector=(0.1 + index, 0.2, 0.3, 0.4), sparse_indices=(index + 1,), sparse_values=(0.8,), payload=IndexPayload( workspace_id=uuid4(), space_id=uuid4(), document_id=uuid4(), document_version_id=uuid4(), chunk_id=point_id, parent_chunk_id=None, source_id=uuid4(), status="published", content_type="text/markdown", language="en", trust_level="source", location={"page": index + 1}, parser_version="text-v1", chunk_profile_version=profile.version, embedding_model_digest=profile.model_digest, text=f"A safe indexed passage {index}.", excerpt=f"A safe indexed passage {index}.", ), ) @pytest.fixture async def qdrant() -> AsyncIterator[AsyncQdrantClient]: client = AsyncQdrantClient( url=os.getenv("RAGCORE_QDRANT_URL", "http://qdrant:6333"), api_key=os.getenv("RAGCORE_QDRANT_API_KEY", "replace-me"), timeout=60, ) try: yield client finally: await client.close() @pytest.fixture async def db_engine() -> AsyncIterator[AsyncEngine]: engine = create_database_engine(Settings(env="test")) try: yield engine finally: await engine.dispose() @pytest.fixture def operations(db_engine: AsyncEngine) -> PostgresIndexOperationStore: return PostgresIndexOperationStore(create_session_factory(db_engine)) async def test_two_concurrent_promotions_only_one_activates( qdrant: AsyncQdrantClient, operations: PostgresIndexOperationStore ) -> None: profile = make_profile() points = (make_point(profile, 0), make_point(profile, 1)) store = QdrantRebuildStore(qdrant) service = IndexBuildService(store, operations, lock_acquire_timeout=_GENEROUS_TIMEOUT) try: results = await asyncio.gather( service.build(profile, FixtureSource(points), PassingSmoke()), service.build(profile, FixtureSource(points), PassingSmoke()), return_exceptions=True, ) succeeded = [result for result in results if not isinstance(result, BaseException)] rejected = [result for result in results if isinstance(result, BaseException)] assert len(succeeded) == 1 assert len(rejected) == 1 # Whichever caller loses the race is safely turned away: either it waits # its turn and then sees the alias already active, or it times out first. # Both prove mutual exclusion held; neither corrupts state. assert isinstance(rejected[0], (IndexBuildError, IndexLockTimeoutError)) assert await store.current_alias_target(profile.alias_name) == profile.collection_name finally: await store.rollback(profile.alias_name, None) await qdrant.delete_collection(profile.collection_name) async def test_promotion_and_rollback_starting_together_are_serialized_not_interleaved( qdrant: AsyncQdrantClient, operations: PostgresIndexOperationStore ) -> None: """A rollback racing a promotion on the same alias must wait, not interleave. The lock is held directly through the primitive for a short, fully controlled window rather than by racing a real build's variable duration, so the outcome is deterministic without depending on Qdrant/container timing. """ profile = make_profile() store = QdrantRebuildStore(qdrant) service = IndexBuildService(store, operations, lock_acquire_timeout=_GENEROUS_TIMEOUT) try: report = await service.build( profile, FixtureSource((make_point(profile, 0),)), PassingSmoke() ) holder_token = uuid4().hex hold_started = datetime.now(UTC) held_operation_id = await operations.begin( alias_name=profile.alias_name, operation="reconcile", collection_name=profile.collection_name, owner_token=holder_token, now=hold_started, lease_duration=timedelta(seconds=30), ) assert held_operation_id is not None async def release_after_delay() -> None: await asyncio.sleep(1.0) await operations.finish( held_operation_id, holder_token, IndexOperationStatus.SUCCEEDED, now=datetime.now(UTC), ) release_task = asyncio.create_task(release_after_delay()) try: started = time.monotonic() await service.rollback(report) elapsed = time.monotonic() - started finally: await release_task # always await: never leak a background task # The rollback could only have completed after the manual hold released. assert elapsed >= 0.9 assert await store.current_alias_target(profile.alias_name) is None finally: await store.rollback(profile.alias_name, None) await qdrant.delete_collection(profile.collection_name) async def test_lock_acquire_timeout_is_bounded_and_raises( qdrant: AsyncQdrantClient, operations: PostgresIndexOperationStore ) -> None: profile = make_profile() store = QdrantRebuildStore(qdrant) holder_token = uuid4().hex operation_id = await operations.begin( alias_name=profile.alias_name, operation="build", collection_name=profile.collection_name, owner_token=holder_token, now=datetime.now(UTC), lease_duration=timedelta(seconds=30), ) assert operation_id is not None service = IndexBuildService( store, operations, lock_acquire_timeout=timedelta(seconds=1), lock_poll_interval_seconds=0.1, ) try: started = time.monotonic() with pytest.raises(IndexLockTimeoutError): await service.build( profile, FixtureSource((make_point(profile, 0),)), PassingSmoke() ) elapsed = time.monotonic() - started # Bounded, not an indefinite hang -- not a tight bound on the nominal # 1s `lock_acquire_timeout` itself. Each of the ~10 poll iterations # this loop takes still does a real Postgres round trip # (`operations.begin()`), and this suite's own real build latency in # this environment can run into the tens of seconds (see # `_GENEROUS_TIMEOUT` above), so a narrow millisecond-scale margin # here flakes under real system load without the timeout logic # itself having failed at all. assert elapsed < _GENEROUS_TIMEOUT.total_seconds() finally: await operations.finish( operation_id, holder_token, IndexOperationStatus.SUCCEEDED, now=datetime.now(UTC), )