Python asyncio.shield() Examples

def init_workers(app: web.Application, conf: WorkersConfig) -> ThreadPoolExecutor:
    n = conf.max_workers
    executor = ThreadPoolExecutor(max_workers=n)

    loop = asyncio.get_event_loop()
    run = loop.run_in_executor
    fs = [run(executor, warm, conf.path_to_model_state) for _ in range(0, n)]
    await asyncio.gather(*fs)

    async def close_executor(app: web.Application) -> None:
        fs = [run(executor, clean) for _ in range(0, n)]
        await asyncio.shield(asyncio.gather(*fs))
        executor.shutdown(wait=True)

    app.on_cleanup.append(close_executor)
    app["executor"] = executor
    return executor 

def _listener(self, subscription, interval, callback):
        while True:
            try:
                # FIXME: This should use an async requester as below
                packet = callback()
                self.__handle_packet(packet, subscription)
                await asyncio.sleep(interval)
            except asyncio.CancelledError:
                ## unsubscribe from all subscriptions
                confirmations = await asyncio.gather(
                    asyncio.shield(self._unsubscribe(subscription)) 
                    for subscription in self.subscriptions)
            except Exception as ex:
                logger.error(ex)
                logger.error(packet)
                raise 

def test_shield_effect(self):
        # Cancelling outer() does not affect inner().
        proof = 0
        waiter = asyncio.Future(loop=self.loop)

        @asyncio.coroutine
        def inner():
            nonlocal proof
            yield from waiter
            proof += 1

        @asyncio.coroutine
        def outer():
            nonlocal proof
            yield from asyncio.shield(inner(), loop=self.loop)
            proof += 100

        f = asyncio.ensure_future(outer(), loop=self.loop)
        test_utils.run_briefly(self.loop)
        f.cancel()
        with self.assertRaises(asyncio.CancelledError):
            self.loop.run_until_complete(f)
        waiter.set_result(None)
        test_utils.run_briefly(self.loop)
        self.assertEqual(proof, 1) 

def test_gather_shield(self):
        child1 = asyncio.Future(loop=self.loop)
        child2 = asyncio.Future(loop=self.loop)
        inner1 = asyncio.shield(child1, loop=self.loop)
        inner2 = asyncio.shield(child2, loop=self.loop)
        parent = asyncio.gather(inner1, inner2, loop=self.loop)
        test_utils.run_briefly(self.loop)
        parent.cancel()
        # This should cancel inner1 and inner2 but bot child1 and child2.
        test_utils.run_briefly(self.loop)
        self.assertIsInstance(parent.exception(), asyncio.CancelledError)
        self.assertTrue(inner1.cancelled())
        self.assertTrue(inner2.cancelled())
        child1.set_result(1)
        child2.set_result(2)
        test_utils.run_briefly(self.loop) 

def send_offsets_to_transaction(self, offsets, group_id):
        self._ensure_transactional()

        if not self._txn_manager.is_in_transaction():
            raise IllegalOperation("Not in the middle of a transaction")

        if not group_id or not isinstance(group_id, str):
            raise ValueError(group_id)

        # validate `offsets` structure
        formatted_offsets = commit_structure_validate(offsets)

        log.debug(
            "Begin adding offsets %s for consumer group %s to transaction",
            formatted_offsets, group_id)
        fut = self._txn_manager.add_offsets_to_txn(formatted_offsets, group_id)
        await asyncio.shield(fut, loop=self._loop) 

def _listener(self):
        await self._connect()
        while True:
            try:
                packet = await self.websocket.recv()
                self.__handle_packet(packet)
            except websockets.exceptions.ConnectionClosed:
                await self._connect()
            except asyncio.CancelledError:
                ## unsubscribe from all subscriptions
                confirmations = await asyncio.gather(
                    asyncio.shield(self._unsubscribe(subscription)) 
                    for subscription in self.subscriptions)
            except Exception as ex:
                logger.error(ex)
                logger.error(packet)
                raise 

