asyncio系列之Lock实现

import types
import select
import time
import socket
import functools
import collections


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:
            self._loop._ready.append(callback)

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

    def __iter__(self):
        print("挂起在yield处")
        yield self
        print("恢复执行")
        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)
        except StopIteration as e:
            self._loop._task_count -= 1
            if self._loop._task_count == 0:
                self._loop.close()
        except Exception as e:
            pass
        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], [], [])
        self._task_count = 0

    def create_task(self, cor):
        task = Task(cor, loop=self)
        self._ready.append(task._step)
        self._task_count += 1
        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))

                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


async def sleep(s, result=None):
    if s <= 0:
        await _sleep()
        return result
    else:
        future = Future(loop=loop)
        future._loop.call_later(s, unless_cancelled, future)
        await future
        return result


def unless_cancelled(future):
    future.set_result(None)


class Lock:
    def __init__(self, *, loop=None):
        self._waiters = collections.deque()
        self._locked = False
        self._loop = loop

    def __repr__(self):
        extra = 'locked' if self._locked else 'unlocked'
        if self._waiters:
            extra = '{},waiters:{}'.format(extra, len(self._waiters))
        return '<[{}]>'.format(extra)

    def locked(self):
        """Return True if lock is acquired."""
        return self._locked

    @types.coroutine
    def acquire(self):
        if not self._locked:
            self._locked = True
            return True

        fut = Future(loop=self._loop)
        self._waiters.append(fut)

        try:
            yield from fut
        finally:
            self._waiters.remove(fut)

        self._locked = True
        return True

    def release(self):
        if self._locked:
            self._locked = False
            self._wake_up_first()
        else:
            raise RuntimeError('Lock is not acquired.')

    def _wake_up_first(self):
        """Wake up the first waiter if it isn't done."""
        try:
            fut = next(iter(self._waiters))
        except StopIteration:
            return

        fut.set_result(True)


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


async def goo(look):
    await look.acquire()
    print(f'enter goo at {time.strftime("%Y-%m-%d %H:%M:%S")}')
    await sleep(1)
    print(f'exit goo at {time.strftime("%Y-%m-%d %H:%M:%S")}')
    look.release()


if __name__ == '__main__':
    loop = Loop()
    look = Lock(loop=loop)
    f = foo(look)
    g = goo(look)
    task1 = loop.create_task(f)
    task2 = loop.create_task(g)
    loop.run_until_complete(task1)