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