go - 为 io.Reader 添加前缀

Question

我编写了一个小型服务器，它以 io.Reader 的形式接收一团数据，添加一个 header 并将结果流式传输回调用者。

我的实现不是特别有效，因为我在内存中缓冲 blob 的数据，以便我可以计算 blob 的长度，这需要构成 header 的一部分。

我看过一些 io.Pipe() 和 io.TeeReader 的例子，但它们更适合拆分 io.Reader 分成两部分，并将它们并行写入。

我正在处理的 blob 大约是 100KB，所以不是很大，但是如果我的服务器变得繁忙，内存很快就会成为一个问题...

有什么想法吗？

func addHeader(in io.Reader) (out io.Reader, err error) {
    buf := new(bytes.Buffer)
    if _, err = io.Copy(buf, in); err != nil {
        return
    }

    header := bytes.NewReader([]byte(fmt.Sprintf("header:%d", buf.Len())))

    return io.MultiReader(header, buf), nil
}

我知道从函数返回接口(interface)不是一个好主意，但这段代码注定不会成为 API，所以我不太关心这一点。

Answer 1

通常，确定 io.Reader 中数据长度的唯一方法是读取到 EOF。有多种方法可以确定特定类型的数据长度。

func addHeader(in io.Reader) (out io.Reader, err error) {
  n := 0
  switch v := in.(type) {
  case *bytes.Buffer:
    n = v.Len()
  case *bytes.Reader:
    n = v.Len()
  case *strings.Reader:
    n = v.Len()
  case io.Seeker:
    cur, err := v.Seek(0, 1)
    if err != nil {
        return nil, err
    }
    end, err := v.Seek(0, 2)
    if err != nil {
        return nil, err
    }
    _, err = v.Seek(cur, 0)
    if err != nil {
        return nil, err
    }
    n = int(end - cur)
  default:
    var buf bytes.Buffer
    if _, err := buf.ReadFrom(in); err != nil {
        return nil, err
    }
    n = buf.Len()
    in = &buf
  }
  header := strings.NewReader(fmt.Sprintf("header:%d", n))
  return io.MultiReader(header, in), nil
}

这类似于 net/http 包 determines the content length of the request body 的方式.