Python asyncio.FIRST_COMPLETED Examples

def wait_for_first_response(tasks, converters):
    """given a list of unawaited tasks and non-coro result parsers to be called on the results,
    this function returns the 1st result that is returned and converted

    if it is possible for 2 tasks to complete at the same time,
    only the 1st result deteremined by asyncio.wait will be returned

    returns None if none successfully complete
    returns 1st error raised if any occur (probably)
    """
    primed = [wait_for_result(t, c) for t, c in zip(tasks, converters)]
    done, pending = await asyncio.wait(primed, return_when=asyncio.FIRST_COMPLETED)
    for p in pending:
        p.cancel()

    try:
        return done.pop().result()
    except NotImplementedError as e:
        raise e
    except:
        return None 

def connect(self, timeout: Optional[float] = None) -> None:
        """Establish a connection with the WebSocket server.

        Returns:
            `(True, <chosen-subprotocol>)` if connection accepted.
            `(False, None)` if connection rejected.

        """
        connected = asyncio.Event()
        self._message_processor = asyncio.ensure_future(
            self._process_messages(connected, timeout or self.TIMEOUT)
        )
        await asyncio.wait(
            [connected.wait(), self._message_processor],
            return_when=asyncio.FIRST_COMPLETED,
        )
        if self._message_processor.done():
            # Make sure to raise an exception from the task.
            self._message_processor.result()
            raise RuntimeError(f"Failed to connect to the server: {self._url}!") 

def _push_error_to_user(self, exc):
        """ Most critical errors are not something we can continue execution
        without user action. Well right now we just drop the Consumer, but
        java client would certainly be ok if we just poll another time, maybe
        it will need to rejoin, but not fail with GroupAuthorizationFailedError
        till the end of days...
        XXX: Research if we can't have the same error several times. For
             example if user gets GroupAuthorizationFailedError and adds
             permission for the group, would Consumer work right away or would
             still raise exception a few times?
        """
        exc = copy.copy(exc)
        self._subscription.abort_waiters(exc)
        self._pending_exception = exc
        self._error_consumed_fut = create_future(loop=self._loop)
        return asyncio.wait(
            [self._error_consumed_fut, self._closing],
            return_when=asyncio.FIRST_COMPLETED,
            loop=self._loop
        ) 

def stop(self):
        """Flush all pending data and close all connections to kafka cluster"""
        if self._closed:
            return
        self._closed = True

        # If the sender task is down there is no way for accumulator to flush
        if self._sender is not None and self._sender.sender_task is not None:
            await asyncio.wait([
                self._message_accumulator.close(),
                self._sender.sender_task],
                return_when=asyncio.FIRST_COMPLETED,
                loop=self._loop)

            await self._sender.close()

        await self.client.close()
        log.debug("The Kafka producer has closed.") 

def onConnect (self, **kwargs):
        self.logger.info ('connect', nick=self.nick, uuid='01f7b138-ea53-4609-88e9-61f3eca3e7e7')

        self.send('NICK', nick=self.nick)
        self.send('USER', user=self.nick, realname='https://github.com/PromyLOPh/crocoite')

        # Don't try to join channels until the server has
        # sent the MOTD, or signaled that there's no MOTD.
        done, pending = await asyncio.wait(
            [self.wait('RPL_ENDOFMOTD'), self.wait('ERR_NOMOTD')],
            loop=self.loop, return_when=asyncio.FIRST_COMPLETED)

        # Cancel whichever waiter's event didn't come in.
        for future in pending:
            future.cancel()

        for c in self.channels:
            self.logger.info ('join', channel=c, uuid='367063a5-9069-4025-907c-65ba88af8593')
            self.send ('JOIN', channel=c)
            # no need for NAMES here, server sends this automatically 

def _push_error_to_user(self, exc):
        """ Most critical errors are not something we can continue execution
        without user action. Well right now we just drop the Consumer, but
        java client would certainly be ok if we just poll another time, maybe
        it will need to rejoin, but not fail with GroupAuthorizationFailedError
        till the end of days...
        XXX: Research if we can't have the same error several times. For
             example if user gets GroupAuthorizationFailedError and adds
             permission for the group, would Consumer work right away or would
             still raise exception a few times?
        """
        exc = copy.copy(exc)
        self._subscription.abort_waiters(exc)
        self._pending_exception = exc
        self._error_consumed_fut = create_future(loop=self._loop)
        return asyncio.wait(
            [self._error_consumed_fut, self._closing],
            return_when=asyncio.FIRST_COMPLETED,
            loop=self._loop
        ) 

