1457433890

Build Type

github-actions

Committed by GitHub

Commit Message

Merge 2cbc251c0 into 5412f89a9

Pull Request Pull Request #119: Feature/worker pool with multiprocessing and RespawningProcessService

Run Details

737 of 999 branches covered (73.77%)

Branch coverage included in aggregate %.

108 of 130 new or added lines in 6 files covered. (83.08%)

28 existing lines in 4 files now uncovered.

3122 of 3710 relevant lines covered (84.15%)

6.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.92

/aiomisc/utils.py

import asyncio
import itertools
import logging.handlers
import socket
import uuid
from functools import wraps
from multiprocessing import cpu_count
from random import getrandbits
from typing import (
    Any, Awaitable, Callable, Iterable, List, Optional, Tuple, TypeVar, Union,
)

from .thread_pool import ThreadPoolExecutor


T = TypeVar("T")
TimeoutType = Union[int, float]

try:
    from time import time_ns
except ImportError:
    from time import time

    def time_ns() -> int:
        return int(time() * 1000000000)

try:
    import uvloop  # type: ignore
    event_loop_policy = uvloop.EventLoopPolicy()
except ImportError:
    event_loop_policy = asyncio.DefaultEventLoopPolicy()


log = logging.getLogger(__name__)


def fast_uuid4() -> uuid.UUID:
    """ Fast UUID4 like identifier """
    return uuid.UUID(int=getrandbits(128), version=4)


__NODE = uuid.getnode()


def fast_uuid1() -> uuid.UUID:
    """ Fast UUID1 like identifier """
    value = time_ns()
    value = (value << 16) + getrandbits(16)
    value = (value << 48) + __NODE
    return uuid.UUID(int=value, version=1)


def chunk_list(iterable: Iterable[T], size: int) -> Iterable[List[T]]:
    """
    Split list or generator by chunks with fixed maximum size.
    """

    iterable = iter(iterable)

    item = list(itertools.islice(iterable, size))
    while item:
        yield item
        item = list(itertools.islice(iterable, size))


OptionsType = Iterable[Tuple[int, int, int]]


if hasattr(socket, "TCP_NODELAY"):
    def _sock_set_nodelay(sock: socket.socket) -> None:
        if sock.proto != socket.IPPROTO_TCP:
            return None
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
else:
    def _sock_set_nodelay(sock: socket.socket) -> None:
        return None

if hasattr(socket, "SO_REUSEPORT"):
    def _sock_set_reuseport(sock: socket.socket, reuse_port: bool) -> None:
        sock.setsockopt(
            socket.SOL_SOCKET, socket.SO_REUSEPORT, int(reuse_port),
        )
else:
    def _sock_set_reuseport(sock: socket.socket, reuse_port: bool) -> None:
        log.debug(
            "SO_REUSEPORT is not implemented by underlying library. Skipping.",
        )


def bind_socket(
    *args: Any, address: str, port: int = 0, options: OptionsType = (),
    reuse_addr: bool = True, reuse_port: bool = False,
    proto_name: Optional[str] = None
) -> socket.socket:
    """

    Bind socket and set ``setblocking(False)`` for just created socket.
    This detects ``address`` format and select socket family automatically.

    :param args: which will be passed to stdlib's socket constructor (optional)
    :param address: bind address
    :param port: bind port
    :param options: Tuple of pairs which contain socket option
                    to set and the option value.
    :param reuse_addr: set socket.SO_REUSEADDR
    :param reuse_port: set socket.SO_REUSEPORT
    :param proto_name: protocol name which will be logged after binding
    :return: socket.socket
    """

    if not args:
        if ":" in address:
            args = (socket.AF_INET6, socket.SOCK_STREAM)
        else:
            args = (socket.AF_INET, socket.SOCK_STREAM)

    sock = socket.socket(*args)
    sock.setblocking(False)

    if not args and ":" in address:
        sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)

    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, reuse_addr.real)

    _sock_set_reuseport(sock, reuse_port)
    _sock_set_nodelay(sock)

    for level, option, value in options:
        sock.setsockopt(level, option, value)

    unix_address_family = getattr(socket, "AF_UNIX", None)
    if sock.family == unix_address_family:
        proto_name = proto_name or "unix"
        sock.bind(address)
    else:
        proto_name = proto_name or "tcp"
        sock.bind((address, port))

    sock_addr = sock.getsockname()
    if not isinstance(sock_addr, str):
        sock_addr = sock_addr[:2]

    if sock.family == socket.AF_INET6:
        log.info("Listening %s://[%s]:%s", proto_name, *sock_addr)
    elif sock.family == unix_address_family:
        log.info("Listening %s://%s", proto_name, sock_addr)
    else:
        log.info("Listening %s://%s:%s", proto_name, *sock_addr)

    return sock


