fix: long-held connections in queue manager (#189)
* fix: long-held connections in queue manager * fix: CR comments. tests * fix (deriver): Limit # of work units claimed based on available workers * fix: detached instance * fix: batch claim work units --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja
GitHub
parent
7a9746ca8b
commit
8533256ca0
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue