使用asyncio/aiohttp+TaskPool批量请求时如何返回结果？

需要发起约50万次HTTP请求，每次获取一个JSON文件，总数据量约8GiB。直接使用asyncio/aiohttp的asyncio.gather会导致内存溢出（OOM），因为该方法会一次性创建所有协程并将所有结果存储在内存中。

原示例代码：

async def get(session, headers, body):
    async with session.get(url='someURL', headers=headers, json=body) as response:
        resp = await response.json()
      
    return resp

async def main(bodies, headers):
    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session:
        data = await asyncio.gather(*[get(session, headers, body) for body in bodies])

    return data

headers = 'some_headers'
bodies = ['list of dictionaries where each dictionary is a body for the request']

data = asyncio.run(main(bodies, headers))

# Other operations with data

尝试过相关并发控制方案但未成功，认为TaskPool类的方案内存占用最优，但无法从TaskPool任务中获取返回结果，仅能打印。以下是针对两种TaskPool实现的修正方案：

1. asyncio-taskpool实现修正 #

原代码返回None的原因是未正确处理结果迭代，且一次性存储所有结果仍会导致OOM。建议边获取结果边写入文件，避免内存堆积：

from asyncio_taskpool import TaskPool 
import asyncio
import aiohttp
import json

async def get_and_process(session, headers, body):
    async with session.get(url='someURL', headers=headers, json=body, timeout=None) as response:
        resp = await response.json()
    # 边获取结果边写入文件，用请求唯一标识命名（假设body包含id字段）
    with open(f"./results/result_{body['id']}.json", 'w', encoding='utf-8') as f:
        json.dump(resp, f, ensure_ascii=False)
    # 仅返回关键标识，不存储完整结果
    return f"请求{body['id']}完成"

async def get_all(bodies, headers, limit_concurrent_tasks):
    pool = TaskPool()
    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session:
        # 用imap迭代获取任务结果，而非一次性收集
        async for status in pool.imap(get_and_process, ((session, headers, body) for body in bodies), num_concurrent=limit_concurrent_tasks):
            print(status)
        await pool.close()

async def main(bodies, headers):
    # 并发数根据服务器承受能力调整，1000可能过高，建议先测试500以内
    limit_concurrent_tasks = 500
    await get_all(bodies, headers, limit_concurrent_tasks)

headers = 'some_headers'
bodies = ['list of dictionaries where each dictionary is a body for the request']

asyncio.run(main(bodies, headers))

2. 自定义TaskPool实现修正 #

原自定义TaskPool缺少结果存储逻辑，且回调函数引用错误。修正后增加结果存储和边处理边写入的逻辑：

import asyncio
import aiohttp
import json

class TaskPool(object):
    def __init__(self, workers):
        self._semaphore = asyncio.Semaphore(workers)
        self._tasks = set()
        self._success_count = 0
        self._error_count = 0

    async def put(self, coro):
        await self._semaphore.acquire()  
        # Python 3.7+推荐用create_task替代ensure_future
        task = asyncio.create_task(coro)
        self._tasks.add(task)
        task.add_done_callback(self._on_task_done)

    def _on_task_done(self, task):
        self._tasks.remove(task)
        self._semaphore.release()
        try:
            resp = task.result()
            # 写入文件，避免内存堆积
            with open(f"./results/result_{resp['id']}.json", 'w', encoding='utf-8') as f:
                json.dump(resp, f, ensure_ascii=False)
            self._success_count += 1
            print(f"完成第{self._success_count}次请求")
        except Exception as e:
            self._error_count += 1
            print(f"请求失败（第{self._error_count}次）：{str(e)}")

    async def join(self):
        # 用return_exceptions=True避免单个任务失败导致整体中断
        await asyncio.gather(*self._tasks, return_exceptions=True)
        return self._success_count, self._error_count

    async def __aenter__(self):
        return self

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

# 并发数建议根据实际测试调整，避免给目标服务器造成过大压力
limit = 500

async def get(session, headers, body):
    async with session.get(url='someURL', headers=headers, json=body, timeout=None) as response:
        resp = await response.json()
    # 给结果添加请求标识，方便文件命名
    resp['id'] = body['id']
    return resp

async def _main(bodies, headers):
    # 确保结果目录存在
    import os
    os.makedirs("./results", exist_ok=True)
    
    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session, TaskPool(workers=limit) as tasks:
        for body in bodies:
            await tasks.put(get(session, headers, body))
        success, error = await tasks.join()
    print(f"全部请求完成：成功{success}次，失败{error}次")

headers = 'some_headers'
bodies = ['list of dictionaries where each dictionary is a body for the request']

asyncio.run(_main(bodies, headers))

核心优化说明 #

1. 控制并发数：通过Semaphore限制同时运行的请求数，避免瞬间创建大量协程占用内存。
2. 内存优化：不将所有结果存储在内存中，而是边获取边写入本地文件，仅保留请求状态统计信息。
3. 异常容错：增加异常捕获逻辑，单个请求失败不会导致整个任务中断，同时统计成功/失败次数。
4. 并发数调整：建议根据目标服务器的限流策略和自身机器性能调整并发数，避免触发目标服务器的反爬或限流机制。

注：所有操作基于Python 3.8环境。

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