python 异步读写锁

import asyncio


class AsyncReadWriteLock:
    def __init__(self):
        self._readers = 0
        self._writers_waiting = 0
        self._writer_active = False
        self._lock = asyncio.Lock()
        self._read_condition = asyncio.Condition(self._lock)
        self._write_condition = asyncio.Condition(self._lock)

    async def acquire_read(self):
        async with self._lock:
            # 当有写操作在等待或正在执行时，读操作需要等待
            while self._writers_waiting > 0 or self._writer_active:
                await self._read_condition.wait()
            self._readers += 1

    async def release_read(self):
        async with self._lock:
            self._readers -= 1
            if self._readers == 0:
                self._write_condition.notify()

    async def acquire_write(self):
        async with self._lock:
            self._writers_waiting += 1
            # 当有读操作正在进行或写操作正在进行时，写操作需要等待
            while self._readers > 0 or self._writer_active:
                await self._write_condition.wait()
            self._writers_waiting -= 1
            self._writer_active = True

    async def release_write(self):
        async with self._lock:
            self._writer_active = False
            # 优先唤醒等待的写操作
            if self._writers_waiting > 0:
                self._write_condition.notify()
            else:
                # 如果没有写操作等待，则唤醒所有读操作
                self._read_condition.notify_all()

    def read_lock(self):
        return _ReadLockContextManager(self)

    def write_lock(self):
        return _WriteLockContextManager(self)


class _ReadLockContextManager:
    def __init__(self, lock):
        self._lock = lock

    async def __aenter__(self):
        await self._lock.acquire_read()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self._lock.release_read()


class _WriteLockContextManager:
    def __init__(self, lock):
        self._lock = lock

    async def __aenter__(self):
        await self._lock.acquire_write()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self._lock.release_write()


# 示例使用
async def reader(lock, id):
    print(f"Reader {id} enters the read function")
    async with lock.read_lock():
        print(f"Reader {id} is reading")
        await asyncio.sleep(1)
        print(f"Reader {id} finished reading")


async def writer(lock, id):
    print(f"Writer {id} enters the write function")
    async with lock.write_lock():
        print(f"Writer {id} is writing")
        await asyncio.sleep(1)
        print(f"Writer {id} finished writing")


async def main():
    lock = AsyncReadWriteLock()
    writers = [asyncio.create_task(writer(lock, i)) for i in range(2)]
    readers = [asyncio.create_task(reader(lock, i)) for i in range(5)]
    await asyncio.gather(*readers, *writers)


if __name__ == "__main__":
    asyncio.run(main())