io.TeeReader与io.MultiWriter的差异解析及SSE流令牌监控场景选型咨询

你的场景需求 #

你想要搭建一个LLM应用，需要添加中间件来监控令牌使用量，而LLM服务商返回的是SSE流。查资料后关注到io.TeeReader和io.MultiWriter，看了一堆对比内容后，除了下游写入器数量的区别，没发现其他差异，还准备了测试代码（不过你贴的代码没写完，咱们先聚焦核心问题聊）。

核心差异远不止写入器数量！ #

别只盯着下游写入器的数量，这俩工具的设计定位和适用场景本质上就不一样：

• io.TeeReader是「读操作时分流」：它是一个io.Reader的包装器，当你从它读取数据时，会自动把读到的内容同步写到指定的io.Writer里。简单说就是，你读一份数据的同时，自动复制一份到目标位置，核心是绑定了读操作和写操作的联动。
• io.MultiWriter是「写操作时分流」：它是一个io.Writer的包装器，当你往它写入数据时，会把这份数据同时写到多个绑定的io.Writer里。核心是把单次写操作扩散到多个目标，它本身完全不涉及读操作。

你的SSE令牌监控场景该选哪个？ #

结合你处理SSE流并监控令牌的需求，我更推荐用io.TeeReader，理由很实在：

1. 完美贴合SSE的流式逻辑：SSE是从LLM的响应流中逐段读取数据，你需要一边把流返回给前端，一边把流内容传给令牌统计逻辑。io.TeeReader刚好能在你读取SSE流的同时，自动把数据复制一份给监控中间件，完全不用手动处理读写同步的问题。
2. 代码更简洁低耦合：你只需要把LLM返回的原始io.Reader包装成io.TeeReader，其中一个Writer指向原响应输出（比如给前端的HTTP响应），另一个指向你的令牌统计处理器。整个流程是线性的，不会打乱原有的流式读取逻辑。

如果用io.MultiWriter的话，你得先把SSE流读到内存缓冲区，再往MultiWriter里写，这直接破坏了SSE的流式特性，还可能带来不必要的内存占用，完全没必要。

给你贴个贴合Go场景的伪代码参考：

func handleLLMRequest(w http.ResponseWriter, r *http.Request) {
    // 调用LLM获取SSE流的Reader
    llmStream, err := llmClient.GetStream(r.Context(), prompt)
    if err != nil {
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    defer llmStream.Close()

    // 实例化自定义的令牌统计Writer
    tokenMonitor := &TokenMonitor{}

    // 用TeeReader包装原流：读数据时同步写给tokenMonitor和响应Writer
    teeStream := io.TeeReader(llmStream, tokenMonitor)

    // 返回流式响应给前端，同时tokenMonitor后台统计令牌
    w.Header().Set("Content-Type", "text/event-stream")
    io.Copy(w, teeStream)

    // 输出最终统计结果
    log.Printf("本次请求令牌使用量：%d", tokenMonitor.TotalTokens())
}

// 自定义TokenMonitor实现io.Writer接口，解析SSE内容统计令牌
type TokenMonitor struct {
    totalTokens int
}

func (tm *TokenMonitor) Write(p []byte) (n int, err error) {
    // 这里实现SSE内容解析和令牌统计逻辑
    // 比如调用tiktoken类库统计字节对应的令牌数
    tokens := countTokens(string(p))
    tm.totalTokens += tokens
    return len(p), nil
}

额外补充 #

如果以后你遇到需要把一份数据同时写到多个存储目标（比如既存本地文件又存数据库）的场景，那io.MultiWriter才是合适的选择，但你的SSE流监控场景，io.TeeReader完全是量身定做的～

备注：内容来源于stack exchange，提问作者lxt