先来看个例子,自己实现的模拟耗时操作

例1

import types
import select
import time
import socket
import functools


class Future:

    def __init__(self, *, loop=None):
        self._result = None
        self._callbacks = []
        self._loop = loop

    def set_result(self, result):
        self._result = result
        callbacks = self._callbacks[:]
        self._callbacks = []
        for callback in callbacks:
            loop._ready.append(callback)

    def add_callback(self, callback):
        self._callbacks.append(callback)

    def __iter__(self):
        print('enter Future ...')
        print('foo 挂起在yield处 ')
        yield self
        print('foo 恢复执行')
        print('exit Future ...')
        return 'future'

    __await__ = __iter__


class Task:

    def __init__(self, cor, *, loop=None):
        self.cor = cor
        self._loop = loop

    def _step(self):
        cor = self.cor
        try:
            result = cor.send(None)
        # 1. cor 协程执行完毕时,会抛出StopIteration,说明cor执行完毕了,这是关闭loop
        except StopIteration as e:
            self._loop.close()
        # 2. 有异常时
        except Exception as e:
            """处理异常逻辑"""
        # 3. result为Future对象时
        else:
            if isinstance(result, Future):
                result.add_callback(self._wakeup)

    def _wakeup(self):
        self._step()


class Loop:

    def __init__(self):
        self._stop = False
        self._ready = []
        self._scheduled = []
        self._time = lambda: time.time()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setblocking(False)
        self._select = functools.partial(select.select, [sock], [], [])

    def create_task(self, cor):
        task = Task(cor, loop=self)
        self._ready.append(task._step)
        return task

    def call_later(self, delay, callback, *args):
        callback._when = delay
        self._scheduled.append((callback, *args))

    def run_until_complete(self, task):
        assert isinstance(task, Task)
        timeout = None

        while not self._stop:

            if self._ready:
                timeout = 0

            if self._scheduled:
                callback, *args = self._scheduled.pop()
                timeout = callback._when
                self._ready.append(functools.partial(callback, *args))

                # 通过select(timeout)来控制阻塞时间
                self._select(timeout)

            n = len(self._ready)
            for i in range(n):
                step = self._ready.pop()
                step()

    def close(self):
        self._stop = True


@types.coroutine
def _sleep():
    yield

# 自己实现一个sleep协程
async def sleep(s, result=None):
    if s <= 0:
        await _sleep()
        return result
    else:
        future = Future(loop=loop)
        future._loop.call_later(s, callback, future)
        await future
        return result

# 延迟回调函数
def callback(future):
    # 时间到了就回调此函数
    future.set_result(None)


async def foo():
    print(f'enter foo at {time.strftime("%Y-%m-%d %H:%M:%S")}')
    await sleep(3)
    print(f'exit foo  at {time.strftime("%Y-%m-%d %H:%M:%S")}')


if __name__ == '__main__':
    f = foo()
    loop = Loop()
    task = loop.create_task(f)
    loop.run_until_complete(task)

执行结果:

enter foo at 2019-07-08 21:09:43
enter Future ...
foo 挂起在yield处
foo 恢复执行
exit Future ...
exit foo  at 2019-07-08 21:09:46

在上一篇文章通过Loop, Task, Future三个类基本上实现了对协程的调度,在此基础上做了一些修改实现了对协程中耗时操作的模拟。

首先我们分析一下async def foo协程中的await sleep(3),这里其实会进入到sleep中 await future这里,再进入到future对象的__await__方法中的yield self,foo协程此时被挂起,上一篇文章中我们分析知道,最终foo还是被这个future对象给分成了part1和part2两部分逻辑。

- foo print('enter foo at ...')
    - sleep
        - future print('enter Future ...')   # 以上是第一次f.send(None)执行的逻辑,命名为part1
        - future yield  self  ---------------------------------------------------------------
        - print('exit Future ...')        #以下是第二次f.send(None)执行的逻辑,命名为part2
    - sleep
- foo print('exit foo at ...')

part1 在 loop 循环的开始就执行了,返回一个 future 对象,把 part2 注册到 future 中,然后挂起了,下半部分 part2 在什么时候执行呢?因为在 sleep 中我们通过注册了一个3秒之后执行的回调函数 callback 到 loop 对象中,loop 对象在执行完 part1 后,会在下一轮的循环中执行 callback 回调函数,由于 loop._scheduled 不为空,timeout 被赋值成3,因此 select(3) 阻塞3秒后就继续往下执行。也就是说 callback 函数的执行时机就是在 select(3) 阻塞3秒后执行,callback 回调函数中又会调用 future.set_result() ,在 set_result 中会把 part2 注册到 loop 中,所以最终又在 loop 的下一轮循环中调用 part2 的逻辑,回到上次 foo 挂起的地方,继续 foo 的流程,直到协程退出。

