如何获取aiohttp.ClientSession收发的原始数据大小？

嘿，这个问题我刚好折腾过！aiohttp本身并没有直接提供获取会话收发原始数据大小的API，但我们可以通过自定义扩展来实现，下面给你两种实用的方案，按需选择：

如果需要统计真正在网络上传输的原始字节数（包括TCP头、SSL加密数据、重传字节等），可以通过自定义TCPConnector来包装底层的socket传输，拦截收发操作并计数。

import aiohttp
from aiohttp import TCPConnector
import asyncio

class TrackingConnector(TCPConnector):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.total_sent = 0
        self.total_received = 0

    async def connect(self, req):
        # 获取原始的transport和protocol
        transport, protocol = await super().connect(req)
        
        # 包装write方法统计发送字节
        original_write = transport.write
        def tracked_write(data):
            nonlocal self
            self.total_sent += len(data)
            original_write(data)
        
        # 包装data_received方法统计接收字节
        original_receive = protocol.data_received
        def tracked_receive(data):
            nonlocal self
            self.total_received += len(data)
            original_receive(data)
        
        # 替换原始方法
        transport.write = tracked_write
        protocol.data_received = tracked_receive
        return transport, protocol

async def main():
    connector = TrackingConnector()
    async with aiohttp.ClientSession(connector=connector) as session:
        # 发送几个请求测试
        async with session.get('https://example.com') as resp:
            await resp.text()
        async with session.post('https://httpbin.org/post', json={"test": "data"}) as resp:
            await resp.json()
        
        print(f"累计发送原始字节: {connector.total_sent}")
        print(f"累计接收原始字节: {connector.total_received}")

asyncio.run(main())

注意：这个方案统计的是TCP层的全部字节，包括握手、加密、重传等，数值会比HTTP层面的统计大。如果服务器返回压缩数据（比如gzip），这里统计的是压缩后的原始传输大小。

如果只需要统计HTTP协议层面的数据量（请求行+请求头+请求体，响应行+响应头+响应体），可以用aiohttp的TraceConfig钩子来实现，更贴近业务层面的统计需求。

import aiohttp
import asyncio

async def main():
    total_sent = 0
    total_received = 0

    # 请求开始时统计请求大小
    async def track_request(session, trace_ctx, params):
        nonlocal total_sent
        # 计算请求行大小（比如GET /path HTTP/1.1\r\n）
        request_line = f"{params.method} {params.url.path_qs} HTTP/1.1\r\n"
        # 计算请求头总大小
        headers_size = sum(len(f"{k}: {v}\r\n") for k, v in params.headers.items())
        # 请求体大小
        body_size = len(params.body) if params.body else 0
        # 加上最后结束的\r\n
        total_sent += len(request_line) + headers_size + body_size + 2

    # 响应结束时统计响应大小
    async def track_response(session, trace_ctx, params):
        nonlocal total_received
        # 计算状态行大小（比如HTTP/1.1 200 OK\r\n）
        status_line = f"HTTP/1.1 {params.response.status} {params.response.reason}\r\n"
        # 响应头总大小
        headers_size = sum(len(f"{k}: {v}\r\n") for k, v in params.response.headers.items())
        # 响应体大小（需要先读取完整响应）
        body_data = await params.response.read()
        body_size = len(body_data)
        # 加上结束的\r\n
        total_received += len(status_line) + headers_size + body_size + 2

    # 创建TraceConfig并注册钩子
    trace_config = aiohttp.TraceConfig()
    trace_config.on_request_start.append(track_request)
    trace_config.on_response_end.append(track_response)

    async with aiohttp.ClientSession(trace_configs=[trace_config]) as session:
        async with session.get('https://example.com') as resp:
            pass  # 钩子会自动处理统计
        async with session.post('https://httpbin.org/post', data={"key": "value"}) as resp:
            pass
        
        print(f"HTTP层面累计发送: {total_sent} 字节")
        print(f"HTTP层面累计接收: {total_received} 字节")

asyncio.run(main())

注意：如果响应是分块编码（Transfer-Encoding: chunked），await params.response.read()会把所有分块数据拼接后统计，适合非流式的场景；如果是流式响应，需要在on_response_chunk_received钩子中逐块累加。

• 如果需要区分单个请求的收发大小，可以在trace_ctx中存储每个请求的统计数据（比如trace_ctx.sent_size = xxx），最后在请求结束时汇总。
• 方案1的TCP统计包含了所有网络开销，适合做网络性能分析；方案2的HTTP统计更贴近业务数据量，适合做接口流量统计。

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