在io.Reader中添加前缀

Question

我编写了一个小服务器，它以io.Reader的形式接收数据io.Reader ，添加一个标头并将结果流回调用方。

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

我已经看到了一些带有io.TeeReader的io.Pipe()示例，但它们更多的是将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
}

我明白从函数返回接口不是一个好主意，但是此代码注定不会成为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包确定请求正文的内容长度的方式 。