def test_maxsize_release(self):
        pool = Pool(url="ws://localhost:8182/",
                    maxsize=2,
                    username="stephen",
                    password="password",
                    future_class=Future)

        async def go():
            c1 = await pool.acquire()
            c2 = await pool.acquire()
            c3 = pool.acquire()
            self.assertIsInstance(c3, Future)
            with self.assertRaises(asyncio.TimeoutError):
                shielded_fut = asyncio.shield(c3)
                await asyncio.wait_for(shielded_fut, 0.1)
            await pool.release(c2)
            c3 = await c3
            self.assertEqual(c2, c3)
            c1.conn.close()
            c2.conn.close()
            c3.conn.close()

        self.loop.run_until_complete(go()) 

def test_maxsize_release(self):
        pool = Pool(url="ws://localhost:8182/",
                    maxsize=2,
                    username="stephen",
                    password="password",
                    future_class=Future)

        @asyncio.coroutine
        def go():
            c1 = yield from pool.acquire()
            c2 = yield from pool.acquire()
            c3 = pool.acquire()
            self.assertIsInstance(c3, Future)
            with self.assertRaises(asyncio.TimeoutError):
                shielded_fut = asyncio.shield(c3)
                yield from asyncio.wait_for(shielded_fut, 0.1)
            yield from pool.release(c2)
            c3 = yield from c3
            self.assertEqual(c2, c3)
            c1.conn.close()
            c2.conn.close()
            c3.conn.close()

        self.loop.run_until_complete(go()) 

def test_gather_shield(self):
        child1 = asyncio.Future(loop=self.loop)
        child2 = asyncio.Future(loop=self.loop)
        inner1 = asyncio.shield(child1, loop=self.loop)
        inner2 = asyncio.shield(child2, loop=self.loop)
        parent = asyncio.gather(inner1, inner2, loop=self.loop)
        test_utils.run_briefly(self.loop)
        parent.cancel()
        # This should cancel inner1 and inner2 but bot child1 and child2.
        test_utils.run_briefly(self.loop)
        self.assertIsInstance(parent.exception(), asyncio.CancelledError)
        self.assertTrue(inner1.cancelled())
        self.assertTrue(inner2.cancelled())
        child1.set_result(1)
        child2.set_result(2)
        test_utils.run_briefly(self.loop) 

def test_shield_effect(self):
        # Cancelling outer() does not affect inner().
        proof = 0
        waiter = asyncio.Future(loop=self.loop)

        @asyncio.coroutine
        def inner():
            nonlocal proof
            yield from waiter
            proof += 1

        @asyncio.coroutine
        def outer():
            nonlocal proof
            yield from asyncio.shield(inner(), loop=self.loop)
            proof += 100

        f = asyncio.ensure_future(outer(), loop=self.loop)
        test_utils.run_briefly(self.loop)
        f.cancel()
        with self.assertRaises(asyncio.CancelledError):
            self.loop.run_until_complete(f)
        waiter.set_result(None)
        test_utils.run_briefly(self.loop)
        self.assertEqual(proof, 1) 

def get_tx_events(self):

        org = 'org1.example.com'
        peer = self.client.get_peer('peer0.' + org)

        org_admin = self.client.get_user(org, 'Admin')
        channel = self.client.get_channel(self.channel_name)
        channel_event_hub = channel.newChannelEventHub(peer, org_admin)
        stream = channel_event_hub.connect(start='oldest',
                                           stop='newest', filtered=False)

        self.txs = {}
        channel_event_hub.registerTxEvent('all', onEvent=self.onTxEvent)

        try:
            await shield(stream)
        except Exception:
            pass

        channel_event_hub.disconnect()

        self.assertEqual(len(self.txs['all']), 4) 

