diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 9f712b2d..802963d2 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -8,8 +8,7 @@ from logging import getLogger import sentry_sdk from dotenv import load_dotenv from sentry_sdk.integrations.asyncio import AsyncioIntegration -from sqlalchemy import delete, insert, select, update -from sqlalchemy.exc import IntegrityError +from sqlalchemy import delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func @@ -117,37 +116,73 @@ class QueueManager: # Polling and Scheduling # ########################## - async def get_available_work_units(self, db: AsyncSession) -> Sequence[str]: + async def get_and_claim_work_units(self) -> Sequence[str]: """ Get available work units that aren't being processed. Returns a list of work unit keys. """ - # Clean up stale work units - five_minutes_ago = datetime.now(timezone.utc) - timedelta( - minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES - ) - await db.execute( - delete(models.ActiveQueueSession).where( - models.ActiveQueueSession.last_updated < five_minutes_ago + claimed_units: list[str] = [] + + async with tracked_db("get_available_work_units") as db: + # Clean up stale work units + five_minutes_ago = datetime.now(timezone.utc) - timedelta( + minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES ) + await db.execute( + delete(models.ActiveQueueSession).where( + models.ActiveQueueSession.last_updated < five_minutes_ago + ) + ) + + # Get number of available workers + limit: int = max(0, self.workers - len(self.owned_work_units)) + + query = ( + select(models.QueueItem.work_unit_key) + .outerjoin( + models.ActiveQueueSession, + models.QueueItem.work_unit_key + == models.ActiveQueueSession.work_unit_key, + ) + .where(~models.QueueItem.processed) + .where(models.QueueItem.work_unit_key.isnot(None)) + .where(models.ActiveQueueSession.work_unit_key.is_(None)) + .distinct() + .limit(limit) + ) + + result = await db.execute(query) + available_units = result.scalars().all() + if not available_units: + await db.commit() + return [] + + claimed_units = await self.claim_work_units(db, available_units) + await db.commit() + + for work_unit in claimed_units: + self.track_work_unit(work_unit) + + return claimed_units + + async def claim_work_units( + self, db: AsyncSession, work_unit_keys: Sequence[str] + ) -> list[str]: + from sqlalchemy.dialects.postgresql import insert + + values = [{"work_unit_key": key} for key in work_unit_keys] + + stmt = ( + insert(models.ActiveQueueSession) + .values(values) + .on_conflict_do_nothing() + .returning(models.ActiveQueueSession.work_unit_key) ) - query = ( - select(models.QueueItem.work_unit_key) - .outerjoin( - models.ActiveQueueSession, - models.QueueItem.work_unit_key - == models.ActiveQueueSession.work_unit_key, - ) - .where(~models.QueueItem.processed) - .where(models.QueueItem.work_unit_key.isnot(None)) - .where(models.ActiveQueueSession.work_unit_key.is_(None)) - .distinct() - .limit(self.workers) - ) - - result = await db.execute(query) - return result.scalars().all() + result = await db.execute(stmt) + claimed_units = result.scalars().all() + logger.debug(f"Claimed {len(claimed_units)} work units") + return list(claimed_units) async def polling_loop(self) -> None: """Main polling loop to find and process new work units""" @@ -166,55 +201,27 @@ class QueueManager: await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) continue - # Use the dependency for transaction safety - async with tracked_db("queue_polling_loop") as db: - try: - new_work_units = await self.get_available_work_units(db) - - if new_work_units and not self.shutdown_event.is_set(): - for work_unit in new_work_units: - try: - # Try to claim the work unit using work_unit_key - await db.execute( - insert(models.ActiveQueueSession).values( - work_unit_key=work_unit - ) - ) - await db.commit() - - # Track this work unit - self.track_work_unit(work_unit) - logger.debug( - f"Claimed work unit {work_unit} for processing" - ) - - # Create a new task for processing this work unit - if not self.shutdown_event.is_set(): - task: Task[None] = asyncio.create_task( - self.process_work_unit(work_unit) - ) - self.add_task(task) - - except IntegrityError: - # Rollback the failed transaction to clear the error state - await db.rollback() - logger.debug( - f"Failed to claim work unit {work_unit}, already owned by another worker" - ) - # If we couldn't claim any work units, avoid tight loop - else: - self.queue_empty_flag.set() - await asyncio.sleep( - settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS - ) - except Exception as e: - logger.exception("Error in polling loop") - if settings.SENTRY.ENABLED: - sentry_sdk.capture_exception(e) - # Note: rollback is handled by tracked_db dependency + try: + new_work_units = await self.get_and_claim_work_units() + if new_work_units: + for work_unit in new_work_units: + # Create a new task for processing this work unit + if not self.shutdown_event.is_set(): + task: Task[None] = asyncio.create_task( + self.process_work_unit(work_unit) + ) + self.add_task(task) + else: + self.queue_empty_flag.set() await asyncio.sleep( settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS ) + except Exception as e: + logger.exception("Error in polling loop") + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(e) + # Note: rollback is handled by tracked_db dependency + await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) finally: logger.info("Polling loop stopped") @@ -227,13 +234,12 @@ class QueueManager: """Process all messages for a specific work unit by routing to the correct handler.""" logger.debug(f"Starting to process work unit {work_unit_key}") async with ( - self.semaphore, - tracked_db("queue_process_work_unit") as db, + self.semaphore ): # Hold the semaphore for the entire work unit duration message_count = 0 try: while not self.shutdown_event.is_set(): - message = await self.get_next_message(db, work_unit_key) + message = await self.get_next_message(work_unit_key) if not message: logger.debug(f"No more messages for work unit {work_unit_key}") break @@ -263,15 +269,22 @@ class QueueManager: sentry_sdk.capture_exception(e) # Prevent malformed messages from stalling queue indefinitely - message.processed = True + async with tracked_db("process_message") as db: + await db.execute( + update(models.QueueItem) + .where(models.QueueItem.id == message.id) + .values(processed=True) + ) - # Update last_updated timestamp to show this work unit is still being processed - await db.execute( - update(models.ActiveQueueSession) - .where(models.ActiveQueueSession.work_unit_key == work_unit_key) - .values(last_updated=func.now()) - ) - await db.commit() + await db.execute( + update(models.ActiveQueueSession) + .where( + models.ActiveQueueSession.work_unit_key == work_unit_key + ) + .values(last_updated=func.now()) + ) + + await db.commit() if self.shutdown_event.is_set(): logger.debug( @@ -286,15 +299,10 @@ class QueueManager: finally: # Remove work unit from active_queue_sessions when done logger.debug(f"Removing work unit {work_unit_key} from active sessions") - delete_result = await db.execute( - delete(models.ActiveQueueSession).where( - models.ActiveQueueSession.work_unit_key == work_unit_key - ) - ) - await db.commit() + removed = await self._cleanup_work_unit(work_unit_key) - # Only publish webhook if we actually removed an active session - if delete_result.rowcount > 0 and message_count > 0: + if removed and message_count > 0: + # Only publish webhook if we actually removed an active session try: from src.deriver.utils import parse_work_unit_key from src.webhooks.events import ( @@ -330,20 +338,32 @@ class QueueManager: self.untrack_work_unit(work_unit_key) @sentry_sdk.trace - async def get_next_message( - self, db: AsyncSession, work_unit_key: str - ) -> QueueItem | None: + async def get_next_message(self, work_unit_key: str) -> QueueItem | None: """Get the next unprocessed message for a specific work unit.""" + async with tracked_db("get_next_message") as db: + query = ( + select(models.QueueItem) + .where(models.QueueItem.work_unit_key == work_unit_key) + .where(~models.QueueItem.processed) + .order_by(models.QueueItem.id) + .limit(1) + ) + result = await db.execute(query) + message = result.scalar_one_or_none() + # Important: commit to avoid tracked_db's rollback expiring the instance + # We rely on expire_on_commit=False to keep attributes accessible post-close + await db.commit() + return message - query = ( - select(models.QueueItem) - .where(models.QueueItem.work_unit_key == work_unit_key) - .where(~models.QueueItem.processed) - .order_by(models.QueueItem.id) - .limit(1) - ) - result = await db.execute(query) - return result.scalar_one_or_none() + async def _cleanup_work_unit(self, work_unit_key: str) -> bool: + async with tracked_db("cleanup_work_unit") as db: + result = await db.execute( + delete(models.ActiveQueueSession).where( + models.ActiveQueueSession.work_unit_key == work_unit_key + ) + ) + await db.commit() + return result.rowcount > 0 async def main(): diff --git a/tests/conftest.py b/tests/conftest.py index 78a5213d..0a00dd78 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -512,7 +512,10 @@ def mock_tracked_db(db_session: AsyncSession): async def mock_tracked_db_context(_: str | None = None): yield db_session - with patch("src.dependencies.tracked_db", mock_tracked_db_context): + with ( + patch("src.dependencies.tracked_db", mock_tracked_db_context), + patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context), + ): yield diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 8f12f5cd..e38571c0 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -1,3 +1,6 @@ +from collections.abc import Callable +from typing import Any + import pytest from sqlalchemy.ext.asyncio import AsyncSession @@ -9,13 +12,12 @@ from src.deriver.queue_manager import QueueManager class TestQueueProcessing: """Test suite for queue processing functionality""" - async def test_get_available_work_units( + async def test_get_and_claim_work_units( self, - db_session: AsyncSession, sample_queue_items: list[models.QueueItem], sample_session_with_peers: tuple[models.Session, list[models.Peer]], - ): - """Test that get_available_work_units correctly identifies unprocessed work""" + ) -> None: + """Test that get_and_claim_work_units correctly identifies unprocessed work""" session, _peers = sample_session_with_peers # pyright: ignore[reportUnusedVariable] # Verify we have queue items from our test setup @@ -25,7 +27,7 @@ class TestQueueProcessing: queue_manager = QueueManager() # Get available work units - work_units = await queue_manager.get_available_work_units(db_session) + work_units = await queue_manager.get_and_claim_work_units() # Should have some work units available (may include items from other tests) assert len(work_units) > 0 @@ -35,7 +37,7 @@ class TestQueueProcessing: assert isinstance(work_unit, str) assert work_unit.split(":")[0] in ["representation", "summary"] - # The test is mainly verifying that get_available_work_units works without errors + # The test is mainly verifying that get_and_claim_work_units works without errors # and returns properly structured work unit key strings async def test_work_unit_claiming( @@ -43,7 +45,7 @@ class TestQueueProcessing: db_session: AsyncSession, sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter] sample_session_with_peers: tuple[models.Session, list[models.Peer]], - ): + ) -> None: """Test that work units can be claimed and are not available to other workers""" _session, _peers = sample_session_with_peers @@ -51,29 +53,141 @@ class TestQueueProcessing: queue_manager = QueueManager() # Get available work units - work_units = await queue_manager.get_available_work_units(db_session) + work_units = await queue_manager.get_and_claim_work_units() assert len(work_units) > 0 - # Claim a work unit by creating an ActiveQueueSession entry + # The API already claimed returned units; verify it's tracked and not returned again + from sqlalchemy import select + work_unit = work_units[0] - active_session = models.ActiveQueueSession( - work_unit_key=work_unit, - ) - db_session.add(active_session) - await db_session.commit() + tracked = ( + await db_session.execute( + select(models.ActiveQueueSession).where( + models.ActiveQueueSession.work_unit_key == work_unit + ) + ) + ).scalar_one_or_none() + assert tracked is not None # Get available work units again - the claimed one should not be available - remaining_work_units = await queue_manager.get_available_work_units(db_session) + remaining_work_units = await queue_manager.get_and_claim_work_units() # The claimed work unit should not be in the remaining list assert work_unit not in remaining_work_units - # We can't assert exact count difference because work units are grouped by unique combinations - async def test_stale_work_unit_cleanup( + @pytest.mark.asyncio + async def test_get_and_claim_excludes_already_claimed( + self, + sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter] + ) -> None: + queue_manager = QueueManager() + first_batch = await queue_manager.get_and_claim_work_units() + assert len(first_batch) > 0 + + # Call again; previously claimed keys should not appear + second_batch = await queue_manager.get_and_claim_work_units() + assert all(k not in second_batch for k in first_batch) + + @pytest.mark.asyncio + async def test_claim_work_unit_conflict_returns_false( + self, + db_session: AsyncSession, + sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter] + ) -> None: + # Pre-create an active session for a key + queue_manager = QueueManager() + claimed = await queue_manager.get_and_claim_work_units() + assert len(claimed) > 0 + key = claimed[0] + + # Trying to claim the same key again via the API should return empty list + claimed_again = await queue_manager.claim_work_units(db_session, [key]) + assert claimed_again == [] + + @pytest.mark.asyncio + async def test_get_next_message_orders_and_filters_simple( self, db_session: AsyncSession, sample_session_with_peers: tuple[models.Session, list[models.Peer]], - ): + create_queue_payload: Callable[..., Any], + add_queue_items: Callable[..., Any], + ) -> None: + from sqlalchemy import select + + session, peers = sample_session_with_peers + peer = peers[0] + + payloads: list[Any] = [] + for i in range(3): + payloads.append( + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=models.Message( + id=i, + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content="hello", + ), # include id for payload builder + task_type="representation", + sender_name=peer.name, + target_name=peer.name, + ) + ) + + items = await add_queue_items(payloads, session.id) + # Determine ascending order by DB id + ordered = ( + ( + await db_session.execute( + select(models.QueueItem) + .where(models.QueueItem.work_unit_key == items[0].work_unit_key) + .order_by(models.QueueItem.id) + ) + ) + .scalars() + .all() + ) + first, second = ordered[0], ordered[1] + + qm = QueueManager() + nxt = await qm.get_next_message(first.work_unit_key) + assert nxt is not None and nxt.id == first.id + + # Mark first processed, next should be the second + first.processed = True + await db_session.commit() + nxt2 = await qm.get_next_message(first.work_unit_key) + assert nxt2 is not None and nxt2.id == second.id + + @pytest.mark.asyncio + async def test_cleanup_work_unit_removes_row( + self, + sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter] + db_session: AsyncSession, + ) -> None: + from sqlalchemy import select + + qm = QueueManager() + claimed = await qm.get_and_claim_work_units() + assert len(claimed) > 0 + key = claimed[0] + + removed = await qm._cleanup_work_unit(key) # pyright: ignore[reportPrivateUsage] + assert removed is True + + remaining = ( + await db_session.execute( + select(models.ActiveQueueSession).where( + models.ActiveQueueSession.work_unit_key == key + ) + ) + ).scalar_one_or_none() + assert remaining is None + + async def test_stale_work_unit_cleanup( + self, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + ) -> None: """Test that stale work units are cleaned up properly""" _session, _peers = sample_session_with_peers @@ -82,19 +196,19 @@ class TestQueueProcessing: # datetime.now(timezone.utc) - timedelta(minutes=10) - # We'll test this by checking that the cleanup logic works in get_available_work_units + # We'll test this by checking that the cleanup logic works in get_and_claim_work_units # which is called by the queue manager during normal operation queue_manager = QueueManager() # Get available work units - this should clean up stale entries - work_units = await queue_manager.get_available_work_units(db_session) + work_units = await queue_manager.get_and_claim_work_units() # This test ensures the cleanup logic doesn't break, though we don't have stale entries yet assert isinstance(work_units, list) async def test_work_unit_key_format( self, sample_session_with_peers: tuple[models.Session, list[models.Peer]] - ): + ) -> None: """Test that work unit keys have the correct format""" session, peers = sample_session_with_peers peer1, peer2, _ = peers