其实所谓的模拟耗时3秒,其实就是在执行完part1后通过 select 函数阻塞3秒,然后再次执行 part2 ,这样就实现了所谓的等待3秒的操作。

要实现这个sleep协程的耗时模拟,主要是有2个关键点:

  • 1.通过 select(timeout) 的 timeout来控制 select 函数的阻塞时间。

      timeout=None    一直阻塞,直到有真实的IO事件到来,如socket的可读可写事件
      timeout=0       无论此时是否有IO事件到来,都立马返回
      timeout=n       阻塞n秒,在这n秒内,只要有IO事件到来,就立马返回,否则阻塞n秒才返回
  • 2.当延迟时间到来时,通过 callback 函数中调用 future.set_result() 方法,来驱动 part2 的执行。

了解到这里之后,我们再来看一下 asyncio 的源码

Loop类

class BaseEventLoop(events.AbstractEventLoop):

    ...

    def __init__(self):
        ...
        # 用来保存包裹task.step方法的handle对象的对端队列
        self._ready = collections.deque()
        # 用来保存包裹延迟回调函数的handle对象的二叉堆,是一个最小二叉堆
        self._scheduled = []
        ...

    def create_task(self, coro):
        """Schedule a coroutine object.

        Return a task object.
        """
        self._check_closed()
        # self._task_factory 默认是None
        if self._task_factory is None:
            # 创建一个task对象
            task = tasks.Task(coro, loop=self)
            if task._source_traceback:
                del task._source_traceback[-1]
        else:
            task = self._task_factory(self, coro)
        # 返回这个task对象
        return task

    def call_soon(self, callback, *args):

        self._check_closed()
        if self._debug:
            self._check_thread()
            self._check_callback(callback, 'call_soon')
        # 关键代码callback就是task._step方法,args是task._step的参数
        handle = self._call_soon(callback, args)
        if handle._source_traceback:
            del handle._source_traceback[-1]
        return handle

    def _call_soon(self, callback, args):
        # 1 handle是一个包裹了task._step方法和args参数的对象
        handle = events.Handle(callback, args, self)
        if handle._source_traceback:
            del handle._source_traceback[-1]
        # 2 关键代码,把handle添加到列表self._ready中
        self._ready.append(handle)
        return handle

    def run_until_complete(self, future):
        ...

        # future就是task对象,下面2句是为了确保future是一个Future类实例对象
        new_task = not futures.isfuture(future)
        future = tasks.ensure_future(future, loop=self)
        if new_task:
            # An exception is raised if the future didn't complete, so there
            # is no need to log the "destroy pending task" message
            future._log_destroy_pending = False

        # 添加回调方法_run_until_complete_cb到当前的task对象的callbacks列表中,_run_until_complete_cb就是最后
        # 把loop的_stop属性设置为ture的,用来结束loop循环的
        future.add_done_callback(_run_until_complete_cb)
        try:
            # 开启无线循环
            self.run_forever()
        except:
            ...
            raise
        finally:
            ...
        # 执行完毕返回cor的返回值
        return future.result()

    def run_forever(self):

        ...

        try:
            events._set_running_loop(self)
            while True:
                # 每次运行一次循环,判断下_stopping是否为true,也就是是否结束循环
                self._run_once()
                if self._stopping:
                    break
        finally:
            ...

    def _run_once(self):

        # loop的_scheduled是一个最小二叉堆,用来存放延迟执行的回调函数,根据延迟的大小,把这些回调函数构成一个最小堆,然后再每次从对顶弹出延迟最小的回调函数放入_ready双端队列中,
        # loop的_ready是双端队列,所有注册到loop的回调函数,最终是要放入到这个队列中,依次取出然后执行的
        # 1. self._scheduled是否为空
        sched_count = len(self._scheduled)
        if (sched_count > _MIN_SCHEDULED_TIMER_HANDLES and
            self._timer_cancelled_count / sched_count >
                _MIN_CANCELLED_TIMER_HANDLES_FRACTION):
            # Remove delayed calls that were cancelled if their number
            # is too high
            new_scheduled = []
            for handle in self._scheduled:
                if handle._cancelled:
                    handle._scheduled = False
                else:
                    new_scheduled.append(handle)

            heapq.heapify(new_scheduled)
            self._scheduled = new_scheduled
            self._timer_cancelled_count = 0
        else:
            # Remove delayed calls that were cancelled from head of queue.
            while self._scheduled and self._scheduled[0]._cancelled:
                self._timer_cancelled_count -= 1
                handle = heapq.heappop(self._scheduled)
                handle._scheduled = False

        # 2. 给timeout赋值,self._scheduled为空,timeout就为None
        timeout = None
        # 只要self._ready和self._scheduled中有一个不为空,timeout就为0
        if self._ready or self._stopping:
            timeout = 0
        # 只要self._scheduled不为空
        elif self._scheduled:
            # Compute the desired timeout.
            # 用堆顶的回调函数的延迟时间作为timeout的等待时间,也就是说用等待时间最短的回调函数的时间作为timeout的等待时间
            when = self._scheduled[0]._when
            timeout = max(0, when - self.time())
        、

        if self._debug and timeout != 0:
            ...
        # 3. 关注else分支,这是关键代码
        else:
            # timeout=None --> 一直阻塞,只要有io事件产生,立马返回event_list事件列表,否则一直阻塞着
            # timeout=0 --> 不阻塞,有io事件产生,就立马返回event_list事件列表,没有也返空列表
            # timeout=2 --> 阻塞等待2s,在这2秒内只要有io事件产生,立马返回event_list事件列表,没有io事件就阻塞2s,然后返回空列表
            event_list = self._selector.select(timeout)

        #  用来处理真正的io事件的函数
        self._process_events(event_list)

        # Handle 'later' callbacks that are ready.
        end_time = self.time() + self._clock_resolution
        # 4. 依次取出堆顶的回调函数handle添加到_ready队列中
        while self._scheduled:
            handle = self._scheduled[0]
            # 当_scheduled[]中有多个延迟回调时,通过handle._when >= end_time来阻止没有到时间的延迟函数被弹出,
            # 也就是说,当有n个延迟回调时,会产生n个timeout,对应n次run_once循环的调用
            if handle._when >= end_time:
                break
            # 从堆中弹出堆顶最小的回调函数,放入 _ready 队列中
            handle = heapq.heappop(self._scheduled)
            handle._scheduled = False
            self._ready.append(handle)

        # 5. 执行self._ready队列中所有的回调函数handle对象
        ntodo = len(self._ready)
        for i in range(ntodo):
            handle = self._ready.popleft()
            if handle._cancelled:
                continue
            if self._debug:
                try:
                    self._current_handle = handle
                    t0 = self.time()
                    handle._run()
                    dt = self.time() - t0
                    if dt >= self.slow_callback_duration:
                        logger.warning('Executing %s took %.3f seconds',
                                       _format_handle(handle), dt)
                finally:
                    self._current_handle = None
            else:
                # handle._run()实际上就是执行task._step(),也就是执行cor.send(None)
                handle._run()
        handle = None  # Needed to break cycles when an exception occurs.