def poll(self, events, timeout=None,
                   return_when=asyncio.FIRST_COMPLETED):
        """Poll for any of a set of event types to be received for this session.
        """
        awaitables = {}
        for name in events:
            awaitables[self.recv(name)] = name
        done, pending = await asyncio.wait(
            awaitables, timeout=timeout, return_when=return_when)

        if done:
            ev_dicts = []
            for fut in done:
                awaitables.pop(fut)
                ev_dicts.append(fut.result())
            return ev_dicts, awaitables.values()
        else:
            raise asyncio.TimeoutError(
                "None of {} was received in {} seconds"
                .format(events, timeout))

    # call control / 'mod_commands' methods
    # TODO: dynamically add @decorated functions to this class
    # and wrap them using functools.update_wrapper ...? 

def stop(self):
        """Flush all pending data and close all connections to kafka cluster"""
        if self._closed:
            return
        self._closed = True

        # If the sender task is down there is no way for accumulator to flush
        if self._sender is not None and self._sender.sender_task is not None:
            await asyncio.wait([
                self._message_accumulator.close(),
                self._sender.sender_task],
                return_when=asyncio.FIRST_COMPLETED,
                loop=self._loop)

            await self._sender.close()

        await self.client.close()
        log.debug("The Kafka producer has closed.") 

def waiter(client):
    async def wait_for(*events, return_when=asyncio.FIRST_COMPLETED):
        if not events:
            return
        done, pending = await asyncio.wait(
            [client.wait(event) for event in events],
            loop=client.loop,
            return_when=return_when)

        # Get the result(s) of the completed task(s).
        ret = [future.result() for future in done]

        # Cancel any events that didn't come in.
        for future in pending:
            future.cancel()

        # Return list of completed event names.
        return ret
    return wait_for


# taken from :ref:`Patterns` 

def _handler(self, websocket: websockets.WebSocketCommonProtocol, *unused_args):
        """Setup the consumer and producer response handlers with asyncio.

        Args:
            websocket: the websocket connection to the client

        """

        asyncio.set_event_loop(self.loop)
        consumer_task = asyncio.ensure_future(self._consumer_handler(websocket))
        producer_task = asyncio.ensure_future(self._producer_handler(websocket))

        done, pending = await asyncio.wait(
            [consumer_task, producer_task], return_when=asyncio.FIRST_COMPLETED
        )

        for task in pending:
            task.cancel() 

def daemon():
    try_init_cgroup()

    async with VJ4Session(config['server_url']) as session:
        while True:
            try:
                await session.login_if_needed(config['uname'], config['password'])
                done, pending = await wait([do_judge(session), do_noop(session)],
                                           return_when=FIRST_COMPLETED)
                for task in pending:
                    task.cancel()
                await gather(*done)
            except Exception as e:
                logger.exception(e)
            logger.info('Retrying after %d seconds', RETRY_DELAY_SEC)
            await sleep(RETRY_DELAY_SEC) 

def ws_handler(self, request):

        ws = web.WebSocketResponse()
        await ws.prepare(request)

        if self.ws is not None:
            await self.ws.close()

        self.ws = ws

        _, unfinished = await asyncio.wait(
            [
                self._websocket_receive(ws),
                self._websocket_send(ws)
            ],
            return_when=asyncio.FIRST_COMPLETED
        )
        for task in unfinished:
            task.cancel()

        return ws 

def mqtt_connected(func):
    """
        MQTTClient coroutines decorator which will wait until connection before calling the decorated method.
        :param func: coroutine to be called once connected
        :return: coroutine result
    """
    @asyncio.coroutine
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if not self._connected_state.is_set():
            base_logger.warning("Client not connected, waiting for it")
            _, pending = yield from asyncio.wait([self._connected_state.wait(), self._no_more_connections.wait()], return_when=asyncio.FIRST_COMPLETED)
            for t in pending:
                t.cancel()
            if self._no_more_connections.is_set():
                raise ClientException("Will not reconnect")
        return (yield from func(self, *args, **kwargs))
    return wrapper 

def _run(self):
        first_completed = asyncio.FIRST_COMPLETED

        if self._get_max_tasks() < 1:
            raise RuntimeError("Executor has no workers")

        try:
            while not self.goal(self.learner):
                futures = self._get_futures()
                done, _ = await asyncio.wait(
                    futures, return_when=first_completed, loop=self.ioloop
                )
                self._process_futures(done)
        finally:
            remaining = self._remove_unfinished()
            if remaining:
                await asyncio.wait(remaining)
            self._cleanup() 

