1import asyncio as I, asyncutils as A
2from asyncutils._internal.py312 import Queue, QueueShutDown
3from asyncutils._internal.helpers import copy_and_clear, fullname, subscriptable
4from asyncutils._internal.submodules import processors_all as __all__
5from _functools import partial
6from time import monotonic
[docs]
7@subscriptable
8class BoundedBatchProcessor:
9 __slots__ = '__b', '__p', '__s'
10 def __init__(self, processor, batch=None, max_concurrent=None): C = A.getcontext(); self.__p, self.__b, self.__s = processor, C.BOUNDED_BATCH_PROCESSOR_DEFAULT_BATCH_SIZE if batch is None else batch, I.Semaphore(C.BOUNDED_BATCH_PROCESSOR_DEFAULT_MAX_CONCURRENT if max_concurrent is None else max_concurrent)
[docs]
11 async def process(self, items):
12 f, p, s = partial(A.collect, A.iter_to_agen(items), self.__b), self.__p, self.__s
13 while b := await f():
14 async with s: x = await p(b)
15 yield x # noqa: RUF070
[docs]
16@subscriptable
17class BatchProcessor(A.LoopContextMixin):
18 __slots__ = '__batch', '__lock', '__lp', '__ms', '__p', '__sleep', '__timer'
19 def __init__(self, processor, *, maxsize=None, maxtime=None, timer=monotonic): C = A.getcontext(); self.__p, self.__ms, self.__sleep, self.__batch, self.__lp, self.__lock, self.__timer = processor, C.BATCH_PROCESSOR_DEFAULT_MAX_SIZE if maxsize is None else maxsize, I.sleep.__get__(C.BATCH_PROCESSOR_DEFAULT_MAX_TIME if maxtime is None else maxtime), [], timer(), I.Lock(), timer
[docs]
20 async def add(self, item):
21 async with self.__lock:
22 (b := self.__batch).append(item)
23 if len(b) >= self.__ms: return await self._process()
24 async def _flush_periodic(self):
25 while True: await self.__sleep(); await self.flush()
26 async def _process(self):
27 if not (b := self.__batch): return
28 b, self.__lp = copy_and_clear(b), self.__timer()
29 await self.__p(b)
[docs]
30 async def flush(self):
31 async with self.__lock:
32 if self.__batch: await self._process()
33 @property
34 def time_since_last_process(self): return self.__timer()-self.__lp
[docs]
35 async def __setup__(self): super().__init__(); self.make(self._flush_periodic())
[docs]
36class Bulkhead(A.LoopContextMixin):
37 __slots__ = '__exc', '__iv', '__mr', '__mt', '__p', '__queue', '__rej', '__sd', '__sem'
38 def __init__(self, max_concurrent, *, max_queue=None, max_rej=None, exc=Exception, processor=None):
39 if max_concurrent <= 0: raise ValueError('asyncutils.processors.Bulkhead: max_concurrent must be positive')
40 C = A.getcontext()
41 if max_queue is None: max_queue = C.BULKHEAD_DEFAULT_MAX_QUEUE
42 if max_queue <= 0: raise ValueError('asyncutils.processors.Bulkhead: max_queue must be positive')
43 if max_rej is None: max_rej = C.BULKHEAD_DEFAULT_MAX_REJ
44 super().__init__(); self.__sem, self.__queue, self.__rej, self.__iv, self.__exc, self.__p, self.__sd, self.__mt, self.__mr = I.Semaphore(max_concurrent), Queue(max_queue), 0, max_concurrent, exc, processor, self.make_fut(), I.Event(), max_rej
[docs]
45 async def execute(self, coro):
46 try: self.__queue.put_nowait(coro)
47 except I.QueueFull as e:
48 if (x := self.__rej) == self.__mr: await self.shutdown(); raise A.BulkheadShutDown(f'{fullname(self)} has been shutdown because too many tasks were rejected') from e
49 self.__rej = x+1; raise A.BulkheadFull(f'{fullname(self)} queue full') from None
50 if self.is_shutdown: raise A.BulkheadShutDown(f'{fullname(self)} is shutting down')
51 async with self.__sem:
52 try: await (await self.__queue.get())
53 except (I.QueueEmpty, QueueShutDown, I.CancelledError): raise A.BulkheadShutDown(f'{fullname(self)} is shutting down') from None
54 except self.__exc as e:
55 if p := self.__p: await p(e)
56 getattr(self.__mt, 'clear' if self.active_tasks else 'set')()
[docs]
57 async def __cleanup__(self): await self.shutdown()
58 @property
59 def available_slots(self): return self.__sem._value
60 @property
61 def active_tasks(self): return self.__iv-self.available_slots
62 @property
63 def curr_qsize(self): return self.__queue.qsize()
64 @property
65 def max_qsize(self): return self.__queue.maxsize
66 @property
67 def available_queue_slots(self): return m-self.curr_qsize if (m := self.max_qsize) > 0 else float('inf')
68 @property
69 def is_shutdown(self): return self.__sd.done()
70 @property
71 def rejected(self): return self.__rej
[docs]
72 async def wait_until_idle(self, timeout=None): await I.wait_for(self.__mt.wait(), timeout)
[docs]
73 def wait_for_shutdown(self, timeout=None): return I.wait_for(self.__sd, timeout)
[docs]
74 async def shutdown(self, timeout=None):
75 self.__sd.set_result(None); (h := (q := self.__queue).shutdown)(); r = []
76 try:
77 async with I.timeout(timeout):
78 await self.__mt.wait(); a = (s := self.__sem).acquire
79 while s._value: await a()
80 except TimeoutError:
81 f, g = r.append, q.get_nowait
82 while True:
83 try: f(g())
84 except: h(True); break # noqa: E722
85 return r