Task类

class Task(futures.Future):

    ...

    def _step(self, exc=None):
        """
        _step方法可以看做是task包装的coroutine对象中的代码的直到yield的前半部分逻辑
        """
        ...
        try:
            if exc is None:

                # 1.关键代码,调用协程
                result = coro.send(None)
            else:
                result = coro.throw(exc)
        # 2. coro执行完毕会抛出StopIteration异常
        except StopIteration as exc:
            if self._must_cancel:
                # Task is cancelled right before coro stops.
                self._must_cancel = False
                self.set_exception(futures.CancelledError())
            else:
                # result为None时,调用task的callbasks列表中的回调方法,在调用loop.run_until_complite,结束loop循环
                self.set_result(exc.value)
        except futures.CancelledError:
            super().cancel()  # I.e., Future.cancel(self).
        except Exception as exc:
            self.set_exception(exc)
        except BaseException as exc:
            self.set_exception(exc)
            raise
        # 3. result = coro.send(None)不抛出异常,说明协程被yield挂起
        else:
            # 4. 查看result是否含有_asyncio_future_blocking属性
            blocking = getattr(result, '_asyncio_future_blocking', None)
            if blocking is not None:
                # Yielded Future must come from Future.__iter__().
                if result._loop is not self._loop:
                    self._loop.call_soon(
                        self._step,
                        RuntimeError(
                            'Task {!r} got Future {!r} attached to a '
                            'different loop'.format(self, result)))

                elif blocking:
                    if result is self:
                        self._loop.call_soon(
                            self._step,
                            RuntimeError(
                                'Task cannot await on itself: {!r}'.format(
                                    self)))
                    # 4.1. 如果result是一个future对象时,blocking会被设置成true
                    else:
                        result._asyncio_future_blocking = False
                        # 把_wakeup回调函数设置到此future对象中,当此future对象调用set_result()方法时,就会调用_wakeup方法
                        result.add_done_callback(self._wakeup)
                        self._fut_waiter = result
                        if self._must_cancel:
                            if self._fut_waiter.cancel():
                                self._must_cancel = False
                else:
                    self._loop.call_soon(
                        self._step,
                        RuntimeError(
                            'yield was used instead of yield from '
                            'in task {!r} with {!r}'.format(self, result)))
            # 5. 如果result是None,则注册task._step到loop对象中去,在下一轮_run_once中被回调
            elif result is None:
                # Bare yield relinquishes control for one event loop iteration.
                self._loop.call_soon(self._step)

            # --------下面的代码可以暂时不关注了--------
            elif inspect.isgenerator(result):
                # Yielding a generator is just wrong.
                self._loop.call_soon(
                    self._step,
                    RuntimeError(
                        'yield was used instead of yield from for '
                        'generator in task {!r} with {}'.format(
                            self, result)))
            else:
                # Yielding something else is an error.
                self._loop.call_soon(
                    self._step,
                    RuntimeError(
                        'Task got bad yield: {!r}'.format(result)))
        finally:
            self.__class__._current_tasks.pop(self._loop)
            self = None  # Needed to break cycles when an exception occurs.

    def _wakeup(self, future):
        try:
            future.result()
        except Exception as exc:
            # This may also be a cancellation.
            self._step(exc)
        else:

            # 这里是关键代码,上次的_step()执行到第一次碰到yield的地方挂住了,此时再次执行_step(),
            # 也就是再次执行 result = coro.send(None) 这句代码,也就是从上次yield的地方继续执行yield后面的逻辑
            self._step()
        self = None  # Needed to break cycles when an exception occurs.

