async def mytask():
nonlocal mysum, detect_cancellation
try:
async with aiotools.aclosing(myiter()) as agen:
async for x in agen:
mysum += x
except asyncio.CancelledError:
detect_cancellation = True
async def detect_manager(self):
log.info('detecting the manager...')
manager_instances = await self.etcd.get_prefix('nodes/manager')
if not manager_instances:
log.warning('watching etcd to wait for the manager being available')
async with aclosing(self.etcd.watch_prefix('nodes/manager')) as agen:
async for ev in agen:
if ev.event == 'put' and ev.value == 'up':
break
log.info('detected at least one manager running')
async def detect_status_update(app):
try:
async with aclosing(
app['shared_config'].watch_manager_status()) as agen:
async for ev in agen:
if ev.event == 'put':
app['shared_config'].get_manager_status.cache_clear()
updated_status = await app['shared_config'
].get_manager_status()
log.debug(
'Process-{0} detected manager status update: {1}',
app['pidx'], updated_status)
except asyncio.CancelledError:
pass
async def get_performance_metric(self) -> FSPerfMetric:
async with self.purity_client as client:
async with aclosing(
client.get_nfs_metric(
self.config["purity_fs_name"]), ) as items:
async for item in items:
return FSPerfMetric(
iops_read=item["reads_per_sec"],
iops_write=item["writes_per_sec"],
io_bytes_read=item["read_bytes_per_sec"],
io_bytes_write=item["write_bytes_per_sec"],
io_usec_read=item["usec_per_read_op"],
io_usec_write=item["usec_per_write_op"],
)
else:
raise RuntimeError(
"no metric found for the configured flashblade filesystem",
)
async def _recv_loop(self) -> None:
if self._connection is None:
raise RuntimeError('consumer is not opened yet.')
while True:
try:
async with aclosing(self._connection.recv_message()) as agen:
async for raw_msg in agen:
if raw_msg is None:
return
msg = StreamMessage.decode(raw_msg, self._deserializer)
for handler in self._handler_registry:
if asyncio.iscoroutinefunction(handler):
asyncio.create_task(
self._task_wrapper(handler, msg))
# TODO: keep weak-refs to tasks and use it to
# cancel running tasks upon shutdown
else:
handler(msg)
except asyncio.CancelledError:
break
except Exception:
log.exception('unexpected error')
async def invoke(self, method: str, body, *,
order_key=None, invoke_timeout=None):
'''
Invoke a remote function via the transport connection.
'''
if invoke_timeout is None:
invoke_timeout = self._invoke_timeout
if order_key is None:
order_key = b''
client_seq_id = self._next_client_seq_id()
try:
request: RPCMessage
server_cancelled = False
with timeout(invoke_timeout):
if callable(body):
# The user is using an upper-layer adaptor.
async with aclosing(body()) as agen:
request = RPCMessage(
None,
RPCMessageTypes.FUNCTION,
method,
order_key,
client_seq_id,
None,
await agen.asend(None),
)
await self._outgoing_queue.put(request)
response = await self._invocation_resolver.wait(
request.request_id)
upper_result = await agen.asend(response.body)
try:
await agen.asend(None)
except StopAsyncIteration:
pass
else:
request = RPCMessage(
None,
RPCMessageTypes.FUNCTION,
method,
order_key,
client_seq_id,
None,
body,
)
await self._outgoing_queue.put(request)
response = await self._invocation_resolver.wait(
request.request_id)
upper_result = response.body
if response.msgtype == RPCMessageTypes.RESULT:
pass
elif response.msgtype == RPCMessageTypes.CANCELLED:
server_cancelled = True
raise asyncio.CancelledError
elif response.msgtype == RPCMessageTypes.FAILURE:
raise RPCUserError(*attr.astuple(response.metadata))
elif response.msgtype == RPCMessageTypes.ERROR:
raise RPCInternalError(*attr.astuple(response.metadata))
return upper_result
except (asyncio.TimeoutError, asyncio.CancelledError):
# propagate cancellation to the connected peer
if not server_cancelled:
cancel_request = RPCMessage.cancel(request)
await self._outgoing_queue.put(cancel_request)
raise
except Exception:
raise
async def _recv_loop(self) -> None:
'''
Receive requests and schedule the request handlers.
'''
if self._connection is None:
raise RuntimeError('consumer is not opened yet.')
func_tasks: Set[asyncio.Task] = set()
while True:
try:
async with aclosing(self._connection.recv_message()) as agen:
async for raw_msg in agen:
# TODO: flow-control in transports or peer queues?
if raw_msg is None:
return
request = RPCMessage.decode(raw_msg, self._deserializer)
client_request_id = request.request_id
server_request_id: Optional[RequestId]
if request.msgtype == RPCMessageTypes.FUNCTION:
server_seq_id = self._next_server_seq_id()
server_request_id = (
client_request_id[0],
client_request_id[1],
server_seq_id,
)
self._req_idmap[(request.peer_id, client_request_id)] = \
server_request_id
func_handler = self._lookup_func(request.method)
task = asyncio.create_task(self._func_task(
server_request_id,
request,
func_handler,
))
func_tasks.add(task)
task.add_done_callback(func_tasks.discard)
task.add_done_callback(
lambda task: self._req_idmap.pop(
(request.peer_id, client_request_id),
None,
)
)
await asyncio.sleep(0)
elif request.msgtype == RPCMessageTypes.CANCEL:
server_request_id = self._req_idmap.pop(
(request.peer_id, client_request_id),
None,
)
if server_request_id is None:
continue
await self._func_scheduler.cancel(server_request_id)
elif request.msgtype in (RPCMessageTypes.RESULT,
RPCMessageTypes.FAILURE,
RPCMessageTypes.CANCELLED,
RPCMessageTypes.ERROR):
self._invocation_resolver.resolve(
client_request_id,
request,
)
await asyncio.sleep(0)
except asyncio.CancelledError:
pending_tasks = []
if func_tasks:
for task in func_tasks:
if not task.done():
task.cancel()
pending_tasks.append(task)
await asyncio.wait(pending_tasks)
await asyncio.sleep(0)
break
except Exception:
log.exception('unexpected error')
async def watch_manager_status(self):
async with aiotools.aclosing(
self.etcd.watch('manager/status')) as agen:
async for ev in agen:
yield ev