From 3d0f96d339a920a265387a81d57b66cf8ad6ab1f Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Wed, 8 Oct 2025 12:44:33 -0400 Subject: [PATCH] Add QueueItem cleanup logic (#228) * feat: start queue cleanup * feat: add Queue cleanup logic * feat: mark messages as errored * fix: rename + commit * fix: rm apscheduler * fix: add default for created_at in migration * fix: consolidate error handling and only throw away first message of an errored batch * fix: move processing to within queue manager * fix: regen migration * fix: CR comment * fix: rename message -> queue_item when dealing with queue_item in the queue manager --- ...ffba56fe8c_add_error_field_to_queueitem.py | 94 ++++++++ src/config.py | 5 + src/deriver/queue_manager.py | 223 ++++++++++++++---- src/main.py | 1 - src/models.py | 4 + tests/deriver/test_queue_processing.py | 16 +- 6 files changed, 290 insertions(+), 53 deletions(-) create mode 100644 migrations/versions/76ffba56fe8c_add_error_field_to_queueitem.py diff --git a/migrations/versions/76ffba56fe8c_add_error_field_to_queueitem.py b/migrations/versions/76ffba56fe8c_add_error_field_to_queueitem.py new file mode 100644 index 00000000..de0ebe80 --- /dev/null +++ b/migrations/versions/76ffba56fe8c_add_error_field_to_queueitem.py @@ -0,0 +1,94 @@ +"""add error field to QueueItem + +Revision ID: 76ffba56fe8c +Revises: 08894082221a +Create Date: 2025-10-08 11:47:17.488301 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from migrations.utils import column_exists, get_schema, index_exists + +# revision identifiers, used by Alembic. +revision: str = "76ffba56fe8c" +down_revision: str | None = "08894082221a" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +schema = get_schema() + + +def upgrade() -> None: + # Add error column + op.add_column( + "queue", + sa.Column("error", sa.TEXT(), nullable=True), + schema=schema, + ) + + op.add_column( + "queue", + sa.Column( + "created_at", + sa.DateTime(timezone=True), + nullable=True, + server_default=sa.func.now(), + ), + schema=schema, + ) + + # Backfill created_at for existing queue items (batched) + op.execute( + sa.text( + f""" + DO $$ + DECLARE + rows_updated INT; + BEGIN + LOOP + UPDATE "{schema}".queue + SET created_at = NOW() + WHERE id IN ( + SELECT id FROM "{schema}".queue + WHERE created_at IS NULL + LIMIT 1000 + ); + + GET DIAGNOSTICS rows_updated = ROW_COUNT; + EXIT WHEN rows_updated = 0; + END LOOP; + END $$; + """ + ) + ) + + # Make created_at non-nullable + op.alter_column("queue", "created_at", nullable=False, schema=schema) + + # Add index on created_at + op.create_index( + op.f("ix_queue_created_at"), + "queue", + ["created_at"], + unique=False, + schema=schema, + ) + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + + if column_exists("queue", "error", inspector): + op.drop_column("queue", "error", schema=schema) + + # Drop index first + if index_exists("queue", "ix_queue_created_at", inspector): + op.drop_index(op.f("ix_queue_created_at"), table_name="queue", schema=schema) + + if column_exists("queue", "created_at", inspector): + op.drop_column("queue", "created_at", schema=schema) diff --git a/src/config.py b/src/config.py index 170bccd5..d50ce340 100644 --- a/src/config.py +++ b/src/config.py @@ -191,6 +191,11 @@ class DeriverSettings(HonchoSettings): ] = 1.0 STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5 + # Retention window (seconds) for keeping errored items in the queue + QUEUE_ERROR_RETENTION_SECONDS: Annotated[ + int, Field(default=30 * 24 * 3600, gt=0) + ] = 30 * 24 * 3600 # 30 days default + PROVIDER: SupportedProviders = "google" MODEL: str = "gemini-2.5-flash-lite" diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index b5706a76..02779b47 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -52,6 +52,7 @@ class QueueManager: self.active_tasks: set[asyncio.Task[None]] = set() self.worker_ownership: dict[str, WorkerOwnership] = {} self.queue_empty_flag: asyncio.Event = asyncio.Event() + self._maintenance_task: asyncio.Task[None] | None = None # Initialize from settings self.workers: int = settings.DERIVER.WORKERS @@ -115,6 +116,12 @@ class QueueManager: ) logger.debug("Signal handlers registered") + # Start background maintenance loop + try: + self._maintenance_task = asyncio.create_task(self._maintenance_loop()) + except Exception: + logger.exception("Failed to start maintenance loop") + # Run the polling loop directly in this task logger.debug("Starting polling loop directly") try: @@ -162,6 +169,14 @@ class QueueManager: finally: self.worker_ownership.clear() + # Cancel maintenance loop if running + if self._maintenance_task is not None: + from contextlib import suppress + + self._maintenance_task.cancel() + with suppress(asyncio.CancelledError): + await self._maintenance_task + ########################## # Polling and Scheduling # ########################## @@ -308,13 +323,99 @@ class QueueManager: ###################### # Queue Worker Logic # ###################### + + async def cleanup_queue_items(self) -> None: + """Delete processed queue items. + Successfully processed queue items are deleted immediately, + while errored queue items are deleted after retention window.""" + async with tracked_db("cleanup_queue_items") as db: + now = datetime.now(timezone.utc) + error_cutoff = now - timedelta( + seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS + ) + + await db.execute( + delete(models.QueueItem).where( + models.QueueItem.processed + & ( + models.QueueItem.error.is_(None) + | ( + models.QueueItem.error.is_not(None) + & (models.QueueItem.created_at < error_cutoff) + ) + ) + ) + ) + await db.commit() + + async def _maintenance_loop(self) -> None: + """Run periodic maintenance tasks on the queue.""" + try: + while not self.shutdown_event.is_set(): + try: + await self.cleanup_queue_items() + except Exception: + logger.exception("Error during maintenance cleanup") + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception() + + # Sleep until interval elapses or shutdown event is set + try: + await asyncio.wait_for( + self.shutdown_event.wait(), + timeout=43200, # 12 hours + ) + break # Shutdown event set + except asyncio.TimeoutError: + # Timeout means it's time for next cleanup + pass + except asyncio.CancelledError: + logger.debug("Maintenance loop cancelled") + raise + + async def _handle_processing_error( + self, + error: Exception, + items: list[QueueItem], + work_unit_key: str, + context: str, + ) -> None: + """ + Handle processing errors by marking queue items as errored, logging, and forwarding to Sentry. + We only mark the first queue item as errored so we don't potentially throw away a batch. This allows us + to incrementally attempt to process the batch while still maintaining progress in a work unit. + + Args: + error: The exception that occurred + items: The queue items that were being processed + work_unit_key: The work unit key for the queue items + context: Context string describing what was being processed (e.g., "processing representation batch") + """ + error_msg = f"{error.__class__.__name__}: {str(error)}" + try: + if items: + await self.mark_queue_item_as_errored( + items[0], work_unit_key, error_msg + ) + except Exception as mark_error: + logger.error( + f"Failed to mark queue items as errored for work unit {work_unit_key}: {mark_error}", + exc_info=True, + ) + + logger.error( + f"Error {context} for work unit {work_unit_key}: {error}", + exc_info=True, + ) + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(error) + async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None: - """Process all messages for a specific work unit by routing to the correct handler.""" + """Process all queue items for a specific work unit by routing to the correct handler.""" logger.debug(f"Starting to process work unit {work_unit_key}") work_unit = parse_work_unit_key(work_unit_key) async with self.semaphore: - message_count = 0 - messages_to_process: list[QueueItem] = [] + queue_item_count = 0 try: while not self.shutdown_event.is_set(): # Get worker ownership info for verification @@ -329,7 +430,7 @@ class QueueManager: ( messages_context, items_to_process, - ) = await self.get_message_batch( + ) = await self.get_queue_item_batch( work_unit.task_type, work_unit_key, ownership.aqs_id ) logger.debug( @@ -341,38 +442,53 @@ class QueueManager: ) break - # Build payloads from the unique messages context window - await process_representation_batch( - messages_context, - observer=work_unit.observer, - observed=work_unit.observed, - ) - - await self.mark_messages_as_processed( - items_to_process, work_unit_key - ) - message_count += len(items_to_process) + try: + await process_representation_batch( + messages_context, + observer=work_unit.observer, + observed=work_unit.observed, + ) + await self.mark_queue_items_as_processed( + items_to_process, work_unit_key + ) + queue_item_count += len(items_to_process) + except Exception as e: + await self._handle_processing_error( + e, + items_to_process, + work_unit_key, + "processing representation batch", + ) else: - messages_to_process = await self.get_next_message( + queue_item = await self.get_next_queue_item( work_unit.task_type, work_unit_key, ownership.aqs_id ) - if not messages_to_process: + if not queue_item: logger.debug( - f"No more messages to process for work unit {work_unit_key} for worker {worker_id}" + f"No more queue items to process for work unit {work_unit_key} for worker {worker_id}" ) break - await process_item( - work_unit.task_type, messages_to_process[0].payload - ) - await self.mark_messages_as_processed( - messages_to_process, work_unit_key - ) - message_count += len(messages_to_process) + + try: + await process_item( + work_unit.task_type, queue_item.payload + ) + await self.mark_queue_items_as_processed( + [queue_item], work_unit_key + ) + queue_item_count += 1 + except Exception as e: + await self._handle_processing_error( + e, + [queue_item], + work_unit_key, + "processing queue item", + ) except Exception as e: logger.error( - f"Error processing tasks for work unit {work_unit_key}: {e}", + f"Error in processing loop for work unit {work_unit_key}: {e}", exc_info=True, ) if settings.SENTRY.ENABLED: @@ -397,7 +513,7 @@ class QueueManager: removed = False self.untrack_worker_work_unit(worker_id, work_unit_key) - if removed and message_count > 0: + if removed and queue_item_count > 0: # Only publish webhook if we actually removed an active session try: if work_unit.task_type in ["representation", "summary"]: @@ -425,22 +541,21 @@ class QueueManager: ) @sentry_sdk.trace - async def get_next_message( + async def get_next_queue_item( self, task_type: str, work_unit_key: str, aqs_id: str - ) -> list[QueueItem]: - """Get the next message to process for a specific work unit.""" + ) -> QueueItem | None: + """Get the next queue item to process for a specific work unit.""" if task_type == "representation": raise ValueError( - "Representation tasks are not supported for get_next_message" + "Representation tasks are not supported for get_next_queue_item" ) - async with tracked_db("get_next_message") as db: + async with tracked_db("get_next_queue_item") as db: # ActiveQueueSession conditions for worker ownership verification aqs_conditions = [ models.ActiveQueueSession.work_unit_key == work_unit_key, models.ActiveQueueSession.id == aqs_id, ] - # For non-representation tasks, just get the next single message. query = ( select(models.QueueItem) .join( @@ -455,15 +570,15 @@ class QueueManager: .limit(1) ) result = await db.execute(query) - messages = result.scalars().all() + queue_item = 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 list(messages) + return queue_item @sentry_sdk.trace - async def get_message_batch( + async def get_queue_item_batch( self, task_type: str, work_unit_key: str, @@ -476,9 +591,9 @@ class QueueManager: """ if task_type != "representation": raise ValueError( - "Non-representation tasks are not supported for get_message_batch" + "Non-representation tasks are not supported for get_queue_item_batch" ) - async with tracked_db("get_message_batch") as db: + async with tracked_db("get_queue_item_batch") as db: # For representation tasks, get a batch based on token limit. # Step 1: Parse work_unit_key to get session context and focused sender parsed_key = parse_work_unit_key(work_unit_key) @@ -592,16 +707,16 @@ class QueueManager: return messages_context, items_to_process - async def mark_messages_as_processed( - self, messages: list[QueueItem], work_unit_key: str + async def mark_queue_items_as_processed( + self, items: list[QueueItem], work_unit_key: str ) -> None: - if not messages: + if not items: return - async with tracked_db("process_message_batch") as db: - message_ids = [msg.id for msg in messages] + async with tracked_db("process_queue_item_batch") as db: + item_ids = [item.id for item in items] await db.execute( update(models.QueueItem) - .where(models.QueueItem.id.in_(message_ids)) + .where(models.QueueItem.id.in_(item_ids)) .where(models.QueueItem.work_unit_key == work_unit_key) .values(processed=True) ) @@ -612,6 +727,26 @@ class QueueManager: ) await db.commit() + async def mark_queue_item_as_errored( + self, item: QueueItem, work_unit_key: str, error: str + ) -> None: + """Mark queue item as processed with an error""" + if not item: + return + async with tracked_db("mark_queue_item_as_errored") as db: + await db.execute( + update(models.QueueItem) + .where(models.QueueItem.id == item.id) + .where(models.QueueItem.work_unit_key == work_unit_key) + .values(processed=True, error=error[:65535]) # Truncate to TEXT limit + ) + await db.execute( + update(models.ActiveQueueSession) + .where(models.ActiveQueueSession.work_unit_key == work_unit_key) + .values(last_updated=func.now()) + ) + await db.commit() + async def _cleanup_work_unit( self, aqs_id: str, diff --git a/src/main.py b/src/main.py index a3e1594e..a881ca21 100644 --- a/src/main.py +++ b/src/main.py @@ -102,7 +102,6 @@ if SENTRY_ENABLED: @asynccontextmanager async def lifespan(_: FastAPI): - # Lifespan events are now handled by the respective services yield await engine.dispose() diff --git a/src/models.py b/src/models.py index a50df97c..9b4e6f1d 100644 --- a/src/models.py +++ b/src/models.py @@ -396,6 +396,10 @@ class QueueItem(Base): task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False) payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False) processed: Mapped[bool] = mapped_column(Boolean, default=False) + error: Mapped[str | None] = mapped_column(TEXT, nullable=True) + created_at: Mapped[datetime.datetime] = mapped_column( + DateTime(timezone=True), index=True, default=func.now() + ) def __repr__(self) -> str: return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed})" diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 3b631a3f..6f1527ad 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -172,7 +172,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(aqs) - _, items_to_process = await qm.get_message_batch( + _, items_to_process = await qm.get_queue_item_batch( task_type="representation", work_unit_key=first.work_unit_key, aqs_id=aqs.id, @@ -183,7 +183,7 @@ class TestQueueProcessing: # Mark first processed, next should be the second first.processed = True await db_session.commit() - _, items_to_process2 = await qm.get_message_batch( + _, items_to_process2 = await qm.get_queue_item_batch( task_type="representation", work_unit_key=first.work_unit_key, aqs_id=aqs.id, @@ -461,7 +461,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(alice_aqs) - alice_messages, alice_items = await qm.get_message_batch( + alice_messages, alice_items = await qm.get_queue_item_batch( task_type="representation", work_unit_key=alice_work_unit_key, aqs_id=alice_aqs.id, @@ -489,7 +489,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(bob_aqs) - bob_messages, bob_items = await qm.get_message_batch( + bob_messages, bob_items = await qm.get_queue_item_batch( task_type="representation", work_unit_key=bob_work_unit_key, aqs_id=bob_aqs.id, @@ -515,7 +515,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(steve_aqs) - steve_messages, steve_items = await qm.get_message_batch( + steve_messages, steve_items = await qm.get_queue_item_batch( task_type="representation", work_unit_key=steve_work_unit_key, aqs_id=steve_aqs.id, @@ -628,7 +628,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(alice_aqs) - alice_messages2, _ = await qm.get_message_batch( + alice_messages2, _ = await qm.get_queue_item_batch( task_type="representation", work_unit_key=alice_work_unit_key, aqs_id=alice_aqs.id, @@ -651,7 +651,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(bob_aqs) - bob_messages2, _ = await qm.get_message_batch( + bob_messages2, _ = await qm.get_queue_item_batch( task_type="representation", work_unit_key=bob_work_unit_key, aqs_id=bob_aqs.id, @@ -670,7 +670,7 @@ class TestQueueProcessing: await db_session.commit() await db_session.refresh(steve_aqs) - steve_messages2, _ = await qm.get_message_batch( + steve_messages2, _ = await qm.get_queue_item_batch( task_type="representation", work_unit_key=steve_work_unit_key, aqs_id=steve_aqs.id,