深入理解Python asyncio模块的实现原理

原创于 2025-06-05 09:08:23 发布 · 319 阅读

4 ·

CC 4.0 BY-SA版权

深入理解Python asyncio模块的实现原理

pysheeet Python Cheat Sheet 项目地址: https://gitcode.com/gh_mirrors/py/pysheeet

前言

Python的asyncio模块是异步编程的核心，它提供了一套完整的异步I/O框架。本文将通过分析pysheeet项目中的asyncio实现示例，深入探讨asyncio背后的工作原理，帮助开发者更好地理解和使用这一强大的异步编程工具。

Task的实现原理

Task是asyncio中用于管理协程执行状态的核心组件。让我们来看一个简化的Task实现：

class Task(Future):
    """Task的简化实现原型"""
    
    def __init__(self, gen, *, loop):
        super().__init__(loop=loop)
        self._gen = gen  # 保存协程对象
        self._loop.call_soon(self._step)  # 立即开始执行
        
    def _step(self, val=None, exc=None):
        try:
            if exc:
                f = self._gen.throw(exc)  # 处理异常
            else:
                f = self._gen.send(val)  # 继续执行协程
        except StopIteration as e:
            self.set_result(e.value)  # 协程完成，设置结果
        except Exception as e:
            self.set_exception(e)  # 协程出错，设置异常
        else:
            f.add_done_callback(self._wakeup)  # 添加回调继续执行
            
    def _wakeup(self, fut):
        try:
            res = fut.result()
            self._step(res, None)  # 正常结果继续执行
        except Exception as e:
            self._step(None, e)  # 异常结果处理异常

关键点解析：

Task继承自Future，可以跟踪协程执行状态
使用生成器的send/throw方法驱动协程执行
通过回调链实现协程的连续执行
自动处理协程完成和异常情况

事件循环的工作原理

事件循环是asyncio的核心调度器，下面是一个简化实现：

class Loop:
    """简化的事件循环原型"""
    
    def __init__(self):
        self._ready = deque()  # 准备执行的任务队列
        self._stopping = False
        
    def run_forever(self):
        """持续运行任务直到停止"""
        try:
            while True:
                self._run_once()  # 执行一轮任务
                if self._stopping:
                    break
        finally:
            self._stopping = False
            
    def _run_once(self):
        """执行一轮任务"""
        ntodo = len(self._ready)
        for i in range(ntodo):
            t, a = self._ready.popleft()
            t(*a)  # 执行任务

事件循环的核心机制：

维护一个就绪任务队列
循环从队列中取出并执行任务
提供停止机制控制循环退出
通过call_soon方法添加任务到队列

asyncio.wait的实现原理

asyncio.wait是常用的等待多个协程完成的工具：

async def wait(fs, loop=None):
    fs = {asyncio.ensure_future(_) for _ in set(fs)}
    waiter = loop.create_future()
    counter = len(fs)
    
    def _on_complete(f):
        nonlocal counter
        counter -= 1
        if counter <= 0 and not waiter.done():
            waiter.set_result(None)
    
    for f in fs:
        f.add_done_callback(_on_complete)
    
    await waiter  # 等待所有任务完成
    
    # 分类已完成和未完成的任务
    done, pending = set(), set()
    for f in fs:
        f.remove_done_callback(_on_complete)
        if f.done():
            done.add(f)
        else:
            pending.add(f)
    return done, pending

关键实现细节：

将所有输入转换为Future对象
创建waiter Future用于等待所有任务完成
为每个任务添加完成回调
当所有任务完成时设置waiter结果
最后分类已完成和未完成的任务

网络I/O的实现原理

asyncio提供了高效的网络I/O操作，下面是socket相关操作的实现：

def sock_accept(self, sock, fut=None, registed=False):
    fd = sock.fileno()
    if fut is None:
        fut = self.create_future()
    if registed:
        self.remove_reader(fd)
    try:
        conn, addr = sock.accept()
        conn.setblocking(False)
    except (BlockingIOError, InterruptedError):
        self.add_reader(fd, self.sock_accept, sock, fut, True)
    except Exception as e:
        fut.set_exception(e)
    else:
        fut.set_result((conn, addr))
    return fut

网络I/O的实现特点：

使用非阻塞socket
当I/O不可立即完成时注册到事件循环
I/O就绪后继续执行并返回结果
通过Future对象跟踪操作状态

创建服务器的实现

create_server是创建TCP服务器的便捷方法：

async def create_server(loop, protocol_factory, host, port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setblocking(False)
    sock.bind((host, port))
    sock.listen(10)
    sockets = [sock]
    server = asyncio.base_events.Server(loop, sockets)
    loop._start_serving(protocol_factory, sock, None, server)
    return server

服务器创建流程：