PyZMQ中asyncio接收端内存泄漏问题解析

问题现象

在使用PyZMQ的asyncio接口时，当接收端持续不断地接收消息时，会出现内存持续增长的问题。具体表现为接收端进程的内存使用量会快速攀升至数GB，而发送端则保持稳定的内存使用。

问题复现

通过以下代码可以稳定复现该问题：

接收端代码：

import zmq.asyncio, gc, asyncio

async def client():
    ctx = zmq.asyncio.Context()
    sock = ctx.socket(zmq.PAIR)
    sock.setsockopt(zmq.RCVBUF, 1024*10)
    sock.setsockopt(zmq.RCVHWM, 1)
    sock.connect("ipc:///tmp/memtest")

    print("client ready")
    i = 0
    while True:
        data = await sock.recv(copy=True)
        print("client recv", i)
        del data
        gc.collect()
        i += 1

asyncio.run(client())

发送端代码：

import zmq.asyncio, asyncio

async def server():
    ctx = zmq.asyncio.Context()
    sock = ctx.socket(zmq.PAIR)
    sock.setsockopt(zmq.SNDBUF, 1024*10)
    sock.setsockopt(zmq.SNDHWM, 1)
    sock.bind("ipc:///tmp/memtest")

    data = b'0'*1000000

    print("server ready")
    i = 0
    while True:
        await sock.send(data, copy=True)
        print("server send", i)
        i += 1

asyncio.run(server())

问题根源分析

这个问题本质上是由asyncio的事件循环机制和PyZMQ的实现方式共同导致的：

1. PyZMQ的优化实现：PyZMQ为了提高性能，直接使用了Future对象而非协程。当socket始终有消息可接收时，这些Future会立即完成。

2. 事件循环的调度特性：在Python的asyncio中，Future的回调是通过call_soon注册的，这意味着它们会在事件循环的下一次迭代中被调用。但是当socket持续有消息时，事件循环永远不会空闲，导致回调无法被执行。

3. 内存泄漏机制：每个完成的Future都会保留对消息数据的引用，而由于回调未被调用，这些引用无法被释放，从而导致内存持续增长。

解决方案

1. 强制事件循环空闲：在接收循环中定期插入asyncio.sleep(0)，给事件循环执行回调的机会：

async def client():
    # ...初始化代码...
    i = 0
    while True:
        data = await sock.recv(copy=True)
        if i % 1000 == 0:
            await asyncio.sleep(0)  # 让事件循环处理回调
        print("client recv", i)
        del data
        i += 1

2. 使用同步上下文：如果不必须使用asyncio，可以改用同步的zmq.Context()，这不会出现内存泄漏问题。

3. PyZMQ内部优化：PyZMQ可以避免注册done回调来减少引用保留（已在PR #1929中实现）。

深入理解

这个问题揭示了asyncio编程中一个重要的概念：事件循环饥饿。当某个协程持续占用事件循环而不主动让出控制权时，其他任务（包括回调处理）就会被"饿死"。

在PyZMQ的特定实现中，由于socket始终有消息可接收，Future会立即完成，导致await操作实际上不会让出控制权。这与常规的I/O操作不同，后者通常会涉及真正的等待，从而自然地让出控制权。

最佳实践建议

1. 在使用PyZMQ的asyncio接口时，特别是高频率消息场景下，应该定期插入短暂的sleep(0)来确保事件循环的健康运行。

2. 对于纯粹的高吞吐量消息处理，考虑使用同步接口可能更为合适，可以避免这类问题。

3. 监控内存使用情况，特别是在长时间运行的消息处理服务中，及时发现潜在的内存问题。

这个问题虽然表现为内存泄漏，但本质上是一个事件循环调度问题，理解这一点对于编写健壮的异步IO程序非常重要。