def _run(self):
        first_completed = concurrent.FIRST_COMPLETED

        if self._get_max_tasks() < 1:
            raise RuntimeError("Executor has no workers")

        try:
            while not self.goal(self.learner):
                futures = self._get_futures()
                done, _ = concurrent.wait(futures, return_when=first_completed)
                self._process_futures(done)
        finally:
            remaining = self._remove_unfinished()
            if remaining:
                concurrent.wait(remaining)
            self._cleanup() 

def cancellable_wait(self, *awaitables: Awaitable[_R], timeout: float = None) -> _R:
        futures = [asyncio.ensure_future(a, loop=self.loop) for a in awaitables + (self.wait(),)]
        try:
            done, pending = await asyncio.wait(
                futures,
                timeout=timeout,
                return_when=asyncio.FIRST_COMPLETED,
                loop=self.loop,
            )
        except CancelledError:
            for future in futures:
                future.cancel()
            raise
        for task in pending:
            task.cancel()
        await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED, loop=self.loop,)
        if not done:
            raise TimeoutError()
        if self.triggered_token is not None:
            for task in done:
                task.exception()
            raise OperationCancelled(f'Cancellation requested by {self.triggered_token} token')
        return done.pop().result() 

def test_wait_first_completed(self):

        def gen():
            when = yield
            self.assertAlmostEqual(10.0, when)
            when = yield 0
            self.assertAlmostEqual(0.1, when)
            yield 0.1

        loop = self.new_test_loop(gen)

        a = asyncio.Task(asyncio.sleep(10.0, loop=loop), loop=loop)
        b = asyncio.Task(asyncio.sleep(0.1, loop=loop), loop=loop)
        task = asyncio.Task(
            asyncio.wait([b, a], return_when=asyncio.FIRST_COMPLETED,
                         loop=loop),
            loop=loop)

        done, pending = loop.run_until_complete(task)
        self.assertEqual({b}, done)
        self.assertEqual({a}, pending)
        self.assertFalse(a.done())
        self.assertTrue(b.done())
        self.assertIsNone(b.result())
        self.assertAlmostEqual(0.1, loop.time())

        # move forward to close generator
        loop.advance_time(10)
        loop.run_until_complete(asyncio.wait([a, b], loop=loop)) 

def wshandler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    recv_task = None
    tick_task = None
    while 1:
        if not recv_task:
            recv_task = asyncio.ensure_future(ws.receive())
        if not tick_task:
            await tick.acquire()
            tick_task = asyncio.ensure_future(tick.wait())

        done, pending = await asyncio.wait(
            [recv_task,
             tick_task],
            return_when=asyncio.FIRST_COMPLETED)

        if recv_task in done:
            msg = recv_task.result()
            if msg.tp == web.MsgType.text:
                print("Got message %s" % msg.data)
                ws.send_str("Pressed key code: {}".format(msg.data))
            elif msg.tp == web.MsgType.close or\
                 msg.tp == web.MsgType.error:
                break
            recv_task = None

        if tick_task in done:
            ws.send_str("game loop ticks")
            tick.release()
            tick_task = None

    return ws 

def wshandler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    recv_task = None
    tick_task = None
    while 1:
        if not recv_task:
            recv_task = asyncio.ensure_future(ws.receive())
        if not tick_task:
            await tick.acquire()
            tick_task = asyncio.ensure_future(tick.wait())

        done, pending = await asyncio.wait(
            [recv_task,
             tick_task],
            return_when=asyncio.FIRST_COMPLETED)

        if recv_task in done:
            msg = recv_task.result()
            if msg.tp == web.MsgType.text:
                print("Got message %s" % msg.data)
                ws.send_str("Pressed key code: {}".format(msg.data))
            elif msg.tp == web.MsgType.close or\
                 msg.tp == web.MsgType.error:
                break
            recv_task = None

        if tick_task in done:
            ws.send_str("game loop ticks")
            tick.release()
            tick_task = None

    return ws 

def test_wait_really_done(self):
        # there is possibility that some tasks in the pending list
        # became done but their callbacks haven't all been called yet

        @asyncio.coroutine
        def coro1():
            yield

        @asyncio.coroutine
        def coro2():
            yield
            yield

        a = asyncio.Task(coro1(), loop=self.loop)
        b = asyncio.Task(coro2(), loop=self.loop)
        task = asyncio.Task(
            asyncio.wait([b, a], return_when=asyncio.FIRST_COMPLETED,
                         loop=self.loop),
            loop=self.loop)

        done, pending = self.loop.run_until_complete(task)
        self.assertEqual({a, b}, done)
        self.assertTrue(a.done())
        self.assertIsNone(a.result())
        self.assertTrue(b.done())
        self.assertIsNone(b.result()) 

