How to use the aiokafka.util.ensure_future function in aiokafka

To help you get started, we’ve selected a few aiokafka examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github aio-libs / aiokafka / tests / test_coordinator.py View on Github external
def start_coordination():
            if coordinator._coordination_task:
                coordinator._coordination_task.cancel()
            coordinator._coordination_task = task = ensure_future(
                coordinator._coordination_routine(), loop=self.loop)
            return task
github aio-libs / aiokafka / tests / test_coordinator.py View on Github external
client, subscription, loop=self.loop,
            group_id='test-my-group', session_timeout_ms=6000,
            heartbeat_interval_ms=1000)
        subscription.subscribe(topics=set(['topic1']))
        await subscription.wait_for_assignment()

        waiter = create_future(loop=self.loop)

        async def commit_offsets(*args, **kw):
            await waiter

        coordinator.commit_offsets = mocked = mock.Mock()
        mocked.side_effect = commit_offsets

        # Close task should call autocommit last time
        close_task = ensure_future(coordinator.close(), loop=self.loop)
        await asyncio.sleep(0.1, loop=self.loop)
        # self.assertFalse(close_task.done())

        # Raising an error should not prevent from closing. Error should be
        # just logged
        waiter.set_exception(Errors.UnknownError())
        await close_task
github aio-libs / aiokafka / tests / test_transactional_producer.py View on Github external
await asyncio.sleep(0)
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

        # test cancel begin_transaction.
        task = ensure_future(producer.begin_transaction())
        await cancel(task)
        self.assertEqual(txn_manager.state, TransactionState.READY)

        # test cancel commit_transaction. Commit should not be cancelled.
        await producer.begin_transaction()
        self.assertEqual(txn_manager.state, TransactionState.IN_TRANSACTION)
        task = ensure_future(producer.commit_transaction())
        await cancel(task)
        self.assertEqual(
            txn_manager.state, TransactionState.COMMITTING_TRANSACTION)
        await asyncio.sleep(0.1)
        self.assertEqual(txn_manager.state, TransactionState.READY)

        # test cancel abort_transaction. Abort should also not be cancelled.
        await producer.begin_transaction()
        self.assertEqual(txn_manager.state, TransactionState.IN_TRANSACTION)
        task = ensure_future(producer.abort_transaction())
        await cancel(task)
        self.assertEqual(
            txn_manager.state, TransactionState.ABORTING_TRANSACTION)
        await asyncio.sleep(0.1)
        self.assertEqual(txn_manager.state, TransactionState.READY)
github aio-libs / aiokafka / aiokafka / producer / sender.py View on Github external
def _maybe_do_transactional_request(self):
        txn_manager = self._txn_manager

        # If we have any new partitions, still not added to the transaction
        # we need to do that before committing
        tps = txn_manager.partitions_to_add()
        if tps:
            return ensure_future(
                self._do_add_partitions_to_txn(tps),
                loop=self._loop)

        # We need to add group to transaction before we can commit the offset
        group_id = txn_manager.consumer_group_to_add()
        if group_id is not None:
            return ensure_future(
                self._do_add_offsets_to_txn(group_id),
                loop=self._loop)

        # Now commit the added group's offset
        commit_data = txn_manager.offsets_to_commit()
        if commit_data is not None:
            offsets, group_id = commit_data
            return ensure_future(
                self._do_txn_offset_commit(offsets, group_id),
github aio-libs / aiokafka / aiokafka / client.py View on Github external
else:
                bootstrap_conn.close()

            log.debug('Received cluster metadata: %s', self.cluster)
            break
        else:
            raise ConnectionError(
                'Unable to bootstrap from {}'.format(self.hosts))

        # detect api version if need
        if self._api_version == 'auto':
            self._api_version = await self.check_version()

        if self._sync_task is None:
            # starting metadata synchronizer task
            self._sync_task = ensure_future(
                self._md_synchronizer(), loop=self._loop)
github aio-libs / aiokafka / aiokafka / consumer / group_coordinator.py View on Github external
def __init__(self, *args, **kw):
        super().__init__(*args, **kw)
        # Reset all committed points, as the GroupCoordinator would
        self._reset_committed_task = ensure_future(
            self._reset_committed_routine(), loop=self._loop)
github aio-libs / aiokafka / aiokafka / consumer / group_coordinator.py View on Github external
def _start_heartbeat_task(self):
        if self._heartbeat_task is None:
            self._heartbeat_task = ensure_future(
                self._heartbeat_routine(), loop=self._loop)
github aio-libs / aiokafka / aiokafka / consumer / group_coordinator.py View on Github external
def maybe_leave_group(self):
        task = ensure_future(self._maybe_leave_group(), loop=self._loop)
        return task
github aio-libs / aiokafka / aiokafka / consumer / fetcher.py View on Github external
def start_pending_task(coro, node_id, self=self):
                task = ensure_future(coro, loop=self._loop)
                self._pending_tasks.add(task)
                self._in_flight.add(node_id)

                def on_done(fut, self=self):
                    self._in_flight.discard(node_id)
                task.add_done_callback(on_done)