def _redirect_async(self, redirect, auth):
        """Redirect the client endpoint using a Link DETACH redirect
        response.

        :param redirect: The Link DETACH redirect details.
        :type redirect: ~uamqp.errors.LinkRedirect
        :param auth: Authentication credentials to the redirected endpoint.
        :type auth: ~uamqp.authentication.common.AMQPAuth
        """
        # pylint: disable=protected-access
        if not self._connection._cbs:
            _logger.info("Closing non-CBS session.")
            await asyncio.shield(self._session.destroy_async(), loop=self.loop)
        self._session = None
        self._auth = auth
        self._hostname = self._remote_address.hostname
        await self._connection.redirect_async(redirect, auth)
        await self._build_session_async() 

def disconnect(self, timeout: Optional[float] = None) -> None:
        """Close the connection gracefully."""
        await self._connection.close(code=1000)
        try:
            await asyncio.wait_for(
                asyncio.shield(self._message_processor), timeout or self.TIMEOUT
            )
            self._message_processor.result()
        except asyncio.TimeoutError:
            pass
        finally:
            if not self._message_processor.done():
                self._message_processor.cancel()
                try:
                    await self._message_processor
                except asyncio.CancelledError:
                    pass 

def downstream(self):
        try:
            self.upstream_buffer_task = \
                    asyncio.ensure_future(self.consume_upstream_buffer())
            async for msg in self.up_conn:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    await self.down_conn.send_str(msg.data)
                    if self.downstream_cb is not None:
                        await asyncio.shield(self.downstream_cb(msg.data))
                if msg.type == aiohttp.WSMsgType.BINARY:
                    await self.down_conn.send_bytes(msg.data)
                    if self.downstream_cb is not None:
                        await asyncio.shield(self.downstream_cb(msg.data))
                elif msg.type == aiohttp.WSMsgType.CLOSED:
                    break
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    break
            # here, server gracefully disconnected
        except asyncio.CancelledError:
            raise
        except Exception:
            log.exception('unexpected error')
        finally:
            await self.close_upstream() 

def send_offsets_to_transaction(self, offsets, group_id):
        self._ensure_transactional()

        if not self._txn_manager.is_in_transaction():
            raise IllegalOperation("Not in the middle of a transaction")

        if not group_id or not isinstance(group_id, str):
            raise ValueError(group_id)

        # validate `offsets` structure
        formatted_offsets = commit_structure_validate(offsets)

        log.debug(
            "Begin adding offsets %s for consumer group %s to transaction",
            formatted_offsets, group_id)
        fut = self._txn_manager.add_offsets_to_txn(formatted_offsets, group_id)
        await asyncio.shield(fut, loop=self._loop) 