def receive(self, *args, **kwargs) -> Optional["Envelope"]:
        """
        Receive an envelope.

        :return: the received envelope, or None if an error occurred.
        """
        if len(self._read_tasks_to_address) == 0:
            logger.warning(
                "Tried to read from the TCP server. However, there is no open connection to read from."
            )
            return None

        try:
            logger.debug("Waiting for incoming messages...")
            done, pending = await asyncio.wait(self._read_tasks_to_address.keys(), return_when=asyncio.FIRST_COMPLETED)  # type: ignore

            # take the first
            task = next(iter(done))
            envelope_bytes = task.result()
            if envelope_bytes is None:  # pragma: no cover
                logger.debug("[{}]: No data received.")
                return None
            envelope = Envelope.decode(envelope_bytes)
            address = self._read_tasks_to_address.pop(task)
            reader = self.connections[address][0]
            new_task = asyncio.ensure_future(self._recv(reader), loop=self._loop)
            self._read_tasks_to_address[new_task] = address
            return envelope
        except asyncio.CancelledError:
            logger.debug("Receiving loop cancelled.")
            return None
        except Exception as e:
            logger.error("Error in the receiving loop: {}".format(str(e)))
            return None 

def _receiving_loop(self) -> None:
        """Process incoming envelopes."""
        logger.debug("Starting receving loop...")
        task_to_connection = {
            asyncio.ensure_future(conn.receive()): conn for conn in self.connections
        }

        while self.connection_status.is_connected and len(task_to_connection) > 0:
            try:
                logger.debug("Waiting for incoming envelopes...")
                done, _pending = await asyncio.wait(
                    task_to_connection.keys(), return_when=asyncio.FIRST_COMPLETED
                )

                # process completed receiving tasks.
                for task in done:
                    envelope = task.result()
                    if envelope is not None:
                        self.in_queue.put_nowait(envelope)

                    # reinstantiate receiving task, but only if the connection is still up.
                    connection = task_to_connection.pop(task)
                    if connection.connection_status.is_connected:
                        new_task = asyncio.ensure_future(connection.receive())
                        task_to_connection[new_task] = connection

            except asyncio.CancelledError:
                logger.debug("Receiving loop cancelled.")
                break
            except Exception as e:  # pylint: disable=broad-except
                logger.error("Error in the receiving loop: {}".format(str(e)))
                break

        # cancel all the receiving tasks.
        for t in task_to_connection.keys():
            t.cancel()
        logger.debug("Receiving loop terminated.") 

def wshandler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    recv_task = None
    tick_task = None
    while 1:
        if not recv_task:
            recv_task = asyncio.ensure_future(ws.receive())
        if not tick_task:
            await tick.acquire()
            tick_task = asyncio.ensure_future(tick.wait())

        done, pending = await asyncio.wait(
            [recv_task,
             tick_task],
            return_when=asyncio.FIRST_COMPLETED)

        if recv_task in done:
            msg = recv_task.result()
            if msg.tp == web.MsgType.text:
                print("Got message %s" % msg.data)
                ws.send_str("Pressed key code: {}".format(msg.data))
            elif msg.tp == web.MsgType.close or\
                 msg.tp == web.MsgType.error:
                break
            recv_task = None

        if tick_task in done:
            ws.send_str("game loop ticks")
            tick.release()
            tick_task = None

    return ws 

def start(self, coro):
        self.server_stopped = asyncio.create_task(coro)
        server_started = asyncio.create_task(self.begin.get())
        done, pending = await asyncio.wait(
            {self.server_stopped, server_started},
            return_when=asyncio.FIRST_COMPLETED,
        )
        # Raise issues if they happened
        for task in done:
            # This branch is only taken if tests fail
            if task is self.server_stopped:  # pragma: no cov
                exception = task.exception()
                if exception is not None:
                    raise exception
        return server_started.result() 

def wait_for_input(loop, abort_sem):
    """Wait for user input (enter) or abort signal."""
    reader = asyncio.StreamReader(loop=loop)
    reader_protocol = asyncio.StreamReaderProtocol(reader)
    await loop.connect_read_pipe(lambda: reader_protocol, sys.stdin)
    await asyncio.wait(
        [reader.readline(), abort_sem.acquire()], return_when=asyncio.FIRST_COMPLETED
    ) 

