async def scan_now(self, polling=True):
"""
Task to scan the whole system.
:param polling: if False, do not add polling tasks
"""
async with anyio.create_task_group() as n:
for s in list(self._servers):
await n.spawn(partial(s.scan_now, polling=polling))
async def test_cancel_scope_in_another_task():
async def child():
nonlocal result, local_scope
async with open_cancel_scope() as local_scope:
await sleep(2)
result = True
local_scope = None
result = False
async with create_task_group() as tg:
await tg.spawn(child)
while local_scope is None:
await sleep(0)
await local_scope.cancel()
assert not result
async def test_manual_acquire(self) -> None:
async def notifier() -> None:
await condition.acquire()
try:
condition.notify_all()
finally:
condition.release()
condition = Condition()
async with create_task_group() as tg:
await condition.acquire()
try:
assert condition.locked()
tg.start_soon(notifier)
await condition.wait()
finally:
condition.release()
async def test_event_cancel(self) -> None:
task_started = event_set = False
async def task() -> None:
nonlocal task_started, event_set
task_started = True
await event.wait()
event_set = True
event = Event()
async with create_task_group() as tg:
tg.start_soon(task)
tg.cancel_scope.cancel()
event.set()
assert task_started
assert not event_set
async def test_alpn_negotiation(self, server_context, client_context):
async def server():
async with await stream_server.accept() as stream:
assert stream.alpn_protocol == 'dummy2'
client_context.set_alpn_protocols(['dummy1', 'dummy2'])
server_context.set_alpn_protocols(['dummy2', 'dummy3'])
async with await create_tcp_server(
interface='localhost',
ssl_context=server_context) as stream_server:
async with create_task_group() as tg:
await tg.spawn(server)
async with await connect_tcp('localhost',
stream_server.port,
ssl_context=client_context,
autostart_tls=True) as client:
assert client.alpn_protocol == 'dummy2'
async def main():
global event
event = anyio.create_event()
logger.info("Starting up the Server.")
# Create a context for the purpose of authenticating clients
context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
# Load the server certificate and private key
context.load_cert_chain(certfile=SERVER_CERT_CHAIN)
logger.debug("Loaded the server certificate.")
async with anyio.create_task_group() as tg:
async with await anyio.create_tcp_server(
9123, ssl_context=context) as server:
logger.info("Started the server on port 9123 (TCP).")
async for client in server.accept_connections():
await tg.spawn(serve, client)
await tg.spawn(download)
async def test_level_cancellation() -> None:
marker = None
async def dummy() -> None:
nonlocal marker
marker = 1
# At this point the task has been cancelled so sleep() will raise an exception
await sleep(0)
# Execution should never get this far
marker = 2
async with create_task_group() as tg:
tg.start_soon(dummy)
assert marker is None
tg.cancel_scope.cancel()
assert marker == 1
async def main():
condition = create_condition()
async with create_task_group() as tg:
for tasknum in range(6):
await tg.spawn(listen, tasknum, condition)
await sleep(1)
async with condition:
await condition.notify(1)
await sleep(1)
async with condition:
await condition.notify(2)
await sleep(1)
async with condition:
await condition.notify_all()
async def test_iterate(self, family):
async def serve():
async for packet in udp2:
await udp2.send(packet[::-1])
async with await create_udp_socket(family=family,
local_host='localhost') as udp1:
host, port = udp1.extra(SocketAttribute.local_address)
async with await create_connected_udp_socket(host, port) as udp2:
host, port = udp2.extra(SocketAttribute.local_address)
async with create_task_group() as tg:
tg.spawn(serve)
await udp1.sendto(b'FOOBAR', host, port)
assert await udp1.receive() == (b'RABOOF', (host, port))
await udp1.sendto(b'123456', host, port)
assert await udp1.receive() == (b'654321', (host, port))
tg.cancel_scope.cancel()
async def test_iterate(self, family):
async def serve():
async for packet, addr in server:
await server.send((packet[::-1], addr))
async with await create_udp_socket(family=family,
local_host='localhost') as server:
host, port = server.extra(SocketAttribute.local_address)
async with await create_udp_socket(
family=family, local_host='localhost') as client:
async with create_task_group() as tg:
tg.spawn(serve)
await client.sendto(b'FOOBAR', host, port)
assert await client.receive() == (b'RABOOF', (host, port))
await client.sendto(b'123456', host, port)
assert await client.receive() == (b'654321', (host, port))
tg.cancel_scope.cancel()
async def _broadcast_loop(self):
async with anyio.create_task_group() as tg:
while True:
broadcast = await self._broadcast_queue.get()
self.logger.debug("broadcasting %r", broadcast)
for k_filter, subscriptions in self._subscriptions.items():
if broadcast['topic'].startswith("$") and (
k_filter.startswith("+")
or k_filter.startswith("#")):
self.logger.debug(
"[MQTT-4.7.2-1] - ignoring brodcasting $ topic to subscriptions starting with + or #"
)
elif self.matches(broadcast['topic'], k_filter):
for (target_session, qos) in subscriptions:
if 'qos' in broadcast:
qos = broadcast['qos']
if target_session.transitions.state == 'connected':
if self.logger.isEnabledFor(logging.DEBUG):
self.logger.debug(
"broadcasting application message from %s on topic '%s' to %s",
format_client_message(
session=broadcast['session']),
broadcast['topic'],
format_client_message(
session=target_session))
handler = self._get_handler(target_session)
await tg.spawn(
partial(handler.mqtt_publish,
broadcast['topic'],
broadcast['data'],
qos,
retain=False))
else:
if self.logger.isEnabledFor(logging.DEBUG):
self.logger.debug(
"retaining application message from %s on topic '%s' to client '%s'",
format_client_message(
session=broadcast['session']),
broadcast['topic'],
format_client_message(
session=target_session))
retained_message = RetainedApplicationMessage(
broadcast['session'], broadcast['topic'],
broadcast['data'], qos)
await target_session.retained_messages.put(
retained_message)
async def test_manual_acquire(self):
async def notifier():
await condition.acquire()
try:
await condition.notify_all()
finally:
await condition.release()
condition = create_condition()
async with create_task_group() as tg:
await condition.acquire()
try:
assert condition.locked()
await tg.spawn(notifier)
await condition.wait()
finally:
await condition.release()
async def test_concurrent_receive(self):
async with await create_connected_udp_socket(
'localhost',
5000,
local_port=5000,
local_host='localhost',
family=socket.AF_INET) as udp:
async with create_task_group() as tg:
await tg.spawn(udp.receive)
await wait_all_tasks_blocked()
try:
with pytest.raises(BusyResourceError) as exc:
await udp.receive()
exc.match('already reading from')
finally:
await tg.cancel_scope.cancel()
async def test_cancel_scope_in_child_task():
async def child():
nonlocal child_scope
async with open_cancel_scope() as child_scope:
await sleep(2)
child_scope = None
host_done = False
async with create_task_group() as tg:
await tg.spawn(child)
await wait_all_tasks_blocked()
await child_scope.cancel()
await sleep(0.1)
host_done = True
assert host_done
assert not tg.cancel_scope.cancel_called
async def test_buffer(self, localhost):
async def server():
async with await stream_server.accept() as stream:
chunks.append(await stream.receive_until(b'\n', 10))
chunks.append(await stream.receive_exactly(4))
chunks.append(await stream.receive_exactly(2))
chunks = []
async with await create_tcp_server(interface=localhost
) as stream_server:
async with create_task_group() as tg:
await tg.spawn(server)
async with await connect_tcp(localhost,
stream_server.port) as client:
await client.send_all(b'blah\nfoobar')
assert chunks == [b'blah', b'foob', b'ar']
async def _run_one(self, val: ValueEvent):
try:
async with anyio.create_task_group() as tg:
self._current_tg = tg
self.stream = await anyio.connect_tcp(self.host, self.port)
ml, self.requests = deque(self.requests), deque()
try:
self._msg_proto = MessageProtocol(self, is_server=False)
e_w = anyio.Event()
e_r = anyio.Event()
tg.spawn(self._writer, e_w)
tg.spawn(self._reader, e_r)
await e_r.wait()
await e_w.wait()
# re-send messages, but skip those that have been cancelled
while ml:
msg = ml.popleft()
if not msg.cancelled:
try:
await self._wqueue_w.send(msg)
except BaseException:
ml.appendleft(msg)
except BaseException:
self.requests = ml
raise
await self.chat(NOPMsg())
if self._scan_args is not None:
tg.spawn(partial(self.start_scan, **self._scan_args))
if val is not None:
val.set(None)
self._backoff = 0.1
pass # wait for tasks
pass # exited tasks
finally:
self._current_tg = None
if self.stream is not None:
with anyio.CancelScope(shield=True):
await self.stream.aclose()
self.stream = None
async def test_run_in_thread_cancelled():
def thread_worker():
nonlocal state
state = 2
async def worker():
nonlocal state
state = 1
await run_sync_in_worker_thread(thread_worker)
state = 3
state = 0
async with create_task_group() as tg:
tg.spawn(worker)
tg.cancel_scope.cancel()
assert state == 1
async def test_send_large_buffer(self, localhost):
async def server():
async with await stream_server.accept() as stream:
await stream.send_all(buffer)
buffer = b'\xff' * 1024 * 1024 # should exceed the maximum kernel send buffer size
async with create_task_group() as tg:
async with await create_tcp_server(interface=localhost
) as stream_server:
await tg.spawn(server)
async with await connect_tcp(localhost,
stream_server.port) as client:
response = await client.receive_exactly(len(buffer))
with pytest.raises(IncompleteRead):
await client.receive_exactly(1)
assert response == buffer
async def test_receive_some_from_cache(self, localhost):
async def server():
async with await stream_server.accept() as stream:
await stream.receive_until(b'a', 10)
request = await stream.receive_some(1)
await stream.send_all(request + b'\n')
async with create_task_group() as tg:
async with await create_tcp_server(interface=localhost
) as stream_server:
await tg.spawn(server)
async with await connect_tcp(localhost,
stream_server.port) as client:
await client.send_all(b'abc')
received = await client.receive_until(b'\n', 3)
assert received == b'b'
async def test_receive_delimited_chunks(self, localhost):
async def server():
async with await stream_server.accept() as stream:
async for chunk in stream.receive_delimited_chunks(b'\r\n', 8):
chunks.append(chunk)
chunks = []
async with await create_tcp_server(interface=localhost
) as stream_server:
async with create_task_group() as tg:
await tg.spawn(server)
async with await connect_tcp(localhost,
stream_server.port) as client:
for chunk in (b'bl', b'ah', b'\r', b'\nfoo', b'bar\r\n'):
await client.send_all(chunk)
assert chunks == [b'blah', b'foobar']
async def as_service(obj=None):
"""
This async context manager provides readiness and keepalive messages to
systemd.
Arguments:
obj: command context. Needs a ``debug`` attribute.
The CM yields a (duck-typed) event whose async ``set`` method will
trigger a ``READY=1`` mesage to systemd.
"""
from systemd.daemon import notify # pylint: disable=no-name-in-module
async def run_keepalive(usec):
usec /= 1_500_000 # 2/3rd of usec ⇒ sec
pid = os.getpid()
while os.getpid() == pid:
notify("WATCHDOG=1")
await anyio.sleep(usec)
def need_keepalive():
pid = os.getpid()
epid = int(os.environ.get("WATCHDOG_PID", pid))
if pid == epid:
return int(os.environ.get("WATCHDOG_USEC", 0))
class RunMsg:
def __init__(self, obj):
self.obj = obj
def set(self):
# TODO: this should be async (set flag and separate thread)
notify("READY=1")
if self.obj is not None and self.obj.debug:
print("Running.")
async with anyio.create_task_group() as tg:
usec = need_keepalive()
if usec:
tg.spawn(run_keepalive, usec)
try:
yield RunMsg(obj)
finally:
with anyio.fail_after(2, shield=True):
tg.cancel_scope.cancel()
async def test_get_running_tasks():
event = create_event()
async with create_task_group() as tg:
existing_tasks = set(get_running_tasks())
await tg.spawn(event.wait, name='task1')
await tg.spawn(event.wait, name='task2')
await wait_all_tasks_blocked()
task_infos = set(get_running_tasks()) - existing_tasks
await event.set()
task_infos = sorted(task_infos, key=lambda info: info.name or '')
assert len(task_infos) == 2
assert task_infos[0].name == 'task1'
assert task_infos[1].name == 'task2'
assert re.fullmatch(r"TaskInfo\(id=\d+, name='task1'\)",
repr(task_infos[0]))
assert re.fullmatch(r"TaskInfo\(id=\d+, name='task2'\)",
repr(task_infos[1]))
async def test_nested_fail_after():
async def killer(scope):
await wait_all_tasks_blocked()
await scope.cancel()
async with create_task_group() as tg:
async with open_cancel_scope() as scope:
async with open_cancel_scope():
await tg.spawn(killer, scope)
async with fail_after(1):
await sleep(2)
pytest.fail('Execution should not reach this point')
pytest.fail('Execution should not reach this point either')
pytest.fail('Execution should also not reach this point')
assert scope.cancel_called
async def test_exception_cancels_siblings():
async def child(fail):
if fail:
raise Exception('foo')
else:
nonlocal sleep_completed
await sleep(1)
sleep_completed = True
sleep_completed = False
with pytest.raises(Exception) as exc:
async with create_task_group() as tg:
await tg.spawn(child, False)
await wait_all_tasks_blocked()
await tg.spawn(child, True)
exc.match('foo')
assert not sleep_completed
async def __aenter__(self):
asynclib = sniffio.current_async_library() or '(unknown)'
if asynclib != 'asyncio':
raise RuntimeError(f'This data store requires asyncio; currently running: {asynclib}')
if self._loans == 0:
self._schedules_event = asyncio.Event()
self._jobs_event = asyncio.Event()
await self._setup()
self._loans += 1
if self._loans == 1 and self._watch_collections:
self._task_group = create_task_group()
await self._task_group.__aenter__()
await self._task_group.spawn(self._watch_collection, self._schedules)
await self._task_group.spawn(self._watch_collection, self._jobs)
return self
async def connect(host="127.0.0.1", port=8282, **kw):
"""A context manager for managing the Resp object"""
async with anyio.create_task_group() as tg:
async with await anyio.connect_tcp(host, port) as st:
s = Resp(st, **kw)
evt = anyio.create_event()
await tg.spawn(s.send, evt)
await evt.wait()
await tg.spawn(reader, s)
try:
yield s
finally:
async with anyio.move_on_after(5, shield=True):
await s.flush()
await anyio.sleep(
0.3) # let's hope that's enough to read errors
await tg.cancel_scope.cancel()
async def _run_with_tg(self, evt: anyio.abc.Event=None):
try:
async with anyio.create_task_group() as tg:
self._tg = tg
await self._task_setup()
if evt is not None:
await evt.set()
await self.run(evt=evt)
except Exception as exc:
self._run_with_exc = exc
if self._run_with_scope is not None:
await self._run_with_scope.cancel()
finally:
self._tg = None
await self._task_teardown()
if self._done is not None:
await self._done.set()
self._done = None
async def _run_async_server(self, raw_socket):
socket = anyio._get_asynclib().Socket(raw_socket)
# TODO: resource usage warning
async with async_exit_stack.AsyncExitStack() as stack:
tcp_server = await stack.enter_async_context(
anyio._networking.SocketStreamServer(socket, None, False,
False))
task_group = await stack.enter_async_context(
anyio.create_task_group())
services_dict = {}
for key, value in self.services.items():
services_dict[key] = await stack.enter_async_context(value)
async for socket in tcp_server.accept_connections():
await task_group.spawn(ConnectionHandler(services_dict),
socket)
async def do_pub(client, arguments):
logger.info("%s Connecting to broker", client.client_id)
await client.connect(uri=arguments['--url'],
cleansession=arguments['--clean-session'],
cafile=arguments['--ca-file'],
capath=arguments['--ca-path'],
cadata=arguments['--ca-data'],
extra_headers=_get_extra_headers(arguments))
try:
qos = _get_qos(arguments)
topic = arguments['-t']
retain = arguments['-r']
async for anyio.create_task_group() as tg:
for message in _get_message(arguments):
logger.info("%s Publishing to '%s'", client.client_id, topic)
await tg.sspawn(client.publish, topic, message, qos, retain)
logger.info("%s Disconnected from broker", client.client_id)
async def test_connect_unix(self, tmp_path_factory, as_path):
async def server():
async with await stream_server.accept() as stream:
command = await stream.receive_some(100)
await stream.send_all(command[::-1])
async with create_task_group() as tg:
path = str(tmp_path_factory.mktemp('unix').joinpath('socket'))
if as_path:
path = Path(path)
async with await create_unix_server(path) as stream_server:
await tg.spawn(server)
async with await connect_unix(path) as client:
await client.send_all(b'blah')
response = await client.receive_some(100)
assert response == b'halb'
async def test_concurrent_read(self, localhost):
async def receive_data():
await client.receive_exactly(1)
async with await create_tcp_server(interface=localhost
) as stream_server:
async with await connect_tcp(localhost,
stream_server.port) as client:
async with create_task_group() as tg:
await tg.spawn(receive_data)
await wait_all_tasks_blocked()
try:
with pytest.raises(ResourceBusyError) as exc:
await client.receive_exactly(1)
exc.match('already reading from')
finally:
await tg.cancel_scope.cancel()
async def OWFS(**kwargs):
async with anyio.create_task_group() as n:
s = Service(n, **kwargs)
async with s:
await yield_(s)
