如何实现处理AsyncIterator当前元素时异步预取下一个元素？

你的问题本质是原异步迭代器的获取和处理是串行的——必须等当前元素处理完，迭代器才会继续取下一个。要实现处理时预取，完全不需要弃用AsyncIterator，只需要调整迭代器的实现逻辑，让它提前启动下一个元素的异步获取任务，和当前元素的处理并行执行。

核心思路 #

把「获取下一个元素」的操作提前：在yield当前元素时，就启动下一个元素的异步获取任务。这样当你处理当前元素的1秒里，下一个元素已经在后台请求了，总耗时会从2n秒降到n+1秒（n为元素数量）。

方案1：修改原有异步生成器 #

直接修改fetch_elements，用asyncio.create_task预取下一个元素，实现最简单的并行预取：

import asyncio
from typing import AsyncIterable

async def _get_next(prev: int) -> int:
    print(f"Fetching element {prev+1}")
    await asyncio.sleep(1)
    return prev + 1

async def fetch_elements(n: int) -> AsyncIterable[int]:
    curr = 0
    # 提前启动第一个元素的获取任务
    next_task = asyncio.create_task(_get_next(curr))
    for _ in range(n):
        # 等待获取完成，拿到当前元素
        curr = await next_task
        print(f"Yielding element {curr}")
        # 非最后一个元素时，预取下一个
        if _ < n - 1:
            next_task = asyncio.create_task(_get_next(curr))
        yield curr

async def process_element(e: int) -> None:
    print(f"Pushing element {e}")
    await asyncio.sleep(1)
    print(f"Element {e} pushed")

async def main():
    n = 3
    async for elem in fetch_elements(n):
        await process_element(elem)

asyncio.run(main())

执行后你会看到输出顺序：

Fetching element 1
Yielding element 1
Pushing element 1
Fetching element 2
Element 1 pushed
Pushing element 2
Fetching element 3
Element 2 pushed
Pushing element 3
Element 3 pushed

明显能看到，处理元素1的同时已经开始获取元素2，完全实现了并行预取的效果。

方案2：更健壮的类式异步迭代器 #

如果需要支持中途终止迭代（比如break），上面的简单版本可能会有任务泄漏。这时可以用类实现AsyncIterator，在迭代结束时自动清理未完成的预取任务：

import asyncio
from typing import AsyncIterator

async def _get_next(prev: int) -> int:
    print(f"Fetching element {prev+1}")
    await asyncio.sleep(1)
    return prev + 1

class FetchElements:
    def __init__(self, n: int):
        self.total = n
        self.fetched = 0
        self.current = 0
        self.next_task = None

    def __aiter__(self) -> AsyncIterator[int]:
        return self

    async def __anext__(self) -> int:
        if self.fetched >= self.total:
            # 清理剩余的预取任务，避免资源泄漏
            if self.next_task:
                self.next_task.cancel()
                try:
                    await self.next_task
                except asyncio.CancelledError:
                    pass
            raise StopAsyncIteration

        # 第一次执行时启动第一个元素的获取任务
        if self.next_task is None:
            self.next_task = asyncio.create_task(_get_next(self.current))
        
        # 等待当前元素获取完成
        self.current = await self.next_task
        self.fetched += 1
        print(f"Yielding element {self.current}")

        # 预取下一个元素（非最后一次迭代）
        if self.fetched < self.total:
            self.next_task = asyncio.create_task(_get_next(self.current))
        else:
            self.next_task = None

        return self.current

async def process_element(e: int) -> None:
    print(f"Pushing element {e}")
    await asyncio.sleep(1)
    print(f"Element {e} pushed")

async def main():
    n = 3
    async for elem in FetchElements(n):
        await process_element(elem)
        # 可以取消注释测试中途终止逻辑
        # if elem == 2:
        #     break

asyncio.run(main())

这个版本在迭代终止（包括正常结束和中途break）时，会自动取消未完成的预取任务，避免不必要的资源占用。

原代码的问题根源 #

原fetch_elements是串行执行逻辑：先await _get_next拿到元素，再yield，此时生成器会暂停，直到下一次迭代（也就是当前元素处理完成）才会继续执行循环，启动下一次_get_next。而修改后的版本把_get_next的启动放在yield之前，让元素获取和处理操作并行进行。

内容的提问来源于stack exchange，提问作者Ayhon Iron