From f4f86e38427bfca1fd788731cf3a35d566d7272e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 30 Jan 2026 11:09:29 +0100 Subject: [PATCH] fix: deadlock in worker polling (#250) * fix: deadlock in worker polling * fix: deadlock in worker polling * fixes --- hindsight-api/hindsight_api/api/http.py | 3 +- hindsight-api/hindsight_api/config.py | 14 +- hindsight-api/hindsight_api/main.py | 3 +- hindsight-api/hindsight_api/worker/main.py | 12 +- hindsight-api/hindsight_api/worker/poller.py | 292 ++++++++++++------ hindsight-api/tests/test_consolidation.py | 22 +- hindsight-api/tests/test_worker.py | 227 +++++++++++++- .../docs/developer/configuration.md | 3 +- 8 files changed, 448 insertions(+), 128 deletions(-) diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index a1b1b9c7..a1a7a566 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -1404,9 +1404,10 @@ def create_app( worker_id=worker_id, executor=memory.execute_task, poll_interval_ms=config.worker_poll_interval_ms, - batch_size=config.worker_batch_size, max_retries=config.worker_max_retries, tenant_extension=getattr(memory, "_tenant_extension", None), + max_slots=config.worker_max_slots, + consolidation_max_slots=config.worker_consolidation_max_slots, ) poller_task = asyncio.create_task(poller.run()) logging.info(f"Worker poller started (worker_id={worker_id})") diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index dec78995..764e5c82 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -143,8 +143,9 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED" ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID" ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS" ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES" -ENV_WORKER_BATCH_SIZE = "HINDSIGHT_API_WORKER_BATCH_SIZE" ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT" +ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS" +ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS" # Reflect agent settings ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS" @@ -229,8 +230,9 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode) DEFAULT_WORKER_ID = None # Will use hostname if not specified DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed -DEFAULT_WORKER_BATCH_SIZE = 10 # Tasks to claim per poll cycle DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health +DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker +DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker # Reflect agent settings DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response @@ -419,8 +421,9 @@ class HindsightConfig: worker_id: str | None worker_poll_interval_ms: int worker_max_retries: int - worker_batch_size: int worker_http_port: int + worker_max_slots: int + worker_consolidation_max_slots: int # Reflect agent settings reflect_max_iterations: int @@ -582,8 +585,11 @@ class HindsightConfig: worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID, worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))), worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))), - worker_batch_size=int(os.getenv(ENV_WORKER_BATCH_SIZE, str(DEFAULT_WORKER_BATCH_SIZE))), worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))), + worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))), + worker_consolidation_max_slots=int( + os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS)) + ), # Reflect agent settings reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))), ) diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 5027e045..d7802b84 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -253,8 +253,9 @@ def main(): worker_id=config.worker_id, worker_poll_interval_ms=config.worker_poll_interval_ms, worker_max_retries=config.worker_max_retries, - worker_batch_size=config.worker_batch_size, worker_http_port=config.worker_http_port, + worker_max_slots=config.worker_max_slots, + worker_consolidation_max_slots=config.worker_consolidation_max_slots, reflect_max_iterations=config.reflect_max_iterations, mental_model_refresh_concurrency=config.mental_model_refresh_concurrency, ) diff --git a/hindsight-api/hindsight_api/worker/main.py b/hindsight-api/hindsight_api/worker/main.py index b5d69466..9b32eee9 100644 --- a/hindsight-api/hindsight_api/worker/main.py +++ b/hindsight-api/hindsight_api/worker/main.py @@ -124,12 +124,6 @@ def main(): default=config.worker_poll_interval_ms, help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)", ) - parser.add_argument( - "--batch-size", - type=int, - default=config.worker_batch_size, - help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)", - ) parser.add_argument( "--max-retries", type=int, @@ -168,8 +162,9 @@ def main(): print(f"Starting Hindsight Worker: {args.worker_id}") print(f" Poll interval: {args.poll_interval}ms") - print(f" Batch size: {args.batch_size}") print(f" Max retries: {args.max_retries}") + print(f" Max slots: {config.worker_max_slots}") + print(f" Consolidation max slots: {config.worker_consolidation_max_slots}") print(f" HTTP server: {args.http_host}:{args.http_port}") print() @@ -213,9 +208,10 @@ def main(): worker_id=args.worker_id, executor=memory.execute_task, poll_interval_ms=args.poll_interval, - batch_size=args.batch_size, max_retries=args.max_retries, tenant_extension=tenant_extension, + max_slots=config.worker_max_slots, + consolidation_max_slots=config.worker_consolidation_max_slots, ) # Create the HTTP app for metrics/health diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py index 3c990c56..bd54c876 100644 --- a/hindsight-api/hindsight_api/worker/poller.py +++ b/hindsight-api/hindsight_api/worker/poller.py @@ -57,10 +57,11 @@ class WorkerPoller: worker_id: str, executor: Callable[[dict[str, Any]], Awaitable[None]], poll_interval_ms: int = 500, - batch_size: int = 10, max_retries: int = 3, schema: str | None = None, tenant_extension: "TenantExtension | None" = None, + max_slots: int = 10, + consolidation_max_slots: int = 2, ): """ Initialize the worker poller. @@ -70,28 +71,32 @@ class WorkerPoller: worker_id: Unique identifier for this worker executor: Async function to execute tasks (typically MemoryEngine.execute_task) poll_interval_ms: Interval between polls when no tasks found (milliseconds) - batch_size: Maximum number of tasks to claim per poll cycle max_retries: Maximum retry attempts before marking task as failed schema: Database schema for single-tenant support (ignored if tenant_extension is set) tenant_extension: Extension for dynamic multi-tenant discovery. If set, list_tenants() is called on each poll cycle to discover schemas dynamically. + max_slots: Maximum concurrent tasks per worker + consolidation_max_slots: Maximum concurrent consolidation tasks per worker """ self._pool = pool self._worker_id = worker_id self._executor = executor self._poll_interval_ms = poll_interval_ms - self._batch_size = batch_size self._max_retries = max_retries self._schema = schema self._tenant_extension = tenant_extension + self._max_slots = max_slots + self._consolidation_max_slots = consolidation_max_slots self._shutdown = asyncio.Event() self._current_tasks: set[asyncio.Task] = set() self._in_flight_count = 0 self._in_flight_lock = asyncio.Lock() self._last_progress_log = 0.0 self._tasks_completed_since_log = 0 - # Track active tasks locally: operation_id -> (op_type, bank_id, schema) - self._active_tasks: dict[str, tuple[str, str, str | None]] = {} + # Track active tasks locally: operation_id -> (op_type, bank_id, schema, asyncio.Task) + self._active_tasks: dict[str, tuple[str, str, str | None, asyncio.Task]] = {} + # Track in-flight tasks by operation type + self._in_flight_by_type: dict[str, int] = {} async def _get_schemas(self) -> list[str | None]: """Get list of schemas to poll. Returns [None] for public schema.""" @@ -102,67 +107,114 @@ class WorkerPoller: # Single schema mode return [self._schema] + async def _get_available_slots(self) -> tuple[int, int]: + """ + Calculate available slots for claiming tasks. + + Returns: + (total_available, consolidation_available) tuple + """ + async with self._in_flight_lock: + total_in_flight = self._in_flight_count + consolidation_in_flight = self._in_flight_by_type.get("consolidation", 0) + + total_available = max(0, self._max_slots - total_in_flight) + consolidation_available = max(0, self._consolidation_max_slots - consolidation_in_flight) + + return total_available, consolidation_available + + async def wait_for_active_tasks(self, timeout: float = 10.0) -> bool: + """ + Wait for all active background tasks to complete (test helper). + + This is a test-only utility that allows tests to synchronize with + fire-and-forget background tasks without using sleep(). + + Args: + timeout: Maximum time to wait in seconds + + Returns: + True if all tasks completed, False if timeout was reached + """ + start_time = asyncio.get_event_loop().time() + while True: + async with self._in_flight_lock: + if self._in_flight_count == 0: + return True + + elapsed = asyncio.get_event_loop().time() - start_time + if elapsed >= timeout: + return False + + # Short sleep to avoid busy-waiting + await asyncio.sleep(0.01) + async def claim_batch(self) -> list[ClaimedTask]: """ - Claim up to batch_size pending tasks atomically across all tenant schemas. + Claim pending tasks atomically across all tenant schemas, + respecting slot limits (total and consolidation). Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers. - For consolidation tasks specifically, skips pending tasks if there's already - a processing consolidation for the same bank (to avoid duplicate work). - - If tenant_extension is configured, dynamically discovers schemas on each call. - Returns: List of ClaimedTask objects containing operation_id, task_dict, and schema """ + # Calculate available slots + total_available, consolidation_available = await self._get_available_slots() + + if total_available <= 0: + return [] + schemas = await self._get_schemas() all_tasks: list[ClaimedTask] = [] - remaining_batch = self._batch_size + remaining_total = total_available + remaining_consolidation = consolidation_available for schema in schemas: - if remaining_batch <= 0: + if remaining_total <= 0: break - tasks = await self._claim_batch_for_schema(schema, remaining_batch) + tasks = await self._claim_batch_for_schema(schema, remaining_total, remaining_consolidation) + + # Update remaining slots based on what was claimed + for task in tasks: + op_type = task.task_dict.get("operation_type", "unknown") + if op_type == "consolidation": + remaining_consolidation -= 1 + all_tasks.extend(tasks) - remaining_batch -= len(tasks) + remaining_total -= len(tasks) return all_tasks - async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]: - """Claim tasks from a specific schema.""" + async def _claim_batch_for_schema( + self, schema: str | None, limit: int, consolidation_limit: int + ) -> list[ClaimedTask]: + """Claim tasks from a specific schema respecting slot limits.""" try: - return await self._claim_batch_for_schema_inner(schema, limit) + return await self._claim_batch_for_schema_inner(schema, limit, consolidation_limit) except Exception as e: logger.warning(f"Worker {self._worker_id} failed to claim tasks for schema {schema or 'public'}: {e}") return [] - async def _claim_batch_for_schema_inner(self, schema: str | None, limit: int) -> list[ClaimedTask]: - """Inner implementation for claiming tasks from a specific schema.""" + async def _claim_batch_for_schema_inner( + self, schema: str | None, limit: int, consolidation_limit: int + ) -> list[ClaimedTask]: + """Inner implementation for claiming tasks from a specific schema with slot limits.""" table = fq_table("async_operations", schema) async with self._pool.acquire() as conn: async with conn.transaction(): - # Select and lock pending tasks - # For consolidation: skip if same bank already has one processing - rows = await conn.fetch( + # Strategy: Claim non-consolidation tasks first, then consolidation up to limit + + # 1. Claim non-consolidation tasks (up to limit) + non_consolidation_rows = await conn.fetch( f""" SELECT operation_id, task_payload - FROM {table} AS pending - WHERE status = 'pending' AND task_payload IS NOT NULL - AND ( - -- Non-consolidation tasks: always claimable - operation_type != 'consolidation' - OR - -- Consolidation: only if no other consolidation processing for same bank - NOT EXISTS ( - SELECT 1 FROM {table} AS processing - WHERE processing.bank_id = pending.bank_id - AND processing.operation_type = 'consolidation' - AND processing.status = 'processing' - ) - ) + FROM {table} + WHERE status = 'pending' + AND task_payload IS NOT NULL + AND operation_type != 'consolidation' ORDER BY created_at LIMIT $1 FOR UPDATE SKIP LOCKED @@ -170,11 +222,39 @@ class WorkerPoller: limit, ) - if not rows: + claimed_count = len(non_consolidation_rows) + remaining_limit = limit - claimed_count + + # 2. Claim consolidation tasks (up to consolidation_limit and remaining_limit) + consolidation_rows = [] + if consolidation_limit > 0 and remaining_limit > 0: + consolidation_rows = await conn.fetch( + f""" + SELECT operation_id, task_payload + FROM {table} AS pending + WHERE status = 'pending' + AND task_payload IS NOT NULL + AND operation_type = 'consolidation' + AND NOT EXISTS ( + SELECT 1 FROM {table} AS processing + WHERE processing.bank_id = pending.bank_id + AND processing.operation_type = 'consolidation' + AND processing.status = 'processing' + ) + ORDER BY created_at + LIMIT $1 + FOR UPDATE SKIP LOCKED + """, + min(consolidation_limit, remaining_limit), + ) + + all_rows = non_consolidation_rows + consolidation_rows + + if not all_rows: return [] # Claim the tasks by updating status and worker_id - operation_ids = [row["operation_id"] for row in rows] + operation_ids = [row["operation_id"] for row in all_rows] await conn.execute( f""" UPDATE {table} @@ -192,7 +272,7 @@ class WorkerPoller: task_dict=json.loads(row["task_payload"]), schema=schema, ) - for row in rows + for row in all_rows ] async def _mark_completed(self, operation_id: str, schema: str | None): @@ -258,18 +338,43 @@ class WorkerPoller: logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})") async def execute_task(self, task: ClaimedTask): - """Execute a single task and update its status.""" + """Execute a single task as a background job (fire-and-forget).""" task_type = task.task_dict.get("type", "unknown") + operation_type = task.task_dict.get("operation_type", "unknown") bank_id = task.task_dict.get("bank_id", "unknown") + # Create background task + bg_task = asyncio.create_task(self._execute_task_inner(task)) + # Track this task as active async with self._in_flight_lock: - self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema) + self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema, bg_task) + self._in_flight_count += 1 + self._in_flight_by_type[operation_type] = self._in_flight_by_type.get(operation_type, 0) + 1 + + # Add cleanup callback + bg_task.add_done_callback(lambda _: asyncio.create_task(self._cleanup_task(task.operation_id, operation_type))) + + async def _cleanup_task(self, operation_id: str, operation_type: str): + """Remove task from tracking after completion.""" + async with self._in_flight_lock: + if operation_id in self._active_tasks: + self._active_tasks.pop(operation_id, None) + self._in_flight_count -= 1 + count = self._in_flight_by_type.get(operation_type, 0) + if count > 0: + self._in_flight_by_type[operation_type] = count - 1 + if self._in_flight_by_type[operation_type] == 0: + del self._in_flight_by_type[operation_type] + + async def _execute_task_inner(self, task: ClaimedTask): + """Inner task execution with error handling.""" + task_type = task.task_dict.get("type", "unknown") + bank_id = task.task_dict.get("bank_id", "unknown") try: schema_info = f", schema={task.schema}" if task.schema else "" logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})") - # Pass schema to executor so it can set the correct context if task.schema: task.task_dict["_schema"] = task.schema await self._executor(task.task_dict) @@ -279,10 +384,6 @@ class WorkerPoller: error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}" logger.error(f"Task {task.operation_id} failed: {e}") await self._retry_or_fail(task.operation_id, error_msg, task.schema) - finally: - # Remove from active tasks - async with self._in_flight_lock: - self._active_tasks.pop(task.operation_id, None) async def recover_own_tasks(self) -> int: """ @@ -325,59 +426,59 @@ class WorkerPoller: async def run(self): """ - Main polling loop. + Main polling loop with fire-and-forget task execution. - Continuously polls for pending tasks, claims them, and executes them - until shutdown is signaled. - - If tenant_extension is configured, dynamically discovers schemas on each poll. + Continuously polls for pending tasks, spawns them as background tasks, + and immediately continues polling (up to slot limits). """ - # Recover any tasks from a previous crash before starting await self.recover_own_tasks() - logger.info(f"Worker {self._worker_id} starting polling loop") + logger.info( + f"Worker {self._worker_id} starting polling loop " + f"(max_slots={self._max_slots}, consolidation_max_slots={self._consolidation_max_slots})" + ) while not self._shutdown.is_set(): try: - # Claim a batch of tasks (across all tenant schemas if configured) + # Claim a batch of tasks (respecting slot limits) tasks = await self.claim_batch() if tasks: # Log batch info task_types: dict[str, int] = {} schemas_seen: set[str | None] = set() + consolidation_count = 0 for task in tasks: t = task.task_dict.get("type", "unknown") + op_type = task.task_dict.get("operation_type", "unknown") task_types[t] = task_types.get(t, 0) + 1 schemas_seen.add(task.schema) + if op_type == "consolidation": + consolidation_count += 1 + types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items()) schemas_str = ", ".join(s or "public" for s in schemas_seen) logger.info( - f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str} (schemas: {schemas_str})" + f"Worker {self._worker_id} claimed {len(tasks)} tasks " + f"({consolidation_count} consolidation): {types_str} (schemas: {schemas_str})" ) - # Track in-flight tasks - async with self._in_flight_lock: - self._in_flight_count += len(tasks) + # Spawn tasks as background jobs (fire-and-forget) + for task in tasks: + await self.execute_task(task) - # Execute tasks concurrently - try: - await asyncio.gather( - *[self.execute_task(task) for task in tasks], - return_exceptions=True, - ) - finally: - async with self._in_flight_lock: - self._in_flight_count -= len(tasks) - else: - # No tasks found, wait before polling again - try: - await asyncio.wait_for( - self._shutdown.wait(), - timeout=self._poll_interval_ms / 1000, - ) - except asyncio.TimeoutError: - pass # Normal timeout, continue polling + # Continue immediately to claim more tasks (if slots available) + continue + + # No tasks claimed (either no pending tasks or slots full) + # Wait before polling again + try: + await asyncio.wait_for( + self._shutdown.wait(), + timeout=self._poll_interval_ms / 1000, + ) + except asyncio.TimeoutError: + pass # Normal timeout, continue polling # Log progress stats periodically await self._log_progress_if_due() @@ -408,15 +509,27 @@ class WorkerPoller: while asyncio.get_event_loop().time() - start_time < timeout: async with self._in_flight_lock: in_flight = self._in_flight_count + active_task_objects = [task_info[3] for task_info in self._active_tasks.values()] if in_flight == 0: logger.info(f"Worker {self._worker_id} graceful shutdown complete") return logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks") - await asyncio.sleep(0.5) - logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s") + # Wait for at least one task to complete + if active_task_objects: + done, _ = await asyncio.wait(active_task_objects, timeout=0.5, return_when=asyncio.FIRST_COMPLETED) + else: + await asyncio.sleep(0.5) + + logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s, cancelling remaining tasks") + + # Cancel remaining tasks + async with self._in_flight_lock: + for operation_id, (_, _, _, bg_task) in list(self._active_tasks.items()): + if not bg_task.done(): + bg_task.cancel() async def _log_progress_if_due(self): """Log progress stats every PROGRESS_LOG_INTERVAL seconds.""" @@ -427,14 +540,19 @@ class WorkerPoller: self._last_progress_log = now try: - # Get local active tasks (this worker only) + # Get local active tasks async with self._in_flight_lock: in_flight = self._in_flight_count - active_tasks = dict(self._active_tasks) # Copy to avoid holding lock + in_flight_by_type = dict(self._in_flight_by_type) + active_tasks = dict(self._active_tasks) - # Build local processing breakdown grouped by (op_type, bank_id) + consolidation_count = in_flight_by_type.get("consolidation", 0) + available_slots = self._max_slots - in_flight + available_consolidation_slots = self._consolidation_max_slots - consolidation_count + + # Build local processing breakdown task_groups: dict[tuple[str, str], int] = {} - for op_type, bank_id, _ in active_tasks.values(): + for op_type, bank_id, _, _ in active_tasks.values(): key = (op_type, bank_id) task_groups[key] = task_groups.get(key, 0) + 1 @@ -443,7 +561,7 @@ class WorkerPoller: if len(processing_info) > 10: processing_str += f" +{len(processing_info) - 10} more" - # Get global stats from DB across all schemas + # Get global stats from DB schemas = await self._get_schemas() global_pending = 0 all_worker_counts: dict[str, int] = {} @@ -455,7 +573,6 @@ class WorkerPoller: row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'") global_pending += row["count"] if row else 0 - # Get processing breakdown by worker worker_rows = await conn.fetch( f""" SELECT worker_id, COUNT(*) as count @@ -468,7 +585,6 @@ class WorkerPoller: wid = wr["worker_id"] or "unknown" all_worker_counts[wid] = all_worker_counts.get(wid, 0) + wr["count"] - # Format other workers' processing counts other_workers = [] for wid, cnt in all_worker_counts.items(): if wid != self._worker_id: @@ -477,7 +593,9 @@ class WorkerPoller: schemas_str = ", ".join(s or "public" for s in schemas) logger.info( - f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | " + f"[WORKER_STATS] worker={self._worker_id} " + f"slots={in_flight}/{self._max_slots} (consolidation={consolidation_count}/{self._consolidation_max_slots}) | " + f"available={available_slots} (consolidation={available_consolidation_slots}) | " f"global: pending={global_pending} (schemas: {schemas_str}) | " f"others: {others_str} | " f"my_active: {processing_str}" diff --git a/hindsight-api/tests/test_consolidation.py b/hindsight-api/tests/test_consolidation.py index 3fe745c5..1ff9e7e2 100644 --- a/hindsight-api/tests/test_consolidation.py +++ b/hindsight-api/tests/test_consolidation.py @@ -346,11 +346,11 @@ class TestConsolidationIntegration: or when one directly updates another (e.g., location change). Given: - - "Nicolò lives in Italy" - - "Nicolò moved to the US recently" (updates the living location) + - "Alex lives in Italy" + - "Alex moved to the US recently" (updates the living location) The second fact should UPDATE the first, not create a separate observation. - But unrelated facts like "Nicolò works at Vectorize" should stay separate. + But unrelated facts like "Alex works at Vectorize" should stay separate. """ bank_id = f"test-consolidation-merge-{uuid.uuid4().hex[:8]}" @@ -360,14 +360,14 @@ class TestConsolidationIntegration: # Retain a memory about living location await memory.retain_async( bank_id=bank_id, - content="Nicolò lives in Italy.", + content="Alex lives in Italy.", request_context=request_context, ) # Retain an unrelated memory (different topic - should NOT merge) await memory.retain_async( bank_id=bank_id, - content="Nicolò works at Vectorize as an engineer.", + content="Alex works at Vectorize as an engineer.", request_context=request_context, ) @@ -384,7 +384,7 @@ class TestConsolidationIntegration: # Add a memory that UPDATES the living location (should merge with first) await memory.retain_async( bank_id=bank_id, - content="Nicolò recently moved to the United States.", + content="Alex recently moved to the United States.", request_context=request_context, ) @@ -485,9 +485,9 @@ class TestConsolidationIntegration: they should be merged into ONE observation that captures the change. Example: - - "Nicolò loves pizza" - - "Nicolò hates pizza" - → Should become: "Nicolò used to love pizza but now hates it" (or similar) + - "Alex loves pizza" + - "Alex hates pizza" + → Should become: "Alex used to love pizza but now hates it" (or similar) """ bank_id = f"test-consolidation-contradict-{uuid.uuid4().hex[:8]}" @@ -497,7 +497,7 @@ class TestConsolidationIntegration: # Add initial fact await memory.retain_async( bank_id=bank_id, - content="Nicolò loves pizza.", + content="Alex loves pizza.", request_context=request_context, ) @@ -515,7 +515,7 @@ class TestConsolidationIntegration: # Add contradicting fact (same person, same topic, opposite sentiment) await memory.retain_async( bank_id=bank_id, - content="Nicolò hates pizza.", + content="Alex hates pizza.", request_context=request_context, ) diff --git a/hindsight-api/tests/test_worker.py b/hindsight-api/tests/test_worker.py index 659554d6..44a8a286 100644 --- a/hindsight-api/tests/test_worker.py +++ b/hindsight-api/tests/test_worker.py @@ -156,7 +156,6 @@ class TestWorkerPoller: pool=pool, worker_id="test-worker-1", executor=mock_executor, - batch_size=10, ) claimed = await poller.claim_batch() @@ -177,8 +176,8 @@ class TestWorkerPoller: assert row["worker_id"] == "test-worker-1" @pytest.mark.asyncio - async def test_claim_batch_respects_batch_size(self, pool, clean_operations): - """Test that claim_batch respects the batch_size limit.""" + async def test_claim_batch_respects_max_slots(self, pool, clean_operations): + """Test that claim_batch respects the max_slots limit.""" from hindsight_api.worker import WorkerPoller # Create 10 pending tasks @@ -196,12 +195,11 @@ class TestWorkerPoller: payload, ) - # Claim with batch_size=3 poller = WorkerPoller( pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=3, + max_slots=3, # Limit to 3 concurrent tasks ) claimed = await poller.claim_batch() @@ -238,11 +236,14 @@ class TestWorkerPoller: executor=mock_executor, ) - # Execute the task + # Execute the task (fire-and-forget) task_dict = json.loads(payload) claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) await poller.execute_task(claimed_task) + # Wait for background task to complete + completed = await poller.wait_for_active_tasks(timeout=5.0) + assert completed, "Task did not complete within timeout" assert len(executed) == 1 # Verify task is marked as completed @@ -283,11 +284,15 @@ class TestWorkerPoller: max_retries=3, ) - # Execute (should fail and retry) + # Execute (should fail and retry) - fire-and-forget task_dict = json.loads(payload) claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) await poller.execute_task(claimed_task) + # Wait for background task to complete + completed = await poller.wait_for_active_tasks(timeout=5.0) + assert completed, "Task did not complete within timeout" + # Verify task is back to pending with incremented retry_count row = await pool.fetchrow( "SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1", @@ -327,11 +332,15 @@ class TestWorkerPoller: max_retries=3, ) - # Execute (should fail permanently) + # Execute (should fail permanently) - fire-and-forget task_dict = json.loads(payload) claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None) await poller.execute_task(claimed_task) + # Wait for background task to complete + completed = await poller.wait_for_active_tasks(timeout=5.0) + assert completed, "Task did not complete within timeout" + # Verify task is marked as failed row = await pool.fetchrow( "SELECT status, error_message FROM async_operations WHERE operation_id = $1", @@ -388,7 +397,6 @@ class TestWorkerPoller: pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=10, ) claimed = await poller.claim_batch() @@ -440,7 +448,6 @@ class TestWorkerPoller: pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=10, ) claimed = await poller.claim_batch() @@ -607,7 +614,6 @@ class TestConcurrentWorkers: pool=pool, worker_id=worker_id, executor=lambda x: None, - batch_size=5, # Each worker tries to claim 5 ) claimed = await poller.claim_batch() workers_claimed[worker_id] = [task.operation_id for task in claimed] @@ -680,7 +686,6 @@ class TestConcurrentWorkers: pool=pool, worker_id="new-worker", executor=lambda x: None, - batch_size=10, ) claimed = await poller.claim_batch() @@ -879,7 +884,6 @@ class TestDynamicTenantDiscovery: pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=10, tenant_extension=mock_extension, ) @@ -946,7 +950,6 @@ class TestDynamicTenantDiscovery: pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=10, tenant_extension=dynamic_extension, ) @@ -1008,7 +1011,6 @@ class TestDynamicTenantDiscovery: pool=pool, worker_id="test-worker-1", executor=lambda x: None, - batch_size=10, ) claimed = await poller.claim_batch() @@ -1017,3 +1019,198 @@ class TestDynamicTenantDiscovery: # All tasks should have schema=None (public) for task in claimed: assert task.schema is None + + +async def test_worker_fire_and_forget_nonblocking(pool, clean_operations): + """ + Test that worker continues polling while tasks run (fire-and-forget pattern). + + This test verifies the FIX: With the old blocking behavior, the worker would + wait for all tasks in a batch to complete before claiming more. This test + would FAIL with the old code because tasks 3-4 wouldn't be claimed until + tasks 1-2 complete. With fire-and-forget, tasks 3-4 are claimed immediately. + """ + from hindsight_api.worker.poller import WorkerPoller + + task_started = {} # operation_id -> Event (set when task starts) + task_canfinish = {} # operation_id -> Event (wait before finishing) + + async def blocking_executor(task_dict: dict): + op_id = task_dict["operation_id"] + # Signal that this task has started + started = asyncio.Event() + task_started[op_id] = started + started.set() + + # Block until we're told to finish + finish = asyncio.Event() + task_canfinish[op_id] = finish + await finish.wait() + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker", + executor=blocking_executor, + poll_interval_ms=50, # Fast polling + max_slots=10, + consolidation_max_slots=2, + ) + + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + + # Submit initial 2 tasks + task_ids = [] + for i in range(2): + op_id = uuid.uuid4() + task_ids.append(str(op_id)) + payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'retain', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + poll_task = asyncio.create_task(poller.run()) + + try: + # Wait for first 2 tasks to start executing (but not finish) + for i in range(100): # Try for up to 1 second + if len(task_started) >= 2: + break + await asyncio.sleep(0.01) + assert len(task_started) == 2, f"Expected 2 tasks started, got {len(task_started)}" + + # Verify tasks are in_flight + async with poller._in_flight_lock: + assert poller._in_flight_count == 2 + + # NOW submit 2 more tasks WHILE the first 2 are still running + for i in range(2): + op_id = uuid.uuid4() + task_ids.append(str(op_id)) + payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'retain', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + # KEY ASSERTION: Worker should claim tasks 3-4 WITHOUT waiting for 1-2 to finish + # This would FAIL with the old blocking behavior + for i in range(100): # Try for up to 1 second + if len(task_started) >= 4: + break + await asyncio.sleep(0.01) + + assert len(task_started) == 4, ( + f"Fire-and-forget FAILED: Expected 4 tasks started, got {len(task_started)}. " + "This means the worker blocked waiting for the first batch to complete." + ) + + # Verify all 4 tasks are in-flight + async with poller._in_flight_lock: + assert poller._in_flight_count == 4 + + # Clean up: allow all tasks to finish + for event in task_canfinish.values(): + event.set() + + finally: + # Ensure cleanup + for event in task_canfinish.values(): + event.set() + await poller.shutdown_graceful(timeout=2.0) + try: + await asyncio.wait_for(poll_task, timeout=1.0) + except asyncio.CancelledError: + pass + + +async def test_worker_slot_limits_enforced(pool, clean_operations): + """Test that worker respects max_slots and won't exceed the limit.""" + from hindsight_api.worker.poller import WorkerPoller + + tasks_started = set() + task_events = {} + + async def controlled_executor(task_dict: dict): + op_id = task_dict["operation_id"] + tasks_started.add(op_id) + event = asyncio.Event() + task_events[op_id] = event + await event.wait() + + poller = WorkerPoller( + pool=pool, + worker_id="test-worker", + executor=controlled_executor, + poll_interval_ms=50, + max_slots=3, # Only allow 3 concurrent tasks + consolidation_max_slots=1, + ) + + # Submit 10 tasks + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + for i in range(10): + op_id = uuid.uuid4() + payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'retain', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + + poll_task = asyncio.create_task(poller.run()) + + try: + # Wait for slots to fill + for i in range(100): + if len(tasks_started) >= 3: + break + await asyncio.sleep(0.01) + + # Should have claimed exactly 3 tasks (slot limit) + assert len(tasks_started) == 3 + + # Wait to ensure no additional tasks are claimed + for i in range(30): + await asyncio.sleep(0.01) + assert len(tasks_started) == 3, "Worker exceeded slot limit!" + + # Release tasks one by one and verify remaining are claimed + completed = 0 + while completed < 10 and len(tasks_started) < 10: + # Release the next batch + events_to_release = list(task_events.values())[completed:completed+3] + for event in events_to_release: + event.set() + completed += len(events_to_release) + + # Wait for new tasks to be claimed + for i in range(100): + if len(tasks_started) >= min(completed + 3, 10): + break + await asyncio.sleep(0.01) + + assert len(tasks_started) == 10 + + finally: + for event in task_events.values(): + event.set() + await poller.shutdown_graceful(timeout=2.0) + try: + await asyncio.wait_for(poll_task, timeout=1.0) + except asyncio.CancelledError: + pass diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index d94f3454..fa0fad81 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -519,9 +519,10 @@ Configuration for background task processing. By default, the API processes task | `HINDSIGHT_API_WORKER_ENABLED` | Enable internal worker in API process | `true` | | `HINDSIGHT_API_WORKER_ID` | Unique worker identifier | hostname | | `HINDSIGHT_API_WORKER_POLL_INTERVAL_MS` | Database polling interval in milliseconds | `500` | -| `HINDSIGHT_API_WORKER_BATCH_SIZE` | Tasks to claim per poll cycle | `10` | | `HINDSIGHT_API_WORKER_MAX_RETRIES` | Max retries before marking task failed | `3` | | `HINDSIGHT_API_WORKER_HTTP_PORT` | HTTP port for worker metrics/health (worker CLI only) | `8889` | +| `HINDSIGHT_API_WORKER_MAX_SLOTS` | Maximum concurrent tasks per worker | `10` | +| `HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS` | Maximum concurrent consolidation tasks per worker | `2` | ### Performance Optimization