def create_default_event_loop(
    pool_size: Optional[int] = None,
    policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
    debug: bool = False,
) -> Tuple[asyncio.AbstractEventLoop, ThreadPoolExecutor]:
    """
    Creates an event loop and thread pool executor

    :param pool_size: thread pool maximal size
    :param policy: event loop policy
    :param debug: set ``loop.set_debug(True)`` if True
    """
    try:
        asyncio.get_event_loop().close()
    except RuntimeError:
        pass  # event loop is not created yet

    asyncio.set_event_loop_policy(policy)

    loop = asyncio.new_event_loop()
    loop.set_debug(debug)
    asyncio.set_event_loop(loop)

    pool_size = pool_size or cpu_count()
    thread_pool = ThreadPoolExecutor(pool_size, statistic_name="default")
    loop.set_default_executor(thread_pool)

    return loop, thread_pool


def new_event_loop(
    pool_size: Optional[int] = None,
    policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
) -> asyncio.AbstractEventLoop:
    loop, thread_pool = create_default_event_loop(pool_size, policy)
    return loop


def shield(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
    """
    Simple and useful decorator for wrap the coroutine to `asyncio.shield`.

    >>> @shield
    ... async def non_cancelable_func():
    ...     await asyncio.sleep(1)

    """

    async def awaiter(future: Awaitable[Any]) -> Any:
        return await future

    @wraps(func)
    def wrap(*args: Any, **kwargs: Any) -> Awaitable[Any]:
        return wraps(func)(awaiter)(asyncio.shield(func(*args, **kwargs)))

    return wrap


class SelectResult:
    def __init__(self, length: int):
        self.length = length
        self.result_idx = None      # type: Optional[int]
        self.is_exception = None    # type: Optional[bool]
        self.value = None           # type: Any

    def set_result(self, idx: int, value: Any, is_exception: bool) -> None:
        if self.result_idx is not None:
            return

        self.value = value
        self.result_idx = idx
        self.is_exception = is_exception

    def result(self) -> T:
        if self.is_exception:
            raise self.value
        return self.value

    def done(self) -> bool:
        return self.result_idx is not None

    def __iter__(self) -> Iterable[Optional[T]]:
        for i in range(self.length):
            if i == self.result_idx:
                yield self.value
            else:
                yield None


def cancel_tasks(tasks: Iterable[asyncio.Future]) -> asyncio.Future:
    """
    All passed tasks will be cancelled and a new task will be returned.

    :param tasks: tasks which will be cancelled
    """

    future = asyncio.get_event_loop().create_future()
    future.set_result(None)

    if not tasks:
        return future

    cancelled_tasks = []
    exc = asyncio.CancelledError()

    for task in tasks:
        if task.done():
            continue

        if isinstance(task, asyncio.Task):
            task.cancel()
            cancelled_tasks.append(task)

        elif isinstance(task, asyncio.Future):
            task.set_exception(exc)

        else:
            log.warning(
                "Skipping object %r because it's not a Task or Future", task,
            )

    if not cancelled_tasks:
        return future

    waiter = asyncio.ensure_future(
        asyncio.gather(
            *cancelled_tasks, return_exceptions=True
        ),
    )

    return waiter


async def _select_waiter(
    idx: int, awaitable: Awaitable[T],
    result: SelectResult,
) -> None:
    try:
        ret = await awaitable
    except asyncio.CancelledError:
        raise
    except Exception as e:
        return result.set_result(idx, e, is_exception=True)

    result.set_result(idx, ret, is_exception=False)


async def select(
    *awaitables: Awaitable[Any],
    return_exceptions: bool = False, cancel: bool = True,
    timeout: Optional[TimeoutType] = None,
    wait: bool = True, loop: Optional[asyncio.AbstractEventLoop] = None
) -> SelectResult:
    """

    :param awaitables: awaitable objects
    :param return_exceptions: if True exception will not be raised
                              just returned as result
    :param cancel: cancel unfinished coroutines (default True)
    :param timeout: execution timeout
    :param wait: when False and ``cancel=True``, unfinished coroutines will
                 be cancelled in the background.
    :param loop: event loop
    """

    loop = loop or asyncio.get_event_loop()
    result = SelectResult(len(awaitables))

    coroutines = [
        loop.create_task(_select_waiter(idx, coroutine, result))
        for idx, coroutine in enumerate(awaitables)
    ]

    _, pending = await loop.create_task(
        asyncio.wait(
            coroutines,
            timeout=timeout,
            return_when=asyncio.FIRST_COMPLETED,
        ),
    )

    if cancel:
        cancelling = cancel_tasks(pending)

        if wait:
            await cancelling

    if result.is_exception and not return_exceptions:
        result.result()

    return result


def awaitable(
    func: Callable[..., Union[T, Awaitable[T]]],
) -> Callable[..., Awaitable[T]]:
    """

    Decorator wraps function and returns a function which returns
    awaitable object. In case than a function returns a future,
    the original future will be returned. In case then the function
    returns a coroutine, the original coroutine will be returned.
    In case than function returns non-awaitable object, it's will
    be wrapped to a new coroutine which just returns this object.
    It's useful when you don't want to check function result before
    use it in ``await`` expression.
    """

    # Avoid python 3.8+ warning
    if asyncio.iscoroutinefunction(func):
        return func     # type: ignore

    async def awaiter(obj: T) -> T:
        return obj

    @wraps(func)
    def wrap(*args: Any, **kwargs: Any) -> Awaitable[T]:
        result = func(*args, **kwargs)

        if hasattr(result, "__await__"):
            return result       # type: ignore
        if asyncio.iscoroutine(result) or asyncio.isfuture(result):
            return result       # type: ignore

        return awaiter(result)  # type: ignore

    return wrap

1	import asyncio	15✔
2	import itertools	15✔
3	import logging.handlers	15✔
4	import socket	15✔
5	import uuid	15✔
6	from functools import wraps	15✔
7	from multiprocessing import cpu_count	15✔
8	from random import getrandbits	15✔
9	from typing import (	15✔
10	Any, Awaitable, Callable, Iterable, List, Optional, Tuple, TypeVar, Union,
11	)
12
13	from .thread_pool import ThreadPoolExecutor	15✔
14
15
16	T = TypeVar("T")	15✔
17	TimeoutType = Union[int, float]	15✔
18
19	try:	15✔
20	from time import time_ns	15✔
NEW 21	except ImportError:	×
NEW 22	from time import time	×
23
NEW 24	def time_ns() -> int:	×
NEW 25	return int(time() * 1000000000)	×
26
27	try:	15✔
28	import uvloop # type: ignore	15✔
29	event_loop_policy = uvloop.EventLoopPolicy()	5✔
30	except ImportError:	10✔
31	event_loop_policy = asyncio.DefaultEventLoopPolicy()	10✔
32
33
34	log = logging.getLogger(__name__)	15✔
35
36
37	def fast_uuid4() -> uuid.UUID:	15✔
38	""" Fast UUID4 like identifier """
39	return uuid.UUID(int=getrandbits(128), version=4)	15✔
40
41
42	__NODE = uuid.getnode()	15✔
43
44
45	def fast_uuid1() -> uuid.UUID:	15✔
46	""" Fast UUID1 like identifier """
NEW 47	value = time_ns()	×
NEW 48	value = (value << 16) + getrandbits(16)	×
NEW 49	value = (value << 48) + __NODE	×
NEW 50	return uuid.UUID(int=value, version=1)	×
51
52
53	def chunk_list(iterable: Iterable[T], size: int) -> Iterable[List[T]]:	15✔
54	"""
55	Split list or generator by chunks with fixed maximum size.
56	"""
57
58	iterable = iter(iterable)	19✔
59
60	item = list(itertools.islice(iterable, size))	19✔
61	while item:	19✔
62	yield item	19✔
63	item = list(itertools.islice(iterable, size))	19✔
64
65
66	OptionsType = Iterable[Tuple[int, int, int]]	15✔
67
68
69	if hasattr(socket, "TCP_NODELAY"):	15!
70	def _sock_set_nodelay(sock: socket.socket) -> None:	15✔
71	if sock.proto != socket.IPPROTO_TCP:	19!
72	return None	19✔
73	sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)	×
74	else:
75	def _sock_set_nodelay(sock: socket.socket) -> None:	×
76	return None	×
77
78	if hasattr(socket, "SO_REUSEPORT"):	15✔
79	def _sock_set_reuseport(sock: socket.socket, reuse_port: bool) -> None:	13✔
80	sock.setsockopt(	17✔
81	socket.SOL_SOCKET, socket.SO_REUSEPORT, int(reuse_port),
82	)
83	else:
84	def _sock_set_reuseport(sock: socket.socket, reuse_port: bool) -> None:	2✔
85	log.debug(	2✔
86	"SO_REUSEPORT is not implemented by underlying library. Skipping.",
87	)
88
89
90	def bind_socket(	15✔
91	*args: Any, address: str, port: int = 0, options: OptionsType = (),
92	reuse_addr: bool = True, reuse_port: bool = False,
93	proto_name: Optional[str] = None
94	) -> socket.socket:
95	"""
96
97	Bind socket and set ``setblocking(False)`` for just created socket.
98	This detects ``address`` format and select socket family automatically.
99
100	:param args: which will be passed to stdlib's socket constructor (optional)
101	:param address: bind address
102	:param port: bind port
103	:param options: Tuple of pairs which contain socket option
104	to set and the option value.
105	:param reuse_addr: set socket.SO_REUSEADDR
106	:param reuse_port: set socket.SO_REUSEPORT
107	:param proto_name: protocol name which will be logged after binding
108	:return: socket.socket
109	"""
110
111	if not args:	19✔
112	if ":" in address:	19✔
113	args = (socket.AF_INET6, socket.SOCK_STREAM)	19✔
114	else:
115	args = (socket.AF_INET, socket.SOCK_STREAM)	19✔
116
117	sock = socket.socket(*args)	19✔
118	sock.setblocking(False)	19✔
119
120	if not args and ":" in address:	19!
121	sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)	×
122
123	sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, reuse_addr.real)	19✔
124
125	_sock_set_reuseport(sock, reuse_port)	19✔
126	_sock_set_nodelay(sock)	19✔
127
128	for level, option, value in options:	19!
129	sock.setsockopt(level, option, value)	×
130
131	unix_address_family = getattr(socket, "AF_UNIX", None)	19✔
132	if sock.family == unix_address_family:	19✔
133	proto_name = proto_name or "unix"	17✔
134	sock.bind(address)	17✔
135	else:
136	proto_name = proto_name or "tcp"	19✔
137	sock.bind((address, port))	19✔
138
139	sock_addr = sock.getsockname()	19✔
140	if not isinstance(sock_addr, str):	19✔
141	sock_addr = sock_addr[:2]	19✔
142
143	if sock.family == socket.AF_INET6:	19✔
144	log.info("Listening %s://[%s]:%s", proto_name, *sock_addr)	19✔
145	elif sock.family == unix_address_family:	19✔
146	log.info("Listening %s://%s", proto_name, sock_addr)	17✔
147	else:
148	log.info("Listening %s://%s:%s", proto_name, *sock_addr)	19✔
149
150	return sock	19✔
151
152
153	def create_default_event_loop(	15✔
154	pool_size: Optional[int] = None,
155	policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
156	debug: bool = False,
157	) -> Tuple[asyncio.AbstractEventLoop, ThreadPoolExecutor]:
158	"""
159	Creates an event loop and thread pool executor
160
161	:param pool_size: thread pool maximal size
162	:param policy: event loop policy
163	:param debug: set ``loop.set_debug(True)`` if True
164	"""
165	try:	19✔
166	asyncio.get_event_loop().close()	19✔
167	except RuntimeError:	×
168	pass # event loop is not created yet	×
169
170	asyncio.set_event_loop_policy(policy)	19✔
171
172	loop = asyncio.new_event_loop()	19✔
173	loop.set_debug(debug)	19✔
174	asyncio.set_event_loop(loop)	19✔
175
176	pool_size = pool_size or cpu_count()	19✔
177	thread_pool = ThreadPoolExecutor(pool_size, statistic_name="default")	19✔
178	loop.set_default_executor(thread_pool)	19✔
179
180	return loop, thread_pool	19✔
181
182
183	def new_event_loop(	15✔
184	pool_size: Optional[int] = None,
185	policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
186	) -> asyncio.AbstractEventLoop:
187	loop, thread_pool = create_default_event_loop(pool_size, policy)	×
188	return loop	×
189
190
191	def shield(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:	15✔
192	"""
193	Simple and useful decorator for wrap the coroutine to `asyncio.shield`.
194
195	>>> @shield
196	... async def non_cancelable_func():
197	... await asyncio.sleep(1)
198
199	"""
200
201	async def awaiter(future: Awaitable[Any]) -> Any:	19✔
202	return await future	19✔
203
204	@wraps(func)	19✔
205	def wrap(args: Any, *kwargs: Any) -> Awaitable[Any]:	19✔
206	return wraps(func)(awaiter)(asyncio.shield(func(args, *kwargs)))	19✔
207
208	return wrap	19✔
209
210
211	class SelectResult:	15✔
212	def __init__(self, length: int):	15✔
213	self.length = length	19✔
214	self.result_idx = None # type: Optional[int]	19✔
215	self.is_exception = None # type: Optional[bool]	19✔
216	self.value = None # type: Any	19✔
217
218	def set_result(self, idx: int, value: Any, is_exception: bool) -> None:	15✔
219	if self.result_idx is not None:	19✔
220	return	19✔
221
222	self.value = value	19✔
223	self.result_idx = idx	19✔
224	self.is_exception = is_exception	19✔
225
226	def result(self) -> T:	15✔
227	if self.is_exception:	19!
228	raise self.value	19✔
229	return self.value	×
230
231	def done(self) -> bool:	15✔
232	return self.result_idx is not None	×
233
234	def __iter__(self) -> Iterable[Optional[T]]:	15✔
235	for i in range(self.length):	19✔
236	if i == self.result_idx:	19✔
237	yield self.value	19✔
238	else:
239	yield None	19✔
240
241
242	def cancel_tasks(tasks: Iterable[asyncio.Future]) -> asyncio.Future:	15✔
243	"""
244	All passed tasks will be cancelled and a new task will be returned.
245
246	:param tasks: tasks which will be cancelled
247	"""
248
249	future = asyncio.get_event_loop().create_future()	19✔
250	future.set_result(None)	19✔
251
252	if not tasks:	19✔
253	return future	19✔
254
255	cancelled_tasks = []	19✔
256	exc = asyncio.CancelledError()	19✔
257
258	for task in tasks:	19✔
259	if task.done():	19✔
260	continue	4✔
261
262	if isinstance(task, asyncio.Task):	19✔
263	task.cancel()	19✔
264	cancelled_tasks.append(task)	19✔
265
266	elif isinstance(task, asyncio.Future):	19!
267	task.set_exception(exc)	19✔
268
269	else:
270	log.warning(	×
271	"Skipping object %r because it's not a Task or Future", task,
272	)
273
274	if not cancelled_tasks:	19✔
275	return future	19✔
276
277	waiter = asyncio.ensure_future(	19✔
278	asyncio.gather(
279	*cancelled_tasks, return_exceptions=True
280	),
281	)
282
283	return waiter	19✔
284
285
286	async def _select_waiter(	15✔
287	idx: int, awaitable: Awaitable[T],
288	result: SelectResult,
289	) -> None:
290	try:	19✔
291	ret = await awaitable	19✔
292	except asyncio.CancelledError:	19✔
293	raise	19✔
294	except Exception as e:	19✔
295	return result.set_result(idx, e, is_exception=True)	19✔
296
297	result.set_result(idx, ret, is_exception=False)	19✔
298
299
300	async def select(	15✔
301	*awaitables: Awaitable[Any],
302	return_exceptions: bool = False, cancel: bool = True,
303	timeout: Optional[TimeoutType] = None,
304	wait: bool = True, loop: Optional[asyncio.AbstractEventLoop] = None
305	) -> SelectResult:
306	"""
307
308	:param awaitables: awaitable objects
309	:param return_exceptions: if True exception will not be raised
310	just returned as result
311	:param cancel: cancel unfinished coroutines (default True)
312	:param timeout: execution timeout
313	:param wait: when False and ``cancel=True``, unfinished coroutines will
314	be cancelled in the background.
315	:param loop: event loop
316	"""
317
318	loop = loop or asyncio.get_event_loop()	19✔
319	result = SelectResult(len(awaitables))	19✔
320
321	coroutines = [	19✔
322	loop.create_task(_select_waiter(idx, coroutine, result))
323	for idx, coroutine in enumerate(awaitables)
324	]
325
326	_, pending = await loop.create_task(	19✔
327	asyncio.wait(
328	coroutines,
329	timeout=timeout,
330	return_when=asyncio.FIRST_COMPLETED,
331	),
332	)
333
334	if cancel:	19✔
335	cancelling = cancel_tasks(pending)	19✔
336
337	if wait:	19!
338	await cancelling	19✔
339
340	if result.is_exception and not return_exceptions:	19✔
341	result.result()	19✔
342
343	return result	19✔
344
345
346	def awaitable(	15✔
347	func: Callable[..., Union[T, Awaitable[T]]],
348	) -> Callable[..., Awaitable[T]]:
349	"""
350
351	Decorator wraps function and returns a function which returns
352	awaitable object. In case than a function returns a future,
353	the original future will be returned. In case then the function
354	returns a coroutine, the original coroutine will be returned.
355	In case than function returns non-awaitable object, it's will
356	be wrapped to a new coroutine which just returns this object.
357	It's useful when you don't want to check function result before
358	use it in ``await`` expression.
359	"""
360
361	# Avoid python 3.8+ warning
362	if asyncio.iscoroutinefunction(func):	19✔
363	return func # type: ignore	19✔
364
365	async def awaiter(obj: T) -> T:	19✔
366	return obj	19✔
367
368	@wraps(func)	19✔
369	def wrap(args: Any, *kwargs: Any) -> Awaitable[T]:	19✔
370	result = func(args, *kwargs)	19✔
371
372	if hasattr(result, "__await__"):	19✔
373	return result # type: ignore	19✔
374	if asyncio.iscoroutine(result) or asyncio.isfuture(result):	19!
375	return result # type: ignore	×
376
377	return awaiter(result) # type: ignore	19✔
378
379	return wrap	19✔

mosquito / aiomisc / 1457433890

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous