Repository: grpc/grpc-go
File: mem/buffer_slice.go (Go)

Code before:
// NewWriter wraps the given BufferSlice and BufferPool to implement the
// io.Writer interface. Every call to Write copies the contents of the given
// buffer into a new Buffer pulled from the given pool and the Buffer is added to
// the given BufferSlice.
func NewWriter(buffers *BufferSlice, pool BufferPool) io.Writer {
	return &writer{buffers: buffers, pool: pool}
}

// ReadAll reads from r until an error or EOF and returns the data it read.
// A successful call returns err == nil, not err == EOF. Because ReadAll is
// defined to read from src until EOF, it does not treat an EOF from Read
// as an error to be reported.
// Make sure to free the returned buffers even if you get an error from this function.
func ReadAll(r io.Reader, pool BufferPool) (BufferSlice, error) {
	var result BufferSlice
	wt, ok := r.(io.WriterTo)
	if ok {
		// This is more optimal since wt knows the size of chunks it wants to write and, hence, we can allocate
		// buffers of an optimal size to fit them. E.g. might be a single big chunk, and we wouldn't chop it into pieces.
		w := NewWriter(&result, pool)
		_, err := wt.WriteTo(w)
		return result, err
	}
	for {
		buf := pool.Get(readAllBufSize)
		// We asked for 32KiB but may have been given a bigger buffer. Use all of it if that's the case.
		*buf = (*buf)[:cap(*buf)]
		usedCap := 0
		for {
			n, err := r.Read((*buf)[usedCap:])
			usedCap += n
			if err != nil {
				if usedCap == 0 {
					// Nothing in this buf, put it back
					pool.Put(buf)
				} else {
					*buf = (*buf)[:usedCap]
					result = append(result, NewBuffer(buf, pool))
				}
				if err == io.EOF {
					err = nil
				}
				return result, err
			}
			if len(*buf) == usedCap {
				result = append(result, NewBuffer(buf, pool))
				break // grab a new buf from pool
			}
		}
	}
}

Review comment (suggestion):
Ditto on this one, if `err != nil`, you should free `result` and return an empty BufferSlice (see [my other comment](https://github.com/grpc/grpc-go/pull/7653/files#r1788364035))

Code after:

// NewWriter wraps the given BufferSlice and BufferPool to implement the
// io.Writer interface. Every call to Write copies the contents of the given
// buffer into a new Buffer pulled from the given pool and the Buffer is
// added to the given BufferSlice.
func NewWriter(buffers *BufferSlice, pool BufferPool) io.Writer {
	return &writer{buffers: buffers, pool: pool}
}

// ReadAll reads from r until an error or EOF and returns the data it read.
// A successful call returns err == nil, not err == EOF. Because ReadAll is
// defined to read from src until EOF, it does not treat an EOF from Read
// as an error to be reported.
//
// Important: A failed call returns a non-nil error and may also return
// partially read buffers. It is the responsibility of the caller to free the
// BufferSlice returned, or its memory will not be reused.
func ReadAll(r io.Reader, pool BufferPool) (BufferSlice, error) {
	var result BufferSlice
	if wt, ok := r.(io.WriterTo); ok {
		// This is more optimal since wt knows the size of chunks it wants to
		// write and, hence, we can allocate buffers of an optimal size to fit
		// them. E.g. might be a single big chunk, and we wouldn't chop it
		// into pieces.
		w := NewWriter(&result, pool)
		_, err := wt.WriteTo(w)
		return result, err
	}
nextBuffer:
	for {
		buf := pool.Get(readAllBufSize)
		// We asked for 32KiB but may have been given a bigger buffer.
		// Use all of it if that's the case.
		*buf = (*buf)[:cap(*buf)]
		usedCap := 0
		for {
			n, err := r.Read((*buf)[usedCap:])
			usedCap += n
			if err != nil {
				if usedCap == 0 {
					// Nothing in this buf, put it back
					pool.Put(buf)
				} else {
					*buf = (*buf)[:usedCap]
					result = append(result, NewBuffer(buf, pool))
				}
				if err == io.EOF {
					err = nil
				}
				return result, err
			}