def _on_page_load_event_fired(self, obj: Dict, *, format: str) -> None:
        if format in ('mhtml', 'pdf'):
            await self._scroll_to_bottom()

        done, pending = await asyncio.wait([
            self._evaluate_prerender_ready(),
            self._wait_responses_ready(),
        ], return_when=asyncio.FIRST_COMPLETED)
        for task in pending:
            task.cancel()
        for task in done:
            task.result()  # To trigger exception if any

        status_code = await self.get_status_code()
        if status_code == 304:
            status_code = 200
        if format == 'html':
            html = await self.get_html()
            self._render_future.set_result((html, status_code))
        elif format == 'mhtml':
            self._render_future.set_result((bytes(self._mhtml), status_code))
        elif format == 'pdf':
            data = await self.print_to_pdf()
            self._render_future.set_result((data, status_code))
        elif format == 'jpeg' or format == 'png':
            data = await self.screenshot(format)
            self._render_future.set_result((data, status_code)) 

def test_wait_first_completed(self):

        def gen():
            when = yield
            self.assertAlmostEqual(10.0, when)
            when = yield 0
            self.assertAlmostEqual(0.1, when)
            yield 0.1

        loop = self.new_test_loop(gen)

        a = asyncio.Task(asyncio.sleep(10.0, loop=loop), loop=loop)
        b = asyncio.Task(asyncio.sleep(0.1, loop=loop), loop=loop)
        task = asyncio.Task(
            asyncio.wait([b, a], return_when=asyncio.FIRST_COMPLETED,
                         loop=loop),
            loop=loop)

        done, pending = loop.run_until_complete(task)
        self.assertEqual({b}, done)
        self.assertEqual({a}, pending)
        self.assertFalse(a.done())
        self.assertTrue(b.done())
        self.assertIsNone(b.result())
        self.assertAlmostEqual(0.1, loop.time())

        # move forward to close generator
        loop.advance_time(10)
        loop.run_until_complete(asyncio.wait([a, b], loop=loop)) 

def test_wait_really_done(self):
        # there is possibility that some tasks in the pending list
        # became done but their callbacks haven't all been called yet

        @asyncio.coroutine
        def coro1():
            yield

        @asyncio.coroutine
        def coro2():
            yield
            yield

        a = asyncio.Task(coro1(), loop=self.loop)
        b = asyncio.Task(coro2(), loop=self.loop)
        task = asyncio.Task(
            asyncio.wait([b, a], return_when=asyncio.FIRST_COMPLETED,
                         loop=self.loop),
            loop=self.loop)

        done, pending = self.loop.run_until_complete(task)
        self.assertEqual({a, b}, done)
        self.assertTrue(a.done())
        self.assertIsNone(a.result())
        self.assertTrue(b.done())
        self.assertIsNone(b.result()) 

def run_forever(self) -> None:
        self._waiting_end = self._loop.create_future()
        while not self._closed:
            self._logger_info(f'正在启动 {self._area_id} 号数据连接')
            if self._waiting_pause is not None:
                self._logger_info(f'暂停启动 {self._area_id} 号数据连接，等待 RESUME 指令')
                await self._waiting_pause
            
            async with self._opening_lock:
                if self._closed:
                    self._logger_info(f'{self._area_id} 号数据连接确认收到关闭信号，正在处理')
                    break
                # 未成功建立数据连接，循环重试
                if await self._prepare_client() and await self._job_open():
                    tasks = [self._loop.create_task(i()) for i in self._funcs_task]

                    self._task_main = self._loop.create_task(self._job_main())
                    tasks.append(self._task_main)

                    task_heartbeat = self._loop.create_task(self._job_heartbeat())
                    tasks.append(task_heartbeat)
                else:
                    continue
                
            _, pending = await asyncio.wait(
                tasks, return_when=asyncio.FIRST_COMPLETED)
            self._logger_info(f'{self._area_id} 号数据连接异常或主动断开，正在处理剩余信息')
            for i in pending:
                if i != self._task_main:
                    i.cancel()
            await self._job_close()
            if pending:
                await asyncio.wait(pending)
            self._logger_info(f'{self._area_id} 号数据连接退出，剩余任务处理完毕')
        await self._conn.clean()
        self._waiting_end.set_result(True) 

def _wait(self) -> None:
        self._timer = asyncio.ensure_future(
            asyncio.sleep(self._send_interval, loop=self._loop),
            loop=self._loop)

        await asyncio.wait([self._timer, self._ender], loop=self._loop,
                           return_when=asyncio.FIRST_COMPLETED)