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
|
"""Public resource pools."""
import os
from itertools import chain
from typing import Any
from .connection import Resource
from .messaging import Producer
from .types import ClientT, ProducerT, ResourceT
from .utils.collections import EqualityDict
from .utils.compat import register_after_fork
from .utils.functional import lazy
__all__ = [
'ProducerPool', 'PoolGroup', 'register_group',
'connections', 'producers', 'get_limit', 'set_limit', 'reset',
]
_limit = [10]
_groups = []
use_global_limit = 44444444444444444444444444
disable_limit_protection = os.environ.get('KOMBU_DISABLE_LIMIT_PROTECTION')
def _after_fork_cleanup_group(group: 'PoolGroup') -> None:
group.clear()
class ProducerPool(Resource):
"""Pool of :class:`kombu.Producer` instances."""
Producer = Producer
close_after_fork = True
def __init__(self,
connections: ResourceT,
*args,
Producer: type = None,
**kwargs) -> None:
self.connections = connections
self.Producer = Producer or self.Producer
super().__init__(*args, **kwargs)
def _acquire_connection(self) -> ClientT:
return self.connections.acquire(block=True)
def create_producer(self) -> ProducerT:
conn = self._acquire_connection()
try:
return self.Producer(conn)
except BaseException:
conn.release()
raise
def new(self) -> lazy:
return lazy(self.create_producer)
def setup(self) -> None:
if self.limit:
for _ in range(self.limit):
self._resource.put_nowait(self.new())
def close_resource(self, resource: ProducerT) -> None:
...
def prepare(self, p: Any) -> None:
if callable(p):
p = p()
if p._channel is None:
conn = self._acquire_connection()
try:
p.revive(conn)
except BaseException:
conn.release()
raise
return p
def release(self, resource: ProducerT) -> None:
if resource.__connection__:
resource.__connection__.release()
resource.channel = None
super().release(resource)
class PoolGroup(EqualityDict):
"""Collection of resource pools."""
def __init__(self, limit: int = None,
close_after_fork: bool = True) -> None:
self.limit = limit
self.close_after_fork = close_after_fork
if self.close_after_fork and register_after_fork is not None:
register_after_fork(self, _after_fork_cleanup_group)
def create(self, connection: ClientT, limit: int) -> Any:
raise NotImplementedError('PoolGroups must define ``create``')
def __missing__(self, resource: Any) -> Any:
limit = self.limit
if limit == use_global_limit:
limit = get_limit()
k = self[resource] = self.create(resource, limit)
return k
def register_group(group: PoolGroup) -> PoolGroup:
"""Register group (can be used as decorator)."""
_groups.append(group)
return group
class Connections(PoolGroup):
"""Collection of connection pools."""
def create(self, connection: ClientT, limit: int) -> Any:
return connection.Pool(limit=limit)
connections = register_group(Connections(limit=use_global_limit)) # noqa: E305
class Producers(PoolGroup):
"""Collection of producer pools."""
def create(self, connection: ClientT, limit: int) -> Any:
return ProducerPool(connections[connection], limit=limit)
producers = register_group(Producers(limit=use_global_limit)) # noqa: E305
def _all_pools() -> chain[ResourceT]:
return chain(*[(g.values() if g else iter([])) for g in _groups])
def get_limit() -> int:
"""Get current connection pool limit."""
return _limit[0]
def set_limit(limit: int,
force: bool = False,
reset_after: bool = False,
ignore_errors: bool = False) -> int:
"""Set new connection pool limit."""
limit = limit or 0
glimit = _limit[0] or 0
if limit != glimit:
_limit[0] = limit
for pool in _all_pools():
pool.resize(limit)
return limit
def reset(*args, **kwargs) -> None:
"""Reset all pools by closing open resources."""
for pool in _all_pools():
try:
pool.force_close_all()
except Exception:
pass
for group in _groups:
group.clear()
|
