在开发基于 Golang 的代理服务器时，处理流式数据可能会遇到一些棘手的问题。例如，当你试图在转发数据给客户端的同时，将其捕获到缓冲区中以便后续使用时，可能会遇到数据丢失或不完整的问题。本文将分享我们在处理这个问题时的经验，并解释为什么手动读取和写入数据流有时比使用 io.TeeReader 和 io.Copy 更加可靠。
问题背景
在构建一个代理服务器时，我们需要将来自上游服务器的流式数据转发给客户端，并同时捕获这些数据以便将其保存到数据库中。最初，我们使用了 io.TeeReader 结合 io.Copy 来实现这一功能。io.TeeReader 是一个便利的工具，它允许我们将读取的数据同时写入两个地方（通常是一个缓冲区和一个目标输出，例如 io.Writer）。
然而，在实际运行中，我们发现捕获的数据存在不完整的情况，尤其是第一行的数据经常丢失。这导致了我们需要重新审视这部分代码，并寻找更可靠的解决方案。
初始解决方案：io.TeeReader 和 io.Copy
我们最初的实现使用了如下代码：
tee := io.TeeReader(response.Body, &buffer)
_, err := io.Copy(writer, tee)
这种方法理论上应该能够将 response.Body 中的数据同时写入 buffer 和 writer。但是，我们发现有时 io.TeeReader 并未如预期般工作，尤其是在处理长行或不规则数据流时，这导致了数据的丢失。
解决方案：手动读取和写入
为了确保数据的完整性，我们决定改用手动读取和写入流式数据的方法。新的实现如下：
proxy.ModifyResponse = func(response *http.Response) error {
log.Println("ModifyResponse started")
wg.Add(1)
go func() {
defer wg.Done()
defer func(writer *io.PipeWriter) {
err := writer.Close()
if err != nil {
log.Println("Error closing pipe writer:", err)
}
}(writer)
buf := make([]byte, 4096) // 定义一个足够大的缓冲区
for {
n, err := response.Body.Read(buf)
if n > 0 {
buffer.Write(buf[:n])
if _, writeErr := writer.Write(buf[:n]); writeErr != nil {
log.Println("Error writing to pipe:", writeErr)
return
}
}
if err != nil {
if err != io.EOF {
log.Println("Error reading from response body:", err)
}
break
}
}
}()
return nil
}
为什么手动读取和写入更可靠？
-
控制缓冲区大小： 手动读取允许我们精确控制每次从
response.Body中读取的数据块大小。通过定义一个合理大小的缓冲区（例如 4096 字节），我们确保每次读取的数据块足够大，可以避免io.TeeReader可能引发的读取块大小不一致问题。 -
更加明确的错误处理： 在手动读取和写入过程中，我们可以更及时、明确地处理可能发生的错误。这包括捕获
io.EOF错误以正确结束读取循环，防止数据丢失或处理不当。 -
避免潜在的竞争条件： 手动管理并发控制（例如使用
sync.WaitGroup和关闭PipeWriter）确保了并发读取和写入操作的顺序执行，避免了在并发环境下可能出现的数据竞争问题。 -
适合流式数据的处理： 流式数据在传输时并不是立即完成的，而是随着时间逐步传递的。手动读取允许我们更好地控制和调试数据的流动，而不像
io.Copy那样隐式处理整个流。这对于确保所有数据都被正确处理特别重要。
总结
通过手动读取 response.Body 的内容并写入缓冲区，我们解决了使用 io.TeeReader 和 io.Copy 时遇到的流式数据处理不完整的问题。这种方法提供了更高的可靠性，特别是在处理复杂的流式数据时。虽然 io.TeeReader 和 io.Copy 是很方便的工具，但在某些特定场景下，手动控制数据流可能是更好的选择。
这种经验教训表明，在开发过程中，我们必须灵活应对工具和技术的局限性，选择最适合当前需求的方法。如果你在 Golang 中处理类似的问题，希望本文提供的思路能够对你有所帮助。