def test_shield_cancel(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        test_utils.run_briefly(self.loop)
        inner.cancel()
        test_utils.run_briefly(self.loop)
        self.assertTrue(outer.cancelled()) 

def add_batch(self, builder, tp, timeout):
        """Add BatchBuilder to queue by topic-partition.

        Arguments:
            builder (BatchBuilder): batch object to enqueue.
            tp (TopicPartition): topic and partition to enqueue this batch for.
            timeout (int): time in seconds to wait for a free slot in the batch
                queue.

        Returns:
            MessageBatch: delivery wrapper around the BatchBuilder object.

        Raises:
            aiokafka.errors.ProducerClosed: the accumulator has already been
                closed and flushed.
            aiokafka.errors.KafkaTimeoutError: the batch could not be added
                within the specified timeout.
        """
        if self._closed:
            raise ProducerClosed()
        if self._exception is not None:
            raise copy.copy(self._exception)

        start = self._loop.time()
        while timeout > 0:
            pending = self._batches.get(tp)
            if pending:
                await pending[-1].wait_drain(timeout=timeout)
                timeout -= self._loop.time() - start
            else:
                batch = self._append_batch(builder, tp)
                return asyncio.shield(batch.future, loop=self._loop)
        raise KafkaTimeoutError() 

def abort_transaction(self):
        self._ensure_transactional()
        log.debug(
            "Aborting transaction for id %s",
            self._txn_manager.transactional_id)
        self._txn_manager.aborting_transaction()
        await asyncio.shield(
            self._txn_manager.wait_for_transaction_end(),
            loop=self._loop
        ) 

def commit_transaction(self):
        self._ensure_transactional()
        log.debug(
            "Committing transaction for id %s",
            self._txn_manager.transactional_id)
        self._txn_manager.committing_transaction()
        await asyncio.shield(
            self._txn_manager.wait_for_transaction_end(),
            loop=self._loop
        ) 

def begin_transaction(self):
        self._ensure_transactional()
        log.debug(
            "Beginning a new transaction for id %s",
            self._txn_manager.transactional_id)
        await asyncio.shield(
            self._txn_manager.wait_for_pid(),
            loop=self._loop
        )
        self._txn_manager.begin_transaction() 

def fetch_next(self, wait=True):
        timeout = Cursor._wait_to_timeout(wait)
        waiter = reusable_waiter(self.conn._io_loop, timeout)
        while len(self.items) == 0 and self.error is None:
            self._maybe_fetch_batch()
            if self.error is not None:
                raise self.error
            with translate_timeout_errors():
                yield from waiter(asyncio.shield(self.new_response))
        # If there is a (non-empty) error to be received, we return True, so the
        # user will receive it on the next `next` call.
        return len(self.items) != 0 or not isinstance(self.error, RqlCursorEmpty) 

def wait_closed(self):
        'wait until pool is closed'
        await asyncio.shield(self._close_waiter, loop=self._loop) 

def try_connect(self):
        while self.active and not self.connected:
            try:
                await self.connect(self._timeout)
            except (TimeoutError, asyncio.TimeoutError):
                logger.error(
                    "[%s] Timeout occurred while connecting to %s",
                    self.name,
                    self.describe_server(),
                )
            except (socket.error, socket.gaierror, OSError, ssl.SSLError):
                logger.error(
                    "[%s] Error occurred while connecting to %s (%s)",
                    self.name,
                    self.describe_server(),
                    traceback.format_exc().splitlines()[-1],
                )
            except Exception as e:
                raise ClientConnectError(self.name, self.describe_server()) from e
            else:
                break

            sleep_time = random.randrange(self._timeout)
            canceller = asyncio.shield(self.cancelled_future)
            try:
                await asyncio.wait_for(canceller, timeout=sleep_time)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pass 

def test_shield_result(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        inner.set_result(42)
        res = self.loop.run_until_complete(outer)
        self.assertEqual(res, 42) 

def test_shield_shortcut(self):
        fut = asyncio.Future(loop=self.loop)
        fut.set_result(42)
        res = self.loop.run_until_complete(asyncio.shield(fut))
        self.assertEqual(res, 42) 

def test_shield_exception(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        test_utils.run_briefly(self.loop)
        exc = RuntimeError('expected')
        inner.set_exception(exc)
        test_utils.run_briefly(self.loop)
        self.assertIs(outer.exception(), exc) 

def test_shield_result(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        inner.set_result(42)
        res = self.loop.run_until_complete(outer)
        self.assertEqual(res, 42) 

def __anext__(self):
        if self._next_future is None:
            if self._iterator is None:
                raise StopAsyncIteration
            self._next_future = self._loop.run_in_executor(self._stream_executor, self._next)
        try:
            return await asyncio.shield(self._next_future, loop=self._loop)
        finally:
            if self._next_future and self._next_future.done():
                self._next_future = None 

def test_shield_exception(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        test_utils.run_briefly(self.loop)
        exc = RuntimeError('expected')
        inner.set_exception(exc)
        test_utils.run_briefly(self.loop)
        self.assertIs(outer.exception(), exc) 

def test_shield_cancel(self):
        inner = asyncio.Future(loop=self.loop)
        outer = asyncio.shield(inner)
        test_utils.run_briefly(self.loop)
        inner.cancel()
        test_utils.run_briefly(self.loop)
        self.assertTrue(outer.cancelled())