1import asyncutils as A, asyncio as I
2from asyncutils._internal.py312 import PriorityQueue
3from asyncutils._internal.helpers import LoopMixinBase, fullname, subscriptable
4from asyncutils._internal.submodules import pools_all as __all__
5from _functools import partial
6from itertools import count, repeat
7from threading import Thread, Lock as TLock
8from time import monotonic
[docs]
9class AdvancedPool(A.LoopContextMixin):
10 __slots__ = '__cnt', '__cur', '__fs', '__kae', '__max', '__min', '__pending', '__q', '__scaling', '__sde', '__shutdown', '__start', '__tl', '__workers', 'completed'
11 @property
12 def __tiebreak(self): return next(self.__cnt)
13 def __init__(self, max_workers=None, min_workers=None, qsize=0, scaling=True, kill_at_exit=False): super().__init__(); C = A.getcontext(); self.__cnt, self.__tl, self.__max, self.__scaling, self.__q, self.__workers, self.__fs, self.__pending, self.__sde, self.__shutdown, self.__start, self.completed, self.__kae = count(), TLock(), C.ADVANCED_POOL_DEFAULT_MAX_WORKERS if max_workers is None else max_workers, scaling, PriorityQueue(qsize), set(), set(), 0, I.Event(), False, monotonic(), 0, kill_at_exit; self.__cur = self.__min = C.ADVANCED_POOL_DEFAULT_MIN_WORKERS if min_workers is None else min_workers
14 def __repr__(self): return f'{fullname(self)}({self.__max}, {self.__min}, {self.__q.maxsize}, {self.__scaling}, {self.__kae})'
15 async def __ts_get(self):
16 with self.__tl: return await self.__q.get()
17 def __ts_task_done(self):
18 with self.__tl: self.__q.task_done()
19 def __wl(self, _):
20 x, g, G = 0, (L := self.loop).call_soon_threadsafe, self.__ts_get
21 try: # noqa: PLW0717
22 while not self.__shutdown:
23 if (F := A.sync_await(G(), loop=L)[2]) is None: self.__ts_task_done(); break
24 f, a, k, F = F
25 try: F.set_result(f(*a, **k))
26 except BaseException as e: F.set_exception(e) # noqa: BLE001
27 finally:
28 self.__ts_task_done(); x += 1
29 with self.__tl: self.completed += 1; self.__pending -= 1
30 except BaseException as e: g(_.set_exception, A.Critical(e) if isinstance(e, A.CRITICAL) else e) # noqa: BLE001
31 else: g(_.set_result, x)
32 def __scale_to(self, new):
33 if (d := new-self.__cur) > 0:
34 a, b, f, g = self.__workers.add, self.__fs.add, self.__wl, self.make_fut
35 for _ in repeat(None, d): (T := Thread(target=f, args=(F := g(),))).start(); a(T); b(F)
36 elif d < 0:
37 f = self.__q.put_nowait
38 for _ in repeat(None, -d): f((0, self.__tiebreak, None))
39 self.__cur = new
40 def __set_adj(self):
41 if not self.__scaling: return
42 C = A.getcontext()
43 if (l := self.__pending/((c := self.__cur) or 1)) > C.ADVANCED_POOL_THRESHOLD_HI and c < (M := self.__max): self.__scale_to(min(M, c+max(1, c>>1)))
44 elif l < C.ADVANCED_POOL_THRESHOLD_LO and c > (m := self.__min): self.__scale_to(max(m, c-max(1, int(c*C.ADVANCED_POOL_FACTOR))))
45 async def wait_for_shutdown(self): return await self.__sde.wait()
[docs]
46 def raise_for_shutdown(self):
47 if self.__shutdown: raise A.PoolShutDown(f'{fullname(self)} is shutting down')
[docs]
48 def submit_nowait(self, f, *a, _priority_=0, **k):
49 self.raise_for_shutdown()
50 if self.full: raise A.PoolFull('asyncutils.pool.AdvancedPool.submit_nowait: task queue full')
51 with self.__tl: self.__pending += 1
52 self.__q.put_nowait((_priority_, self.__tiebreak, (f, a, k, F := self.make_fut()))); self.__set_adj(); return F
53 async def _kill_helper(self):
54 f, g = (q := self.__q).get_nowait, q.task_done
55 with self.__tl, A.ignore_qempty:
56 while True:
57 if (F := f()[2]) is not None: await A.safe_cancel(F[-1])
58 g()
[docs]
59 async def complete(self, *a, **k): return await (await self.submit(*a, **k))
[docs]
60 async def submit(self, f, *a, _priority_=0, **k):
61 self.raise_for_shutdown()
62 with self.__tl: self.__pending += 1
63 await self.__q.put((_priority_, self.__tiebreak, (f, a, k, F := self.make_fut()))); self.__set_adj(); return F
[docs]
64 async def shutdown(self, cancel_pending=False, idle_timeout=None):
65 if self.__shutdown: return await self.wait_for_shutdown()
66 if cancel_pending: await self._kill_helper()
67 else:
68 try: await self.wait_for_slot(idle_timeout)
69 except A.PoolFull: await self._kill_helper()
70 p = (q := self.__q).put
71 with A.ignore_qshutdown:
72 for _ in repeat(None, self.__cur): await p((0, self.__tiebreak, None))
73 with self.__tl: q.shutdown(True)
74 await self.join(); self.__sde.set(); return self.uptime
[docs]
75 async def join(self): return await I.gather(*self.__fs, return_exceptions=True)
[docs]
76 async def map(self, f, /, *i, priority=0, strict=False): return await A.agather(A.amap(partial(self.complete, f, _priority_=priority), *i, strict=strict))
[docs]
77 async def starmap(self, f, /, it, priority=0): return await A.agather(A.astarmap(partial(self.complete, f, _priority_=priority), it))
[docs]
78 async def double_starmap(self, f, /, it, priority=0): return await A.agather(A.adouble_starmap(partial(self.complete, f, _priority_=priority), it))
[docs]
79 async def starmap_with_kwds(self, f, /, it, priority=0): return await A.agather(A.astarmap_with_kwds(partial(self.complete, f, _priority_=priority), it))
[docs]
80 async def resize(self, min_workers, max_workers): M = max(max_workers, m := max(1, min_workers)); self.__scale_to(min(max(self.__cur, m), M)); self.__min, self.__max = m, M
[docs]
81 def drain(self): return self.__q.join()
[docs]
82 async def wait_for_slot(self, timeout=None):
83 self.raise_for_shutdown()
84 if not self.full: return 0.0
85 try: t = monotonic(); await I.wait_for(self.drain(), timeout); return monotonic()-t
86 except TimeoutError: raise A.PoolFull('asyncutils.pools.AdvancedPool: timeout waiting for queue space') from None
[docs]
87 async def __cleanup__(self): await self.shutdown(self.__kae)
[docs]
88 def __del__(self):
89 if self.loop.is_running(): self.make(self.shutdown(True, 0.03))
90 @property
91 def full(self): return self.__q.full()
92 @property
93 def empty(self): return self.__q.empty()
94 @property
95 def qsize(self): return self.__q.qsize()
96 @property
97 def idle(self): return self.empty and not self.__pending
98 @property
99 def uptime(self): return monotonic()-self.__start
[docs]
100@subscriptable
101class ConnectionPool(LoopMixinBase):
102 __slots__ = '__av', '__clean', '__cts', '__factory', '__hc', '__in_use', '__lock', '__mtr', '__pool', 'maxlife', 'maxsize', 'minsize'
103 def __init__(self, factory, maxsize=None, minsize=None, maxlife=None, healthchecker=None, cleaner=None): C = A.getcontext(); self.__factory, self.maxsize, self.minsize, self.maxlife, self.__hc, self.__clean, self.__pool, self.__in_use, self.__cts, self.__lock, self.__av, self.__mtr = factory, C.CONNECTION_POOL_DEFAULT_MAX_SIZE if maxsize is None else maxsize, C.CONNECTION_POOL_DEFAULT_MIN_SIZE if minsize is None else minsize, C.CONNECTION_POOL_DEFAULT_MAX_LIFE if maxlife is None else maxlife, healthchecker or (lambda _: True), cleaner or (lambda _: None), [], set(), {}, I.Lock(), I.Event(), None
104 def _is_healthy(self, conn, /): return self.__hc(conn) and not ((t := self.__cts.get(id(conn))) and monotonic()-t > self.maxlife)
[docs]
105 async def create_connection(self, *a, _executor_=None, **k): self.__cts[id(c := await self.loop.run_in_executor(_executor_, partial(self.__factory, *a, **k)))] = monotonic(); return c
[docs]
106 async def acquire(self, *a, **k):
107 p = self.__pool
108 async with self.__lock:
109 while p:
110 if self._is_healthy(c := p.pop()): self.__in_use.add(c); return c
111 self.__clean(c)
112 if self.cursize < self.maxsize: self.__in_use.add(c := await self.create_connection(*a, **k)); return c
113 await self.__av.wait(); return await self.acquire(*a, **k)
[docs]
114 def release(self, c, /, *a, **k):
115 self.__in_use.discard(c)
116 if self._is_healthy(c) and len(self.__pool) < self.maxsize: self.__pool.append(c); self.__av.set(); self.__av.clear()
117 else:
118 self.__clean(c)
119 if self.cursize < self.minsize: self.make(self.create_connection(*a, **k))
120 async def _maintain(self):
121 f = I.sleep.__get__(A.getcontext().CONNECTION_POOL_MAINTENANCE_INTERVAL)
122 while True:
123 await f(); n, g = [], self.create_connection
124 for c in self.__pool: (n.append if self.__hc(c) else self.__clean)(c)
125 async with self.__lock: n.extend(await I.gather(*(g() for _ in repeat(None, self.minsize-self.cursize)))); self.__pool = n
[docs]
126 async def start(self, akgen=None, executor=None):
127 f = self.create_connection
128 if akgen is None:
129 for _ in repeat(None, self.minsize): await f(_executor_=executor)
130 else:
131 async for a, k in A.take(akgen, self.minsize, ((), {})): await f(*a, _executor_=executor, **k)
132 self.__mtr = self.make(self._maintain())
[docs]
133 async def stop(self):
134 if m := self.__mtr: await A.safe_cancel(m)
135 p, c, f = (P := self.__pool).pop, self.__clean, (u := self.__in_use).pop
136 while P: c(p())
137 while u: c(f())
[docs]
138 async def __aenter__(self): await self.start(); return self
[docs]
139 async def __aexit__(self, /, *_): await self.stop()
140 @property
141 def cursize(self): return self.available+self.in_use
142 @property
143 def available(self): return len(self.__pool)
144 @property
145 def in_use(self): return len(self.__in_use)