Repository: grpc/grpc-go

Issue:
mem.NewBuffer() should look at slice capacity?
### What version of gRPC are you using?

1.66.2

### What version of Go are you using (`go version`)?

N/A

### What operating system (Linux, Windows, …) and version?

N/A

### What did you do?

Roughly this:

```go
buf := pool.Get(32*1024)

_, _ = reader.Read(buf)

buffer := mem.NewBuffer(buf, pool)

buffer.Free()
```

I'm getting a ~big buffer from the pool and reading into it.

### What did you expect to see?

I expect to get a "normal" buffer for the `buf` slice. That way it can be put back into the pool when it's no longer needed.

### What did you see instead?

`mem.NewBuffer()` looks at the length of the `buf` and, some times, it's under the threshold. In that case the `buf` is wrapped into `SliceBuffer` and returned as is. This means `buf` will not be returned into the pool, but become garbage. This defeats the point of buffer pooling, obviously.

Perhaps `mem.NewBuffer()` should look at slice capacity (`cap(buf)`) instead of length (`len(buf)`)?

---

Current (undocumented) behavior contradicts the function documentation too:

```
// NewBuffer creates a new Buffer from the given data, initializing the reference
// counter to 1. The data will then be returned to the given pool when all
// references to the returned Buffer are released. As a special case to avoid
// additional allocations, if the given buffer pool is nil, the returned buffer
// will be a "no-op" Buffer where invoking Buffer.Free() does nothing and the
// underlying data is never freed.
//
// Note that the backing array of the given data is not copied.
```

The buffer will not be returned to the pool if it's length is too small.

Hints from issue comments:
@PapaCharlie: I think you're totally right on this one, there's definitely an edge case here. That being said, it's only an edge case in the specific code path you mentioned, i.e. some buffer is pulled from the pool, but only partially filled. If I did this:
```
buf := make([]byte, 32*1034)

_, _ = reader.Read(buf)

buffer := mem.NewBuffer(buf, pool)

buffer.Free()
```
Then the buffer would get put into the pool. I guess this is far less likely, in this case, the `pool` argument should be `nil`, making this really a non-issue.

Stepping back for a bit though, is there a reason you are requesting large buffers but reading very little data in them? Is it because you have no way of knowing ahead of time how much data you will be reading from the reader? Is it not possible to request smaller buffers?

