From 5ca1131ba3a08ffef999398b8a2e71b3f7424b27 Mon Sep 17 00:00:00 2001 From: busya Date: Thu, 4 Jun 2026 20:30:43 +0300 Subject: [PATCH] 032: Phase 5 (T036-T039) + Phase 6 (T040-T044) completed T036: superset_compilation_adapter fully async T037: fileio.py async wrappers (aiofiles+run_blocking) T038-T039: tests for plugins fileio concurrency T040: providers async (aiosmtplib, httpx.AsyncClient) T041: dispatch_report parallel via asyncio.gather T042: TaskManager ThreadPoolExecutor->asyncio.create_task T043: EventBus asyncio.Queue(maxsize=10000) T044: lifecycle async context manager Remaining: T045 git/_base.py, T046-T048 tests, T005-T020-T021 tests --- backend/src/core/task_manager/event_bus.py | 204 ++++++++++-------- backend/src/core/task_manager/lifecycle.py | 81 ++++--- backend/src/core/task_manager/manager.py | 97 +++++---- backend/src/core/utils/fileio.py | 105 ++++++++- .../utils/superset_compilation_adapter.py | 26 +-- .../src/services/notifications/providers.py | 36 ++-- backend/src/services/notifications/service.py | 24 ++- backend/tests/core/test_fileio_async.py | 135 ++++++++++++ backend/tests/plugins/test_plugins_async.py | 95 ++++++++ 9 files changed, 595 insertions(+), 208 deletions(-) create mode 100644 backend/tests/core/test_fileio_async.py create mode 100644 backend/tests/plugins/test_plugins_async.py diff --git a/backend/src/core/task_manager/event_bus.py b/backend/src/core/task_manager/event_bus.py index 86962c14..c3bfbf8e 100644 --- a/backend/src/core/task_manager/event_bus.py +++ b/backend/src/core/task_manager/event_bus.py @@ -1,110 +1,123 @@ # #region EventBusModule [C:5] [TYPE Module] [SEMANTICS event,bus,log,buffer,subscriber,flush] # @BRIEF Task-scoped log buffering, persistence flushes, and subscriber fan-out for real-time -# WebSocket observers — extracted from the monolithic TaskManager. +# WebSocket observers — fully async with asyncio.Queue fan-out. # @LAYER Core # @RELATION DEPENDS_ON -> [LogEntry] # @RELATION DEPENDS_ON -> [TaskLogPersistenceService] # @RELATION DEPENDS_ON -> [should_log_task_level] -# @RATIONALE Extracted from TaskManager to satisfy INV_7. EventBus owns log buffering, batch -# persistence flushing, and pub/sub subscription management — a cohesive unit with -# independent threading and I/O patterns. -# @REJECTED Keeping log/event concerns inside TaskManager was rejected — the background flusher -# thread, buffer locks, and subscriber queues are architecturally distinct from task -# lifecycle state machines and the in-memory task registry. +# @RATIONALE Extracted from TaskManager to satisfy INV_7. Migrated from threading to asyncio +# for T043 — replaces queue.Queue + call_soon_threadsafe with asyncio.Queue +# and async subscriber generator pattern. +# @REJECTED Keeping threading-based EventBus was rejected — call_soon_threadsafe is fragile +# under high load and cannot be tested without a running event loop. # @PRE TaskLogPersistenceService is initialized. # @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers. # @SIDE_EFFECT Writes task logs to persistence, mutates in-memory buffers, pushes log entries -# into subscriber queues, runs background flusher thread. +# into subscriber queues via async put. # @INVARIANT Buffered logs are retried on persistence failure; every subscriber receives only # task-scoped events. # @DATA_CONTRACT Input: LogEntry -> Output: persisted log + subscriber notification import asyncio -import threading from typing import Any -from ..cot_logger import seed_trace_id -from ..logger import belief_scope, logger, should_log_task_level -from .models import LogEntry, LogFilter, LogStats -from .persistence import TaskLogPersistenceService +from src.core.cot_logger import seed_trace_id +from src.core.logger import belief_scope, logger, should_log_task_level +from src.core.task_manager.models import LogEntry, LogFilter, LogStats +from src.core.task_manager.persistence import TaskLogPersistenceService + +QUEUE_MAXSIZE = 10000 # #region EventBus [C:5] [TYPE Class] [SEMANTICS event,bus,log,buffer,subscriber,flush] # @BRIEF Coordinates task-scoped log buffering, persistence flushes, and subscriber fan-out -# for real-time observers. +# for real-time observers. Fully async — uses asyncio.Queue for fan-out delivery. # @RELATION DEPENDS_ON -> [LogEntry] # @RELATION DEPENDS_ON -> [TaskLogPersistenceService] # @RELATION DEPENDS_ON -> [should_log_task_level] -# @PRE asyncio loop is available for thread-safe queue operations. +# @PRE TaskLogPersistenceService is initialized. # @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers. # @INVARIANT Buffered logs are retried on persistence failure. -# @RATIONALE Extracted as standalone pub/sub event bus to isolate threading (background flusher -# thread, buffer locks) and subscription management from task lifecycle and registry. -# @REJECTED Keeping log buffering inside TaskManager was rejected — the background flusher -# thread and subscriber queues are architecturally distinct from task state machines -# and should be independently stoppable/testable. +# @RATIONALE Fully async event bus with asyncio.Queue(maxsize=10000) for each subscriber. +# Broadcast methods use await queue.put() instead of call_soon_threadsafe. +# @REJECTED Keeping threading-based EventBus with call_soon_threadsafe — fragile under load, +# hard to test, and incompatible with fully async task execution. class EventBus: - """Task-scoped log buffering, persistence flushes, and subscriber fan-out.""" + """Task-scoped log buffering, persistence flushes, and subscriber fan-out — fully async.""" LOG_FLUSH_INTERVAL = 2.0 # #region __init__ [TYPE Function] - # @BRIEF Initialize empty log buffer, subscriber registry, status subscribers, and start background flusher. - def __init__(self, log_persistence_service: TaskLogPersistenceService, loop: asyncio.AbstractEventLoop): + # @BRIEF Initialize log buffer, subscriber registries, and async flusher. + # @POST EventBus is ready to accept logs. Subscriber dicts are empty. + def __init__(self, log_persistence_service: TaskLogPersistenceService): self.log_persistence_service = log_persistence_service - self.loop = loop self.subscribers: dict[str, list[asyncio.Queue]] = {} - # Per-task status subscribers (same pattern as log subscribers) self._status_subscribers: dict[str, list[asyncio.Queue]] = {} - # Global task event subscribers (all tasks) self._task_event_subscribers: list[asyncio.Queue] = [] - # Maintenance event subscribers self._maintenance_subscribers: list[asyncio.Queue] = [] self._log_buffer: dict[str, list[LogEntry]] = {} - self._log_buffer_lock = threading.Lock() - self._flusher_stop_event = threading.Event() - self._flusher_thread = threading.Thread( - target=self._flusher_loop, daemon=True - ) - self._flusher_thread.start() + self._flusher_stop_event = asyncio.Event() + self._flusher_task: asyncio.Task | None = None # #endregion __init__ + # #region start [TYPE Function] + # @BRIEF Start the background async flusher task. + # @POST Flusher task is created and scheduled in the event loop. + def start(self) -> None: + if self._flusher_task is None or self._flusher_task.done(): + self._flusher_task = asyncio.create_task(self.async_flusher_loop()) + # #endregion start + # #region stop [TYPE Function] - # @BRIEF Signal the flusher thread to stop and wait for it. - def stop(self) -> None: + # @BRIEF Signal the flusher to stop and wait for it. + # @POST Flusher task is cancelled and awaited. + async def stop(self) -> None: self._flusher_stop_event.set() - self._flusher_thread.join(timeout=2) + if self._flusher_task and not self._flusher_task.done(): + self._flusher_task.cancel() + from contextlib import suppress + with suppress(asyncio.CancelledError, Exception): + await self._flusher_task + self._flusher_task = None # #endregion stop - # #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,thread] - # @BRIEF Background thread that periodically flushes log buffer to database. - def _flusher_loop(self): + # #region async_flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async] + # @BRIEF Background async task that periodically flushes log buffer to database. + async def async_flusher_loop(self): seed_trace_id() - while not self._flusher_stop_event.is_set(): - self._flush_logs() - self._flusher_stop_event.wait(self.LOG_FLUSH_INTERVAL) - # #endregion _flusher_loop + try: + while not self._flusher_stop_event.is_set(): + await self._flush_logs() + try: + await asyncio.wait_for( + self._flusher_stop_event.wait(), + timeout=self.LOG_FLUSH_INTERVAL, + ) + break # Event was set + except TimeoutError: + continue # Timeout is normal — flush again + except asyncio.CancelledError: + pass + # #endregion async_flusher_loop # #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence] # @BRIEF Flush all buffered logs to the database. # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs] - def _flush_logs(self): + async def _flush_logs(self): seed_trace_id() - with self._log_buffer_lock: - task_ids = list(self._log_buffer.keys()) + task_ids = list(self._log_buffer.keys()) for task_id in task_ids: - with self._log_buffer_lock: - logs = self._log_buffer.pop(task_id, []) + logs = self._log_buffer.pop(task_id, []) if logs: try: self.log_persistence_service.add_logs(task_id, logs) logger.debug(f"Flushed {len(logs)} logs for task {task_id}") except Exception as e: logger.error(f"Failed to flush logs for task {task_id}: {e}") - with self._log_buffer_lock: - if task_id not in self._log_buffer: - self._log_buffer[task_id] = [] - self._log_buffer[task_id].extend(logs) + if task_id not in self._log_buffer: + self._log_buffer[task_id] = [] + self._log_buffer[task_id].extend(logs) # #endregion _flush_logs # #region flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence] @@ -112,10 +125,9 @@ class EventBus: # @PRE task_id exists. # @POST Task's buffered logs are written to database. # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs] - def flush_task_logs(self, task_id: str): + async def flush_task_logs(self, task_id: str): with belief_scope("EventBus.flush_task_logs"): - with self._log_buffer_lock: - logs = self._log_buffer.pop(task_id, []) + logs = self._log_buffer.pop(task_id, []) if logs: try: self.log_persistence_service.add_logs(task_id, logs) @@ -124,11 +136,11 @@ class EventBus: # #endregion flush_task_logs # #region add_log [C:3] [TYPE Function] [SEMANTICS log,buffer,notify,subscriber] - # @BRIEF Adds a log entry to a task buffer and notifies subscribers. + # @BRIEF Adds a log entry to a task buffer and notifies subscribers via async put. # @PRE Task exists. - # @POST Log added to buffer and pushed to queues (if level meets task_log_level filter). + # @POST Log added to buffer and pushed to subscriber queues (if level passes filter). # @RELATION CALLS -> [should_log_task_level] - def add_log( + async def add_log( self, task_id: str, level: str, @@ -148,25 +160,31 @@ class EventBus: metadata=metadata, context=context, ) - # Add to in-memory logs (for backward compatibility with legacy JSON field) + # Add to in-memory logs (backward compatibility) if task_logs_list is not None: task_logs_list.append(log_entry) # Add to buffer for batch persistence - with self._log_buffer_lock: - if task_id not in self._log_buffer: - self._log_buffer[task_id] = [] - self._log_buffer[task_id].append(log_entry) - # Notify subscribers (for real-time WebSocket updates) + if task_id not in self._log_buffer: + self._log_buffer[task_id] = [] + self._log_buffer[task_id].append(log_entry) + # Notify subscribers via async put (fan-out) if task_id in self.subscribers: for queue in self.subscribers[task_id]: - self.loop.call_soon_threadsafe(queue.put_nowait, log_entry) + try: + await queue.put(log_entry) + except asyncio.QueueFull: + logger.warning( + f"Subscriber queue full for task {task_id}, dropping log entry" + ) # #endregion add_log + # ── Log Subscribers ── + # #region subscribe_logs [C:2] [TYPE Function] [SEMANTICS subscribe,queue,realtime] # @BRIEF Subscribes to real-time logs for a task. - # @POST Returns an asyncio.Queue for log entries. + # @POST Returns an asyncio.Queue (maxsize=10000) for log entries. async def subscribe_logs(self, task_id: str) -> asyncio.Queue: - queue = asyncio.Queue() + queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE) if task_id not in self.subscribers: self.subscribers[task_id] = [] self.subscribers[task_id].append(queue) @@ -188,9 +206,9 @@ class EventBus: # #region subscribe_status [C:2] [TYPE Function] [SEMANTICS subscribe,status,realtime] # @BRIEF Subscribes to real-time status updates for a specific task. - # @POST Returns an asyncio.Queue for Task status dicts. + # @POST Returns an asyncio.Queue (maxsize=10000) for Task status dicts. async def subscribe_status(self, task_id: str) -> asyncio.Queue: - queue = asyncio.Queue() + queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE) if task_id not in self._status_subscribers: self._status_subscribers[task_id] = [] self._status_subscribers[task_id].append(queue) @@ -198,7 +216,7 @@ class EventBus: # #endregion subscribe_status # #region unsubscribe_status [C:2] [TYPE Function] [SEMANTICS unsubscribe,status,remove] - # @BRIEF Unsubscribes from real-time status updates for a task. + # @BRIEF Unsubscribes from status updates for a task. # @POST Queue removed from status subscribers. def unsubscribe_status(self, task_id: str, queue: asyncio.Queue): if task_id in self._status_subscribers: @@ -210,32 +228,35 @@ class EventBus: # #region broadcast_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,notify] # @BRIEF Broadcast a task status update to all per-task and global subscribers. - # @POST Status event pushed to all matching queues. - def broadcast_status(self, task_id: str, task_dict: dict): - from ..logger import logger as _lg + # @POST Status event pushed to all matching queues via async put. + async def broadcast_status(self, task_id: str, task_dict: dict): event = {"type": "task_status", "task_id": task_id, "task": task_dict} # Per-task subscribers if task_id in self._status_subscribers: for queue in self._status_subscribers[task_id]: try: - self.loop.call_soon_threadsafe(queue.put_nowait, event) - except Exception as _e: - _lg.warning("Failed to send status to per-task subscriber", extra={"task_id": task_id, "error": str(_e)}) + await queue.put(event) + except asyncio.QueueFull: + logger.warning( + f"Status subscriber queue full for task {task_id}, dropping event" + ) # Global task event subscribers for queue in self._task_event_subscribers: try: - self.loop.call_soon_threadsafe(queue.put_nowait, event) - except Exception as _e: - _lg.warning("Failed to send status to global subscriber", extra={"task_id": task_id, "error": str(_e)}) + await queue.put(event) + except asyncio.QueueFull: + logger.warning( + "Global task event subscriber queue full, dropping event" + ) # #endregion broadcast_status # ── Maintenance Event Subscribers ── # #region subscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS subscribe,maintenance,events] # @BRIEF Subscribes to maintenance event updates (created/ended/banner changes). - # @POST Returns an asyncio.Queue for maintenance event dicts. + # @POST Returns an asyncio.Queue (maxsize=10000) for maintenance event dicts. async def subscribe_maintenance_events(self) -> asyncio.Queue: - queue = asyncio.Queue() + queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE) self._maintenance_subscribers.append(queue) return queue # #endregion subscribe_maintenance_events @@ -250,23 +271,24 @@ class EventBus: # #region broadcast_maintenance_event [C:2] [TYPE Function] [SEMANTICS broadcast,maintenance,notify] # @BRIEF Broadcast a maintenance event to all subscribers. - # @POST Event pushed to all maintenance subscriber queues. - def broadcast_maintenance_event(self, event: dict): - from ..logger import logger as _lg + # @POST Event pushed to all maintenance subscriber queues via async put. + async def broadcast_maintenance_event(self, event: dict): for queue in self._maintenance_subscribers: try: - self.loop.call_soon_threadsafe(queue.put_nowait, event) - except Exception as _e: - _lg.warning("Failed to send maintenance event to subscriber", extra={"event_type": event.get("type"), "error": str(_e)}) + await queue.put(event) + except asyncio.QueueFull: + logger.warning( + "Maintenance subscriber queue full, dropping event" + ) # #endregion broadcast_maintenance_event # ── Global Task Event Subscribers ── # #region subscribe_task_events [C:2] [TYPE Function] [SEMANTICS subscribe,global,events] # @BRIEF Subscribes to ALL task events (status changes) globally. - # @POST Returns an asyncio.Queue for task event dicts. + # @POST Returns an asyncio.Queue (maxsize=10000) for task event dicts. async def subscribe_task_events(self) -> asyncio.Queue: - queue = asyncio.Queue() + queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE) self._task_event_subscribers.append(queue) return queue # #endregion subscribe_task_events @@ -279,13 +301,15 @@ class EventBus: self._task_event_subscribers.remove(queue) # #endregion unsubscribe_task_events + # ── Log retrieval methods (sync, delegates to persistence) ── + # #region get_task_logs [C:3] [TYPE Function] [SEMANTICS logs,read,persistence,backfill] # @BRIEF Retrieves logs for a specific task (from memory for running, persistence for completed). # @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_logs] def get_task_logs( self, task_id: str, log_filter: LogFilter | None = None, task_status=None, task_logs: list | None = None ) -> list[LogEntry]: - from .models import TaskStatus + from src.core.task_manager.models import TaskStatus # For completed tasks, fetch from persistence if task_status in [TaskStatus.SUCCESS, TaskStatus.FAILED]: if log_filter is None: diff --git a/backend/src/core/task_manager/lifecycle.py b/backend/src/core/task_manager/lifecycle.py index 636feab5..8515cbb6 100644 --- a/backend/src/core/task_manager/lifecycle.py +++ b/backend/src/core/task_manager/lifecycle.py @@ -23,16 +23,15 @@ # a new wait future is created. import asyncio -from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime import inspect from typing import Any -from ..cot_logger import seed_trace_id -from ..logger import belief_scope, logger -from .context import TaskContext -from .graph import TaskGraph -from .models import Task, TaskStatus +from src.core.cot_logger import seed_trace_id +from src.core.logger import belief_scope, logger +from src.core.task_manager.context import TaskContext +from src.core.task_manager.graph import TaskGraph +from src.core.task_manager.models import Task, TaskStatus def _task_to_dict(task: Task) -> dict: """Serialize a Task to a dict for WebSocket broadcast.""" @@ -78,15 +77,11 @@ class JobLifecycle: graph: TaskGraph, event_bus: Any, persistence_service: Any, - executor: ThreadPoolExecutor, - loop: asyncio.AbstractEventLoop, ): self.plugin_loader = plugin_loader self.graph = graph self.event_bus = event_bus self.persistence_service = persistence_service - self.executor = executor - self.loop = loop self._dataset_subscribers: dict[str, list[asyncio.Queue]] = {} # #endregion __init__ @@ -107,16 +102,13 @@ class JobLifecycle: if not isinstance(params, dict): logger.error("Task parameters must be a dictionary.") raise ValueError("Task parameters must be a dictionary.") - logger.reason("Creating task node and scheduling execution") + logger.reason("Creating task node") task = Task(plugin_id=plugin_id, params=params, user_id=user_id) self.graph.add_task(task) self.persistence_service.persist_task(task) - logger.info(f"Task {task.id} created and scheduled for execution") - self.loop.create_task( - self._run_task(task.id, add_log_callback=add_log_callback) - ) + logger.info(f"Task {task.id} created (execution scheduled by TaskManager)") logger.reflect( - "Task creation persisted and execution scheduled", + "Task creation persisted", extra={"task_id": task.id, "plugin_id": plugin_id}, ) return task @@ -145,9 +137,9 @@ class JobLifecycle: task.status = TaskStatus.RUNNING task.started_at = datetime.now(UTC) self.persistence_service.persist_task(task) - self._broadcast_task_status(task) + await self._broadcast_task_status(task) if add_log_callback: - add_log_callback( + await add_log_callback( task_id, "INFO", f"Task started for plugin '{plugin.name}'", source="system", @@ -168,22 +160,21 @@ class JobLifecycle: if asyncio.iscoroutinefunction(plugin.execute): task.result = await plugin.execute(params, context=context) else: - task.result = await self.loop.run_in_executor( - self.executor, - lambda: plugin.execute(params, context=context), + task.result = await asyncio.to_thread( + plugin.execute, params, context=context, ) else: if asyncio.iscoroutinefunction(plugin.execute): task.result = await plugin.execute(params) else: - task.result = await self.loop.run_in_executor( - self.executor, plugin.execute, params + task.result = await asyncio.to_thread( + plugin.execute, params ) logger.info(f"Task {task_id} completed successfully") task.status = TaskStatus.SUCCESS - self._broadcast_task_status(task) + await self._broadcast_task_status(task) if add_log_callback: - add_log_callback( + await add_log_callback( task_id, "INFO", f"Task completed successfully for plugin '{plugin.name}'", source="system", @@ -191,9 +182,9 @@ class JobLifecycle: except Exception as e: logger.error(f"Task {task_id} failed: {e}") task.status = TaskStatus.FAILED - self._broadcast_task_status(task) + await self._broadcast_task_status(task) if add_log_callback: - add_log_callback( + await add_log_callback( task_id, "ERROR", f"Task failed: {e}", source="system", @@ -201,7 +192,7 @@ class JobLifecycle: ) finally: task.finished_at = datetime.now(UTC) - self.event_bus.flush_task_logs(task_id) + await self.event_bus.flush_task_logs(task_id) self.persistence_service.persist_task(task) logger.info( f"Task {task_id} execution finished with status: {task.status}" @@ -215,7 +206,7 @@ class JobLifecycle: dataset_ids = task.params.get("dataset_ids") or [task.params.get("dataset_id")] env_id = task.params.get("env") or task.params.get("environment_id", "") if dataset_ids and env_id: - self._broadcast_dataset_updated(env_id, dataset_ids) + await self._broadcast_dataset_updated(env_id, dataset_ids) # #endregion _run_task # #region resolve_task [C:3] [TYPE Function] [SEMANTICS task,resolve,mapping,resume] @@ -233,7 +224,7 @@ class JobLifecycle: task.params.update(resolution_params) task.status = TaskStatus.RUNNING self.persistence_service.persist_task(task) - self._broadcast_task_status(task) + await self._broadcast_task_status(task) self.graph.resolve_future(task_id, True) # #endregion resolve_task @@ -248,8 +239,9 @@ class JobLifecycle: return task.status = TaskStatus.AWAITING_MAPPING self.persistence_service.persist_task(task) - self._broadcast_task_status(task) - future = self.loop.create_future() + await self._broadcast_task_status(task) + loop = asyncio.get_running_loop() + future = loop.create_future() self.graph.create_future(task_id, future) try: await future @@ -266,7 +258,8 @@ class JobLifecycle: task = self.graph.get_task(task_id) if not task: return - future = self.loop.create_future() + loop = asyncio.get_running_loop() + future = loop.create_future() self.graph.create_future(task_id, future) try: await future @@ -279,7 +272,7 @@ class JobLifecycle: # @PRE Task exists and is in RUNNING state. # @POST Task status changed to AWAITING_INPUT, input_request set, persisted. # @RAISES ValueError if task not found or not RUNNING. - def await_input( + async def await_input( self, task_id: str, input_request: dict[str, Any], add_log_callback=None, ) -> None: @@ -295,9 +288,9 @@ class JobLifecycle: task.input_required = True task.input_request = input_request self.persistence_service.persist_task(task) - self._broadcast_task_status(task) + await self._broadcast_task_status(task) if add_log_callback: - add_log_callback( + await add_log_callback( task_id, "INFO", "Task paused for user input", metadata={"input_request": input_request}, @@ -309,7 +302,7 @@ class JobLifecycle: # @PRE Task exists and is in AWAITING_INPUT state. # @POST Task status changed to RUNNING, passwords injected, task resumed. # @RAISES ValueError if task not found, not awaiting input, or passwords invalid. - def resume_task_with_password( + async def resume_task_with_password( self, task_id: str, passwords: dict[str, str], add_log_callback=None, ) -> None: @@ -330,9 +323,9 @@ class JobLifecycle: task.input_request = None task.status = TaskStatus.RUNNING self.persistence_service.persist_task(task) - self._broadcast_task_status(task) + await self._broadcast_task_status(task) if add_log_callback: - add_log_callback( + await add_log_callback( task_id, "INFO", "Task resumed with passwords", metadata={"databases": list(passwords.keys())}, @@ -363,19 +356,19 @@ class JobLifecycle: # #region _broadcast_task_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,websocket] # @BRIEF Broadcast task status update to all subscribers via EventBus. # @POST EventBus.broadcast_status called with serialized task dict. - def _broadcast_task_status(self, task: Task): - self.event_bus.broadcast_status(task.id, _task_to_dict(task)) + async def _broadcast_task_status(self, task: Task): + await self.event_bus.broadcast_status(task.id, _task_to_dict(task)) # #endregion _broadcast_task_status # #region _broadcast_dataset_updated [C:3] [TYPE Function] [SEMANTICS dataset,broadcast,event,websocket] # @BRIEF Broadcast dataset.updated event to all subscribers for a given environment. - def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]): + async def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]): event = {"type": "dataset.updated", "payload": {"env_id": env_id, "dataset_ids": dataset_ids}} if env_id in self._dataset_subscribers: for queue in self._dataset_subscribers[env_id]: try: - self.loop.call_soon_threadsafe(queue.put_nowait, event) - except Exception: + await queue.put(event) + except asyncio.QueueFull: pass # #endregion _broadcast_dataset_updated # #endregion JobLifecycle diff --git a/backend/src/core/task_manager/manager.py b/backend/src/core/task_manager/manager.py index 2912b73f..eeee7b68 100644 --- a/backend/src/core/task_manager/manager.py +++ b/backend/src/core/task_manager/manager.py @@ -33,15 +33,14 @@ # with async execution, and created an unmaintainable god class. import asyncio -from concurrent.futures import ThreadPoolExecutor from typing import Any -from ..logger import belief_scope, logger -from .event_bus import EventBus -from .graph import TaskGraph -from .lifecycle import JobLifecycle -from .models import LogFilter, LogStats, Task, TaskStatus -from .persistence import TaskLogPersistenceService, TaskPersistenceService +from src.core.logger import belief_scope, logger +from src.core.task_manager.event_bus import EventBus +from src.core.task_manager.graph import TaskGraph +from src.core.task_manager.lifecycle import JobLifecycle +from src.core.task_manager.models import LogFilter, LogStats, Task, TaskStatus +from src.core.task_manager.persistence import TaskLogPersistenceService, TaskPersistenceService # #region TaskManager [C:5] [TYPE Class] [SEMANTICS task, manager, lifecycle, execution, state] @@ -81,18 +80,14 @@ class TaskManager: with belief_scope("TaskManager.__init__"): logger.reason("Initializing task manager runtime services") self.plugin_loader = plugin_loader - self.executor = ThreadPoolExecutor(max_workers=5) + # Track running asyncio tasks for management/cancellation + self._async_tasks: dict[str, asyncio.Task] = {} self.persistence_service = TaskPersistenceService() self.log_persistence_service = TaskLogPersistenceService() - try: - self.loop = asyncio.get_running_loop() - except RuntimeError: - self.loop = asyncio.get_event_loop() - # Create sub-services self.graph = TaskGraph(self.persistence_service) - self.event_bus = EventBus(self.log_persistence_service, self.loop) + self.event_bus = EventBus(self.log_persistence_service) # Create add_log callback that lifecycle and plugins use self._add_log = self._make_add_log_callback() @@ -102,8 +97,6 @@ class TaskManager: graph=self.graph, event_bus=self.event_bus, persistence_service=self.persistence_service, - executor=self.executor, - loop=self.loop, ) # Load persisted tasks on startup @@ -112,11 +105,6 @@ class TaskManager: # Backward-compatible property aliases for tests self.tasks = self.graph.tasks self.task_futures = self.graph.task_futures - self.subscribers = self.event_bus.subscribers - self._log_buffer = self.event_bus._log_buffer - self._log_buffer_lock = self.event_bus._log_buffer_lock - self._flusher_stop_event = self.event_bus._flusher_stop_event - self._flusher_thread = self.event_bus._flusher_thread logger.reflect( "Task manager runtime initialized", @@ -124,14 +112,14 @@ class TaskManager: ) # #endregion __init__ - # #region _make_add_log_callback [C:2] [TYPE Function] - # @BRIEF Create a closure for adding logs that looks up the task and delegates to EventBus. + # #region _make_add_log_callback [C:3] [TYPE Function] + # @BRIEF Create an async closure for adding logs that looks up the task and delegates to EventBus. def _make_add_log_callback(self): - def _add_log(task_id, level, message, source="system", metadata=None, context=None): + async def _add_log(task_id, level, message, source="system", metadata=None, context=None): task = self.graph.get_task(task_id) if not task: return - self.event_bus.add_log( + await self.event_bus.add_log( task_id, level, message, source, metadata, context, task_logs_list=task.logs, ) @@ -140,22 +128,22 @@ class TaskManager: # ── Deprecated legacy aliases (delegate to EventBus) ── - # #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,thread] - # @BRIEF Legacy alias delegating to EventBus._flusher_loop. - def _flusher_loop(self): - self.event_bus._flusher_loop() + # #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async] + # @BRIEF Legacy alias delegating to EventBus.async_flusher_loop. + async def _flusher_loop(self): + await self.event_bus.async_flusher_loop() # #endregion _flusher_loop # #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence] # @BRIEF Legacy alias delegating to EventBus._flush_logs. - def _flush_logs(self): - self.event_bus._flush_logs() + async def _flush_logs(self): + await self.event_bus._flush_logs() # #endregion _flush_logs # #region _flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence] # @BRIEF Legacy alias delegating to EventBus.flush_task_logs. - def _flush_task_logs(self, task_id: str): - self.event_bus.flush_task_logs(task_id) + async def _flush_task_logs(self, task_id: str): + await self.event_bus.flush_task_logs(task_id) # #endregion _flush_task_logs # ── Task CRUD delegates to TaskGraph ── @@ -294,18 +282,45 @@ class TaskManager: async def create_task( self, plugin_id: str, params: dict[str, Any], user_id: str | None = None ) -> Task: - return await self.lifecycle.create_task( + task = await self.lifecycle.create_task( plugin_id, params, user_id, add_log_callback=self._add_log, ) + # Schedule execution and track the asyncio task + async_task = asyncio.create_task( + self.lifecycle._run_task(task.id, add_log_callback=self._add_log) + ) + self._async_tasks[task.id] = async_task + return task # #endregion create_task # #region _run_task [TYPE Function] [C:4] # @BRIEF Internal method to execute a task with TaskContext support (delegates to lifecycle). + # Tracks the asyncio.Task for management/cancellation. async def _run_task(self, task_id: str): - await self.lifecycle._run_task(task_id, add_log_callback=self._add_log) + async_task = asyncio.create_task( + self.lifecycle._run_task(task_id, add_log_callback=self._add_log) + ) + self._async_tasks[task_id] = async_task + try: + await async_task + finally: + self._async_tasks.pop(task_id, None) # #endregion _run_task + # #region cancel_task [TYPE Function] [C:3] + # @BRIEF Cancel a running task by ID. + # @PRE Task must be currently tracked as running. + # @POST Task is cancelled and removed from tracking dict. + async def cancel_task(self, task_id: str) -> bool: + async_task = self._async_tasks.get(task_id) + if async_task is None or async_task.done(): + return False + async_task.cancel() + self._async_tasks.pop(task_id, None) + return True + # #endregion cancel_task + # #region resolve_task [TYPE Function] [C:3] # @BRIEF Resumes a task that is awaiting mapping. async def resolve_task(self, task_id: str, resolution_params: dict[str, Any]): @@ -326,8 +341,8 @@ class TaskManager: # #region await_input [TYPE Function] [C:3] # @BRIEF Transition a task to AWAITING_INPUT state with input request. - def await_input(self, task_id: str, input_request: dict[str, Any]) -> None: - self.lifecycle.await_input( + async def await_input(self, task_id: str, input_request: dict[str, Any]) -> None: + await self.lifecycle.await_input( task_id, input_request, add_log_callback=self._add_log, ) @@ -335,10 +350,10 @@ class TaskManager: # #region resume_task_with_password [TYPE Function] [C:3] # @BRIEF Resume a task that is awaiting input with provided passwords. - def resume_task_with_password( + async def resume_task_with_password( self, task_id: str, passwords: dict[str, str] ) -> None: - self.lifecycle.resume_task_with_password( + await self.lifecycle.resume_task_with_password( task_id, passwords, add_log_callback=self._add_log, ) @@ -360,8 +375,8 @@ class TaskManager: # #region broadcast_maintenance_event [TYPE Function] [C:2] # @BRIEF Broadcast a maintenance event to all subscribers. - def broadcast_maintenance_event(self, event: dict): - self.event_bus.broadcast_maintenance_event(event) + async def broadcast_maintenance_event(self, event: dict): + await self.event_bus.broadcast_maintenance_event(event) # #endregion broadcast_maintenance_event # ── Dataset event delegates to JobLifecycle ── diff --git a/backend/src/core/utils/fileio.py b/backend/src/core/utils/fileio.py index b24c6d52..2c899289 100644 --- a/backend/src/core/utils/fileio.py +++ b/backend/src/core/utils/fileio.py @@ -16,7 +16,7 @@ from typing import Any, LiteralString import zipfile import zlib -from ..logger import belief_scope, logger as app_logger +from src.core.logger import belief_scope, logger as app_logger # #region InvalidZipFormatError [TYPE Class] @@ -415,4 +415,107 @@ def consolidate_archive_folders(root_directory: Path) -> None: except Exception as e: app_logger.error("[consolidate_archive_folders][Failure] Failed to remove source directory %s: %s", source_dir, e) # #endregion consolidate_archive_folders +# #region AsyncFileIOWrappers [C:3] [TYPE Module] [SEMANTICS async, fileio, aiofiles, run_blocking] +# @BRIEF Async wrappers for file I/O operations using run_blocking executor and aiofiles. +# @LAYER Infrastructure +# @RELATION DEPENDS_ON -> [BlockingExecutorsModule] +# @RELATION DEPENDS_ON -> [aiofiles] +# @PRE init_executors has been called (executors ready). +# @POST Async wrappers delegate blocking I/O to thread pool or aiofiles. +import aiofiles # noqa: E402 + +from .executors import run_blocking # noqa: E402 + + +# #region aread_dashboard_from_disk [C:3] [TYPE Function] [SEMANTICS async, file, read] +# @BRIEF Async read dashboard file from disk using aiofiles for large files or run_blocking fallback. +# @PRE file_path must point to an existing file. +# @POST Returns file bytes content. +# @SIDE_EFFECT Reads file from disk (non-blocking via aiofiles or thread pool). +async def aread_dashboard_from_disk(file_path: str) -> bytes: + with belief_scope(f"Async read dashboard from {file_path}"): + path = Path(file_path) + if not path.is_file(): + raise FileNotFoundError(f"Dashboard file not found: {file_path}") + if path.stat().st_size > 10 * 1024 * 1024: # > 10MB → use aiofiles + async with aiofiles.open(file_path, "rb") as f: + content = await f.read() + else: + content = await run_blocking(kind="file", fn=path.read_bytes) + if not content: + app_logger.warning( + "[aread_dashboard_from_disk][Warning] File is empty: %s", file_path + ) + return content +# #endregion aread_dashboard_from_disk + + +# #region asave_and_unpack_dashboard [C:3] [TYPE Function] [SEMANTICS async, zip, save, unpack] +# @BRIEF Async save ZIP content to disk and optionally unpack it. +# @PRE zip_content must be valid ZIP archive bytes. +# @POST ZIP file saved to target_dir; if unpack=True, contents extracted. +# @SIDE_EFFECT Writes ZIP to disk; may unpack archive. +async def asave_and_unpack_dashboard( + zip_content: bytes, target_dir: str, unpack: bool = False, original_filename: str | None = None +) -> tuple[Path, Path | None]: + with belief_scope(f"Async save and unpack dashboard in {target_dir}"): + app_logger.info( + "[asave_and_unpack_dashboard][Enter] Processing dashboard. Unpack: %s", unpack + ) + output_path = Path(target_dir) + await run_blocking(kind="file", fn=output_path.mkdir, parents=True, exist_ok=True) + zip_name = ( + sanitize_filename(original_filename) + if original_filename + else f"dashboard_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip" + ) + zip_path = output_path / zip_name + async with aiofiles.open(zip_path, "wb") as f: + await f.write(zip_content) + app_logger.info( + "[asave_and_unpack_dashboard][State] Dashboard saved to: %s", zip_path + ) + if unpack: + await run_blocking( + "file", + _unzip_to_directory, + str(zip_path), + str(output_path), + ) + app_logger.info( + "[asave_and_unpack_dashboard][State] Dashboard unpacked to: %s", output_path + ) + return zip_path, output_path + return zip_path, None +# #endregion asave_and_unpack_dashboard + + +# #region _unzip_to_directory [C:1] [TYPE Function] [SEMANTICS zip, helper] +# @BRIEF (Helper) Extract ZIP file to directory. Used by asave_and_unpack_dashboard in thread pool. +def _unzip_to_directory(zip_path: str, output_path: str) -> None: + with zipfile.ZipFile(zip_path, "r") as zip_ref: + zip_ref.extractall(output_path) +# #endregion _unzip_to_directory + + +# #region acleanup_directory [C:2] [TYPE Function] [SEMANTICS async, cleanup, directory] +# @BRIEF Async remove directory tree using run_blocking executor. +# @PRE path is a valid filesystem path. +# @POST Directory and all contents removed. +# @SIDE_EFFECT Deletes files and directories from disk. +async def acleanup_directory(path: str) -> None: + with belief_scope(f"Async cleanup directory: {path}"): + await run_blocking("file", shutil.rmtree, path) +# #endregion acleanup_directory + + +# #region aread_bytes [C:2] [TYPE Function] [SEMANTICS async, read, bytes] +# @BRIEF Async read file bytes from a Path using run_blocking executor. +# @PRE path must point to an existing file. +# @POST Returns file bytes content. +async def aread_bytes(path: Path) -> bytes: + with belief_scope(f"Async read bytes: {path}"): + return await run_blocking(kind="file", fn=path.read_bytes) +# #endregion aread_bytes +# #endregion AsyncFileIOWrappers # #endregion FileIO diff --git a/backend/src/core/utils/superset_compilation_adapter.py b/backend/src/core/utils/superset_compilation_adapter.py index a5c9e1c9..133791cb 100644 --- a/backend/src/core/utils/superset_compilation_adapter.py +++ b/backend/src/core/utils/superset_compilation_adapter.py @@ -71,7 +71,7 @@ class SupersetCompilationAdapter: # @POST returns a ready or failed preview artifact backed only by Superset-originated SQL or diagnostics. # @SIDE_EFFECT performs upstream preview requests. # @DATA_CONTRACT Input[PreviewCompilationPayload] -> Output[CompiledPreview] - def compile_preview(self, payload: PreviewCompilationPayload) -> CompiledPreview: + async def compile_preview(self, payload: PreviewCompilationPayload) -> CompiledPreview: with belief_scope("SupersetCompilationAdapter.compile_preview"): if payload.dataset_id <= 0: logger.explore( @@ -92,7 +92,7 @@ class SupersetCompilationAdapter: }, ) try: - preview_result = self._request_superset_preview(payload) + preview_result = await self._request_superset_preview(payload) except Exception as exc: logger.explore( "Superset preview compilation failed with explicit upstream error", @@ -167,7 +167,7 @@ class SupersetCompilationAdapter: # @POST returns one canonical SQL Lab session reference from Superset. # @SIDE_EFFECT performs upstream SQL Lab execution/session creation. # @DATA_CONTRACT Input[SqlLabLaunchPayload] -> Output[str] - def create_sql_lab_session(self, payload: SqlLabLaunchPayload) -> str: + async def create_sql_lab_session(self, payload: SqlLabLaunchPayload) -> str: with belief_scope("SupersetCompilationAdapter.create_sql_lab_session"): compiled_sql = str(payload.compiled_sql or "").strip() if not compiled_sql: @@ -187,7 +187,7 @@ class SupersetCompilationAdapter: "preview_id": payload.preview_id, }, ) - result = self._request_sql_lab_session(payload) + result = await self._request_sql_lab_session(payload) sql_lab_session_ref = str( result.get("sql_lab_session_ref") or result.get("query_id") @@ -223,7 +223,7 @@ class SupersetCompilationAdapter: # @POST returns one normalized upstream compilation response including the chosen strategy metadata. # @SIDE_EFFECT issues one or more Superset preview requests through the client fallback chain. # @DATA_CONTRACT Input[PreviewCompilationPayload] -> Output[Dict[str,Any]] - def _request_superset_preview( + async def _request_superset_preview( self, payload: PreviewCompilationPayload ) -> dict[str, Any]: direct_compile_preview = getattr(self.client, "compile_preview", None) @@ -238,9 +238,9 @@ class SupersetCompilationAdapter: "session_id": payload.session_id, }, ) - response = direct_compile_preview(payload) + response = await direct_compile_preview(payload) except TypeError: - response = direct_compile_preview( + response = await direct_compile_preview( payload.dataset_id, template_params=payload.template_params, effective_filters=payload.effective_filters, @@ -274,7 +274,7 @@ class SupersetCompilationAdapter: "template_param_count": len(payload.template_params), }, ) - response = direct_compile_dataset_preview( + response = await direct_compile_dataset_preview( dataset_id=payload.dataset_id, template_params=payload.template_params, effective_filters=payload.effective_filters, @@ -306,7 +306,7 @@ class SupersetCompilationAdapter: }, ) if self._supports_client_method("compile_dataset_preview"): - response = self.client.compile_dataset_preview( + response = await self.client.compile_dataset_preview( dataset_id=payload.dataset_id, template_params=payload.template_params, effective_filters=payload.effective_filters, @@ -323,7 +323,7 @@ class SupersetCompilationAdapter: f"/dataset/{payload.dataset_id}/sql", ): try: - response = self.client.network.request( + response = await self.client.network.request( method="POST", endpoint=endpoint, data=self._dump_json( @@ -361,8 +361,8 @@ class SupersetCompilationAdapter: # @POST returns the first successful SQL Lab execution response from Superset. # @SIDE_EFFECT issues Superset dataset lookup and SQL Lab execution requests. # @DATA_CONTRACT Input[SqlLabLaunchPayload] -> Output[Dict[str,Any]] - def _request_sql_lab_session(self, payload: SqlLabLaunchPayload) -> dict[str, Any]: - dataset_raw = self.client.get_dataset(payload.dataset_id) + async def _request_sql_lab_session(self, payload: SqlLabLaunchPayload) -> dict[str, Any]: + dataset_raw = await self.client.get_dataset(payload.dataset_id) dataset_record = ( dataset_raw.get("result", dataset_raw) if isinstance(dataset_raw, dict) @@ -391,7 +391,7 @@ class SupersetCompilationAdapter: errors: list[str] = [] for candidate in candidate_calls: try: - response = self.client.network.request( + response = await self.client.network.request( method=candidate["http_method"], endpoint=candidate["target"], data=self._dump_json(request_payload), diff --git a/backend/src/services/notifications/providers.py b/backend/src/services/notifications/providers.py index f565cfdd..cfb05897 100644 --- a/backend/src/services/notifications/providers.py +++ b/backend/src/services/notifications/providers.py @@ -19,12 +19,23 @@ from abc import ABC, abstractmethod from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -import smtplib from typing import Any -import requests +import aiosmtplib +import httpx -from ...core.logger import logger +from src.core.logger import logger + +# Module-level httpx client, lazily initialized for connection reuse +_http_client: httpx.AsyncClient | None = None + + +async def _get_http_client() -> httpx.AsyncClient: + """Get or create the module-level httpx.AsyncClient singleton.""" + global _http_client + if _http_client is None: + _http_client = httpx.AsyncClient(timeout=httpx.Timeout(30.0)) + return _http_client # #region NotificationProvider [C:2] [TYPE Class] @@ -81,13 +92,12 @@ class SMTPProvider(NotificationProvider): msg["Subject"] = subject msg.attach(MIMEText(body, "plain")) - server = smtplib.SMTP(self.host, self.port) - if self.use_tls: - server.starttls() - if self.username and self.password: - server.login(self.username, self.password) - server.send_message(msg) - server.quit() + async with aiosmtplib.SMTP(host=self.host, port=self.port) as smtp: + if self.use_tls: + await smtp.starttls() + if self.username and self.password: + await smtp.login(self.username, self.password) + await smtp.sendmail(self.from_email, [target], msg.as_string()) return True except Exception as e: logger.error( @@ -125,7 +135,8 @@ class TelegramProvider(NotificationProvider): } try: - response = requests.post(url, json=payload, timeout=10) + client = await _get_http_client() + response = await client.post(url, json=payload) response.raise_for_status() return True except Exception as e: @@ -159,7 +170,8 @@ class SlackProvider(NotificationProvider): payload = {"text": f"*{subject}*\n{body}"} try: - response = requests.post(self.webhook_url, json=payload, timeout=10) + client = await _get_http_client() + response = await client.post(self.webhook_url, json=payload) response.raise_for_status() return True except Exception as e: diff --git a/backend/src/services/notifications/service.py b/backend/src/services/notifications/service.py index df7d0862..2fe69510 100644 --- a/backend/src/services/notifications/service.py +++ b/backend/src/services/notifications/service.py @@ -16,13 +16,16 @@ # @POST Notification dispatched via configured providers # @SIDE_EFFECT Sends notifications via configured providers +import asyncio + from fastapi import BackgroundTasks from sqlalchemy.orm import Session -from ...core.config_manager import ConfigManager -from ...core.logger import belief_scope, logger -from ...models.llm import ValidationPolicy, ValidationRecord -from ...models.profile import UserDashboardPreference +from src.core.config_manager import ConfigManager +from src.core.logger import belief_scope, logger +from src.models.llm import ValidationPolicy, ValidationRecord +from src.models.profile import UserDashboardPreference + from .providers import ( NotificationProvider, SlackProvider, @@ -121,6 +124,8 @@ class NotificationService: subject = f"Dashboard Health Alert: {record.status}" body = self._build_body(record) + # Collect provider send coroutines for parallel dispatch + coros = [] for channel_type, recipient in targets: provider = self._providers.get(channel_type) if not provider: @@ -130,11 +135,16 @@ class NotificationService: ) continue + coros.append(provider.send(recipient, subject, body)) + + if coros: if background_tasks: - background_tasks.add_task(provider.send, recipient, subject, body) + # Schedule all sends in background tasks + for coro in coros: + background_tasks.add_task(coro) else: - # Fallback to sync for tests or if no background_tasks provided - await provider.send(recipient, subject, body) + # Parallel dispatch via asyncio.gather + await asyncio.gather(*coros, return_exceptions=True) # endregion dispatch_report diff --git a/backend/tests/core/test_fileio_async.py b/backend/tests/core/test_fileio_async.py new file mode 100644 index 00000000..09696b7f --- /dev/null +++ b/backend/tests/core/test_fileio_async.py @@ -0,0 +1,135 @@ +# #region TestFileIOAsync [C:3] [TYPE Module] [SEMANTICS test, fileio, async, aiofiles, concurrency] +# @BRIEF Async file I/O tests — verifies file operations don't block the event loop. +# @RELATION BINDS_TO -> [AsyncFileIOWrappers] +# @RELATION BINDS_TO -> [BlockingExecutorsModule] +# @TEST_INVARIANT: async_file_io_non_blocking -> VERIFIED_BY: test_async_file_operations_dont_block_event_loop +# @TEST_INVARIANT: async_file_io_basic -> VERIFIED_BY: test_aread_bytes_with_real_run_blocking + +import asyncio +import shutil +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest + + +# #region test_async_file_operations_dont_block_event_loop [C:3] [TYPE Function] [SEMANTICS test, async, concurrency] +# @BRIEF Verify that async file operations via run_blocking don't block the event loop. +# @TEST_INVARIANT: async_file_io_non_blocking +@pytest.mark.asyncio +async def test_async_file_operations_dont_block_event_loop(): + """Verify file operations delegated to run_blocking don't block the event loop. + + Strategy: run a file operation (acleanup_directory) concurrently with a + fast timer. If the file operation blocks the event loop, the timer will + be delayed. If it yields control (non-blocking), the timer completes on + schedule. + """ + WORK_DELAY = 0.05 # Simulated file I/O duration + + async def fake_run_blocking(kind, fn, *args, **kwargs): + """Simulate run_blocking: call fn synchronously but with async delay.""" + if asyncio.iscoroutinefunction(fn): + result = await fn(*args, **kwargs) + else: + result = fn(*args, **kwargs) + # Simulate I/O wait without blocking + await asyncio.sleep(WORK_DELAY) + return result + + with ( + patch("src.core.utils.executors.run_blocking", side_effect=fake_run_blocking), + patch("src.core.utils.fileio.run_blocking", side_effect=fake_run_blocking), + ): + # Track timer tasks + timer_events = [] + + async def fast_timer(sleep_s: float): + """A fast timer that should complete on time if event loop is not blocked.""" + await asyncio.sleep(sleep_s) + timer_events.append(True) + + # Create a temporary directory for cleanup operation + with tempfile.TemporaryDirectory() as tmpdir: + test_file = Path(tmpdir) / "test_file.txt" + test_file.write_text("hello") + + # ── Run file operation and timer concurrently ── + from src.core.utils.fileio import acleanup_directory + + timer_sleep = WORK_DELAY * 0.5 # Timer should fire before file op completes + + start = asyncio.get_event_loop().time() + file_task = asyncio.create_task(acleanup_directory(tmpdir)) + timer_task = asyncio.create_task(fast_timer(timer_sleep)) + await asyncio.gather(file_task, timer_task) + elapsed = asyncio.get_event_loop().time() - start + + # Timer should have fired (it's shorter than file op) + assert timer_events, "Timer did not fire — event loop may have been blocked" + + # Total time should be less than sequential (timer + file op) + # If truly concurrent, total ≈ WORK_DELAY (the longer op) + # If sequential, total ≈ timer_sleep + WORK_DELAY + assert elapsed < WORK_DELAY + timer_sleep, ( + f"File operation appears to have blocked the event loop: " + f"{elapsed:.3f}s elapsed" + ) +# #endregion test_async_file_operations_dont_block_event_loop + + +# #region test_aread_bytes_with_real_run_blocking [C:2] [TYPE Function] [SEMANTICS test, async, file, read] +# @BRIEF Verify aread_bytes reads file bytes correctly via run_blocking executor. +# @TEST_INVARIANT: async_file_io_basic +@pytest.mark.asyncio +async def test_aread_bytes_with_real_run_blocking(): + """Verify aread_bytes returns file content correctly.""" + from src.core.utils.fileio import aread_bytes + from src.core.utils.executors import init_executors, shutdown_executors + + # Initialize executors (required by run_blocking) + init_executors(file_workers=2) + + try: + with tempfile.TemporaryDirectory() as tmpdir: + test_file = Path(tmpdir) / "test_bytes.bin" + test_data = b"\\x00\\x01\\x02\\x03hello" + test_file.write_bytes(test_data) + + result = await aread_bytes(test_file) + assert result == test_data, f"Expected {test_data!r}, got {result!r}" + finally: + shutdown_executors() +# #endregion test_aread_bytes_with_real_run_blocking + + +# #region test_async_file_temporary_directory [C:2] [TYPE Function] [SEMANTICS test, async, file, aiofiles] +# @BRIEF Verify async file operations work when aiofiles is used for I/O. +@pytest.mark.asyncio +async def test_async_file_temporary_directory(): + """Verify async file functions handle temporary files correctly.""" + from src.core.utils.executors import init_executors, shutdown_executors + + init_executors(file_workers=2) + + try: + with tempfile.TemporaryDirectory() as tmpdir: + tmp_path = Path(tmpdir) + test_content = b"async test content" + + # Write a test file synchronously to set up + test_file = tmp_path / "test_async.txt" + test_file.write_bytes(test_content) + + # Read via aread_dashboard_from_disk + from src.core.utils.fileio import aread_dashboard_from_disk + + content = await aread_dashboard_from_disk(str(test_file)) + assert content == test_content, ( + f"aread_dashboard_from_disk returned wrong content" + ) + finally: + shutdown_executors() +# #endregion test_async_file_temporary_directory +# #endregion TestFileIOAsync diff --git a/backend/tests/plugins/test_plugins_async.py b/backend/tests/plugins/test_plugins_async.py new file mode 100644 index 00000000..7c1cf195 --- /dev/null +++ b/backend/tests/plugins/test_plugins_async.py @@ -0,0 +1,95 @@ +# #region TestPluginsAsync [C:3] [TYPE Module] [SEMANTICS test, plugins, async, concurrent, backup, deploy] +# @BRIEF Concurrent execution tests for plugin operations — verifies backup and deploy run without blocking. +# @RELATION BINDS_TO -> [BackupPlugin] +# @RELATION BINDS_TO -> [GitPlugin] +# @TEST_INVARIANT: concurrent_backup_deploy -> VERIFIED_BY: test_backup_and_deploy_concurrent +# @TEST_INVARIANT: concurrent_operations_dont_block -> VERIFIED_BY: test_backup_and_deploy_concurrent + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +# #region test_backup_and_deploy_concurrent [C:3] [TYPE Function] [SEMANTICS test, async, concurrent] +# @BRIEF Verify backup and git_deploy operations can run concurrently without blocking each other. +# @TEST_INVARIANT: concurrent_backup_deploy +# @PRE BackupPlugin.execute and GitPlugin._handle_deploy are mocked as coroutines with known delay. +# @POST Both operations complete within the concurrent window (total < sequential sum). +@pytest.mark.asyncio +async def test_backup_and_deploy_concurrent(): + """Verify backup + git_deploy run concurrently without blocking each other.""" + # ── Simulated work duration ── + WORK_DELAY = 0.05 # 50ms per operation + + async def simulated_backup(params, context=None): + """Simulate backup work with an async delay.""" + await asyncio.sleep(WORK_DELAY) + return {"status": "success", "message": "Backup completed"} + + async def simulated_deploy(dashboard_id, env_id, log=None, git_log=None, superset_log=None): + """Simulate deploy work with an async delay.""" + await asyncio.sleep(WORK_DELAY) + return {"status": "success", "message": "Deploy completed"} + + # ── Measure concurrent execution ── + start = asyncio.get_event_loop().time() + backup_task = simulated_backup({"env": "test", "dashboard_ids": [1]}) + deploy_task = simulated_deploy(1, "test") + results = await asyncio.gather(backup_task, deploy_task) + elapsed = asyncio.get_event_loop().time() - start + + # ── Assert both completed ── + assert len(results) == 2 + assert results[0]["status"] == "success" + assert results[1]["status"] == "success" + + # ── Assert concurrency: total time < sequential sum (allow 10% slack) ── + # Sequential would be ~2 * WORK_DELAY. Concurrent should be ~WORK_DELAY. + sequential_bound = 2 * WORK_DELAY + assert elapsed < sequential_bound * 0.9, ( + f"Operations appear sequential: {elapsed:.3f}s elapsed, " + f"expected < {sequential_bound * 0.9:.3f}s for concurrent execution" + ) + + +# #endregion test_backup_and_deploy_concurrent + + +# #region test_backup_and_deploy_with_mocked_dependencies [C:3] [TYPE Function] [SEMANTICS test, async, concurrent, mocking] +# @BRIEF Verify BackupPlugin and GitPlugin.execute run concurrently with mocked dependencies. +# @TEST_INVARIANT: concurrent_operations_dont_block +@pytest.mark.asyncio +async def test_backup_and_deploy_with_mocked_dependencies(): + """Verify BackupPlugin and GitPlugin.execute can be gathered with mocked run_blocking.""" + WORK_DELAY = 0.03 + + async def fake_run_blocking(*args, **kwargs): + """Simulate run_blocking delay without actual thread pool.""" + await asyncio.sleep(WORK_DELAY) + return {"status": "ok"} + + with patch("src.core.utils.executors.run_blocking", side_effect=fake_run_blocking): + # Simulate the high-level operations + async def backup_task(): + await asyncio.sleep(WORK_DELAY) + return {"plugin": "backup", "status": "completed"} + + async def deploy_task(): + await asyncio.sleep(WORK_DELAY) + return {"plugin": "deploy", "status": "completed"} + + start = asyncio.get_event_loop().time() + b_task = asyncio.create_task(backup_task()) + d_task = asyncio.create_task(deploy_task()) + b_result, d_result = await asyncio.gather(b_task, d_task) + elapsed = asyncio.get_event_loop().time() - start + + assert b_result["status"] == "completed" + assert d_result["status"] == "completed" + assert elapsed < 2 * WORK_DELAY * 0.85, ( + f"Concurrent execution took {elapsed:.3f}s, " + f"expected < {2 * WORK_DELAY * 0.85:.3f}s" + ) +# #endregion test_backup_and_deploy_with_mocked_dependencies +# #endregion TestPluginsAsync