async def test_nested_move_on_after():
sleep_completed = inner_scope_completed = False
async with move_on_after(0.1) as outer_scope:
assert await current_effective_deadline() == outer_scope.deadline
async with move_on_after(1) as inner_scope:
assert await current_effective_deadline() == outer_scope.deadline
await sleep(2)
sleep_completed = True
inner_scope_completed = True
assert not sleep_completed
assert not inner_scope_completed
assert outer_scope.cancel_called
assert not inner_scope.cancel_called
async def acquire_jobs(self,
worker_id: str,
limit: Optional[int] = None) -> List[Job]:
while True:
jobs: List[Job] = []
async with self.pool.acquire() as conn, conn.transaction():
now = datetime.now(timezone.utc)
acquired_until = datetime.fromtimestamp(
now.timestamp() + self.lock_expiration_delay, timezone.utc)
records = await conn.fetch(
f"""
WITH job_ids AS (
SELECT id FROM {self.schema}.jobs
WHERE acquired_until IS NULL OR acquired_until < $1
ORDER BY created_at
FOR NO KEY UPDATE SKIP LOCKED
FETCH FIRST $2 ROWS ONLY
)
UPDATE {self.schema}.jobs SET acquired_by = $3, acquired_until = $4
WHERE id IN (SELECT id FROM job_ids)
RETURNING serialized_data
""", now, limit, worker_id, acquired_until)
for record in records:
job = self.serializer.deserialize(record['serialized_data'])
jobs.append(job)
if jobs:
return jobs
async with move_on_after(self.max_poll_time):
await self._jobs_event.wait()
self._jobs_event.clear()
async def acquire_schedules(self, scheduler_id: str,
limit: int) -> List[Schedule]:
while True:
schedules: List[Schedule] = []
async with self.pool.acquire() as conn, conn.transaction():
acquired_until = datetime.fromtimestamp(
datetime.now(timezone.utc).timestamp() +
self.lock_expiration_delay, timezone.utc)
records = await conn.fetch(
f"""
WITH schedule_ids AS (
SELECT id FROM {self.schema}.schedules
WHERE next_fire_time IS NOT NULL AND next_fire_time <= $1
AND (acquired_until IS NULL OR $1 > acquired_until)
ORDER BY next_fire_time
FOR NO KEY UPDATE SKIP LOCKED
FETCH FIRST $2 ROWS ONLY
)
UPDATE {self.schema}.schedules SET acquired_by = $3, acquired_until = $4
WHERE id IN (SELECT id FROM schedule_ids)
RETURNING serialized_data
""", datetime.now(timezone.utc), limit, scheduler_id,
acquired_until)
for record in records:
schedule = self.serializer.deserialize(
record['serialized_data'])
schedules.append(schedule)
if schedules:
return schedules
async with move_on_after(self.max_poll_time):
await self._schedules_event.wait()
async def debouncer(self, chan, q):
while True:
e1 = await q.get()
logger.debug("new %s %s", self.name, e1)
val = not self.value
# something happens. Send an event immediately,
# no matter the event's value.
await self.handle_event(val, chan)
e2 = None
while True:
async with anyio.move_on_after(abs(self.debounce)) as skip:
e2 = await q.get()
logger.debug("and %s %s", self.name, e2)
if skip.cancel_called:
break
logger.debug("done %s %s", self.name, e2)
if self.debounce < 0:
if e2 is None or e1.value == e2.value:
await self.handle_event(e1, chan)
else:
await self.handle_event(e1, chan)
if e2 is not None and e1.value != e2.value:
await self.handle_event(e2, chan)
async def multi_sleep(self, durations: List[float], timeout: float) -> bool:
async with anyio.move_on_after(timeout) as cancel_scope:
async with anyio.create_task_group() as task_group:
for duration in durations:
await task_group.spawn(self.sleep, duration)
return not cancel_scope.cancel_called
async def __aexit__(self, *tb):
self._q = None
try:
async with anyio.move_on_after(2, shield=True):
await self.client._unsubscribe(self)
except ClientException:
pass
async def test_streaming_response_stops_if_receiving_http_disconnect():
streamed = 0
disconnected = anyio.Event()
async def receive_disconnect():
await disconnected.wait()
return {"type": "http.disconnect"}
async def send(message):
nonlocal streamed
if message["type"] == "http.response.body":
streamed += len(message.get("body", b""))
# Simulate disconnection after download has started
if streamed >= 16:
disconnected.set()
async def stream_indefinitely():
while True:
# Need a sleep for the event loop to switch to another task
await anyio.sleep(0)
yield b"chunk "
response = StreamingResponse(content=stream_indefinitely())
with anyio.move_on_after(1) as cancel_scope:
await response({}, receive_disconnect, send)
assert not cancel_scope.cancel_called, "Content streaming should stop itself."
async def _next_to_send(self):
"""
Returns the next message on the heap
"""
while True:
while self._heap:
if self._heap_large is not None and len(
self._heap) < self._heap_max / 2:
await self._heap_large.set()
self._heap_large = None
if self._ending.is_set(
) or self._heap[0].time <= self._t - self._delay:
return heapq.heappop(self._heap)
self._t = time.time()
if self._heap[0].time <= self._t - self._delay:
return heapq.heappop(self._heap)
async with anyio.move_on_after(
max(self._delay + self._heap[0].time - self._t, 0)):
await self._ending.wait()
await self.flush_buf()
if self._done is not None:
await self._done.set()
self._heap_item = anyio.create_event()
await self._heap_item.wait()
async def test_acquire_schedules_lock_timeout(store, schedules, events,
freezer):
"""
Test that a scheduler can acquire schedules that were acquired by another scheduler but not
released within the lock timeout period.
"""
# First, one scheduler acquires the first available schedule
await store.add_schedule(schedules[0], ConflictPolicy.exception)
acquired = await store.acquire_schedules('dummy-id1', 1)
assert len(acquired) == 1
assert acquired[0].id == 's1'
# Try to acquire the schedule just at the threshold (now == acquired_until).
# This should not yield any schedules.
freezer.tick(30)
async with move_on_after(0.2):
await store.acquire_schedules('dummy-id2', 1)
pytest.fail('The call should have timed out')
# Right after that, the schedule should be available
freezer.tick(1)
acquired = await store.acquire_schedules('dummy-id2', 1)
assert len(acquired) == 1
assert acquired[0].id == 's1'
async def teardown(self, hangup_reason="normal"):
"""Removes all channels from the bridge and destroys it.
All remaining channels are hung up.
This method is typically called when leaving the bridge's context
manager. If you want to keep it online, e.g. for being able to
cleanly restart a PBX without downtime, you may override this --
but you're then responsible for recovering state after restarting,
and you still need to clean up bridges that are no longer needed.
"""
if self.bridge is None:
return
async with anyio.move_on_after(2, shield=True) as s:
log.info("TEARDOWN %s %s", self, self.bridge.channels)
for ch in list(self.bridge.channels) + list(self.calls):
try:
await ch.hang_up(reason=hangup_reason)
except Exception as exc:
log.info("%s gone: %s", ch, exc)
try:
await self.bridge.removeChannel(channel=ch.id)
except Exception as exc:
log.info("%s detached: %s", ch, exc)
await self.bridge.destroy()
async def self_cancel(*, task_status: TaskStatus) -> None:
nonlocal done
task_status.started()
with move_on_after(-1):
await Event().wait()
done = True
async def acquire_schedules(self, scheduler_id: str, limit: int) -> List[Schedule]:
while True:
schedules: List[Schedule] = []
async with await self.client.start_session() as s, s.start_transaction():
cursor = self._schedules.find(
{'next_fire_time': {'$ne': None},
'$or': [{'acquired_until': {'$exists': False}},
{'acquired_until': {'$lt': datetime.now(timezone.utc)}}]
},
projection=['serialized_data']
).sort('next_fire_time')
for document in await cursor.to_list(length=limit):
schedule = self.serializer.deserialize(document['serialized_data'])
schedules.append(schedule)
if schedules:
now = datetime.now(timezone.utc)
acquired_until = datetime.fromtimestamp(
now.timestamp() + self.lock_expiration_delay, now.tzinfo)
filters = {'_id': {'$in': [schedule.id for schedule in schedules]}}
update = {'$set': {'acquired_by': scheduler_id,
'acquired_until': acquired_until}}
await self._schedules.update_many(filters, update)
return schedules
async with move_on_after(self.max_poll_time):
await self._schedules_event.wait()
async def _sender_loop(self, evt):
keepalive_timeout = self.session.keep_alive
if keepalive_timeout <= 0:
keepalive_timeout = None
try:
async with anyio.open_cancel_scope() as scope:
self._sender_task = scope
await evt.set()
while True:
packet = None
async with anyio.move_on_after(keepalive_timeout):
packet = await self._send_q.get()
if packet is None: # closing
break
if packet is None: # timeout
await self.handle_write_timeout()
continue
# self.logger.debug("%s > %r",'B' if 'Broker' in type(self).__name__ else 'C', packet)
await packet.to_stream(self.stream)
await self.plugins_manager.fire_event(
EVENT_MQTT_PACKET_SENT,
packet=packet,
session=self.session)
except ConnectionResetError:
await self.handle_connection_closed()
except anyio.get_cancelled_exc_class():
raise
except BaseException as e:
self.logger.warning("Unhandled exception", exc_info=e)
raise
finally:
async with anyio.fail_after(2, shield=True):
await self._sender_stopped.set()
self._sender_task = None
async def acquire_schedules(self, scheduler_id: str,
limit: int) -> List[Schedule]:
while True:
now = datetime.now(timezone.utc)
schedules: List[Schedule] = []
wait_time = None
for state in self._schedules:
if state.acquired_by is not None and state.acquired_until >= now:
continue
elif state.next_fire_time is None:
break
elif state.next_fire_time > now:
wait_time = state.next_fire_time.timestamp(
) - now.timestamp()
break
schedules.append(state.schedule)
state.acquired_by = scheduler_id
state.acquired_until = datetime.fromtimestamp(
now.timestamp() + self.lock_expiration_delay, now.tzinfo)
if len(schedules) == limit:
break
if schedules:
return schedules
# Wait until reaching the next known fire time, or a schedule is added or updated
async with move_on_after(wait_time):
await self._schedules_event.wait()
async def test_acquire_release_jobs(store, jobs, events):
for job in jobs:
await store.add_job(job)
events.clear()
# The first worker gets the first job in the queue
jobs1 = await store.acquire_jobs('dummy-id1', 1)
assert len(jobs1) == 1
assert jobs1[0].id == jobs[0].id
# The second worker gets the second job
jobs2 = await store.acquire_jobs('dummy-id2', 1)
assert len(jobs2) == 1
assert jobs2[0].id == jobs[1].id
# The third worker gets nothing
async with move_on_after(0.2):
await store.acquire_jobs('dummy-id3', 1)
pytest.fail('The call should have timed out')
# All the jobs should still be returned
visible_jobs = await store.get_jobs()
assert len(visible_jobs) == 2
await store.release_jobs('dummy-id1', jobs1)
await store.release_jobs('dummy-id2', jobs2)
# All the jobs should be gone
visible_jobs = await store.get_jobs()
assert len(visible_jobs) == 0
async def test_escaping_cancelled_error_from_cancelled_task():
"""Regression test for issue #88. No CancelledError should escape the outer scope."""
async with open_cancel_scope() as scope:
async with move_on_after(0.1):
await sleep(1)
await scope.cancel()
async def mqtt_publish(self, topic, data, qos, retain, ack_timeout=None):
"""
Sends a MQTT publish message and manages messages flows.
This methods doesn't return until the message has been acknowledged by receiver or timeout occur
:param topic: MQTT topic to publish
:param data: data to send on topic
:param qos: quality of service to use for message flow. Can be QOS_0, QOS_1 or QOS_2
:param retain: retain message flag
:param ack_timeout: acknowledge timeout. If set, this method will return a TimeOut error if the acknowledgment
is not completed before ack_timeout second
:return: ApplicationMessage used during inflight operations
"""
if qos in (QOS_1, QOS_2):
packet_id = self.session.next_packet_id
if packet_id in self.session.inflight_out:
raise HBMQTTException(
"A message with the same packet ID '%d' is already in flight"
% packet_id)
else:
packet_id = None
message = OutgoingApplicationMessage(packet_id, topic, qos, data,
retain)
# Handle message flow
if ack_timeout is not None and ack_timeout > 0:
async with anyio.move_on_after(ack_timeout):
await self._handle_message_flow(message)
else:
await self._handle_message_flow(message)
return message
async def acquire_jobs(self, worker_id: str, limit: Optional[int] = None) -> List[Job]:
jobs: List[Job] = []
while True:
async with await self.client.start_session() as s, s.start_transaction():
cursor = self._jobs.find(
{'$or': [{'acquired_until': {'$exists': False}},
{'acquired_until': {'$lt': datetime.now(timezone.utc)}}]
},
projection=['serialized_data'],
sort=[('created_at', ASCENDING)]
)
for document in await cursor.to_list(length=limit):
job = self.serializer.deserialize(document['serialized_data'])
jobs.append(job)
if jobs:
now = datetime.now(timezone.utc)
acquired_until = datetime.fromtimestamp(
now.timestamp() + self.lock_expiration_delay, timezone.utc)
filters = {'_id': {'$in': [job.id for job in jobs]}}
update = {'$set': {'acquired_by': worker_id,
'acquired_until': acquired_until}}
await self._jobs.update_many(filters, update)
return jobs
async with move_on_after(self.max_poll_time):
await self._jobs_event.wait()
self._jobs_event.clear()
async def test_escaping_cancelled_error_from_cancelled_task() -> None:
"""Regression test for issue #88. No CancelledError should escape the outer scope."""
with CancelScope() as scope:
with move_on_after(0.1):
await sleep(1)
scope.cancel()
async def _handle_disconnect(self, disconnect, wait=True): # pylint: disable=arguments-differ
self.logger.debug("Client disconnecting")
self.clean_disconnect = False # depending on 'disconnect' (if set)
async with anyio.fail_after(2, shield=True):
if wait:
async with anyio.move_on_after(self.session.keep_alive):
await self._reader_stopped.wait()
await self.stop()
async def _handle_disconnect(self, disconnect, wait=True):
self.logger.debug("Client disconnecting")
self.clean_disconnect = False # depending on 'disconnect' (if set)
async with anyio.open_cancel_scope(shield=True):
if wait:
async with anyio.move_on_after(self.session.keep_alive):
await self._reader_stopped.wait()
await self.stop()
async def test_move_on_after(delay: float) -> None:
result = False
with move_on_after(delay) as scope:
await sleep(1)
result = True
assert not result
assert scope.cancel_called
async def test_should_work_when_using_task_done_and_join_methods(self):
async with Queue(size=1) as queue:
await queue.put(2)
await queue.get()
assert 1 == queue._tasks_in_progress
with anyio.move_on_after(1) as cancel_scope:
await queue.join()
assert cancel_scope.cancel_called
queue.task_done()
assert 0 == queue._tasks_in_progress
with anyio.move_on_after(1) as cancel_scope:
await queue.join()
assert not cancel_scope.cancel_called
async def __aexit__(self, *tb):
if not self.channel.is_open:
return
async with anyio.move_on_after(2, shield=True):
try:
await self.channel.close()
except exceptions.AmqpClosedConnection:
pass
async def socketSend(ws):
async with capture_continuous() as agen:
with anyio.move_on_after(20):
async for frame, msg in agen:
print("sending message")
await ws.send(msg)
print("waiting for response")
print(await ws.recv())
async def serve(self) -> None:
await self._master.register_service(self.service_name)
try:
await anyio.sleep_forever()
except anyio.get_cancelled_exc_class():
with anyio.CancelScope(shield=True), anyio.move_on_after(1):
await self._master.unregister_service(self.service_name)
raise
async def test_move_on_after():
result = False
async with move_on_after(0.1) as scope:
await sleep(1)
result = True
assert not result
assert scope.cancel_called
async def test_move_on_after_no_timeout():
result = False
async with move_on_after(None) as scope:
assert scope.deadline == float('inf')
await sleep(0.1)
result = True
assert result
assert not scope.cancel_called
async def process_one(message):
async with anyio.move_on_after(10):
try:
await self._handle_message_flow(message)
except CancelledError:
pass
nonlocal done
done += 1
async def process_one(message):
async with anyio.move_on_after(10):
try:
await self._handle_message_flow(message)
except CancelledError:
pass
nonlocal done
done += 1 # pylint: disable=undefined-variable ## fixed in 2.5