Future类

class Future:

    ...

    def add_done_callback(self, fn, *, context=None):
        if self._state != _PENDING:
            self._loop.call_soon(fn, self, context=context)
        else:
            if context is None:
                context = contextvars.copy_context()
            self._callbacks.append((fn, context))

    def set_result(self, result):

        if self._state != _PENDING:
            raise InvalidStateError('{}: {!r}'.format(self._state, self))
        self._result = result
        self._state = _FINISHED
    self.__schedule_callbacks()

    def __iter__(self):
        # self.done()返回False,
        if not self.done():
            self._asyncio_future_blocking = True
            # 把Future对象自己返回出去
            yield self  # This tells Task to wait for completion.
        assert self.done(), "yield from wasn't used with future"
        return self.result()  # May raise too.

    if compat.PY35:
        __await__ = __iter__ # make compatible with 'await' expression

sleep协程

#延迟回调函数,里面调用fut.set_result
def _set_result_unless_cancelled(fut, result):
    if fut.cancelled():
        return
    # 关键是这一步,驱动协程从上次挂起的地方继续执行
    fut.set_result(result)

@types.coroutine
def __sleep0():

    yield

async def sleep(delay, result=None, *, loop=None):
    """Coroutine that completes after a given time (in seconds)."""
    if delay <= 0:
        await __sleep0()
        return result

    if loop is None:
        loop = events.get_event_loop()
    # 创建一个future对象
    future = loop.create_future()
    # 注册一个延迟回调函数到loop对象中
    h = loop.call_later(delay, futures._set_result_unless_cancelled, future,  result)
    try:
        return await future
    finally:
        h.cancel()

关键地方我都写了注释,如果能耐着性子细心看下来,你会发现例1中的实现,就是模仿asyncio中的这几个类去实现的。

asyncio的sleep中的延迟回调函数是_set_result_unless_cancelled与我写的callback对应,关键都是要回调future.set_result方法,这样才能驱动协程从上次挂起的地方开始继续执行。

对于使用asyncio.sleep的例子

import asyncio


async def cor():
    print('enter cor ...')
    await asyncio.sleep(2)
    print('exit cor ...')

    return 'cor'

loop = asyncio.get_event_loop()
task = loop.create_task(cor())
rst = loop.run_until_complete(task)
print(rst)

await asyncio.sleep(2) 这句代码同样是把cor协程分为如下两个部分:

- cor print('enter cor ...')
    - sleep
        - future print('enter Future ...')   # 以上是第一次cor.send(None)执行的逻辑,命名为part1
        - future yield  self  ---------------------------------------------------------------
        - future print('exit Future ...')    # 以下是第二次cor.send(None)执行的逻辑,命名为part2
    - sleep
- cor print('exit foo ...')

总之,只要有要耗时的地方,就必须要有一个 future 用来 await future,然后协程就被分成了part1和part2,part1和part2就被分别封装到了task._step和task._wakeup中,然后在loop循环中先调用part1,再通过select函数阻塞n秒之后,再执行part2,最后,协程执行完毕。

07-09 08:19