I think either way, fixing the size check [here](https://github.com/grpc/grpc-go/blob/master/mem/buffers.go#L95C20-L95C49) makes sense to me@ash2k: > is there a reason you are requesting large buffers but reading very little data in them? Is it because you have no way of knowing ahead of time how much data you will be reading from the reader?

Yes, I don't know how much data there will be. It depends. Can be nothing, can be up to a few MB. It's a somewhat generic piece of code.

> Is it not possible to request smaller buffers?

It is but then the overhead is bigger. I think it should be cheaper to use a ~single buffer size across many such places in the codebase. That makes the chance of reusing the buffer much higher since it's of the same capacity. A lot of places in my code use 32KiB, so this codepath also does that.

Solution patch:
diff --git a/balancer/weightedtarget/weightedtarget.go b/balancer/weightedtarget/weightedtarget.go
index c004e112..a617f6a6 100644
--- a/balancer/weightedtarget/weightedtarget.go
+++ b/balancer/weightedtarget/weightedtarget.go
@@ -108,6 +108,7 @@ func (b *weightedTargetBalancer) UpdateClientConnState(s balancer.ClientConnStat
 		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
 	}
 	addressesSplit := hierarchy.Group(s.ResolverState.Addresses)
+	endpointsSplit := hierarchy.GroupEndpoints(s.ResolverState.Endpoints)
 
 	b.stateAggregator.PauseStateUpdates()
 	defer b.stateAggregator.ResumeStateUpdates()
@@ -155,6 +156,7 @@ func (b *weightedTargetBalancer) UpdateClientConnState(s balancer.ClientConnStat
 		_ = b.bg.UpdateClientConnState(name, balancer.ClientConnState{
 			ResolverState: resolver.State{
 				Addresses:     addressesSplit[name],
+				Endpoints:     endpointsSplit[name],
 				ServiceConfig: s.ResolverState.ServiceConfig,
 				Attributes:    s.ResolverState.Attributes.WithValue(localityKey, name),
 			},
diff --git a/internal/hierarchy/hierarchy.go b/internal/hierarchy/hierarchy.go
index c3baac36..362c05fa 100644
--- a/internal/hierarchy/hierarchy.go
+++ b/internal/hierarchy/hierarchy.go
@@ -48,6 +48,18 @@ func (p pathValue) Equal(o any) bool {
 	return true
 }
 
+// FromEndpoint returns the hierarchical path of endpoint.
+func FromEndpoint(endpoint resolver.Endpoint) []string {
+	path, _ := endpoint.Attributes.Value(pathKey).(pathValue)
+	return path
+}
+
+// SetInEndpoint overrides the hierarchical path in endpoint with path.
+func SetInEndpoint(endpoint resolver.Endpoint, path []string) resolver.Endpoint {
+	endpoint.Attributes = endpoint.Attributes.WithValue(pathKey, pathValue(path))
+	return endpoint
+}
+
 // Get returns the hierarchical path of addr.
 func Get(addr resolver.Address) []string {
 	attrs := addr.BalancerAttributes
@@ -110,3 +122,50 @@ func Group(addrs []resolver.Address) map[string][]resolver.Address {
 	}
 	return ret
 }
+
+// GroupEndpoints splits a slice of endpoints into groups based on
+// the first hierarchy path. The first hierarchy path will be removed from the
+// result.
+//
+// Input:
+// [
+//
+//	{endpoint0, path: [p0, wt0]}
+//	{endpoint1, path: [p0, wt1]}
+//	{endpoint2, path: [p1, wt2]}
+//	{endpoint3, path: [p1, wt3]}
+//
+// ]
+//
+// Endpoints will be split into p0/p1, and the p0/p1 will be removed from the
+// path.
+//
+// Output:
+//
+//	{
+//	  p0: [
+//	    {endpoint0, path: [wt0]},
+//	    {endpoint1, path: [wt1]},
+//	  ],
+//	  p1: [
+//	    {endpoint2, path: [wt2]},
+//	    {endpoint3, path: [wt3]},
+//	  ],
+//	}
+//
+// If hierarchical path is not set, or has no path in it, the endpoint is
+// dropped.
+func GroupEndpoints(endpoints []resolver.Endpoint) map[string][]resolver.Endpoint {
+	ret := make(map[string][]resolver.Endpoint)
+	for _, endpoint := range endpoints {
+		oldPath := FromEndpoint(endpoint)
+		if len(oldPath) == 0 {
+			continue
+		}
+		curPath := oldPath[0]
+		newPath := oldPath[1:]
+		newEndpoint := SetInEndpoint(endpoint, newPath)
+		ret[curPath] = append(ret[curPath], newEndpoint)
+	}
+	return ret
+}
diff --git a/mem/buffer_slice.go b/mem/buffer_slice.go
index 228e9c2f..65002e2c 100644
--- a/mem/buffer_slice.go
+++ b/mem/buffer_slice.go
@@ -22,6 +22,11 @@ import (
 	"io"
 )
 
+const (
+	// 32 KiB is what io.Copy uses.
+	readAllBufSize = 32 * 1024
+)
+
 // BufferSlice offers a means to represent data that spans one or more Buffer
 // instances. A BufferSlice is meant to be immutable after creation, and methods
 // like Ref create and return copies of the slice. This is why all methods have
@@ -219,8 +224,58 @@ func (w *writer) Write(p []byte) (n int, err error) {
 
 // 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.
+// 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
+			}
+			if len(*buf) == usedCap {
+				result = append(result, NewBuffer(buf, pool))
+				continue nextBuffer
+			}
+		}
+	}
+}
diff --git a/rpc_util.go b/rpc_util.go
index 033ffdc1..06c1f1b2 100644
--- a/rpc_util.go
+++ b/rpc_util.go
@@ -899,8 +899,7 @@ func decompress(compressor encoding.Compressor, d mem.BufferSlice, maxReceiveMes
 	//	}
 	//}
 
-	var out mem.BufferSlice
-	_, err = io.Copy(mem.NewWriter(&out, pool), io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
+	out, err := mem.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1), pool)
 	if err != nil {
 		out.Free()
 		return nil, 0, err
diff --git a/server.go b/server.go
index 23a16600..05f0c6e2 100644
--- a/server.go
+++ b/server.go
@@ -87,12 +87,13 @@ func init() {
 var statusOK = status.New(codes.OK, "")
 var logger = grpclog.Component("core")
 
-type methodHandler func(srv any, ctx context.Context, dec func(any) error, interceptor UnaryServerInterceptor) (any, error)
+// MethodHandler is a function type that processes a unary RPC method call.
+type MethodHandler func(srv any, ctx context.Context, dec func(any) error, interceptor UnaryServerInterceptor) (any, error)
 
 // MethodDesc represents an RPC service's method specification.
 type MethodDesc struct {
 	MethodName string
-	Handler    methodHandler
+	Handler    MethodHandler
 }
 
 // ServiceDesc represents an RPC service's specification.
diff --git a/stream.go b/stream.go
index 34c846a4..722209cc 100644
--- a/stream.go
+++ b/stream.go
@@ -113,7 +113,9 @@ type ClientStream interface {
 	// SendMsg is generally called by generated code. On error, SendMsg aborts
 	// the stream. If the error was generated by the client, the status is
 	// returned directly; otherwise, io.EOF is returned and the status of
-	// the stream may be discovered using RecvMsg.
+	// the stream may be discovered using RecvMsg. For unary or server-streaming
+	// RPCs (StreamDesc.ClientStreams is false), a nil error is returned
+	// unconditionally.
 	//
 	// SendMsg blocks until:
 	//   - There is sufficient flow control to schedule m with the transport, or
diff --git a/xds/internal/balancer/clustermanager/clustermanager.go b/xds/internal/balancer/clustermanager/clustermanager.go
index ef5b34ea..24ad2399 100644
--- a/xds/internal/balancer/clustermanager/clustermanager.go
+++ b/xds/internal/balancer/clustermanager/clustermanager.go
@@ -87,6 +87,7 @@ func (b *bal) updateChildren(s balancer.ClientConnState, newConfig *lbConfig) er
 	// TODO: Get rid of handling hierarchy in addresses. This LB policy never
 	// gets addresses from the resolver.
 	addressesSplit := hierarchy.Group(s.ResolverState.Addresses)
+	endpointsSplit := hierarchy.GroupEndpoints(s.ResolverState.Endpoints)
 
 	// Remove sub-balancers that are not in the new list from the aggregator and
 	// balancergroup.
@@ -139,6 +140,7 @@ func (b *bal) updateChildren(s balancer.ClientConnState, newConfig *lbConfig) er
 		if err := b.bg.UpdateClientConnState(childName, balancer.ClientConnState{
 			ResolverState: resolver.State{
 				Addresses:     addressesSplit[childName],
+				Endpoints:     endpointsSplit[childName],
 				ServiceConfig: s.ResolverState.ServiceConfig,
 				Attributes:    s.ResolverState.Attributes,
 			},
diff --git a/xds/internal/balancer/clusterresolver/clusterresolver.go b/xds/internal/balancer/clusterresolver/clusterresolver.go
index 3b996989..ae2c5fe9 100644
--- a/xds/internal/balancer/clusterresolver/clusterresolver.go
+++ b/xds/internal/balancer/clusterresolver/clusterresolver.go
@@ -252,7 +252,6 @@ func (b *clusterResolverBalancer) updateChildConfig() {
 	for i, a := range addrs {
 		endpoints[i].Attributes = a.BalancerAttributes
 		endpoints[i].Addresses = []resolver.Address{a}
-		endpoints[i].Addresses[0].BalancerAttributes = nil
 	}
 	if err := b.child.UpdateClientConnState(balancer.ClientConnState{
 		ResolverState: resolver.State{
diff --git a/xds/internal/balancer/priority/balancer.go b/xds/internal/balancer/priority/balancer.go
index c17c62f2..ba3fe52e 100644
--- a/xds/internal/balancer/priority/balancer.go
+++ b/xds/internal/balancer/priority/balancer.go
@@ -123,6 +123,7 @@ func (b *priorityBalancer) UpdateClientConnState(s balancer.ClientConnState) err
 		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
 	}
 	addressesSplit := hierarchy.Group(s.ResolverState.Addresses)
+	endpointsSplit := hierarchy.GroupEndpoints(s.ResolverState.Endpoints)
 
 	b.mu.Lock()
 	// Create and remove children, since we know all children from the config
@@ -142,6 +143,7 @@ func (b *priorityBalancer) UpdateClientConnState(s balancer.ClientConnState) err
 			cb := newChildBalancer(name, b, bb.Name(), b.cc)
 			cb.updateConfig(newSubConfig, resolver.State{
 				Addresses:     addressesSplit[name],
+				Endpoints:     endpointsSplit[name],
 				ServiceConfig: s.ResolverState.ServiceConfig,
 				Attributes:    s.ResolverState.Attributes,
 			})
@@ -163,6 +165,7 @@ func (b *priorityBalancer) UpdateClientConnState(s balancer.ClientConnState) err
 		// be built, if it's a low priority).
 		currentChild.updateConfig(newSubConfig, resolver.State{
 			Addresses:     addressesSplit[name],
+			Endpoints:     endpointsSplit[name],
 			ServiceConfig: s.ResolverState.ServiceConfig,
 			Attributes:    s.ResolverState.Attributes,
 		})



Test patch:
diff --git a/mem/buffer_slice_test.go b/mem/buffer_slice_test.go
index bb438443..bb9303f0 100644
--- a/mem/buffer_slice_test.go
+++ b/mem/buffer_slice_test.go
@@ -20,6 +20,8 @@ package mem_test
 
 import (
 	"bytes"
+	"crypto/rand"
+	"errors"
 	"fmt"
 	"io"
 	"testing"
@@ -27,6 +29,12 @@ import (
 	"google.golang.org/grpc/mem"
 )
 
+const (
+	minReadSize = 1
+	// Should match the constant in buffer_slice.go (another package)
+	readAllBufSize = 32 * 1024 // 32 KiB
+)
+
 func newBuffer(data []byte, pool mem.BufferPool) mem.Buffer {
 	return mem.NewBuffer(&data, pool)
 }
@@ -156,6 +164,252 @@ func (s) TestBufferSlice_Reader(t *testing.T) {
 	}
 }
 
+// TestBufferSlice_ReadAll_Reads exercises ReadAll by allowing it to read
+// various combinations of data, empty data, EOF.
+func (s) TestBufferSlice_ReadAll_Reads(t *testing.T) {
+	testcases := []struct {
+		name     string
+		reads    []readStep
+		wantErr  string
+		wantBufs int
+	}{
+		{
+			name: "EOF",
+			reads: []readStep{
+				{
+					err: io.EOF,
+				},
+			},
+		},
+		{
+			name: "data,EOF",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data+EOF",
+			reads: []readStep{
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "0,data+EOF",
+			reads: []readStep{
+				{},
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "0,data,EOF",
+			reads: []readStep{
+				{},
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+EOF",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "error",
+			reads: []readStep{
+				{
+					err: errors.New("boom"),
+				},
+			},
+			wantErr: "boom",
+		},
+		{
+			name: "data+error",
+			reads: []readStep{
+				{
+					n:   minReadSize,
+					err: errors.New("boom"),
+				},
+			},
+			wantErr:  "boom",
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+error",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   minReadSize,
+					err: errors.New("boom"),
+				},
+			},
+			wantErr:  "boom",
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+EOF - whole buf",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   readAllBufSize - minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data,EOF - whole buf",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n: readAllBufSize - minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data,EOF - 2 bufs",
+			reads: []readStep{
+				{
+					n: readAllBufSize,
+				},
+				{
+					n: minReadSize,
+				},
+				{
+					n: readAllBufSize - minReadSize,
+				},
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 3,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			pool := &testPool{
+				allocated: make(map[*[]byte]struct{}),
+			}
+			r := &stepReader{
+				reads: tc.reads,
+			}
+			data, err := mem.ReadAll(r, pool)
+			if tc.wantErr != "" {
+				if err == nil || err.Error() != tc.wantErr {
+					t.Fatalf("ReadAll() returned err %v, wanted %q", err, tc.wantErr)
+				}
+			} else {
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			gotData := data.Materialize()
+			if !bytes.Equal(r.read, gotData) {
+				t.Fatalf("ReadAll() returned data %q, wanted %q", gotData, r.read)
+			}
+			if len(data) != tc.wantBufs {
+				t.Fatalf("ReadAll() returned %d bufs, wanted %d bufs", len(data), tc.wantBufs)
+			}
+			// all but last should be full buffers
+			for i := 0; i < len(data)-1; i++ {
+				if data[i].Len() != readAllBufSize {
+					t.Fatalf("ReadAll() returned data length %d, wanted %d", data[i].Len(), readAllBufSize)
+				}
+			}
+			data.Free()
+			if len(pool.allocated) > 0 {
+				t.Fatalf("got %d allocated buffers, wanted none", len(pool.allocated))
+			}
+		})
+	}
+}
+
+func (s) TestBufferSlice_ReadAll_WriteTo(t *testing.T) {
+	testcases := []struct {
+		name string
+		size int
+	}{
+		{
+			name: "small",
+			size: minReadSize,
+		},
+		{
+			name: "exact size",
+			size: readAllBufSize,
+		},
+		{
+			name: "big",
+			size: readAllBufSize * 3,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			pool := &testPool{
+				allocated: make(map[*[]byte]struct{}),
+			}
+			buf := make([]byte, tc.size)
+			_, err := rand.Read(buf)
+			if err != nil {
+				t.Fatal(err)
+			}
+			r := bytes.NewBuffer(buf)
+			data, err := mem.ReadAll(r, pool)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			gotData := data.Materialize()
+			if !bytes.Equal(buf, gotData) {
+				t.Fatalf("ReadAll() = %q, wanted %q", gotData, buf)
+			}
+			data.Free()
+			if len(pool.allocated) > 0 {
+				t.Fatalf("wanted no allocated buffers, got %d", len(pool.allocated))
+			}
+		})
+	}
+}
+
 func ExampleNewWriter() {
 	var bs mem.BufferSlice
 	pool := mem.DefaultBufferPool()
@@ -176,3 +430,57 @@ func ExampleNewWriter() {
 	// Wrote 4 bytes, err: <nil>
 	// abcdabcdabcd
 }
+
+var (
+	_ io.Reader      = (*stepReader)(nil)
+	_ mem.BufferPool = (*testPool)(nil)
+)
+
+// readStep describes what a single stepReader.Read should do - how much data
+// to return and what error to return.
+type readStep struct {
+	n   int
+	err error
+}
+
+// stepReader implements io.Reader that reads specified amount of data and/or
+// returns the specified error in specified steps.
+// The read data is accumulated in the read field.
+type stepReader struct {
+	reads []readStep
+	read  []byte
+}
+
+func (s *stepReader) Read(buf []byte) (int, error) {
+	if len(s.reads) == 0 {
+		panic("unexpected Read() call")
+	}
+	read := s.reads[0]
+	s.reads = s.reads[1:]
+	_, err := rand.Read(buf[:read.n])
+	if err != nil {
+		panic(err)
+	}
+	s.read = append(s.read, buf[:read.n]...)
+	return read.n, read.err
+}
+
+// testPool is an implementation of BufferPool that allows to ensure that:
+// - there are matching Put calls for all Get calls.
+// - there are no unexpected Put calls.
+type testPool struct {
+	allocated map[*[]byte]struct{}
+}
+
+func (t *testPool) Get(length int) *[]byte {
+	buf := make([]byte, length)
+	t.allocated[&buf] = struct{}{}
+	return &buf
+}
+
+func (t *testPool) Put(buf *[]byte) {
+	if _, ok := t.allocated[buf]; !ok {
+		panic("unexpected put")
+	}
+	delete(t.allocated, buf)
+}