在io.Reader中添加前缀

Question

I've written a little server which receives a blob of data in the form of an io.Reader, adds a header and streams the result back to the caller.

My implementation isn't particularly efficient as I'm buffering the blob's data in-memory so that I can calculate the blob's length, which needs to form part of the header.

I've seen some examples of io.Pipe() with io.TeeReader but they're more for splitting an io.Reader into two, and writing them away in parallel.

The blobs I'm dealing with are around 100KB, so not huge but if my server gets busy, memory's going to quickly become an issue...

Any ideas?

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
}

I appreciate it's not a good idea to return interfaces from functions but this code isn't destined to become an API, so I'm not too concerned with that bit.

Answer 1

In general, the only way to determine the length of data in an io.Reader is to read until EOF. There are ways to determine the length of the data for specific types.

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
}

This is similar to how the net/http package determines the content length of the request body.