1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
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
|
