diff options
Diffstat (limited to 'src/services/api/background.py')
| -rw-r--r-- | src/services/api/background.py | 179 |
1 files changed, 179 insertions, 0 deletions
diff --git a/src/services/api/background.py b/src/services/api/background.py new file mode 100644 index 000000000..2b25af307 --- /dev/null +++ b/src/services/api/background.py @@ -0,0 +1,179 @@ +# Copyright VyOS maintainers and contributors <maintainers@vyos.io> +# +# This library is free software; you can redistribute it and/or +# modify it under the terms of the GNU Lesser General Public +# License as published by the Free Software Foundation; either +# version 2.1 of the License, or (at your option) any later version. +# +# This library is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with this library. If not, see <http://www.gnu.org/licenses/>. + +import time +import functools +from collections import deque +from enum import Enum +from threading import Lock +from typing import Any +from typing import Callable +from typing import Optional +from uuid import uuid4 + +from fastapi import BackgroundTasks +from pydantic import BaseModel +from pydantic import StrictStr +from pydantic import StrictInt + + +def _ts(): + """Return current Unix timestamp (seconds since epoch)""" + return int(time.time()) + + +class BackgroundOpStatus(str, Enum): + queued = 'queued' + running = 'running' + succeeded = 'succeeded' + failed = 'failed' + + @property + def is_completed(self): + """True if the operation is in a terminal state (succeeded/failed)""" + return self in (BackgroundOpStatus.succeeded, BackgroundOpStatus.failed) + + +class BackgroundOpRecord(BaseModel): + """Metadata and outcome for a single background operation""" + + op_id: StrictStr + created_at: StrictInt + started_at: Optional[StrictInt] = None + finished_at: Optional[StrictInt] = None + status: BackgroundOpStatus = BackgroundOpStatus.queued + result: Optional[Any] = None + error: Optional[StrictStr] = None + + +class BackgroundOpError(Exception): + """Raised when a background operation cannot be enqueued/executed""" + + pass + + +class BackgroundOpManager: + """ + In-memory FIFO operation queue. + + Uses BackgroundTasks to schedule a `drain()` call after the response, + so `enqueue()` is fast and non-blocking for the client. + """ + + DEFAULT_MAX_QUEUE_SIZE = 128 + + def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE): + # max number of queued (pending) operations allowed at a time + self._max_queue_size = max_queue_size + + # FIFO queue of operation IDs waiting to be executed + self._queue = deque() + self._jobs = {} + self._workers = {} + + # protects _queue/_jobs/_workers/_drain_scheduled from concurrent access + self._mx = Lock() + + # whether a drain task has already been scheduled via BackgroundTasks + self._drain_scheduled = False + + def enqueue( + self, + background_tasks: BackgroundTasks, + func: Callable, + *args, + **kwargs, + ) -> BackgroundOpRecord: + """Enqueue a function for background execution and return its record""" + + assert isinstance(background_tasks, BackgroundTasks) + assert callable(func), '`func` argument should be function or lambda' + + with self._mx: + if len(self._queue) >= self._max_queue_size: + raise BackgroundOpError( + f'Background operation queue is full ({self._max_queue_size})' + ) + + op_id = str(uuid4()) + record = BackgroundOpRecord(op_id=op_id, created_at=_ts()) + + self._jobs[op_id] = record + # store the callable for later execution (outside the lock) + self._workers[op_id] = functools.partial(func, *args, **kwargs) + self._queue.append(op_id) + + if not self._drain_scheduled: + # schedule a single drain() call after the current response + background_tasks.add_task(self.drain) + self._drain_scheduled = True + + # Best-effort pruning: keep history bounded by dropping oldest completed records + if len(self._jobs) > self._max_queue_size: + oldest = min(self._jobs.values(), key=lambda record: record.created_at) + if oldest.status.is_completed: + del self._jobs[oldest.op_id] + + return record + + def drain(self): + """Run queued operations sequentially until the queue is empty""" + + while True: + with self._mx: + if not self._queue: + # allow future enqueue() calls to schedule the next drain() + self._drain_scheduled = False + return + + op_id = self._queue.popleft() + record = self._jobs[op_id] + func = self._workers.pop(op_id) + + record.status = BackgroundOpStatus.running + record.started_at = _ts() + + # execute outside the lock to avoid blocking enqueues/status reads + result = error = status = None + try: + result = func() + except Exception as e: # noqa: BLE001 + status = BackgroundOpStatus.failed + error = str(e) + else: + status = BackgroundOpStatus.succeeded + + with self._mx: + record.result = result + record.error = error + record.status = status + record.finished_at = _ts() + + def get_record(self, op_id: str) -> BackgroundOpRecord | None: + """Return a deep copy of a single record""" + + with self._mx: + record = self._jobs.get(op_id) + return record.copy(deep=True) if record else None + + def get_records(self) -> list: + """Return deep copies of all records, sorted oldest-first by created_at""" + + with self._mx: + records = [record.copy(deep=True) for record in self._jobs.values()] + + # stable-ish ordering (oldest first) + records.sort(key=lambda record: record.created_at) + return records |