RAGcore-Public/tests/integration/indexes/test_concurrent_operations.py

"""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),
        )