Постигая asyncio. Часть первая — сопрограммы, неблокирующий ввод-вывод и асинхронный таймер.

Как же быть в том случае, если мы хотим, чтобы корутина поспала полсекунды и продолжила свою работу не блокируя остальные? Для этого существует асинхронный sleep. Реализация его не такая уж и сложная, давайте разбираться.

Асинхронный таймер.

Концептуально это будет выглядеть следующим образом: создаем дополнительную очередь задач в нашем планировщике, причем не простую, а очередь с приоритетом, где значением будет  Task, а весом — временная метка в миллисекундах + время задержки. Там эти таски будут храниться в очереди до того момента, пока их временная метка не станет меньше текущего времени, тогда их можно будет переместить в основную очередь, где они продолжат свое выполнение.

Такая реализация довольно примитивна и имеет ряд недостатков, но вполне себе рабочая и может быть использована как пример.

import time
from collections import deque
from heapq import heappush, heappop  # очередь с приоритетом


def get_timestamp(offset=0):
    return time.time() * 1000 + (offset * 1000)


def async_sleep(t):
    '''
    Наш асинхронный sleep,
    ничего особенного он не делает,
    только дает планировщику понять,
    что надо приостановить таску
    '''
    return 'sleep', t


class Task:
    def __init__(self, coroutine):
        self._coroutine = coroutine
        self._is_ready = False
        self._is_finished = False

    def run(self):
        try:
            # инициализируем корутину
            if not self._is_ready:
                self._is_ready = True
                return self._coroutine.send(None)
            ret = self._coroutine.send(None)
            return ret
        except StopIteration:
            self._is_finished = True

    @property
    def is_finished(self):
        return self._is_finished


class TimeQueue:
    def __init__(self):
        self._q = []

    def append(self, task, ts):
        heappush(self._q, (ts, task))

    def expired(self, ts):
        while self._q:
            ret_ts, task = heappop(self._q)
            if ret_ts <= ts:
                yield task
            else:
                heappush(self._q, (ret_ts, task))
                break

    @property
    def is_empty(self):
        return not bool(self._q)


class Scheduler:
    def __init__(self):
        self._task_queue = deque([])
        self._time_queue = TimeQueue()

    def new_task(self, coroutine):
        task = Task(coroutine)
        self.schedule(task)

    def sleep(self, task, ts):
        self._time_queue.append(task, ts)
        
    def schedule(self, task):
        self._task_queue.append(task)

    def run(self):
        while True:
            if self._task_queue:
                task = self._task_queue.popleft()
                ret = task.run()
                # проверяем, не надо ли таске немного поспать?
                if ret and len(ret) == 2 and ret[0] == 'sleep':
                    self.sleep(task, get_timestamp(ret[1]))
                else:
                    if not task.is_finished:
                        self.schedule(task)

            for task in self._time_queue.expired(get_timestamp()):
                task.run()
                if not task.is_finished:
                    self.schedule(task)

            if not self._task_queue and self._time_queue.is_empty:
                break


def async_socket_emu(total_bytes=4096, chunk=1024, _id=1):
    while total_bytes > 0:
        print(f'Read {chunk} bytes from socket {_id}')
        yield async_sleep(0.5)  # теперь "сокет" спит в асинхронном режиме
        total_bytes -= chunk
        yield
    print(f'Finish with socket {_id}')


scheduler = Scheduler()
scheduler.new_task(async_socket_emu(_id=1))
scheduler.new_task(async_socket_emu(chunk=2048, _id=2))
scheduler.new_task(async_socket_emu(chunk=512, _id=3))

start = time.time()

scheduler.run()

print('Total time:', time.time()-start, 'seconds')

Read 1024 bytes from socket 1
Read 2048 bytes from socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Read 2048 bytes from socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Finish with socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Read 512 bytes from socket 3
Finish with socket 1
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Finish with socket 3
Total time: 4.002506971359253 seconds

import asyncio


async def async_socket_emu(total_bytes=4096, chunk=1024, _id=1):
    while total_bytes > 0:
        print(f'Read {chunk} bytes from socket {_id}')
        await asyncio.sleep(0.5)
        total_bytes -= chunk
    print(f'Finish with socket {_id}')


loop = asyncio.get_event_loop()

loop.create_task(async_socket_emu(_id=1))
loop.create_task(async_socket_emu(chunk=2048, _id=2))
loop.create_task(async_socket_emu(chunk=512, _id=3))
pending = asyncio.Task.all_tasks(loop=loop)
group = asyncio.gather(*pending, return_exceptions=True)

start = time.time()

loop.run_until_complete(group)

print('Total time:', time.time()-start, 'seconds')

Read 1024 bytes from socket 1
Read 2048 bytes from socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Read 2048 bytes from socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Finish with socket 2
Read 512 bytes from socket 3
Read 1024 bytes from socket 1
Read 512 bytes from socket 3
Finish with socket 1
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Read 512 bytes from socket 3
Finish with socket 3
Total time: 4.018957853317261 seconds