grpc
/
grpc-go
mirror of https://github.com/grpc/grpc-go.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Implement new Codec that uses `mem.BufferSlice` instead of `[]byte` (#7356)

This commit is contained in:
Paul Chesnais 2024-08-21 17:11:39 -04:00 committed by GitHub
parent 7e12068baf
commit 9ab8b62505
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
39 changed files with 1844 additions and 1067 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
.github/workflows/testing.yml vendored
View File

@ -55,6 +55,10 @@ jobs:
goversion: '1.22'
testflags: -race
- type: tests
goversion: '1.22'
testflags: '-race -tags=buffer_pooling'
- type: tests
goversion: '1.22'
goarch: 386

34
benchmark/benchmain/main.go
View File

@ -66,11 +66,11 @@ import (
"google.golang.org/grpc/benchmark/stats"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/encoding/gzip"
"google.golang.org/grpc/experimental"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/channelz"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/test/bufconn"
@ -153,6 +153,33 @@ const (
warmuptime = time.Second
)
var useNopBufferPool atomic.Bool
type swappableBufferPool struct {
mem.BufferPool
}
func (p swappableBufferPool) Get(length int) *[]byte {
var pool mem.BufferPool
if useNopBufferPool.Load() {
pool = mem.NopBufferPool{}
} else {
pool = p.BufferPool
}
return pool.Get(length)
}
func (p swappableBufferPool) Put(i *[]byte) {
if useNopBufferPool.Load() {
return
}
p.BufferPool.Put(i)
}
func init() {
internal.SetDefaultBufferPoolForTesting.(func(mem.BufferPool))(swappableBufferPool{mem.DefaultBufferPool()})
}
var (
allWorkloads = []string{workloadsUnary, workloadsStreaming, workloadsUnconstrained, workloadsAll}
allCompModes = []string{compModeOff, compModeGzip, compModeNop, compModeAll}
@ -343,10 +370,9 @@ func makeClients(bf stats.Features) ([]testgrpc.BenchmarkServiceClient, func())
}
switch bf.RecvBufferPool {
case recvBufferPoolNil:
// Do nothing.
useNopBufferPool.Store(true)
case recvBufferPoolSimple:
opts = append(opts, experimental.WithRecvBufferPool(grpc.NewSharedBufferPool()))
sopts = append(sopts, experimental.RecvBufferPool(grpc.NewSharedBufferPool()))
// Do nothing as buffering is enabled by default.
default:
logger.Fatalf("Unknown shared recv buffer pool type: %v", bf.RecvBufferPool)
}

75
codec.go
View File

@ -21,18 +21,79 @@ package grpc
import (
"google.golang.org/grpc/encoding"
_ "google.golang.org/grpc/encoding/proto" // to register the Codec for "proto"
"google.golang.org/grpc/mem"
)
// baseCodec contains the functionality of both Codec and encoding.Codec, but
// omits the name/string, which vary between the two and are not needed for
// anything besides the registry in the encoding package.
// baseCodec captures the new encoding.CodecV2 interface without the Name
// function, allowing it to be implemented by older Codec and encoding.Codec
// implementations. The omitted Name function is only needed for the register in
// the encoding package and is not part of the core functionality.
type baseCodec interface {
Marshal(v any) ([]byte, error)
Unmarshal(data []byte, v any) error
Marshal(v any) (mem.BufferSlice, error)
Unmarshal(data mem.BufferSlice, v any) error
}
var _ baseCodec = Codec(nil)
var _ baseCodec = encoding.Codec(nil)
// getCodec returns an encoding.CodecV2 for the codec of the given name (if
// registered). Initially checks the V2 registry with encoding.GetCodecV2 and
// returns the V2 codec if it is registered. Otherwise, it checks the V1 registry
// with encoding.GetCodec and if it is registered wraps it with newCodecV1Bridge
// to turn it into an encoding.CodecV2. Returns nil otherwise.
func getCodec(name string) encoding.CodecV2 {
codecV2 := encoding.GetCodecV2(name)
if codecV2 != nil {
return codecV2
}
codecV1 := encoding.GetCodec(name)
if codecV1 != nil {
return newCodecV1Bridge(codecV1)
}
return nil
}
func newCodecV0Bridge(c Codec) baseCodec {
return codecV0Bridge{codec: c}
}
func newCodecV1Bridge(c encoding.Codec) encoding.CodecV2 {
return codecV1Bridge{
codecV0Bridge: codecV0Bridge{codec: c},
name: c.Name(),
}
}
var _ baseCodec = codecV0Bridge{}
type codecV0Bridge struct {
codec interface {
Marshal(v any) ([]byte, error)
Unmarshal(data []byte, v any) error
}
}
func (c codecV0Bridge) Marshal(v any) (mem.BufferSlice, error) {
data, err := c.codec.Marshal(v)
if err != nil {
return nil, err
}
return mem.BufferSlice{mem.NewBuffer(&data, nil)}, nil
}
func (c codecV0Bridge) Unmarshal(data mem.BufferSlice, v any) (err error) {
return c.codec.Unmarshal(data.Materialize(), v)
}
var _ encoding.CodecV2 = codecV1Bridge{}
type codecV1Bridge struct {
codecV0Bridge
name string
}
func (c codecV1Bridge) Name() string {
return c.name
}
// Codec defines the interface gRPC uses to encode and decode messages.
// Note that implementations of this interface must be thread safe;

27
dialoptions.go
View File

@ -33,6 +33,7 @@ import (
"google.golang.org/grpc/internal/binarylog"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/stats"
)
@ -60,7 +61,7 @@ func init() {
internal.WithBinaryLogger = withBinaryLogger
internal.JoinDialOptions = newJoinDialOption
internal.DisableGlobalDialOptions = newDisableGlobalDialOptions
internal.WithRecvBufferPool = withRecvBufferPool
internal.WithBufferPool = withBufferPool
}
// dialOptions configure a Dial call. dialOptions are set by the DialOption
@ -92,7 +93,6 @@ type dialOptions struct {
defaultServiceConfigRawJSON *string
resolvers []resolver.Builder
idleTimeout time.Duration
recvBufferPool SharedBufferPool
defaultScheme string
maxCallAttempts int
}
@ -679,11 +679,11 @@ func defaultDialOptions() dialOptions {
WriteBufferSize: defaultWriteBufSize,
UseProxy: true,
UserAgent: grpcUA,
BufferPool: mem.DefaultBufferPool(),
},
bs: internalbackoff.DefaultExponential,
healthCheckFunc: internal.HealthCheckFunc,
idleTimeout: 30 * time.Minute,
recvBufferPool: nopBufferPool{},
defaultScheme: "dns",
maxCallAttempts: defaultMaxCallAttempts,
}
@ -760,25 +760,8 @@ func WithMaxCallAttempts(n int) DialOption {
})
}
// WithRecvBufferPool returns a DialOption that configures the ClientConn
// to use the provided shared buffer pool for parsing incoming messages. Depending
// on the application's workload, this could result in reduced memory allocation.
//
// If you are unsure about how to implement a memory pool but want to utilize one,
// begin with grpc.NewSharedBufferPool.
//
// Note: The shared buffer pool feature will not be active if any of the following
// options are used: WithStatsHandler, EnableTracing, or binary logging. In such
// cases, the shared buffer pool will be ignored.
//
// Deprecated: use experimental.WithRecvBufferPool instead. Will be deleted in
// v1.60.0 or later.
func WithRecvBufferPool(bufferPool SharedBufferPool) DialOption {
return withRecvBufferPool(bufferPool)
}
func withRecvBufferPool(bufferPool SharedBufferPool) DialOption {
func withBufferPool(bufferPool mem.BufferPool) DialOption {
return newFuncDialOption(func(o *dialOptions) {
o.recvBufferPool = bufferPool
o.copts.BufferPool = bufferPool
})
}

82
encoding/encoding_v2.go Normal file
View File

@ -0,0 +1,82 @@
/*
*
* Copyright 2024 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package encoding
import (
"strings"
"google.golang.org/grpc/mem"
)
// CodecV2 defines the interface gRPC uses to encode and decode messages. Note
// that implementations of this interface must be thread safe; a CodecV2's
// methods can be called from concurrent goroutines.
type CodecV2 interface {
// Marshal returns the wire format of v. The buffers in the returned
// [mem.BufferSlice] must have at least one reference each, which will be freed
// by gRPC when they are no longer needed.
Marshal(v any) (out mem.BufferSlice, err error)
// Unmarshal parses the wire format into v. Note that data will be freed as soon
// as this function returns. If the codec wishes to guarantee access to the data
// after this function, it must take its own reference that it frees when it is
// no longer needed.
Unmarshal(data mem.BufferSlice, v any) error
// Name returns the name of the Codec implementation. The returned string
// will be used as part of content type in transmission. The result must be
// static; the result cannot change between calls.
Name() string
}
var registeredV2Codecs = make(map[string]CodecV2)
// RegisterCodecV2 registers the provided CodecV2 for use with all gRPC clients and
// servers.
//
// The CodecV2 will be stored and looked up by result of its Name() method, which
// should match the content-subtype of the encoding handled by the CodecV2. This
// is case-insensitive, and is stored and looked up as lowercase. If the
// result of calling Name() is an empty string, RegisterCodecV2 will panic. See
// Content-Type on
// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
// more details.
//
// If both a Codec and CodecV2 are registered with the same name, the CodecV2
// will be used.
//
// NOTE: this function must only be called during initialization time (i.e. in
// an init() function), and is not thread-safe. If multiple Codecs are
// registered with the same name, the one registered last will take effect.
func RegisterCodecV2(codec CodecV2) {
if codec == nil {
panic("cannot register a nil CodecV2")
}
if codec.Name() == "" {
panic("cannot register CodecV2 with empty string result for Name()")
}
contentSubtype := strings.ToLower(codec.Name())
registeredV2Codecs[contentSubtype] = codec
}
// GetCodecV2 gets a registered CodecV2 by content-subtype, or nil if no CodecV2 is
// registered for the content-subtype.
//
// The content-subtype is expected to be lowercase.
func GetCodecV2(contentSubtype string) CodecV2 {
return registeredV2Codecs[contentSubtype]
}

81
encoding/proto/proto_v2.go Normal file
View File

@ -0,0 +1,81 @@
/*
*
* Copyright 2024 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package proto
import (
"fmt"
"google.golang.org/grpc/encoding"
"google.golang.org/grpc/mem"
"google.golang.org/protobuf/proto"
)
func init() {
encoding.RegisterCodecV2(&codecV2{})
}
// codec is a CodecV2 implementation with protobuf. It is the default codec for
// gRPC.
type codecV2 struct{}
var _ encoding.CodecV2 = (*codecV2)(nil)
func (c *codecV2) Marshal(v any) (data mem.BufferSlice, err error) {
vv := messageV2Of(v)
if vv == nil {
return nil, fmt.Errorf("proto: failed to marshal, message is %T, want proto.Message", v)
}
size := proto.Size(vv)
if mem.IsBelowBufferPoolingThreshold(size) {
buf, err := proto.Marshal(vv)
if err != nil {
return nil, err
}
data = append(data, mem.SliceBuffer(buf))
} else {
pool := mem.DefaultBufferPool()
buf := pool.Get(size)
if _, err := (proto.MarshalOptions{}).MarshalAppend((*buf)[:0], vv); err != nil {
pool.Put(buf)
return nil, err
}
data = append(data, mem.NewBuffer(buf, pool))
}
return data, nil
}
func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) (err error) {
vv := messageV2Of(v)
if vv == nil {
return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
}
buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
defer buf.Free()
// TODO: Upgrade proto.Unmarshal to support mem.BufferSlice. Right now, it's not
// really possible without a major overhaul of the proto package, but the
// vtprotobuf library may be able to support this.
return proto.Unmarshal(buf.ReadOnlyData(), vv)
}
func (c *codecV2) Name() string {
return Name
}

55
experimental/experimental.go
View File

@ -28,38 +28,37 @@ package experimental
import (
"google.golang.org/grpc"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/mem"
)
// WithRecvBufferPool returns a grpc.DialOption that configures the use of
// bufferPool for parsing incoming messages on a grpc.ClientConn. Depending on
// the application's workload, this could result in reduced memory allocation.
// WithBufferPool returns a grpc.DialOption that configures the use of bufferPool
// for parsing incoming messages on a grpc.ClientConn, and for temporary buffers
// when marshaling outgoing messages. By default, mem.DefaultBufferPool is used,
// and this option only exists to provide alternative buffer pool implementations
// to the client, such as more optimized size allocations etc. However, the
// default buffer pool is already tuned to account for many different use-cases.
//
// If you are unsure about how to implement a memory pool but want to utilize
// one, begin with grpc.NewSharedBufferPool.
//
// Note: The shared buffer pool feature will not be active if any of the
// following options are used: WithStatsHandler, EnableTracing, or binary
// logging. In such cases, the shared buffer pool will be ignored.
//
// Note: It is not recommended to use the shared buffer pool when compression is
// enabled.
func WithRecvBufferPool(bufferPool grpc.SharedBufferPool) grpc.DialOption {
return internal.WithRecvBufferPool.(func(grpc.SharedBufferPool) grpc.DialOption)(bufferPool)
// Note: The following options will interfere with the buffer pool because they
// require a fully materialized buffer instead of a sequence of buffers:
// EnableTracing, and binary logging. In such cases, materializing the buffer
// will generate a lot of garbage, reducing the overall benefit from using a
// pool.
func WithBufferPool(bufferPool mem.BufferPool) grpc.DialOption {
return internal.WithBufferPool.(func(mem.BufferPool) grpc.DialOption)(bufferPool)
}
// RecvBufferPool returns a grpc.ServerOption that configures the server to use
// the provided shared buffer pool for parsing incoming messages. Depending on
// the application's workload, this could result in reduced memory allocation.
// BufferPool returns a grpc.ServerOption that configures the server to use the
// provided buffer pool for parsing incoming messages and for temporary buffers
// when marshaling outgoing messages. By default, mem.DefaultBufferPool is used,
// and this option only exists to provide alternative buffer pool implementations
// to the server, such as more optimized size allocations etc. However, the
// default buffer pool is already tuned to account for many different use-cases.
//
// If you are unsure about how to implement a memory pool but want to utilize
// one, begin with grpc.NewSharedBufferPool.
//
// Note: The shared buffer pool feature will not be active if any of the
// following options are used: StatsHandler, EnableTracing, or binary logging.
// In such cases, the shared buffer pool will be ignored.
//
// Note: It is not recommended to use the shared buffer pool when compression is
// enabled.
func RecvBufferPool(bufferPool grpc.SharedBufferPool) grpc.ServerOption {
return internal.RecvBufferPool.(func(grpc.SharedBufferPool) grpc.ServerOption)(bufferPool)
// Note: The following options will interfere with the buffer pool because they
// require a fully materialized buffer instead of a sequence of buffers:
// EnableTracing, and binary logging. In such cases, materializing the buffer
// will generate a lot of garbage, reducing the overall benefit from using a
// pool.
func BufferPool(bufferPool mem.BufferPool) grpc.ServerOption {
return internal.BufferPool.(func(mem.BufferPool) grpc.ServerOption)(bufferPool)
}

18
experimental/shared_buffer_pool_test.go
View File

@ -46,6 +46,9 @@ func Test(t *testing.T) {
const defaultTestTimeout = 10 * time.Second
func (s) TestRecvBufferPoolStream(t *testing.T) {
// TODO: How much of this test can be preserved now that buffer reuse happens at
// the codec and HTTP/2 level?
t.SkipNow()
tcs := []struct {
name string
callOpts []grpc.CallOption
@ -83,8 +86,8 @@ func (s) TestRecvBufferPoolStream(t *testing.T) {
}
pool := &checkBufferPool{}
sopts := []grpc.ServerOption{experimental.RecvBufferPool(pool)}
dopts := []grpc.DialOption{experimental.WithRecvBufferPool(pool)}
sopts := []grpc.ServerOption{experimental.BufferPool(pool)}
dopts := []grpc.DialOption{experimental.WithBufferPool(pool)}
if err := ss.Start(sopts, dopts...); err != nil {
t.Fatalf("Error starting endpoint server: %v", err)
}
@ -129,6 +132,8 @@ func (s) TestRecvBufferPoolStream(t *testing.T) {
}
func (s) TestRecvBufferPoolUnary(t *testing.T) {
// TODO: See above
t.SkipNow()
tcs := []struct {
name string
callOpts []grpc.CallOption
@ -159,8 +164,8 @@ func (s) TestRecvBufferPoolUnary(t *testing.T) {
}
pool := &checkBufferPool{}
sopts := []grpc.ServerOption{experimental.RecvBufferPool(pool)}
dopts := []grpc.DialOption{experimental.WithRecvBufferPool(pool)}
sopts := []grpc.ServerOption{experimental.BufferPool(pool)}
dopts := []grpc.DialOption{experimental.WithBufferPool(pool)}
if err := ss.Start(sopts, dopts...); err != nil {
t.Fatalf("Error starting endpoint server: %v", err)
}
@ -196,8 +201,9 @@ type checkBufferPool struct {
puts [][]byte
}
func (p *checkBufferPool) Get(size int) []byte {
return make([]byte, size)
func (p *checkBufferPool) Get(size int) *[]byte {
b := make([]byte, size)
return &b
}
func (p *checkBufferPool) Put(bs *[]byte) {

8
gcp/observability/logging_test.go
View File

@ -204,7 +204,7 @@ func (s) TestClientRPCEventsLogAll(t *testing.T) {
SequenceID: 2,
Authority: ss.Address,
Payload: payload{
Message: []uint8{},
Message: nil,
},
},
{
@ -285,7 +285,7 @@ func (s) TestClientRPCEventsLogAll(t *testing.T) {
SequenceID: 2,
Authority: ss.Address,
Payload: payload{
Message: []uint8{},
Message: nil,
},
},
{
@ -512,7 +512,7 @@ func (s) TestServerRPCEventsLogAll(t *testing.T) {
SequenceID: 4,
Authority: ss.Address,
Payload: payload{
Message: []uint8{},
Message: nil,
},
},
{
@ -870,7 +870,7 @@ func (s) TestPrecedenceOrderingInConfiguration(t *testing.T) {
SequenceID: 2,
Authority: ss.Address,
Payload: payload{
Message: []uint8{},
Message: nil,
},
},
{

8
internal/experimental.go
View File

@ -18,11 +18,11 @@
package internal
var (
// WithRecvBufferPool is implemented by the grpc package and returns a dial
// WithBufferPool is implemented by the grpc package and returns a dial
// option to configure a shared buffer pool for a grpc.ClientConn.
WithRecvBufferPool any // func (grpc.SharedBufferPool) grpc.DialOption
WithBufferPool any // func (grpc.SharedBufferPool) grpc.DialOption
// RecvBufferPool is implemented by the grpc package and returns a server
// BufferPool is implemented by the grpc package and returns a server
// option to configure a shared buffer pool for a grpc.Server.
RecvBufferPool any // func (grpc.SharedBufferPool) grpc.ServerOption
BufferPool any // func (grpc.SharedBufferPool) grpc.ServerOption
)

19
internal/grpctest/grpctest.go
View File

@ -24,17 +24,22 @@ import (
"strings"
"sync/atomic"
"testing"
"time"
"google.golang.org/grpc/internal/leakcheck"
)
var lcFailed uint32
type errorer struct {
type logger struct {
t *testing.T
}
func (e errorer) Errorf(format string, args ...any) {
func (e logger) Logf(format string, args ...any) {
e.t.Logf(format, args...)
}
func (e logger) Errorf(format string, args ...any) {
atomic.StoreUint32(&lcFailed, 1)
e.t.Errorf(format, args...)
}
@ -48,16 +53,22 @@ type Tester struct{}
// Setup updates the tlogger.
func (Tester) Setup(t *testing.T) {
TLogger.Update(t)
// TODO: There is one final leak around closing connections without completely
// draining the recvBuffer that has yet to be resolved. All other leaks have been
// completely addressed, and this can be turned back on as soon as this issue is
// fixed.
leakcheck.SetTrackingBufferPool(logger{t: t})
}
// Teardown performs a leak check.
func (Tester) Teardown(t *testing.T) {
leakcheck.CheckTrackingBufferPool()
if atomic.LoadUint32(&lcFailed) == 1 {
return
}
leakcheck.Check(errorer{t: t})
leakcheck.CheckGoroutines(logger{t: t}, 10*time.Second)
if atomic.LoadUint32(&lcFailed) == 1 {
t.Log("Leak check disabled for future tests")
t.Log("Goroutine leak check disabled for future tests")
}
TLogger.EndTest(t)
}

4
internal/internal.go
View File

@ -225,6 +225,10 @@ var (
// SetDefaultBufferPoolForTesting updates the default buffer pool, for
// testing purposes.
SetDefaultBufferPoolForTesting any // func(mem.BufferPool)
// SetBufferPoolingThresholdForTesting updates the buffer pooling threshold, for
// testing purposes.
SetBufferPoolingThresholdForTesting any // func(int)
)
// HealthChecker defines the signature of the client-side LB channel health

200
internal/leakcheck/leakcheck.go
View File

@ -16,18 +16,171 @@
*
*/
// Package leakcheck contains functions to check leaked goroutines.
// Package leakcheck contains functions to check leaked goroutines and buffers.
//
// Call "defer leakcheck.Check(t)" at the beginning of tests.
// Call the following at the beginning of test:
//
// defer leakcheck.NewLeakChecker(t).Check()
package leakcheck
import (
"runtime"
"runtime/debug"
"slices"
"sort"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/mem"
)
// failTestsOnLeakedBuffers is a special flag that will cause tests to fail if
// leaked buffers are detected, instead of simply logging them as an
// informational failure. This can be enabled with the "checkbuffers" compile
// flag, e.g.:
//
// go test -tags=checkbuffers
var failTestsOnLeakedBuffers = false
func init() {
defaultPool := mem.DefaultBufferPool()
globalPool.Store(&defaultPool)
(internal.SetDefaultBufferPoolForTesting.(func(mem.BufferPool)))(&globalPool)
}
var globalPool swappableBufferPool
type swappableBufferPool struct {
atomic.Pointer[mem.BufferPool]
}
func (b *swappableBufferPool) Get(length int) *[]byte {
return (*b.Load()).Get(length)
}
func (b *swappableBufferPool) Put(buf *[]byte) {
(*b.Load()).Put(buf)
}
// SetTrackingBufferPool replaces the default buffer pool in the mem package to
// one that tracks where buffers are allocated. CheckTrackingBufferPool should
// then be invoked at the end of the test to validate that all buffers pulled
// from the pool were returned.
func SetTrackingBufferPool(logger Logger) {
newPool := mem.BufferPool(&trackingBufferPool{
pool: *globalPool.Load(),
logger: logger,
allocatedBuffers: make(map[*[]byte][]uintptr),
})
globalPool.Store(&newPool)
}
// CheckTrackingBufferPool undoes the effects of SetTrackingBufferPool, and fails
// unit tests if not all buffers were returned. It is invalid to invoke this
// method without previously having invoked SetTrackingBufferPool.
func CheckTrackingBufferPool() {
p := (*globalPool.Load()).(*trackingBufferPool)
p.lock.Lock()
defer p.lock.Unlock()
globalPool.Store(&p.pool)
type uniqueTrace struct {
stack []uintptr
count int
}
var totalLeakedBuffers int
var uniqueTraces []uniqueTrace
for _, stack := range p.allocatedBuffers {
idx, ok := slices.BinarySearchFunc(uniqueTraces, stack, func(trace uniqueTrace, stack []uintptr) int {
return slices.Compare(trace.stack, stack)
})
if !ok {
uniqueTraces = slices.Insert(uniqueTraces, idx, uniqueTrace{stack: stack})
}
uniqueTraces[idx].count++
totalLeakedBuffers++
}
for _, ut := range uniqueTraces {
frames := runtime.CallersFrames(ut.stack)
var trace strings.Builder
for {
f, ok := frames.Next()
if !ok {
break
}
trace.WriteString(f.Function)
trace.WriteString("\n\t")
trace.WriteString(f.File)
trace.WriteString(":")
trace.WriteString(strconv.Itoa(f.Line))
trace.WriteString("\n")
}
format := "%d allocated buffers never freed:\n%s"
args := []any{ut.count, trace.String()}
if failTestsOnLeakedBuffers {
p.logger.Errorf(format, args...)
} else {
p.logger.Logf("WARNING "+format, args...)
}
}
if totalLeakedBuffers > 0 {
p.logger.Logf("%g%% of buffers never freed", float64(totalLeakedBuffers)/float64(p.bufferCount))
}
}
type trackingBufferPool struct {
pool mem.BufferPool
logger Logger
lock sync.Mutex
bufferCount int
allocatedBuffers map[*[]byte][]uintptr
}
func (p *trackingBufferPool) Get(length int) *[]byte {
p.lock.Lock()
defer p.lock.Unlock()
p.bufferCount++
buf := p.pool.Get(length)
var stackBuf [16]uintptr
var stack []uintptr
skip := 2
for {
n := runtime.Callers(skip, stackBuf[:])
stack = append(stack, stackBuf[:n]...)
if n < len(stackBuf) {
break
}
skip += len(stackBuf)
}
p.allocatedBuffers[buf] = stack
return buf
}
func (p *trackingBufferPool) Put(buf *[]byte) {
p.lock.Lock()
defer p.lock.Unlock()
if _, ok := p.allocatedBuffers[buf]; !ok {
p.logger.Errorf("Unknown buffer freed:\n%s", string(debug.Stack()))
} else {
delete(p.allocatedBuffers, buf)
}
p.pool.Put(buf)
}
var goroutinesToIgnore = []string{
"testing.Main(",
"testing.tRunner(",
@ -94,13 +247,17 @@ func interestingGoroutines() (gs []string) {
return
}
// Errorfer is the interface that wraps the Errorf method. It's a subset of
// testing.TB to make it easy to use Check.
type Errorfer interface {
// Logger is the interface that wraps the Logf and Errorf method. It's a subset
// of testing.TB to make it easy to use this package.
type Logger interface {
Logf(format string, args ...any)
Errorf(format string, args ...any)
}
func check(efer Errorfer, timeout time.Duration) {
// CheckGoroutines looks at the currently-running goroutines and checks if there
// are any interesting (created by gRPC) goroutines leaked. It waits up to 10
// seconds in the error cases.
func CheckGoroutines(logger Logger, timeout time.Duration) {
// Loop, waiting for goroutines to shut down.
// Wait up to timeout, but finish as quickly as possible.
deadline := time.Now().Add(timeout)
@ -112,13 +269,32 @@ func check(efer Errorfer, timeout time.Duration) {
time.Sleep(50 * time.Millisecond)
}
for _, g := range leaked {
efer.Errorf("Leaked goroutine: %v", g)
logger.Errorf("Leaked goroutine: %v", g)
}
}
// Check looks at the currently-running goroutines and checks if there are any
// interesting (created by gRPC) goroutines leaked. It waits up to 10 seconds
// in the error cases.
func Check(efer Errorfer) {
check(efer, 10*time.Second)
// LeakChecker captures an Logger and is returned by NewLeakChecker as a
// convenient method to set up leak check tests in a unit test.
type LeakChecker struct {
logger Logger
}
// Check executes the leak check tests, failing the unit test if any buffer or
// goroutine leaks are detected.
func (lc *LeakChecker) Check() {
CheckTrackingBufferPool()
CheckGoroutines(lc.logger, 10*time.Second)
}
// NewLeakChecker offers a convenient way to set up the leak checks for a
// specific unit test. It can be used as follows, at the beginning of tests:
//
// defer leakcheck.NewLeakChecker(t).Check()
//
// It initially invokes SetTrackingBufferPool to set up buffer tracking, then the
// deferred LeakChecker.Check call will invoke CheckTrackingBufferPool and
// CheckGoroutines with a default timeout of 10 seconds.
func NewLeakChecker(logger Logger) *LeakChecker {
SetTrackingBufferPool(logger)
return &LeakChecker{logger: logger}
}

35
shared_buffer_pool_test.go → internal/leakcheck/leakcheck_enabled.go
View File

@ -1,6 +1,8 @@
//go:build checkbuffers
/*
*
* Copyright 2023 gRPC authors.
* Copyright 2017 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@ -16,33 +18,8 @@
*
*/
package grpc
package leakcheck
import "testing"
func (s) TestSharedBufferPool(t *testing.T) {
pools := []SharedBufferPool{
nopBufferPool{},
NewSharedBufferPool(),
}
lengths := []int{
level4PoolMaxSize + 1,
level4PoolMaxSize,
level3PoolMaxSize,
level2PoolMaxSize,
level1PoolMaxSize,
level0PoolMaxSize,
}
for _, p := range pools {
for _, l := range lengths {
bs := p.Get(l)
if len(bs) != l {
t.Fatalf("Expected buffer of length %d, got %d", l, len(bs))
}
p.Put(&bs)
}
}
func init() {
failTestsOnLeakedBuffers = true
}

23
internal/leakcheck/leakcheck_test.go
View File

@ -25,12 +25,15 @@ import (
"time"
)
type testErrorfer struct {
type testLogger struct {
errorCount int
errors []string
}
func (e *testErrorfer) Errorf(format string, args ...any) {
func (e *testLogger) Logf(format string, args ...any) {
}
func (e *testLogger) Errorf(format string, args ...any) {
e.errors = append(e.errors, fmt.Sprintf(format, args...))
e.errorCount++
}
@ -43,13 +46,13 @@ func TestCheck(t *testing.T) {
if ig := interestingGoroutines(); len(ig) == 0 {
t.Error("blah")
}
e := &testErrorfer{}
check(e, time.Second)
e := &testLogger{}
CheckGoroutines(e, time.Second)
if e.errorCount != leakCount {
t.Errorf("check found %v leaks, want %v leaks", e.errorCount, leakCount)
t.Errorf("CheckGoroutines found %v leaks, want %v leaks", e.errorCount, leakCount)
t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
}
check(t, 3*time.Second)
CheckGoroutines(t, 3*time.Second)
}
func ignoredTestingLeak(d time.Duration) {
@ -66,11 +69,11 @@ func TestCheckRegisterIgnore(t *testing.T) {
if ig := interestingGoroutines(); len(ig) == 0 {
t.Error("blah")
}
e := &testErrorfer{}
check(e, time.Second)
e := &testLogger{}
CheckGoroutines(e, time.Second)
if e.errorCount != leakCount {
t.Errorf("check found %v leaks, want %v leaks", e.errorCount, leakCount)
t.Errorf("CheckGoroutines found %v leaks, want %v leaks", e.errorCount, leakCount)
t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
}
check(t, 3*time.Second)
CheckGoroutines(t, 3*time.Second)
}

88
internal/transport/controlbuf.go
View File

@ -32,6 +32,7 @@ import (
"golang.org/x/net/http2/hpack"
"google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/internal/grpcutil"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/status"
)
@ -148,9 +149,9 @@ type dataFrame struct {
streamID uint32
endStream bool
h []byte
d []byte
reader mem.Reader
// onEachWrite is called every time
// a part of d is written out.
// a part of data is written out.
onEachWrite func()
}
@ -454,12 +455,13 @@ func (c *controlBuffer) finish() {
// These streams need to be cleaned out since the transport
// is still not aware of these yet.
for head := c.list.dequeueAll(); head != nil; head = head.next {
hdr, ok := head.it.(*headerFrame)
if !ok {
continue
}
if hdr.onOrphaned != nil { // It will be nil on the server-side.
hdr.onOrphaned(ErrConnClosing)
switch v := head.it.(type) {
case *headerFrame:
if v.onOrphaned != nil { // It will be nil on the server-side.
v.onOrphaned(ErrConnClosing)
}
case *dataFrame:
_ = v.reader.Close()
}
}
@ -509,12 +511,13 @@ type loopyWriter struct {
draining bool
conn net.Conn
logger *grpclog.PrefixLogger
bufferPool mem.BufferPool
// Side-specific handlers
ssGoAwayHandler func(*goAway) (bool, error)
}
func newLoopyWriter(s side, fr *framer, cbuf *controlBuffer, bdpEst *bdpEstimator, conn net.Conn, logger *grpclog.PrefixLogger, goAwayHandler func(*goAway) (bool, error)) *loopyWriter {
func newLoopyWriter(s side, fr *framer, cbuf *controlBuffer, bdpEst *bdpEstimator, conn net.Conn, logger *grpclog.PrefixLogger, goAwayHandler func(*goAway) (bool, error), bufferPool mem.BufferPool) *loopyWriter {
var buf bytes.Buffer
l := &loopyWriter{
side: s,
@ -530,6 +533,7 @@ func newLoopyWriter(s side, fr *framer, cbuf *controlBuffer, bdpEst *bdpEstimato
conn: conn,
logger: logger,
ssGoAwayHandler: goAwayHandler,
bufferPool: bufferPool,
}
return l
}
@ -787,6 +791,11 @@ func (l *loopyWriter) cleanupStreamHandler(c *cleanupStream) error {
// not be established yet.
delete(l.estdStreams, c.streamID)
str.deleteSelf()
for head := str.itl.dequeueAll(); head != nil; head = head.next {
if df, ok := head.it.(*dataFrame); ok {
_ = df.reader.Close()
}
}
}
if c.rst { // If RST_STREAM needs to be sent.
if err := l.framer.fr.WriteRSTStream(c.streamID, c.rstCode); err != nil {
@ -922,16 +931,18 @@ func (l *loopyWriter) processData() (bool, error) {
dataItem := str.itl.peek().(*dataFrame) // Peek at the first data item this stream.
// A data item is represented by a dataFrame, since it later translates into
// multiple HTTP2 data frames.
// Every dataFrame has two buffers; h that keeps grpc-message header and d that is actual data.
// As an optimization to keep wire traffic low, data from d is copied to h to make as big as the
// maximum possible HTTP2 frame size.
// Every dataFrame has two buffers; h that keeps grpc-message header and data
// that is the actual message. As an optimization to keep wire traffic low, data
// from data is copied to h to make as big as the maximum possible HTTP2 frame
// size.
if len(dataItem.h) == 0 && len(dataItem.d) == 0 { // Empty data frame
if len(dataItem.h) == 0 && dataItem.reader.Remaining() == 0 { // Empty data frame
// Client sends out empty data frame with endStream = true
if err := l.framer.fr.WriteData(dataItem.streamID, dataItem.endStream, nil); err != nil {
return false, err
}
str.itl.dequeue() // remove the empty data item from stream
_ = dataItem.reader.Close()
if str.itl.isEmpty() {
str.state = empty
} else if trailer, ok := str.itl.peek().(*headerFrame); ok { // the next item is trailers.
@ -946,9 +957,7 @@ func (l *loopyWriter) processData() (bool, error) {
}
return false, nil
}
var (
buf []byte
)
// Figure out the maximum size we can send
maxSize := http2MaxFrameLen
if strQuota := int(l.oiws) - str.bytesOutStanding; strQuota <= 0 { // stream-level flow control.
@ -962,43 +971,50 @@ func (l *loopyWriter) processData() (bool, error) {
}
// Compute how much of the header and data we can send within quota and max frame length
hSize := min(maxSize, len(dataItem.h))
dSize := min(maxSize-hSize, len(dataItem.d))
if hSize != 0 {
if dSize == 0 {
buf = dataItem.h
} else {
// We can add some data to grpc message header to distribute bytes more equally across frames.
// Copy on the stack to avoid generating garbage
var localBuf [http2MaxFrameLen]byte
copy(localBuf[:hSize], dataItem.h)
copy(localBuf[hSize:], dataItem.d[:dSize])
buf = localBuf[:hSize+dSize]
}
} else {
buf = dataItem.d
}
dSize := min(maxSize-hSize, dataItem.reader.Remaining())
remainingBytes := len(dataItem.h) + dataItem.reader.Remaining() - hSize - dSize
size := hSize + dSize
var buf *[]byte
if hSize != 0 && dSize == 0 {
buf = &dataItem.h
} else {
// Note: this is only necessary because the http2.Framer does not support
// partially writing a frame, so the sequence must be materialized into a buffer.
// TODO: Revisit once https://github.com/golang/go/issues/66655 is addressed.
pool := l.bufferPool
if pool == nil {
// Note that this is only supposed to be nil in tests. Otherwise, stream is
// always initialized with a BufferPool.
pool = mem.DefaultBufferPool()
}
buf = pool.Get(size)
defer pool.Put(buf)
copy((*buf)[:hSize], dataItem.h)
_, _ = dataItem.reader.Read((*buf)[hSize:])
}
// Now that outgoing flow controls are checked we can replenish str's write quota
str.wq.replenish(size)
var endStream bool
// If this is the last data message on this stream and all of it can be written in this iteration.
if dataItem.endStream && len(dataItem.h)+len(dataItem.d) <= size {
if dataItem.endStream && remainingBytes == 0 {
endStream = true
}
if dataItem.onEachWrite != nil {
dataItem.onEachWrite()
}
if err := l.framer.fr.WriteData(dataItem.streamID, endStream, buf[:size]); err != nil {
if err := l.framer.fr.WriteData(dataItem.streamID, endStream, (*buf)[:size]); err != nil {
return false, err
}
str.bytesOutStanding += size
l.sendQuota -= uint32(size)
dataItem.h = dataItem.h[hSize:]
dataItem.d = dataItem.d[dSize:]
if len(dataItem.h) == 0 && len(dataItem.d) == 0 { // All the data from that message was written out.
if remainingBytes == 0 { // All the data from that message was written out.
_ = dataItem.reader.Close()
str.itl.dequeue()
}
if str.itl.isEmpty() {

24
internal/transport/grpchttp2/http2bridge.go
View File

@ -85,10 +85,10 @@ func (fr *FramerBridge) ReadFrame() (Frame, error) {
switch f := f.(type) {
case *http2.DataFrame:
buf := fr.pool.Get(int(hdr.Size))
copy(buf, f.Data())
copy(*buf, f.Data())
return &DataFrame{
hdr: hdr,
Data: buf,
Data: *buf,
free: func() { fr.pool.Put(buf) },
}, nil
case *http2.RSTStreamFrame:
@ -111,21 +111,21 @@ func (fr *FramerBridge) ReadFrame() (Frame, error) {
}, nil
case *http2.PingFrame:
buf := fr.pool.Get(int(hdr.Size))
copy(buf, f.Data[:])
copy(*buf, f.Data[:])
return &PingFrame{
hdr: hdr,
Data: buf,
Data: *buf,
free: func() { fr.pool.Put(buf) },
}, nil
case *http2.GoAwayFrame:
// Size of the frame minus the code and lastStreamID
buf := fr.pool.Get(int(hdr.Size) - 8)
copy(buf, f.DebugData())
copy(*buf, f.DebugData())
return &GoAwayFrame{
hdr: hdr,
LastStreamID: f.LastStreamID,
Code: ErrCode(f.ErrCode),
DebugData: buf,
DebugData: *buf,
free: func() { fr.pool.Put(buf) },
}, nil
case *http2.WindowUpdateFrame:
@ -141,10 +141,10 @@ func (fr *FramerBridge) ReadFrame() (Frame, error) {
default:
buf := fr.pool.Get(int(hdr.Size))
uf := f.(*http2.UnknownFrame)
copy(buf, uf.Payload())
copy(*buf, uf.Payload())
return &UnknownFrame{
hdr: hdr,
Payload: buf,
Payload: *buf,
free: func() { fr.pool.Put(buf) },
}, nil
}
@ -156,19 +156,19 @@ func (fr *FramerBridge) WriteData(streamID uint32, endStream bool, data ...[]byt
return fr.framer.WriteData(streamID, endStream, data[0])
}
var buf []byte
tl := 0
for _, s := range data {
tl += len(s)
}
buf = fr.pool.Get(tl)[:0]
buf := fr.pool.Get(tl)
*buf = (*buf)[:0]
defer fr.pool.Put(buf)
for _, s := range data {
buf = append(buf, s...)
*buf = append(*buf, s...)
}
return fr.framer.WriteData(streamID, endStream, buf)
return fr.framer.WriteData(streamID, endStream, *buf)
}
// WriteHeaders writes a Headers Frame into the underlying writer.

43
internal/transport/handler_server.go
View File

@ -24,7 +24,6 @@
package transport
import (
"bytes"
"context"
"errors"
"fmt"
@ -40,6 +39,7 @@ import (
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/internal/grpcutil"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/stats"
@ -50,7 +50,7 @@ import (
// NewServerHandlerTransport returns a ServerTransport handling gRPC from
// inside an http.Handler, or writes an HTTP error to w and returns an error.
// It requires that the http Server supports HTTP/2.
func NewServerHandlerTransport(w http.ResponseWriter, r *http.Request, stats []stats.Handler) (ServerTransport, error) {
func NewServerHandlerTransport(w http.ResponseWriter, r *http.Request, stats []stats.Handler, bufferPool mem.BufferPool) (ServerTransport, error) {
if r.Method != http.MethodPost {
w.Header().Set("Allow", http.MethodPost)
msg := fmt.Sprintf("invalid gRPC request method %q", r.Method)
@ -98,6 +98,7 @@ func NewServerHandlerTransport(w http.ResponseWriter, r *http.Request, stats []s
contentType: contentType,
contentSubtype: contentSubtype,
stats: stats,
bufferPool: bufferPool,
}
st.logger = prefixLoggerForServerHandlerTransport(st)
@ -171,6 +172,8 @@ type serverHandlerTransport struct {
stats []stats.Handler
logger *grpclog.PrefixLogger
bufferPool mem.BufferPool
}
func (ht *serverHandlerTransport) Close(err error) {
@ -330,16 +333,28 @@ func (ht *serverHandlerTransport) writeCustomHeaders(s *Stream) {
s.hdrMu.Unlock()
}
func (ht *serverHandlerTransport) Write(s *Stream, hdr []byte, data []byte, opts *Options) error {
func (ht *serverHandlerTransport) Write(s *Stream, hdr []byte, data mem.BufferSlice, opts *Options) error {
// Always take a reference because otherwise there is no guarantee the data will
// be available after this function returns. This is what callers to Write
// expect.
data.Ref()
headersWritten := s.updateHeaderSent()
return ht.do(func() {
err := ht.do(func() {
defer data.Free()
if !headersWritten {
ht.writePendingHeaders(s)
}
ht.rw.Write(hdr)
ht.rw.Write(data)
for _, b := range data {
_, _ = ht.rw.Write(b.ReadOnlyData())
}
ht.rw.(http.Flusher).Flush()
})
if err != nil {
data.Free()
return err
}
return nil
}
func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
@ -406,7 +421,7 @@ func (ht *serverHandlerTransport) HandleStreams(ctx context.Context, startStream
headerWireLength: 0, // won't have access to header wire length until golang/go#18997.
}
s.trReader = &transportReader{
reader: &recvBufferReader{ctx: s.ctx, ctxDone: s.ctx.Done(), recv: s.buf, freeBuffer: func(*bytes.Buffer) {}},
reader: &recvBufferReader{ctx: s.ctx, ctxDone: s.ctx.Done(), recv: s.buf},
windowHandler: func(int) {},
}
@ -415,21 +430,19 @@ func (ht *serverHandlerTransport) HandleStreams(ctx context.Context, startStream
go func() {
defer close(readerDone)
// TODO: minimize garbage, optimize recvBuffer code/ownership
const readSize = 8196
for buf := make([]byte, readSize); ; {
n, err := req.Body.Read(buf)
for {
buf := ht.bufferPool.Get(http2MaxFrameLen)
n, err := req.Body.Read(*buf)
if n > 0 {
s.buf.put(recvMsg{buffer: bytes.NewBuffer(buf[:n:n])})
buf = buf[n:]
*buf = (*buf)[:n]
s.buf.put(recvMsg{buffer: mem.NewBuffer(buf, ht.bufferPool)})
} else {
ht.bufferPool.Put(buf)
}
if err != nil {
s.buf.put(recvMsg{err: mapRecvMsgError(err)})
return
}
if len(buf) == 0 {
buf = make([]byte, readSize)
}
}
}()

9
internal/transport/handler_server_test.go
View File

@ -33,6 +33,7 @@ import (
epb "google.golang.org/genproto/googleapis/rpc/errdetails"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
@ -203,7 +204,7 @@ func (s) TestHandlerTransport_NewServerHandlerTransport(t *testing.T) {
if tt.modrw != nil {
rw = tt.modrw(rw)
}
got, gotErr := NewServerHandlerTransport(rw, tt.req, nil)
got, gotErr := NewServerHandlerTransport(rw, tt.req, nil, mem.DefaultBufferPool())
if (gotErr != nil) != (tt.wantErr != "") || (gotErr != nil && gotErr.Error() != tt.wantErr) {
t.Errorf("%s: error = %q; want %q", tt.name, gotErr.Error(), tt.wantErr)
continue
@ -259,7 +260,7 @@ func newHandleStreamTest(t *testing.T) *handleStreamTest {
Body: bodyr,
}
rw := newTestHandlerResponseWriter().(testHandlerResponseWriter)
ht, err := NewServerHandlerTransport(rw, req, nil)
ht, err := NewServerHandlerTransport(rw, req, nil, mem.DefaultBufferPool())
if err != nil {
t.Fatal(err)
}
@ -374,7 +375,7 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
Body: bodyr,
}
rw := newTestHandlerResponseWriter().(testHandlerResponseWriter)
ht, err := NewServerHandlerTransport(rw, req, nil)
ht, err := NewServerHandlerTransport(rw, req, nil, mem.DefaultBufferPool())
if err != nil {
t.Fatal(err)
}
@ -439,7 +440,7 @@ func (s) TestHandlerTransport_HandleStreams_WriteStatusWrite(t *testing.T) {
st.bodyw.Close() // no body
st.ht.WriteStatus(s, status.New(codes.OK, ""))
st.ht.Write(s, []byte("hdr"), []byte("data"), &Options{})
st.ht.Write(s, []byte("hdr"), newBufferSlice([]byte("data")), &Options{})
})
}

38
internal/transport/http2_client.go
View File

@ -47,6 +47,7 @@ import (
isyscall "google.golang.org/grpc/internal/syscall"
"google.golang.org/grpc/internal/transport/networktype"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
@ -146,7 +147,7 @@ type http2Client struct {
onClose func(GoAwayReason)
bufferPool *bufferPool
bufferPool mem.BufferPool
connectionID uint64
logger *grpclog.PrefixLogger
@ -348,7 +349,7 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
streamQuota: defaultMaxStreamsClient,
streamsQuotaAvailable: make(chan struct{}, 1),
keepaliveEnabled: keepaliveEnabled,
bufferPool: newBufferPool(),
bufferPool: opts.BufferPool,
onClose: onClose,
}
var czSecurity credentials.ChannelzSecurityValue
@ -465,7 +466,7 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
return nil, err
}
go func() {
t.loopy = newLoopyWriter(clientSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger, t.outgoingGoAwayHandler)
t.loopy = newLoopyWriter(clientSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger, t.outgoingGoAwayHandler, t.bufferPool)
if err := t.loopy.run(); !isIOError(err) {
// Immediately close the connection, as the loopy writer returns
// when there are no more active streams and we were draining (the
@ -506,7 +507,6 @@ func (t *http2Client) newStream(ctx context.Context, callHdr *CallHdr) *Stream {
closeStream: func(err error) {
t.CloseStream(s, err)
},
freeBuffer: t.bufferPool.put,
},
windowHandler: func(n int) {
t.updateWindow(s, uint32(n))
@ -1078,27 +1078,36 @@ func (t *http2Client) GracefulClose() {
// Write formats the data into HTTP2 data frame(s) and sends it out. The caller
// should proceed only if Write returns nil.
func (t *http2Client) Write(s *Stream, hdr []byte, data []byte, opts *Options) error {
func (t *http2Client) Write(s *Stream, hdr []byte, data mem.BufferSlice, opts *Options) error {
reader := data.Reader()
if opts.Last {
// If it's the last message, update stream state.
if !s.compareAndSwapState(streamActive, streamWriteDone) {
_ = reader.Close()
return errStreamDone
}
} else if s.getState() != streamActive {
_ = reader.Close()
return errStreamDone
}
df := &dataFrame{
streamID: s.id,
endStream: opts.Last,
h: hdr,
d: data,
reader: reader,
}
if hdr != nil || data != nil { // If it's not an empty data frame, check quota.
if err := s.wq.get(int32(len(hdr) + len(data))); err != nil {
if hdr != nil || df.reader.Remaining() != 0 { // If it's not an empty data frame, check quota.
if err := s.wq.get(int32(len(hdr) + df.reader.Remaining())); err != nil {
_ = reader.Close()
return err
}
}
return t.controlBuf.put(df)
if err := t.controlBuf.put(df); err != nil {
_ = reader.Close()
return err
}
return nil
}
func (t *http2Client) getStream(f http2.Frame) *Stream {
@ -1203,10 +1212,13 @@ func (t *http2Client) handleData(f *http2.DataFrame) {
// guarantee f.Data() is consumed before the arrival of next frame.
// Can this copy be eliminated?
if len(f.Data()) > 0 {
buffer := t.bufferPool.get()
buffer.Reset()
buffer.Write(f.Data())
s.write(recvMsg{buffer: buffer})
pool := t.bufferPool
if pool == nil {
// Note that this is only supposed to be nil in tests. Otherwise, stream is
// always initialized with a BufferPool.
pool = mem.DefaultBufferPool()
}
s.write(recvMsg{buffer: mem.Copy(f.Data(), pool)})
}
}
// The server has closed the stream without sending trailers. Record that

43
internal/transport/http2_server.go
View File

@ -39,6 +39,7 @@ import (
"google.golang.org/grpc/internal/grpcutil"
"google.golang.org/grpc/internal/pretty"
"google.golang.org/grpc/internal/syscall"
"google.golang.org/grpc/mem"
"google.golang.org/protobuf/proto"
"google.golang.org/grpc/codes"
@ -119,7 +120,7 @@ type http2Server struct {
// Fields below are for channelz metric collection.
channelz *channelz.Socket
bufferPool *bufferPool
bufferPool mem.BufferPool
connectionID uint64
@ -261,7 +262,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
idle: time.Now(),
kep: kep,
initialWindowSize: iwz,
bufferPool: newBufferPool(),
bufferPool: config.BufferPool,
}
var czSecurity credentials.ChannelzSecurityValue
if au, ok := authInfo.(credentials.ChannelzSecurityInfo); ok {
@ -330,7 +331,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
t.handleSettings(sf)
go func() {
t.loopy = newLoopyWriter(serverSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger, t.outgoingGoAwayHandler)
t.loopy = newLoopyWriter(serverSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger, t.outgoingGoAwayHandler, t.bufferPool)
err := t.loopy.run()
close(t.loopyWriterDone)
if !isIOError(err) {
@ -613,10 +614,9 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
s.wq = newWriteQuota(defaultWriteQuota, s.ctxDone)
s.trReader = &transportReader{
reader: &recvBufferReader{
ctx: s.ctx,
ctxDone: s.ctxDone,
recv: s.buf,
freeBuffer: t.bufferPool.put,
ctx: s.ctx,
ctxDone: s.ctxDone,
recv: s.buf,
},
windowHandler: func(n int) {
t.updateWindow(s, uint32(n))
@ -813,10 +813,13 @@ func (t *http2Server) handleData(f *http2.DataFrame) {
// guarantee f.Data() is consumed before the arrival of next frame.
// Can this copy be eliminated?
if len(f.Data()) > 0 {
buffer := t.bufferPool.get()
buffer.Reset()
buffer.Write(f.Data())
s.write(recvMsg{buffer: buffer})
pool := t.bufferPool
if pool == nil {
// Note that this is only supposed to be nil in tests. Otherwise, stream is
// always initialized with a BufferPool.
pool = mem.DefaultBufferPool()
}
s.write(recvMsg{buffer: mem.Copy(f.Data(), pool)})
}
}
if f.StreamEnded() {
@ -1114,27 +1117,37 @@ func (t *http2Server) WriteStatus(s *Stream, st *status.Status) error {
// Write converts the data into HTTP2 data frame and sends it out. Non-nil error
// is returns if it fails (e.g., framing error, transport error).
func (t *http2Server) Write(s *Stream, hdr []byte, data []byte, opts *Options) error {
func (t *http2Server) Write(s *Stream, hdr []byte, data mem.BufferSlice, opts *Options) error {
reader := data.Reader()
if !s.isHeaderSent() { // Headers haven't been written yet.
if err := t.WriteHeader(s, nil); err != nil {
_ = reader.Close()
return err
}
} else {
// Writing headers checks for this condition.
if s.getState() == streamDone {
_ = reader.Close()
return t.streamContextErr(s)
}
}
df := &dataFrame{
streamID: s.id,
h: hdr,
d: data,
reader: reader,
onEachWrite: t.setResetPingStrikes,
}
if err := s.wq.get(int32(len(hdr) + len(data))); err != nil {
if err := s.wq.get(int32(len(hdr) + df.reader.Remaining())); err != nil {
_ = reader.Close()
return t.streamContextErr(s)
}
return t.controlBuf.put(df)
if err := t.controlBuf.put(df); err != nil {
_ = reader.Close()
return err
}
return nil
}
// keepalive running in a separate goroutine does the following:

236
internal/transport/transport.go
View File

@ -22,7 +22,6 @@
package transport
import (
"bytes"
"context"
"errors"
"fmt"
@ -37,6 +36,7 @@ import (
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/internal/channelz"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
@ -47,32 +47,10 @@ import (
const logLevel = 2
type bufferPool struct {
pool sync.Pool
}
func newBufferPool() *bufferPool {
return &bufferPool{
pool: sync.Pool{
New: func() any {
return new(bytes.Buffer)
},
},
}
}
func (p *bufferPool) get() *bytes.Buffer {
return p.pool.Get().(*bytes.Buffer)
}
func (p *bufferPool) put(b *bytes.Buffer) {
p.pool.Put(b)
}
// recvMsg represents the received msg from the transport. All transport
// protocol specific info has been removed.
type recvMsg struct {
buffer *bytes.Buffer
buffer mem.Buffer
// nil: received some data
// io.EOF: stream is completed. data is nil.
// other non-nil error: transport failure. data is nil.
@ -102,6 +80,9 @@ func newRecvBuffer() *recvBuffer {
func (b *recvBuffer) put(r recvMsg) {
b.mu.Lock()
if b.err != nil {
// drop the buffer on the floor. Since b.err is not nil, any subsequent reads
// will always return an error, making this buffer inaccessible.
r.buffer.Free()
b.mu.Unlock()
// An error had occurred earlier, don't accept more
// data or errors.
@ -148,45 +129,70 @@ type recvBufferReader struct {
ctx context.Context
ctxDone <-chan struct{} // cache of ctx.Done() (for performance).
recv *recvBuffer
last *bytes.Buffer // Stores the remaining data in the previous calls.
last mem.Buffer // Stores the remaining data in the previous calls.
err error
freeBuffer func(*bytes.Buffer)
}
// Read reads the next len(p) bytes from last. If last is drained, it tries to
// read additional data from recv. It blocks if there no additional data available
// in recv. If Read returns any non-nil error, it will continue to return that error.
func (r *recvBufferReader) Read(p []byte) (n int, err error) {
func (r *recvBufferReader) ReadHeader(header []byte) (n int, err error) {
if r.err != nil {
return 0, r.err
}
if r.last != nil {
// Read remaining data left in last call.
copied, _ := r.last.Read(p)
if r.last.Len() == 0 {
r.freeBuffer(r.last)
r.last = nil
}
return copied, nil
n, r.last = mem.ReadUnsafe(header, r.last)
return n, nil
}
if r.closeStream != nil {
n, r.err = r.readClient(p)
n, r.err = r.readHeaderClient(header)
} else {
n, r.err = r.read(p)
n, r.err = r.readHeader(header)
}
return n, r.err
}
func (r *recvBufferReader) read(p []byte) (n int, err error) {
// Read reads the next n bytes from last. If last is drained, it tries to read
// additional data from recv. It blocks if there no additional data available in
// recv. If Read returns any non-nil error, it will continue to return that
// error.
func (r *recvBufferReader) Read(n int) (buf mem.Buffer, err error) {
if r.err != nil {
return nil, r.err
}
if r.last != nil {
buf = r.last
if r.last.Len() > n {
buf, r.last = mem.SplitUnsafe(buf, n)
} else {
r.last = nil
}
return buf, nil
}
if r.closeStream != nil {
buf, r.err = r.readClient(n)
} else {
buf, r.err = r.read(n)
}
return buf, r.err
}
func (r *recvBufferReader) readHeader(header []byte) (n int, err error) {
select {
case <-r.ctxDone:
return 0, ContextErr(r.ctx.Err())
case m := <-r.recv.get():
return r.readAdditional(m, p)
return r.readHeaderAdditional(m, header)
}
}
func (r *recvBufferReader) readClient(p []byte) (n int, err error) {
func (r *recvBufferReader) read(n int) (buf mem.Buffer, err error) {
select {
case <-r.ctxDone:
return nil, ContextErr(r.ctx.Err())
case m := <-r.recv.get():
return r.readAdditional(m, n)
}
}
func (r *recvBufferReader) readHeaderClient(header []byte) (n int, err error) {
// If the context is canceled, then closes the stream with nil metadata.
// closeStream writes its error parameter to r.recv as a recvMsg.
// r.readAdditional acts on that message and returns the necessary error.
@ -207,25 +213,67 @@ func (r *recvBufferReader) readClient(p []byte) (n int, err error) {
// faster.
r.closeStream(ContextErr(r.ctx.Err()))
m := <-r.recv.get()
return r.readAdditional(m, p)
return r.readHeaderAdditional(m, header)
case m := <-r.recv.get():
return r.readAdditional(m, p)
return r.readHeaderAdditional(m, header)
}
}
func (r *recvBufferReader) readAdditional(m recvMsg, p []byte) (n int, err error) {
func (r *recvBufferReader) readClient(n int) (buf mem.Buffer, err error) {
// If the context is canceled, then closes the stream with nil metadata.
// closeStream writes its error parameter to r.recv as a recvMsg.
// r.readAdditional acts on that message and returns the necessary error.
select {
case <-r.ctxDone:
// Note that this adds the ctx error to the end of recv buffer, and
// reads from the head. This will delay the error until recv buffer is
// empty, thus will delay ctx cancellation in Recv().
//
// It's done this way to fix a race between ctx cancel and trailer. The
// race was, stream.Recv() may return ctx error if ctxDone wins the
// race, but stream.Trailer() may return a non-nil md because the stream
// was not marked as done when trailer is received. This closeStream
// call will mark stream as done, thus fix the race.
//
// TODO: delaying ctx error seems like a unnecessary side effect. What
// we really want is to mark the stream as done, and return ctx error
// faster.
r.closeStream(ContextErr(r.ctx.Err()))
m := <-r.recv.get()
return r.readAdditional(m, n)
case m := <-r.recv.get():
return r.readAdditional(m, n)
}
}
func (r *recvBufferReader) readHeaderAdditional(m recvMsg, header []byte) (n int, err error) {
r.recv.load()
if m.err != nil {
if m.buffer != nil {
m.buffer.Free()
}
return 0, m.err
}
copied, _ := m.buffer.Read(p)
if m.buffer.Len() == 0 {
r.freeBuffer(m.buffer)
r.last = nil
} else {
r.last = m.buffer
n, r.last = mem.ReadUnsafe(header, m.buffer)
return n, nil
}
func (r *recvBufferReader) readAdditional(m recvMsg, n int) (b mem.Buffer, err error) {
r.recv.load()
if m.err != nil {
if m.buffer != nil {
m.buffer.Free()
}
return nil, m.err
}
return copied, nil
if m.buffer.Len() > n {
m.buffer, r.last = mem.SplitUnsafe(m.buffer, n)
}
return m.buffer, nil
}
type streamState uint32
@ -251,7 +299,7 @@ type Stream struct {
recvCompress string
sendCompress string
buf *recvBuffer
trReader io.Reader
trReader *transportReader
fc *inFlow
wq *writeQuota
@ -499,14 +547,55 @@ func (s *Stream) write(m recvMsg) {
s.buf.put(m)
}
// Read reads all p bytes from the wire for this stream.
func (s *Stream) Read(p []byte) (n int, err error) {
func (s *Stream) ReadHeader(header []byte) (err error) {
// Don't request a read if there was an error earlier
if er := s.trReader.(*transportReader).er; er != nil {
return 0, er
if er := s.trReader.er; er != nil {
return er
}
s.requestRead(len(p))
return io.ReadFull(s.trReader, p)
s.requestRead(len(header))
for len(header) != 0 {
n, err := s.trReader.ReadHeader(header)
header = header[n:]
if len(header) == 0 {
err = nil
}
if err != nil {
if n > 0 && err == io.EOF {
err = io.ErrUnexpectedEOF
}
return err
}
}
return nil
}
// Read reads n bytes from the wire for this stream.
func (s *Stream) Read(n int) (data mem.BufferSlice, err error) {
// Don't request a read if there was an error earlier
if er := s.trReader.er; er != nil {
return nil, er
}
s.requestRead(n)
for n != 0 {
buf, err := s.trReader.Read(n)
var bufLen int
if buf != nil {
bufLen = buf.Len()
}
n -= bufLen
if n == 0 {
err = nil
}
if err != nil {
if bufLen > 0 && err == io.EOF {
err = io.ErrUnexpectedEOF
}
data.Free()
return nil, err
}
data = append(data, buf)
}
return data, nil
}
// transportReader reads all the data available for this Stream from the transport and
@ -514,21 +603,31 @@ func (s *Stream) Read(p []byte) (n int, err error) {
// The error is io.EOF when the stream is done or another non-nil error if
// the stream broke.
type transportReader struct {
reader io.Reader
reader *recvBufferReader
// The handler to control the window update procedure for both this
// particular stream and the associated transport.
windowHandler func(int)
er error
}
func (t *transportReader) Read(p []byte) (n int, err error) {
n, err = t.reader.Read(p)
func (t *transportReader) ReadHeader(header []byte) (int, error) {
n, err := t.reader.ReadHeader(header)
if err != nil {
t.er = err
return
return 0, err
}
t.windowHandler(n)
return
t.windowHandler(len(header))
return n, nil
}
func (t *transportReader) Read(n int) (mem.Buffer, error) {
buf, err := t.reader.Read(n)
if err != nil {
t.er = err
return buf, err
}
t.windowHandler(buf.Len())
return buf, nil
}
// BytesReceived indicates whether any bytes have been received on this stream.
@ -574,6 +673,7 @@ type ServerConfig struct {
ChannelzParent *channelz.Server
MaxHeaderListSize *uint32
HeaderTableSize *uint32
BufferPool mem.BufferPool
}
// ConnectOptions covers all relevant options for communicating with the server.
@ -612,6 +712,8 @@ type ConnectOptions struct {
MaxHeaderListSize *uint32
// UseProxy specifies if a proxy should be used.
UseProxy bool
// The mem.BufferPool to use when reading/writing to the wire.
BufferPool mem.BufferPool
}
// NewClientTransport establishes the transport with the required ConnectOptions
@ -673,7 +775,7 @@ type ClientTransport interface {
// Write sends the data for the given stream. A nil stream indicates
// the write is to be performed on the transport as a whole.
Write(s *Stream, hdr []byte, data []byte, opts *Options) error
Write(s *Stream, hdr []byte, data mem.BufferSlice, opts *Options) error
// NewStream creates a Stream for an RPC.
NewStream(ctx context.Context, callHdr *CallHdr) (*Stream, error)
@ -725,7 +827,7 @@ type ServerTransport interface {
// Write sends the data for the given stream.
// Write may not be called on all streams.
Write(s *Stream, hdr []byte, data []byte, opts *Options) error
Write(s *Stream, hdr []byte, data mem.BufferSlice, opts *Options) error
// WriteStatus sends the status of a stream to the client. WriteStatus is
// the final call made on a stream and always occurs.

161
internal/transport/transport_test.go
View File

@ -46,6 +46,7 @@ import (
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/leakcheck"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/status"
@ -74,6 +75,29 @@ func init() {
expectedResponseLarge[len(expectedResponseLarge)-1] = 'c'
}
func newBufferSlice(b []byte) mem.BufferSlice {
return mem.BufferSlice{mem.NewBuffer(&b, nil)}
}
func (s *Stream) readTo(p []byte) (int, error) {
data, err := s.Read(len(p))
defer data.Free()
if err != nil {
return 0, err
}
if data.Len() != len(p) {
if err == nil {
err = io.ErrUnexpectedEOF
}
return 0, err
}
data.CopyTo(p)
return len(p), nil
}
type testStreamHandler struct {
t *http2Server
notify chan struct{}
@ -114,7 +138,7 @@ func (h *testStreamHandler) handleStream(t *testing.T, s *Stream) {
resp = expectedResponseLarge
}
p := make([]byte, len(req))
_, err := s.Read(p)
_, err := s.readTo(p)
if err != nil {
return
}
@ -124,7 +148,7 @@ func (h *testStreamHandler) handleStream(t *testing.T, s *Stream) {
return
}
// send a response back to the client.
h.t.Write(s, nil, resp, &Options{})
h.t.Write(s, nil, newBufferSlice(resp), &Options{})
// send the trailer to end the stream.
h.t.WriteStatus(s, status.New(codes.OK, ""))
}
@ -132,7 +156,7 @@ func (h *testStreamHandler) handleStream(t *testing.T, s *Stream) {
func (h *testStreamHandler) handleStreamPingPong(t *testing.T, s *Stream) {
header := make([]byte, 5)
for {
if _, err := s.Read(header); err != nil {
if _, err := s.readTo(header); err != nil {
if err == io.EOF {
h.t.WriteStatus(s, status.New(codes.OK, ""))
return
@ -143,7 +167,7 @@ func (h *testStreamHandler) handleStreamPingPong(t *testing.T, s *Stream) {
}
sz := binary.BigEndian.Uint32(header[1:])
msg := make([]byte, int(sz))
if _, err := s.Read(msg); err != nil {
if _, err := s.readTo(msg); err != nil {
t.Errorf("Error on server while reading message: %v", err)
h.t.WriteStatus(s, status.New(codes.Internal, "panic"))
return
@ -152,7 +176,7 @@ func (h *testStreamHandler) handleStreamPingPong(t *testing.T, s *Stream) {
buf[0] = byte(0)
binary.BigEndian.PutUint32(buf[1:], uint32(sz))
copy(buf[5:], msg)
h.t.Write(s, nil, buf, &Options{})
h.t.Write(s, nil, newBufferSlice(buf), &Options{})
}
}
@ -178,10 +202,11 @@ func (h *testStreamHandler) handleStreamMisbehave(t *testing.T, s *Stream) {
p = make([]byte, n+1)
}
}
data := newBufferSlice(p)
conn.controlBuf.put(&dataFrame{
streamID: s.id,
h: nil,
d: p,
reader: data.Reader(),
onEachWrite: func() {},
})
sent += len(p)
@ -191,6 +216,8 @@ func (h *testStreamHandler) handleStreamMisbehave(t *testing.T, s *Stream) {
func (h *testStreamHandler) handleStreamEncodingRequiredStatus(s *Stream) {
// raw newline is not accepted by http2 framer so it must be encoded.
h.t.WriteStatus(s, encodingTestStatus)
// Drain any remaining buffers from the stream since it was closed early.
s.Read(math.MaxInt)
}
func (h *testStreamHandler) handleStreamInvalidHeaderField(s *Stream) {
@ -260,7 +287,7 @@ func (h *testStreamHandler) handleStreamDelayRead(t *testing.T, s *Stream) {
t.Errorf("Server timed-out.")
return
}
_, err := s.Read(p)
_, err := s.readTo(p)
if err != nil {
t.Errorf("s.Read(_) = _, %v, want _, <nil>", err)
return
@ -273,14 +300,14 @@ func (h *testStreamHandler) handleStreamDelayRead(t *testing.T, s *Stream) {
// This write will cause server to run out of stream level,
// flow control and the other side won't send a window update
// until that happens.
if err := h.t.Write(s, nil, resp, &Options{}); err != nil {
if err := h.t.Write(s, nil, newBufferSlice(resp), &Options{}); err != nil {
t.Errorf("server Write got %v, want <nil>", err)
return
}
// Read one more time to ensure that everything remains fine and
// that the goroutine, that we launched earlier to signal client
// to read, gets enough time to process.
_, err = s.Read(p)
_, err = s.readTo(p)
if err != nil {
t.Errorf("s.Read(_) = _, %v, want _, nil", err)
return
@ -502,7 +529,7 @@ func (s) TestInflightStreamClosing(t *testing.T) {
serr := status.Error(codes.Internal, "client connection is closing")
go func() {
defer close(donec)
if _, err := stream.Read(make([]byte, defaultWindowSize)); err != serr {
if _, err := stream.readTo(make([]byte, defaultWindowSize)); err != serr {
t.Errorf("unexpected Stream error %v, expected %v", err, serr)
}
}()
@ -592,15 +619,15 @@ func (s) TestClientSendAndReceive(t *testing.T) {
t.Fatalf("wrong stream id: %d", s2.id)
}
opts := Options{Last: true}
if err := ct.Write(s1, nil, expectedRequest, &opts); err != nil && err != io.EOF {
if err := ct.Write(s1, nil, newBufferSlice(expectedRequest), &opts); err != nil && err != io.EOF {
t.Fatalf("failed to send data: %v", err)
}
p := make([]byte, len(expectedResponse))
_, recvErr := s1.Read(p)
_, recvErr := s1.readTo(p)
if recvErr != nil || !bytes.Equal(p, expectedResponse) {
t.Fatalf("Error: %v, want <nil>; Result: %v, want %v", recvErr, p, expectedResponse)
}
_, recvErr = s1.Read(p)
_, recvErr = s1.readTo(p)
if recvErr != io.EOF {
t.Fatalf("Error: %v; want <EOF>", recvErr)
}
@ -629,16 +656,16 @@ func performOneRPC(ct ClientTransport) {
return
}
opts := Options{Last: true}
if err := ct.Write(s, []byte{}, expectedRequest, &opts); err == nil || err == io.EOF {
if err := ct.Write(s, []byte{}, newBufferSlice(expectedRequest), &opts); err == nil || err == io.EOF {
time.Sleep(5 * time.Millisecond)
// The following s.Recv()'s could error out because the
// underlying transport is gone.
//
// Read response
p := make([]byte, len(expectedResponse))
s.Read(p)
s.readTo(p)
// Read io.EOF
s.Read(p)
s.readTo(p)
}
}
@ -674,14 +701,14 @@ func (s) TestLargeMessage(t *testing.T) {
if err != nil {
t.Errorf("%v.NewStream(_, _) = _, %v, want _, <nil>", ct, err)
}
if err := ct.Write(s, []byte{}, expectedRequestLarge, &Options{Last: true}); err != nil && err != io.EOF {
if err := ct.Write(s, []byte{}, newBufferSlice(expectedRequestLarge), &Options{Last: true}); err != nil && err != io.EOF {
t.Errorf("%v.Write(_, _, _) = %v, want <nil>", ct, err)
}
p := make([]byte, len(expectedResponseLarge))
if _, err := s.Read(p); err != nil || !bytes.Equal(p, expectedResponseLarge) {
if _, err := s.readTo(p); err != nil || !bytes.Equal(p, expectedResponseLarge) {
t.Errorf("s.Read(%v) = _, %v, want %v, <nil>", err, p, expectedResponse)
}
if _, err = s.Read(p); err != io.EOF {
if _, err = s.readTo(p); err != io.EOF {
t.Errorf("Failed to complete the stream %v; want <EOF>", err)
}
}()
@ -765,7 +792,7 @@ func (s) TestLargeMessageWithDelayRead(t *testing.T) {
// This write will cause client to run out of stream level,
// flow control and the other side won't send a window update
// until that happens.
if err := ct.Write(s, []byte{}, expectedRequestLarge, &Options{}); err != nil {
if err := ct.Write(s, []byte{}, newBufferSlice(expectedRequestLarge), &Options{}); err != nil {
t.Fatalf("write(_, _, _) = %v, want <nil>", err)
}
p := make([]byte, len(expectedResponseLarge))
@ -777,13 +804,13 @@ func (s) TestLargeMessageWithDelayRead(t *testing.T) {
case <-ctx.Done():
t.Fatalf("Client timed out")
}
if _, err := s.Read(p); err != nil || !bytes.Equal(p, expectedResponseLarge) {
if _, err := s.readTo(p); err != nil || !bytes.Equal(p, expectedResponseLarge) {
t.Fatalf("s.Read(_) = _, %v, want _, <nil>", err)
}
if err := ct.Write(s, []byte{}, expectedRequestLarge, &Options{Last: true}); err != nil {
if err := ct.Write(s, []byte{}, newBufferSlice(expectedRequestLarge), &Options{Last: true}); err != nil {
t.Fatalf("Write(_, _, _) = %v, want <nil>", err)
}
if _, err = s.Read(p); err != io.EOF {
if _, err = s.readTo(p); err != io.EOF {
t.Fatalf("Failed to complete the stream %v; want <EOF>", err)
}
}
@ -792,6 +819,7 @@ func (s) TestLargeMessageWithDelayRead(t *testing.T) {
// proceed until they complete naturally, while not allowing creation of new
// streams during this window.
func (s) TestGracefulClose(t *testing.T) {
leakcheck.SetTrackingBufferPool(t)
server, ct, cancel := setUp(t, 0, pingpong)
defer cancel()
defer func() {
@ -800,7 +828,8 @@ func (s) TestGracefulClose(t *testing.T) {
server.lis.Close()
// Check for goroutine leaks (i.e. GracefulClose with an active stream
// doesn't eventually close the connection when that stream completes).
leakcheck.Check(t)
leakcheck.CheckGoroutines(t, 10*time.Second)
leakcheck.CheckTrackingBufferPool()
// Correctly clean up the server
server.stop()
}()
@ -818,15 +847,15 @@ func (s) TestGracefulClose(t *testing.T) {
outgoingHeader[0] = byte(0)
binary.BigEndian.PutUint32(outgoingHeader[1:], uint32(len(msg)))
incomingHeader := make([]byte, 5)
if err := ct.Write(s, outgoingHeader, msg, &Options{}); err != nil {
if err := ct.Write(s, outgoingHeader, newBufferSlice(msg), &Options{}); err != nil {
t.Fatalf("Error while writing: %v", err)
}
if _, err := s.Read(incomingHeader); err != nil {
if _, err := s.readTo(incomingHeader); err != nil {
t.Fatalf("Error while reading: %v", err)
}
sz := binary.BigEndian.Uint32(incomingHeader[1:])
recvMsg := make([]byte, int(sz))
if _, err := s.Read(recvMsg); err != nil {
if _, err := s.readTo(recvMsg); err != nil {
t.Fatalf("Error while reading: %v", err)
}
@ -851,7 +880,7 @@ func (s) TestGracefulClose(t *testing.T) {
// Confirm the existing stream still functions as expected.
ct.Write(s, nil, nil, &Options{Last: true})
if _, err := s.Read(incomingHeader); err != io.EOF {
if _, err := s.readTo(incomingHeader); err != io.EOF {
t.Fatalf("Client expected EOF from the server. Got: %v", err)
}
wg.Wait()
@ -879,13 +908,13 @@ func (s) TestLargeMessageSuspension(t *testing.T) {
}()
// Write should not be done successfully due to flow control.
msg := make([]byte, initialWindowSize*8)
ct.Write(s, nil, msg, &Options{})
err = ct.Write(s, nil, msg, &Options{Last: true})
ct.Write(s, nil, newBufferSlice(msg), &Options{})
err = ct.Write(s, nil, newBufferSlice(msg), &Options{Last: true})
if err != errStreamDone {
t.Fatalf("Write got %v, want io.EOF", err)
}
expectedErr := status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error())
if _, err := s.Read(make([]byte, 8)); err.Error() != expectedErr.Error() {
if _, err := s.readTo(make([]byte, 8)); err.Error() != expectedErr.Error() {
t.Fatalf("Read got %v of type %T, want %v", err, err, expectedErr)
}
ct.Close(fmt.Errorf("closed manually by test"))
@ -997,11 +1026,12 @@ func (s) TestServerContextCanceledOnClosedConnection(t *testing.T) {
if err != nil {
t.Fatalf("Failed to open stream: %v", err)
}
d := newBufferSlice(make([]byte, http2MaxFrameLen))
ct.controlBuf.put(&dataFrame{
streamID: s.id,
endStream: false,
h: nil,
d: make([]byte, http2MaxFrameLen),
reader: d.Reader(),
onEachWrite: func() {},
})
// Loop until the server side stream is created.
@ -1078,7 +1108,7 @@ func (s) TestClientConnDecoupledFromApplicationRead(t *testing.T) {
t.Fatalf("Didn't find stream corresponding to client cstream.id: %v on the server", cstream1.id)
}
// Exhaust client's connection window.
if err := st.Write(sstream1, []byte{}, make([]byte, defaultWindowSize), &Options{}); err != nil {
if err := st.Write(sstream1, []byte{}, newBufferSlice(make([]byte, defaultWindowSize)), &Options{}); err != nil {
t.Fatalf("Server failed to write data. Err: %v", err)
}
notifyChan = make(chan struct{})
@ -1103,17 +1133,17 @@ func (s) TestClientConnDecoupledFromApplicationRead(t *testing.T) {
t.Fatalf("Didn't find stream corresponding to client cstream.id: %v on the server", cstream2.id)
}
// Server should be able to send data on the new stream, even though the client hasn't read anything on the first stream.
if err := st.Write(sstream2, []byte{}, make([]byte, defaultWindowSize), &Options{}); err != nil {
if err := st.Write(sstream2, []byte{}, newBufferSlice(make([]byte, defaultWindowSize)), &Options{}); err != nil {
t.Fatalf("Server failed to write data. Err: %v", err)
}
// Client should be able to read data on second stream.
if _, err := cstream2.Read(make([]byte, defaultWindowSize)); err != nil {
if _, err := cstream2.readTo(make([]byte, defaultWindowSize)); err != nil {
t.Fatalf("_.Read(_) = _, %v, want _, <nil>", err)
}
// Client should be able to read data on first stream.
if _, err := cstream1.Read(make([]byte, defaultWindowSize)); err != nil {
if _, err := cstream1.readTo(make([]byte, defaultWindowSize)); err != nil {
t.Fatalf("_.Read(_) = _, %v, want _, <nil>", err)
}
}
@ -1149,7 +1179,7 @@ func (s) TestServerConnDecoupledFromApplicationRead(t *testing.T) {
t.Fatalf("Failed to create 1st stream. Err: %v", err)
}
// Exhaust server's connection window.
if err := client.Write(cstream1, nil, make([]byte, defaultWindowSize), &Options{Last: true}); err != nil {
if err := client.Write(cstream1, nil, newBufferSlice(make([]byte, defaultWindowSize)), &Options{Last: true}); err != nil {
t.Fatalf("Client failed to write data. Err: %v", err)
}
//Client should be able to create another stream and send data on it.
@ -1157,7 +1187,7 @@ func (s) TestServerConnDecoupledFromApplicationRead(t *testing.T) {
if err != nil {
t.Fatalf("Failed to create 2nd stream. Err: %v", err)
}
if err := client.Write(cstream2, nil, make([]byte, defaultWindowSize), &Options{}); err != nil {
if err := client.Write(cstream2, nil, newBufferSlice(make([]byte, defaultWindowSize)), &Options{}); err != nil {
t.Fatalf("Client failed to write data. Err: %v", err)
}
// Get the streams on server.
@ -1179,11 +1209,11 @@ func (s) TestServerConnDecoupledFromApplicationRead(t *testing.T) {
}
st.mu.Unlock()
// Reading from the stream on server should succeed.
if _, err := sstream1.Read(make([]byte, defaultWindowSize)); err != nil {
if _, err := sstream1.readTo(make([]byte, defaultWindowSize)); err != nil {
t.Fatalf("_.Read(_) = %v, want <nil>", err)
}
if _, err := sstream1.Read(make([]byte, 1)); err != io.EOF {
if _, err := sstream1.readTo(make([]byte, 1)); err != io.EOF {
t.Fatalf("_.Read(_) = %v, want io.EOF", err)
}
@ -1435,6 +1465,9 @@ func (s) TestClientWithMisbehavedServer(t *testing.T) {
t.Fatalf("Test timed-out.")
case <-success:
}
// Drain the remaining buffers in the stream by reading until an error is
// encountered.
str.Read(math.MaxInt)
}
var encodingTestStatus = status.New(codes.Internal, "\n")
@ -1453,11 +1486,11 @@ func (s) TestEncodingRequiredStatus(t *testing.T) {
return
}
opts := Options{Last: true}
if err := ct.Write(s, nil, expectedRequest, &opts); err != nil && err != errStreamDone {
if err := ct.Write(s, nil, newBufferSlice(expectedRequest), &opts); err != nil && err != errStreamDone {
t.Fatalf("Failed to write the request: %v", err)
}
p := make([]byte, http2MaxFrameLen)
if _, err := s.trReader.(*transportReader).Read(p); err != io.EOF {
if _, err := s.readTo(p); err != io.EOF {
t.Fatalf("Read got error %v, want %v", err, io.EOF)
}
if !testutils.StatusErrEqual(s.Status().Err(), encodingTestStatus.Err()) {
@ -1465,6 +1498,8 @@ func (s) TestEncodingRequiredStatus(t *testing.T) {
}
ct.Close(fmt.Errorf("closed manually by test"))
server.stop()
// Drain any remaining buffers from the stream since it was closed early.
s.Read(math.MaxInt)
}
func (s) TestInvalidHeaderField(t *testing.T) {
@ -1481,7 +1516,7 @@ func (s) TestInvalidHeaderField(t *testing.T) {
return
}
p := make([]byte, http2MaxFrameLen)
_, err = s.trReader.(*transportReader).Read(p)
_, err = s.readTo(p)
if se, ok := status.FromError(err); !ok || se.Code() != codes.Internal || !strings.Contains(err.Error(), expectedInvalidHeaderField) {
t.Fatalf("Read got error %v, want error with code %s and contains %q", err, codes.Internal, expectedInvalidHeaderField)
}
@ -1639,17 +1674,17 @@ func testFlowControlAccountCheck(t *testing.T, msgSize int, wc windowSizeConfig)
opts := Options{}
header := make([]byte, 5)
for i := 1; i <= 5; i++ {
if err := client.Write(stream, nil, buf, &opts); err != nil {
if err := client.Write(stream, nil, newBufferSlice(buf), &opts); err != nil {
t.Errorf("Error on client while writing message %v on stream %v: %v", i, stream.id, err)
return
}
if _, err := stream.Read(header); err != nil {
if _, err := stream.readTo(header); err != nil {
t.Errorf("Error on client while reading data frame header %v on stream %v: %v", i, stream.id, err)
return
}
sz := binary.BigEndian.Uint32(header[1:])
recvMsg := make([]byte, int(sz))
if _, err := stream.Read(recvMsg); err != nil {
if _, err := stream.readTo(recvMsg); err != nil {
t.Errorf("Error on client while reading data %v on stream %v: %v", i, stream.id, err)
return
}
@ -1680,7 +1715,7 @@ func testFlowControlAccountCheck(t *testing.T, msgSize int, wc windowSizeConfig)
// Close all streams
for _, stream := range clientStreams {
client.Write(stream, nil, nil, &Options{Last: true})
if _, err := stream.Read(make([]byte, 5)); err != io.EOF {
if _, err := stream.readTo(make([]byte, 5)); err != io.EOF {
t.Fatalf("Client expected an EOF from the server. Got: %v", err)
}
}
@ -1752,21 +1787,19 @@ func (s) TestReadGivesSameErrorAfterAnyErrorOccurs(t *testing.T) {
}
s.trReader = &transportReader{
reader: &recvBufferReader{
ctx: s.ctx,
ctxDone: s.ctx.Done(),
recv: s.buf,
freeBuffer: func(*bytes.Buffer) {},
ctx: s.ctx,
ctxDone: s.ctx.Done(),
recv: s.buf,
},
windowHandler: func(int) {},
}
testData := make([]byte, 1)
testData[0] = 5
testBuffer := bytes.NewBuffer(testData)
testErr := errors.New("test error")
s.write(recvMsg{buffer: testBuffer, err: testErr})
s.write(recvMsg{buffer: mem.NewBuffer(&testData, nil), err: testErr})
inBuf := make([]byte, 1)
actualCount, actualErr := s.Read(inBuf)
actualCount, actualErr := s.readTo(inBuf)
if actualCount != 0 {
t.Errorf("actualCount, _ := s.Read(_) differs; want 0; got %v", actualCount)
}
@ -1774,12 +1807,12 @@ func (s) TestReadGivesSameErrorAfterAnyErrorOccurs(t *testing.T) {
t.Errorf("_ , actualErr := s.Read(_) differs; want actualErr.Error() to be %v; got %v", testErr.Error(), actualErr.Error())
}
s.write(recvMsg{buffer: testBuffer, err: nil})
s.write(recvMsg{buffer: testBuffer, err: errors.New("different error from first")})
s.write(recvMsg{buffer: mem.NewBuffer(&testData, nil), err: nil})
s.write(recvMsg{buffer: mem.NewBuffer(&testData, nil), err: errors.New("different error from first")})
for i := 0; i < 2; i++ {
inBuf := make([]byte, 1)
actualCount, actualErr := s.Read(inBuf)
actualCount, actualErr := s.readTo(inBuf)
if actualCount != 0 {
t.Errorf("actualCount, _ := s.Read(_) differs; want %v; got %v", 0, actualCount)
}
@ -2208,11 +2241,11 @@ func (s) TestPingPong1B(t *testing.T) {
runPingPongTest(t, 1)
}
func (s) TestPingPong1KB(t *testing.T) {
func TestPingPong1KB(t *testing.T) {
runPingPongTest(t, 1024)
}
func (s) TestPingPong64KB(t *testing.T) {
func TestPingPong64KB(t *testing.T) {
runPingPongTest(t, 65536)
}
@ -2247,24 +2280,24 @@ func runPingPongTest(t *testing.T, msgSize int) {
opts := &Options{}
incomingHeader := make([]byte, 5)
ctx, cancel = context.WithTimeout(ctx, time.Second)
ctx, cancel = context.WithTimeout(ctx, 10*time.Millisecond)
defer cancel()
for ctx.Err() == nil {
if err := client.Write(stream, outgoingHeader, msg, opts); err != nil {
if err := client.Write(stream, outgoingHeader, newBufferSlice(msg), opts); err != nil {
t.Fatalf("Error on client while writing message. Err: %v", err)
}
if _, err := stream.Read(incomingHeader); err != nil {
if _, err := stream.readTo(incomingHeader); err != nil {
t.Fatalf("Error on client while reading data header. Err: %v", err)
}
sz := binary.BigEndian.Uint32(incomingHeader[1:])
recvMsg := make([]byte, int(sz))
if _, err := stream.Read(recvMsg); err != nil {
if _, err := stream.readTo(recvMsg); err != nil {
t.Fatalf("Error on client while reading data. Err: %v", err)
}
}
client.Write(stream, nil, nil, &Options{Last: true})
if _, err := stream.Read(incomingHeader); err != io.EOF {
if _, err := stream.readTo(incomingHeader); err != io.EOF {
t.Fatalf("Client expected EOF from the server. Got: %v", err)
}
}

56
mem/buffer_pool.go
View File

@ -29,10 +29,10 @@ import (
// decreased memory allocation.
type BufferPool interface {
// Get returns a buffer with specified length from the pool.
Get(length int) []byte
Get(length int) *[]byte
// Put returns a buffer to the pool.
Put([]byte)
Put(*[]byte)
}
var defaultBufferPoolSizes = []int{
@ -48,7 +48,13 @@ var defaultBufferPool BufferPool
func init() {
defaultBufferPool = NewTieredBufferPool(defaultBufferPoolSizes...)
internal.SetDefaultBufferPoolForTesting = func(pool BufferPool) { defaultBufferPool = pool }
internal.SetDefaultBufferPoolForTesting = func(pool BufferPool) {
defaultBufferPool = pool
}
internal.SetBufferPoolingThresholdForTesting = func(threshold int) {
bufferPoolingThreshold = threshold
}
}
// DefaultBufferPool returns the current default buffer pool. It is a BufferPool
@ -78,12 +84,12 @@ type tieredBufferPool struct {
fallbackPool simpleBufferPool
}
func (p *tieredBufferPool) Get(size int) []byte {
func (p *tieredBufferPool) Get(size int) *[]byte {
return p.getPool(size).Get(size)
}
func (p *tieredBufferPool) Put(buf []byte) {
p.getPool(cap(buf)).Put(buf)
func (p *tieredBufferPool) Put(buf *[]byte) {
p.getPool(cap(*buf)).Put(buf)
}
func (p *tieredBufferPool) getPool(size int) BufferPool {
@ -111,21 +117,22 @@ type sizedBufferPool struct {
defaultSize int
}
func (p *sizedBufferPool) Get(size int) []byte {
bs := *p.pool.Get().(*[]byte)
return bs[:size]
func (p *sizedBufferPool) Get(size int) *[]byte {
buf := p.pool.Get().(*[]byte)
b := *buf
clear(b[:cap(b)])
*buf = b[:size]
return buf
}
func (p *sizedBufferPool) Put(buf []byte) {
if cap(buf) < p.defaultSize {
func (p *sizedBufferPool) Put(buf *[]byte) {
if cap(*buf) < p.defaultSize {
// Ignore buffers that are too small to fit in the pool. Otherwise, when
// Get is called it will panic as it tries to index outside the bounds
// of the buffer.
return
}
buf = buf[:cap(buf)]
clear(buf)
p.pool.Put(&buf)
p.pool.Put(buf)
}
func newSizedBufferPool(size int) *sizedBufferPool {
@ -150,10 +157,11 @@ type simpleBufferPool struct {
pool sync.Pool
}
func (p *simpleBufferPool) Get(size int) []byte {
func (p *simpleBufferPool) Get(size int) *[]byte {
bs, ok := p.pool.Get().(*[]byte)
if ok && cap(*bs) >= size {
return (*bs)[:size]
*bs = (*bs)[:size]
return bs
}
// A buffer was pulled from the pool, but it is too small. Put it back in
@ -162,13 +170,12 @@ func (p *simpleBufferPool) Get(size int) []byte {
p.pool.Put(bs)
}
return make([]byte, size)
b := make([]byte, size)
return &b
}
func (p *simpleBufferPool) Put(buf []byte) {
buf = buf[:cap(buf)]
clear(buf)
p.pool.Put(&buf)
func (p *simpleBufferPool) Put(buf *[]byte) {
p.pool.Put(buf)
}
var _ BufferPool = NopBufferPool{}
@ -177,10 +184,11 @@ var _ BufferPool = NopBufferPool{}
type NopBufferPool struct{}
// Get returns a buffer with specified length from the pool.
func (NopBufferPool) Get(length int) []byte {
return make([]byte, length)
func (NopBufferPool) Get(length int) *[]byte {
b := make([]byte, length)
return &b
}
// Put returns a buffer to the pool.
func (NopBufferPool) Put([]byte) {
func (NopBufferPool) Put(*[]byte) {
}

36
mem/buffer_pool_test.go
View File

@ -19,9 +19,10 @@
package mem_test
import (
"bytes"
"testing"
"unsafe"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/mem"
)
@ -38,8 +39,8 @@ func (s) TestBufferPool(t *testing.T) {
for _, p := range pools {
for _, l := range testSizes {
bs := p.Get(l)
if len(bs) != l {
t.Fatalf("Get(%d) returned buffer of length %d, want %d", l, len(bs), l)
if len(*bs) != l {
t.Fatalf("Get(%d) returned buffer of length %d, want %d", l, len(*bs), l)
}
p.Put(bs)
@ -50,24 +51,37 @@ func (s) TestBufferPool(t *testing.T) {
func (s) TestBufferPoolClears(t *testing.T) {
pool := mem.NewTieredBufferPool(4)
buf := pool.Get(4)
copy(buf, "1234")
pool.Put(buf)
for {
buf1 := pool.Get(4)
copy(*buf1, "1234")
pool.Put(buf1)
if !cmp.Equal(buf, make([]byte, 4)) {
t.Fatalf("buffer not cleared")
buf2 := pool.Get(4)
if unsafe.SliceData(*buf1) != unsafe.SliceData(*buf2) {
pool.Put(buf2)
// This test is only relevant if a buffer is reused, otherwise try again. This
// can happen if a GC pause happens between putting the buffer back in the pool
// and getting a new one.
continue
}
if !bytes.Equal(*buf1, make([]byte, 4)) {
t.Fatalf("buffer not cleared")
}
break
}
}
func (s) TestBufferPoolIgnoresShortBuffers(t *testing.T) {
pool := mem.NewTieredBufferPool(10, 20)
buf := pool.Get(1)
if cap(buf) != 10 {
t.Fatalf("Get(1) returned buffer with capacity: %d, want 10", cap(buf))
if cap(*buf) != 10 {
t.Fatalf("Get(1) returned buffer with capacity: %d, want 10", cap(*buf))
}
// Insert a short buffer into the pool, which is currently empty.
pool.Put(make([]byte, 1))
short := make([]byte, 1)
pool.Put(&short)
// Then immediately request a buffer that would be pulled from the pool where the
// short buffer would have been returned. If the short buffer is pulled from the
// pool, it could cause a panic.

102
mem/buffer_slice.go
View File

@ -19,6 +19,7 @@
package mem
import (
"compress/flate"
"io"
)
@ -36,7 +37,7 @@ import (
// By convention, any APIs that return (mem.BufferSlice, error) should reduce
// the burden on the caller by never returning a mem.BufferSlice that needs to
// be freed if the error is non-nil, unless explicitly stated.
type BufferSlice []*Buffer
type BufferSlice []Buffer
// Len returns the sum of the length of all the Buffers in this slice.
//
@ -52,14 +53,11 @@ func (s BufferSlice) Len() int {
return length
}
// Ref returns a new BufferSlice containing a new reference of each Buffer in the
// input slice.
func (s BufferSlice) Ref() BufferSlice {
out := make(BufferSlice, len(s))
for i, b := range s {
out[i] = b.Ref()
// Ref invokes Ref on each buffer in the slice.
func (s BufferSlice) Ref() {
for _, b := range s {
b.Ref()
}
return out
}
// Free invokes Buffer.Free() on each Buffer in the slice.
@ -97,54 +95,73 @@ func (s BufferSlice) Materialize() []byte {
// to a single Buffer pulled from the given BufferPool. As a special case, if the
// input BufferSlice only actually has one Buffer, this function has nothing to
// do and simply returns said Buffer.
func (s BufferSlice) MaterializeToBuffer(pool BufferPool) *Buffer {
func (s BufferSlice) MaterializeToBuffer(pool BufferPool) Buffer {
if len(s) == 1 {
return s[0].Ref()
s[0].Ref()
return s[0]
}
buf := pool.Get(s.Len())
s.CopyTo(buf)
return NewBuffer(buf, pool.Put)
sLen := s.Len()
if sLen == 0 {
return emptyBuffer{}
}
buf := pool.Get(sLen)
s.CopyTo(*buf)
return NewBuffer(buf, pool)
}
// Reader returns a new Reader for the input slice after taking references to
// each underlying buffer.
func (s BufferSlice) Reader() *Reader {
return &Reader{
data: s.Ref(),
func (s BufferSlice) Reader() Reader {
s.Ref()
return &sliceReader{
data: s,
len: s.Len(),
}
}
var _ io.ReadCloser = (*Reader)(nil)
// Reader exposes a BufferSlice's data as an io.Reader, allowing it to interface
// with other parts systems. It also provides an additional convenience method
// Remaining(), which returns the number of unread bytes remaining in the slice.
//
// Note that reading data from the reader does not free the underlying buffers!
// Only calling Close once all data is read will free the buffers.
type Reader struct {
// Buffers will be freed as they are read.
type Reader interface {
flate.Reader
// Close frees the underlying BufferSlice and never returns an error. Subsequent
// calls to Read will return (0, io.EOF).
Close() error
// Remaining returns the number of unread bytes remaining in the slice.
Remaining() int
}
type sliceReader struct {
data BufferSlice
len int
// The index into data[0].ReadOnlyData().
bufferIdx int
}
// Remaining returns the number of unread bytes remaining in the slice.
func (r *Reader) Remaining() int {
func (r *sliceReader) Remaining() int {
return r.len
}
// Close frees the underlying BufferSlice and never returns an error. Subsequent
// calls to Read will return (0, io.EOF).
func (r *Reader) Close() error {
func (r *sliceReader) Close() error {
r.data.Free()
r.data = nil
r.len = 0
return nil
}
func (r *Reader) Read(buf []byte) (n int, _ error) {
func (r *sliceReader) freeFirstBufferIfEmpty() bool {
if len(r.data) == 0 || r.bufferIdx != len(r.data[0].ReadOnlyData()) {
return false
}
r.data[0].Free()
r.data = r.data[1:]
r.bufferIdx = 0
return true
}
func (r *sliceReader) Read(buf []byte) (n int, _ error) {
if r.len == 0 {
return 0, io.EOF
}
@ -159,19 +176,32 @@ func (r *Reader) Read(buf []byte) (n int, _ error) {
n += copied // Increment the total number of bytes read.
buf = buf[copied:] // Shrink the given byte slice.
// If we have copied all of the data from the first Buffer, free it and
// advance to the next in the slice.
if r.bufferIdx == len(data) {
oldBuffer := r.data[0]
oldBuffer.Free()
r.data = r.data[1:]
r.bufferIdx = 0
}
// If we have copied all the data from the first Buffer, free it and advance to
// the next in the slice.
r.freeFirstBufferIfEmpty()
}
return n, nil
}
func (r *sliceReader) ReadByte() (byte, error) {
if r.len == 0 {
return 0, io.EOF
}
// There may be any number of empty buffers in the slice, clear them all until a
// non-empty buffer is reached. This is guaranteed to exit since r.len is not 0.
for r.freeFirstBufferIfEmpty() {
}
b := r.data[0].ReadOnlyData()[r.bufferIdx]
r.len--
r.bufferIdx++
// Free the first buffer in the slice if the last byte was read
r.freeFirstBufferIfEmpty()
return b, nil
}
var _ io.Writer = (*writer)(nil)
type writer struct {

35
mem/buffer_slice_test.go
View File

@ -27,6 +27,10 @@ import (
"google.golang.org/grpc/mem"
)
func newBuffer(data []byte, pool mem.BufferPool) mem.Buffer {
return mem.NewBuffer(&data, pool)
}
func (s) TestBufferSlice_Len(t *testing.T) {
tests := []struct {
name string
@ -40,15 +44,15 @@ func (s) TestBufferSlice_Len(t *testing.T) {
},
{
name: "single",
in: mem.BufferSlice{mem.NewBuffer([]byte("abcd"), nil)},
in: mem.BufferSlice{newBuffer([]byte("abcd"), nil)},
want: 4,
},
{
name: "multiple",
in: mem.BufferSlice{
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
},
want: 12,
},
@ -65,15 +69,15 @@ func (s) TestBufferSlice_Len(t *testing.T) {
func (s) TestBufferSlice_Ref(t *testing.T) {
// Create a new buffer slice and a reference to it.
bs := mem.BufferSlice{
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
}
bsRef := bs.Ref()
bs.Ref()
// Free the original buffer slice and verify that the reference can still
// read data from it.
bs.Free()
got := bsRef.Materialize()
got := bs.Materialize()
want := []byte("abcdabcd")
if !bytes.Equal(got, want) {
t.Errorf("BufferSlice.Materialize() = %s, want %s", string(got), string(want))
@ -89,16 +93,16 @@ func (s) TestBufferSlice_MaterializeToBuffer(t *testing.T) {
}{
{
name: "single",
in: mem.BufferSlice{mem.NewBuffer([]byte("abcd"), nil)},
in: mem.BufferSlice{newBuffer([]byte("abcd"), nil)},
pool: nil, // MaterializeToBuffer should not use the pool in this case.
wantData: []byte("abcd"),
},
{
name: "multiple",
in: mem.BufferSlice{
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
},
pool: mem.DefaultBufferPool(),
wantData: []byte("abcdabcdabcd"),
@ -106,6 +110,7 @@ func (s) TestBufferSlice_MaterializeToBuffer(t *testing.T) {
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
defer tt.in.Free()
got := tt.in.MaterializeToBuffer(tt.pool)
defer got.Free()
if !bytes.Equal(got.ReadOnlyData(), tt.wantData) {
@ -117,9 +122,9 @@ func (s) TestBufferSlice_MaterializeToBuffer(t *testing.T) {
func (s) TestBufferSlice_Reader(t *testing.T) {
bs := mem.BufferSlice{
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
mem.NewBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
newBuffer([]byte("abcd"), nil),
}
wantData := []byte("abcdabcdabcd")

217
mem/buffers.go
View File

@ -27,13 +27,14 @@ package mem
import (
"fmt"
"sync"
"sync/atomic"
)
// A Buffer represents a reference counted piece of data (in bytes) that can be
// acquired by a call to NewBuffer() or Copy(). A reference to a Buffer may be
// released by calling Free(), which invokes the given free function only after
// all references are released.
// released by calling Free(), which invokes the free function given at creation
// only after all references are released.
//
// Note that a Buffer is not safe for concurrent access and instead each
// goroutine should use its own reference to the data, which can be acquired via
@ -41,23 +42,61 @@ import (
//
// Attempts to access the underlying data after releasing the reference to the
// Buffer will panic.
type Buffer struct {
data []byte
refs *atomic.Int32
free func()
freed bool
type Buffer interface {
// ReadOnlyData returns the underlying byte slice. Note that it is undefined
// behavior to modify the contents of this slice in any way.
ReadOnlyData() []byte
// Ref increases the reference counter for this Buffer.
Ref()
// Free decrements this Buffer's reference counter and frees the underlying
// byte slice if the counter reaches 0 as a result of this call.
Free()
// Len returns the Buffer's size.
Len() int
split(n int) (left, right Buffer)
read(buf []byte) (int, Buffer)
}
// NewBuffer creates a new Buffer from the given data, initializing the
// reference counter to 1. The given free function is called when all references
// to the returned Buffer are released.
var (
bufferPoolingThreshold = 1 << 10
bufferObjectPool = sync.Pool{New: func() any { return new(buffer) }}
refObjectPool = sync.Pool{New: func() any { return new(atomic.Int32) }}
)
func IsBelowBufferPoolingThreshold(size int) bool {
return size <= bufferPoolingThreshold
}
type buffer struct {
origData *[]byte
data []byte
refs *atomic.Int32
pool BufferPool
}
func newBuffer() *buffer {
return bufferObjectPool.Get().(*buffer)
}
// 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.
func NewBuffer(data []byte, onFree func([]byte)) *Buffer {
b := &Buffer{data: data, refs: new(atomic.Int32)}
if onFree != nil {
b.free = func() { onFree(data) }
func NewBuffer(data *[]byte, pool BufferPool) Buffer {
if pool == nil || IsBelowBufferPoolingThreshold(len(*data)) {
return (SliceBuffer)(*data)
}
b := newBuffer()
b.origData = data
b.data = *data
b.pool = pool
b.refs = refObjectPool.Get().(*atomic.Int32)
b.refs.Add(1)
return b
}
@ -68,82 +107,146 @@ func NewBuffer(data []byte, onFree func([]byte)) *Buffer {
// It acquires a []byte from the given pool and copies over the backing array
// of the given data. The []byte acquired from the pool is returned to the
// pool when all references to the returned Buffer are released.
func Copy(data []byte, pool BufferPool) *Buffer {
func Copy(data []byte, pool BufferPool) Buffer {
if IsBelowBufferPoolingThreshold(len(data)) {
buf := make(SliceBuffer, len(data))
copy(buf, data)
return buf
}
buf := pool.Get(len(data))
copy(buf, data)
return NewBuffer(buf, pool.Put)
copy(*buf, data)
return NewBuffer(buf, pool)
}
// ReadOnlyData returns the underlying byte slice. Note that it is undefined
// behavior to modify the contents of this slice in any way.
func (b *Buffer) ReadOnlyData() []byte {
if b.freed {
func (b *buffer) ReadOnlyData() []byte {
if b.refs == nil {
panic("Cannot read freed buffer")
}
return b.data
}
// Ref returns a new reference to this Buffer's underlying byte slice.
func (b *Buffer) Ref() *Buffer {
if b.freed {
func (b *buffer) Ref() {
if b.refs == nil {
panic("Cannot ref freed buffer")
}
b.refs.Add(1)
return &Buffer{
data: b.data,
refs: b.refs,
free: b.free,
}
}
// Free decrements this Buffer's reference counter and frees the underlying
// byte slice if the counter reaches 0 as a result of this call.
func (b *Buffer) Free() {
if b.freed {
return
func (b *buffer) Free() {
if b.refs == nil {
panic("Cannot free freed buffer")
}
b.freed = true
refs := b.refs.Add(-1)
if refs == 0 && b.free != nil {
b.free()
switch {
case refs > 0:
return
case refs == 0:
if b.pool != nil {
b.pool.Put(b.origData)
}
refObjectPool.Put(b.refs)
b.origData = nil
b.data = nil
b.refs = nil
b.pool = nil
bufferObjectPool.Put(b)
default:
panic("Cannot free freed buffer")
}
b.data = nil
}
// Len returns the Buffer's size.
func (b *Buffer) Len() int {
// Convenience: io.Reader returns (n int, err error), and n is often checked
// before err is checked. To mimic this, Len() should work on nil Buffers.
if b == nil {
return 0
}
func (b *buffer) Len() int {
return len(b.ReadOnlyData())
}
// Split modifies the receiver to point to the first n bytes while it returns a
// new reference to the remaining bytes. The returned Buffer functions just like
// a normal reference acquired using Ref().
func (b *Buffer) Split(n int) *Buffer {
if b.freed {
func (b *buffer) split(n int) (Buffer, Buffer) {
if b.refs == nil {
panic("Cannot split freed buffer")
}
b.refs.Add(1)
split := newBuffer()
split.origData = b.origData
split.data = b.data[n:]
split.refs = b.refs
split.pool = b.pool
split := &Buffer{
refs: b.refs,
free: b.free,
b.data = b.data[:n]
return b, split
}
func (b *buffer) read(buf []byte) (int, Buffer) {
if b.refs == nil {
panic("Cannot read freed buffer")
}
b.data, split.data = b.data[:n], b.data[n:]
n := copy(buf, b.data)
if n == len(b.data) {
b.Free()
return n, nil
}
return split
b.data = b.data[n:]
return n, b
}
// String returns a string representation of the buffer. May be used for
// debugging purposes.
func (b *Buffer) String() string {
func (b *buffer) String() string {
return fmt.Sprintf("mem.Buffer(%p, data: %p, length: %d)", b, b.ReadOnlyData(), len(b.ReadOnlyData()))
}
func ReadUnsafe(dst []byte, buf Buffer) (int, Buffer) {
return buf.read(dst)
}
// SplitUnsafe modifies the receiver to point to the first n bytes while it
// returns a new reference to the remaining bytes. The returned Buffer functions
// just like a normal reference acquired using Ref().
func SplitUnsafe(buf Buffer, n int) (left, right Buffer) {
return buf.split(n)
}
type emptyBuffer struct{}
func (e emptyBuffer) ReadOnlyData() []byte {
return nil
}
func (e emptyBuffer) Ref() {}
func (e emptyBuffer) Free() {}
func (e emptyBuffer) Len() int {
return 0
}
func (e emptyBuffer) split(n int) (left, right Buffer) {
return e, e
}
func (e emptyBuffer) read(buf []byte) (int, Buffer) {
return 0, e
}
type SliceBuffer []byte
func (s SliceBuffer) ReadOnlyData() []byte { return s }
func (s SliceBuffer) Ref() {}
func (s SliceBuffer) Free() {}
func (s SliceBuffer) Len() int { return len(s) }
func (s SliceBuffer) split(n int) (left, right Buffer) {
return s[:n], s[n:]
}
func (s SliceBuffer) read(buf []byte) (int, Buffer) {
n := copy(buf, s)
if n == len(s) {
return n, nil
}
return n, s[n:]
}

259
mem/buffers_test.go
View File

@ -20,24 +20,20 @@ package mem_test
import (
"bytes"
"fmt"
"testing"
"time"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/mem"
)
const (
defaultTestTimeout = 5 * time.Second
defaultTestShortTimeout = 100 * time.Millisecond
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
internal.SetBufferPoolingThresholdForTesting.(func(int))(0)
grpctest.RunSubTests(t, s{})
}
@ -45,29 +41,23 @@ func Test(t *testing.T) {
// the free function with the correct data.
func (s) TestBuffer_NewBufferAndFree(t *testing.T) {
data := "abcd"
errCh := make(chan error, 1)
freeF := func(got []byte) {
if !bytes.Equal(got, []byte(data)) {
errCh <- fmt.Errorf("Free function called with bytes %s, want %s", string(got), string(data))
return
freed := false
freeF := poolFunc(func(got *[]byte) {
if !bytes.Equal(*got, []byte(data)) {
t.Fatalf("Free function called with bytes %s, want %s", string(*got), data)
}
errCh <- nil
}
freed = true
})
buf := mem.NewBuffer([]byte(data), freeF)
buf := newBuffer([]byte(data), freeF)
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
}
// Verify that the free function is invoked when all references are freed.
buf.Free()
select {
case err := <-errCh:
if err != nil {
t.Fatal(err)
}
case <-time.After(defaultTestTimeout):
t.Fatalf("Timeout waiting for Buffer to be freed")
if !freed {
t.Fatalf("Buffer not freed")
}
}
@ -76,84 +66,87 @@ func (s) TestBuffer_NewBufferAndFree(t *testing.T) {
// correct data, but only after all references are released.
func (s) TestBuffer_NewBufferRefAndFree(t *testing.T) {
data := "abcd"
errCh := make(chan error, 1)
freeF := func(got []byte) {
if !bytes.Equal(got, []byte(data)) {
errCh <- fmt.Errorf("Free function called with bytes %s, want %s", string(got), string(data))
return
freed := false
freeF := poolFunc(func(got *[]byte) {
if !bytes.Equal(*got, []byte(data)) {
t.Fatalf("Free function called with bytes %s, want %s", string(*got), string(data))
}
errCh <- nil
}
freed = true
})
buf := mem.NewBuffer([]byte(data), freeF)
buf := newBuffer([]byte(data), freeF)
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
}
bufRef := buf.Ref()
if got := bufRef.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
buf.Ref()
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
t.Fatalf("New reference to the Buffer contains data %s, want %s", string(got), string(data))
}
// Verify that the free function is not invoked when all references are yet
// to be freed.
buf.Free()
select {
case <-errCh:
if freed {
t.Fatalf("Free function called before all references freed")
case <-time.After(defaultTestShortTimeout):
}
// Verify that the free function is invoked when all references are freed.
bufRef.Free()
select {
case err := <-errCh:
if err != nil {
t.Fatal(err)
}
case <-time.After(defaultTestTimeout):
t.Fatalf("Timeout waiting for Buffer to be freed")
buf.Free()
if !freed {
t.Fatalf("Buffer not freed")
}
}
// testBufferPool is a buffer pool that makes new buffer without pooling, and
// notifies on a channel that a buffer was returned to the pool.
type testBufferPool struct {
putCh chan []byte
func (s) TestBuffer_FreeAfterFree(t *testing.T) {
buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
if buf.Len() != 4 {
t.Fatalf("Buffer length is %d, want 4", buf.Len())
}
// Ensure that a double free does panic.
buf.Free()
defer checkForPanic(t, "Cannot free freed buffer")
buf.Free()
}
func (t *testBufferPool) Get(length int) []byte {
return make([]byte, length)
type singleBufferPool struct {
t *testing.T
data *[]byte
}
func (t *testBufferPool) Put(data []byte) {
t.putCh <- data
func (s *singleBufferPool) Get(length int) *[]byte {
if len(*s.data) != length {
s.t.Fatalf("Invalid requested length, got %d want %d", length, len(*s.data))
}
return s.data
}
func newTestBufferPool() *testBufferPool {
return &testBufferPool{putCh: make(chan []byte, 1)}
func (s *singleBufferPool) Put(b *[]byte) {
if s.data != b {
s.t.Fatalf("Wrong buffer returned to pool, got %p want %p", b, s.data)
}
s.data = nil
}
// Tests that a buffer created with Copy, which when later freed, returns the underlying
// byte slice to the buffer pool.
func (s) TestBuffer_CopyAndFree(t *testing.T) {
data := "abcd"
testPool := newTestBufferPool()
data := []byte("abcd")
testPool := &singleBufferPool{
t: t,
data: &data,
}
buf := mem.Copy([]byte(data), testPool)
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
buf := mem.Copy(data, testPool)
if got := buf.ReadOnlyData(); !bytes.Equal(got, data) {
t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
}
// Verify that the free function is invoked when all references are freed.
buf.Free()
select {
case got := <-testPool.putCh:
if !bytes.Equal(got, []byte(data)) {
t.Fatalf("Free function called with bytes %s, want %s", string(got), string(data))
}
case <-time.After(defaultTestTimeout):
t.Fatalf("Timeout waiting for Buffer to be freed")
if testPool.data != nil {
t.Fatalf("Buffer not freed")
}
}
@ -161,68 +154,103 @@ func (s) TestBuffer_CopyAndFree(t *testing.T) {
// acquired, which when later freed, returns the underlying byte slice to the
// buffer pool.
func (s) TestBuffer_CopyRefAndFree(t *testing.T) {
data := "abcd"
testPool := newTestBufferPool()
data := []byte("abcd")
testPool := &singleBufferPool{
t: t,
data: &data,
}
buf := mem.Copy([]byte(data), testPool)
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
buf := mem.Copy(data, testPool)
if got := buf.ReadOnlyData(); !bytes.Equal(got, data) {
t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
}
bufRef := buf.Ref()
if got := bufRef.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
buf.Ref()
if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
t.Fatalf("New reference to the Buffer contains data %s, want %s", string(got), string(data))
}
// Verify that the free function is not invoked when all references are yet
// to be freed.
buf.Free()
select {
case <-testPool.putCh:
if testPool.data == nil {
t.Fatalf("Free function called before all references freed")
case <-time.After(defaultTestShortTimeout):
}
// Verify that the free function is invoked when all references are freed.
bufRef.Free()
select {
case got := <-testPool.putCh:
if !bytes.Equal(got, []byte(data)) {
t.Fatalf("Free function called with bytes %s, want %s", string(got), string(data))
}
case <-time.After(defaultTestTimeout):
t.Fatalf("Timeout waiting for Buffer to be freed")
buf.Free()
if testPool.data != nil {
t.Fatalf("Free never called")
}
}
func (s) TestBuffer_ReadOnlyDataAfterFree(t *testing.T) {
// Verify that reading before freeing does not panic.
buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
buf.ReadOnlyData()
buf.Free()
defer checkForPanic(t, "Cannot read freed buffer")
buf.ReadOnlyData()
}
func (s) TestBuffer_RefAfterFree(t *testing.T) {
// Verify that acquiring a ref before freeing does not panic.
buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
buf.Ref()
// This first call should not panc and bring the ref counter down to 1
buf.Free()
// This second call actually frees the buffer
buf.Free()
defer checkForPanic(t, "Cannot ref freed buffer")
buf.Ref()
}
func (s) TestBuffer_SplitAfterFree(t *testing.T) {
// Verify that splitting before freeing does not panic.
buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
buf, bufSplit := mem.SplitUnsafe(buf, 2)
bufSplit.Free()
buf.Free()
defer checkForPanic(t, "Cannot split freed buffer")
mem.SplitUnsafe(buf, 2)
}
type poolFunc func(*[]byte)
func (p poolFunc) Get(length int) *[]byte {
panic("Get should never be called")
}
func (p poolFunc) Put(i *[]byte) {
p(i)
}
func (s) TestBuffer_Split(t *testing.T) {
ready := false
freed := false
data := []byte{1, 2, 3, 4}
buf := mem.NewBuffer(data, func(bytes []byte) {
buf := mem.NewBuffer(&data, poolFunc(func(bytes *[]byte) {
if !ready {
t.Fatalf("Freed too early")
}
freed = true
})
checkBufData := func(b *mem.Buffer, expected []byte) {
}))
checkBufData := func(b mem.Buffer, expected []byte) {
t.Helper()
if !bytes.Equal(b.ReadOnlyData(), expected) {
t.Fatalf("Buffer did not contain expected data %v, got %v", expected, b.ReadOnlyData())
}
}
// Take a ref of the original buffer
ref1 := buf.Ref()
split1 := buf.Split(2)
buf, split1 := mem.SplitUnsafe(buf, 2)
checkBufData(buf, data[:2])
checkBufData(split1, data[2:])
// Check that even though buf was split, the reference wasn't modified
checkBufData(ref1, data)
ref1.Free()
// Check that splitting the buffer more than once works as intended.
split2 := split1.Split(1)
split1, split2 := mem.SplitUnsafe(split1, 1)
checkBufData(split1, data[2:3])
checkBufData(split2, data[3:])
@ -242,52 +270,9 @@ func checkForPanic(t *testing.T, wantErr string) {
t.Helper()
r := recover()
if r == nil {
t.Fatalf("Use after free dit not panic")
t.Fatalf("Use after free did not panic")
}
if r.(string) != wantErr {
if msg, ok := r.(string); !ok || msg != wantErr {
t.Fatalf("panic called with %v, want %s", r, wantErr)
}
}
func (s) TestBuffer_ReadOnlyDataAfterFree(t *testing.T) {
// Verify that reading before freeing does not panic.
buf := mem.NewBuffer([]byte("abcd"), nil)
buf.ReadOnlyData()
buf.Free()
defer checkForPanic(t, "Cannot read freed buffer")
buf.ReadOnlyData()
}
func (s) TestBuffer_RefAfterFree(t *testing.T) {
// Verify that acquiring a ref before freeing does not panic.
buf := mem.NewBuffer([]byte("abcd"), nil)
bufRef := buf.Ref()
defer bufRef.Free()
buf.Free()
defer checkForPanic(t, "Cannot ref freed buffer")
buf.Ref()
}
func (s) TestBuffer_SplitAfterFree(t *testing.T) {
// Verify that splitting before freeing does not panic.
buf := mem.NewBuffer([]byte("abcd"), nil)
bufSplit := buf.Split(2)
defer bufSplit.Free()
buf.Free()
defer checkForPanic(t, "Cannot split freed buffer")
buf.Split(1)
}
func (s) TestBuffer_FreeAfterFree(t *testing.T) {
buf := mem.NewBuffer([]byte("abcd"), nil)
if buf.Len() != 4 {
t.Fatalf("Buffer length is %d, want 4", buf.Len())
}
// Ensure that a double free does not panic.
buf.Free()
buf.Free()
}

28
preloader.go
View File

@ -20,6 +20,7 @@ package grpc
import (
"google.golang.org/grpc/codes"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/status"
)
@ -31,9 +32,10 @@ import (
// later release.
type PreparedMsg struct {
// Struct for preparing msg before sending them
encodedData []byte
encodedData mem.BufferSlice
hdr []byte
payload []byte
payload mem.BufferSlice
pf payloadFormat
}
// Encode marshalls and compresses the message using the codec and compressor for the stream.
@ -57,11 +59,27 @@ func (p *PreparedMsg) Encode(s Stream, msg any) error {
if err != nil {
return err
}
p.encodedData = data
compData, err := compress(data, rpcInfo.preloaderInfo.cp, rpcInfo.preloaderInfo.comp)
materializedData := data.Materialize()
data.Free()
p.encodedData = mem.BufferSlice{mem.NewBuffer(&materializedData, nil)}
// TODO: it should be possible to grab the bufferPool from the underlying
// stream implementation with a type cast to its actual type (such as
// addrConnStream) and accessing the buffer pool directly.
var compData mem.BufferSlice
compData, p.pf, err = compress(p.encodedData, rpcInfo.preloaderInfo.cp, rpcInfo.preloaderInfo.comp, mem.DefaultBufferPool())
if err != nil {
return err
}
p.hdr, p.payload = msgHeader(data, compData)
if p.pf.isCompressed() {
materializedCompData := compData.Materialize()
compData.Free()
compData = mem.BufferSlice{mem.NewBuffer(&materializedCompData, nil)}
}
p.hdr, p.payload = msgHeader(p.encodedData, compData, p.pf)
return nil
}

267
rpc_util.go
View File

@ -19,7 +19,6 @@
package grpc
import (
"bytes"
"compress/gzip"
"context"
"encoding/binary"
@ -35,6 +34,7 @@ import (
"google.golang.org/grpc/encoding"
"google.golang.org/grpc/encoding/proto"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/stats"
@ -511,11 +511,51 @@ type ForceCodecCallOption struct {
}
func (o ForceCodecCallOption) before(c *callInfo) error {
c.codec = o.Codec
c.codec = newCodecV1Bridge(o.Codec)
return nil
}
func (o ForceCodecCallOption) after(c *callInfo, attempt *csAttempt) {}
// ForceCodecV2 returns a CallOption that will set codec to be used for all
// request and response messages for a call. The result of calling Name() will
// be used as the content-subtype after converting to lowercase, unless
// CallContentSubtype is also used.
//
// See Content-Type on
// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
// more details. Also see the documentation on RegisterCodec and
// CallContentSubtype for more details on the interaction between Codec and
// content-subtype.
//
// This function is provided for advanced users; prefer to use only
// CallContentSubtype to select a registered codec instead.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func ForceCodecV2(codec encoding.CodecV2) CallOption {
return ForceCodecV2CallOption{CodecV2: codec}
}
// ForceCodecV2CallOption is a CallOption that indicates the codec used for
// marshaling messages.
//
// # Experimental
//
// Notice: This type is EXPERIMENTAL and may be changed or removed in a
// later release.
type ForceCodecV2CallOption struct {
CodecV2 encoding.CodecV2
}
func (o ForceCodecV2CallOption) before(c *callInfo) error {
c.codec = o.CodecV2
return nil
}
func (o ForceCodecV2CallOption) after(c *callInfo, attempt *csAttempt) {}
// CallCustomCodec behaves like ForceCodec, but accepts a grpc.Codec instead of
// an encoding.Codec.
//
@ -536,7 +576,7 @@ type CustomCodecCallOption struct {
}
func (o CustomCodecCallOption) before(c *callInfo) error {
c.codec = o.Codec
c.codec = newCodecV0Bridge(o.Codec)
return nil
}
func (o CustomCodecCallOption) after(c *callInfo, attempt *csAttempt) {}
@ -577,19 +617,28 @@ const (
compressionMade payloadFormat = 1 // compressed
)
func (pf payloadFormat) isCompressed() bool {
return pf == compressionMade
}
type streamReader interface {
ReadHeader(header []byte) error
Read(n int) (mem.BufferSlice, error)
}
// parser reads complete gRPC messages from the underlying reader.
type parser struct {
// r is the underlying reader.
// See the comment on recvMsg for the permissible
// error types.
r io.Reader
r streamReader
// The header of a gRPC message. Find more detail at
// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
header [5]byte
// recvBufferPool is the pool of shared receive buffers.
recvBufferPool SharedBufferPool
// bufferPool is the pool of shared receive buffers.
bufferPool mem.BufferPool
}
// recvMsg reads a complete gRPC message from the stream.
@ -604,14 +653,15 @@ type parser struct {
// - an error from the status package
//
// No other error values or types must be returned, which also means
// that the underlying io.Reader must not return an incompatible
// that the underlying streamReader must not return an incompatible
// error.
func (p *parser) recvMsg(maxReceiveMessageSize int) (pf payloadFormat, msg []byte, err error) {
if _, err := p.r.Read(p.header[:]); err != nil {
func (p *parser) recvMsg(maxReceiveMessageSize int) (payloadFormat, mem.BufferSlice, error) {
err := p.r.ReadHeader(p.header[:])
if err != nil {
return 0, nil, err
}
pf = payloadFormat(p.header[0])
pf := payloadFormat(p.header[0])
length := binary.BigEndian.Uint32(p.header[1:])
if length == 0 {
@ -623,20 +673,21 @@ func (p *parser) recvMsg(maxReceiveMessageSize int) (pf payloadFormat, msg []byt
if int(length) > maxReceiveMessageSize {
return 0, nil, status.Errorf(codes.ResourceExhausted, "grpc: received message larger than max (%d vs. %d)", length, maxReceiveMessageSize)
}
msg = p.recvBufferPool.Get(int(length))
if _, err := p.r.Read(msg); err != nil {
data, err := p.r.Read(int(length))
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return 0, nil, err
}
return pf, msg, nil
return pf, data, nil
}
// encode serializes msg and returns a buffer containing the message, or an
// error if it is too large to be transmitted by grpc. If msg is nil, it
// generates an empty message.
func encode(c baseCodec, msg any) ([]byte, error) {
func encode(c baseCodec, msg any) (mem.BufferSlice, error) {
if msg == nil { // NOTE: typed nils will not be caught by this check
return nil, nil
}
@ -644,7 +695,8 @@ func encode(c baseCodec, msg any) ([]byte, error) {
if err != nil {
return nil, status.Errorf(codes.Internal, "grpc: error while marshaling: %v", err.Error())
}
if uint(len(b)) > math.MaxUint32 {
if uint(b.Len()) > math.MaxUint32 {
b.Free()
return nil, status.Errorf(codes.ResourceExhausted, "grpc: message too large (%d bytes)", len(b))
}
return b, nil
@ -655,34 +707,41 @@ func encode(c baseCodec, msg any) ([]byte, error) {
// indicating no compression was done.
//
// TODO(dfawley): eliminate cp parameter by wrapping Compressor in an encoding.Compressor.
func compress(in []byte, cp Compressor, compressor encoding.Compressor) ([]byte, error) {
if compressor == nil && cp == nil {
return nil, nil
}
if len(in) == 0 {
return nil, nil
func compress(in mem.BufferSlice, cp Compressor, compressor encoding.Compressor, pool mem.BufferPool) (mem.BufferSlice, payloadFormat, error) {
if (compressor == nil && cp == nil) || in.Len() == 0 {
return nil, compressionNone, nil
}
var out mem.BufferSlice
w := mem.NewWriter(&out, pool)
wrapErr := func(err error) error {
out.Free()
return status.Errorf(codes.Internal, "grpc: error while compressing: %v", err.Error())
}
cbuf := &bytes.Buffer{}
if compressor != nil {
z, err := compressor.Compress(cbuf)
z, err := compressor.Compress(w)
if err != nil {
return nil, wrapErr(err)
return nil, 0, wrapErr(err)
}
if _, err := z.Write(in); err != nil {
return nil, wrapErr(err)
for _, b := range in {
if _, err := z.Write(b.ReadOnlyData()); err != nil {
return nil, 0, wrapErr(err)
}
}
if err := z.Close(); err != nil {
return nil, wrapErr(err)
return nil, 0, wrapErr(err)
}
} else {
if err := cp.Do(cbuf, in); err != nil {
return nil, wrapErr(err)
// This is obviously really inefficient since it fully materializes the data, but
// there is no way around this with the old Compressor API. At least it attempts
// to return the buffer to the provider, in the hopes it can be reused (maybe
// even by a subsequent call to this very function).
buf := in.MaterializeToBuffer(pool)
defer buf.Free()
if err := cp.Do(w, buf.ReadOnlyData()); err != nil {
return nil, 0, wrapErr(err)
}
}
return cbuf.Bytes(), nil
return out, compressionMade, nil
}
const (
@ -693,28 +752,31 @@ const (
// msgHeader returns a 5-byte header for the message being transmitted and the
// payload, which is compData if non-nil or data otherwise.
func msgHeader(data, compData []byte) (hdr []byte, payload []byte) {
func msgHeader(data, compData mem.BufferSlice, pf payloadFormat) (hdr []byte, payload mem.BufferSlice) {
hdr = make([]byte, headerLen)
if compData != nil {
hdr[0] = byte(compressionMade)
data = compData
hdr[0] = byte(pf)
var length uint32
if pf.isCompressed() {
length = uint32(compData.Len())
payload = compData
} else {
hdr[0] = byte(compressionNone)
length = uint32(data.Len())
payload = data
}
// Write length of payload into buf
binary.BigEndian.PutUint32(hdr[payloadLen:], uint32(len(data)))
return hdr, data
binary.BigEndian.PutUint32(hdr[payloadLen:], length)
return hdr, payload
}
func outPayload(client bool, msg any, data, payload []byte, t time.Time) *stats.OutPayload {
func outPayload(client bool, msg any, dataLength, payloadLength int, t time.Time) *stats.OutPayload {
return &stats.OutPayload{
Client: client,
Payload: msg,
Data: data,
Length: len(data),
WireLength: len(payload) + headerLen,
CompressedLength: len(payload),
Length: dataLength,
WireLength: payloadLength + headerLen,
CompressedLength: payloadLength,
SentTime: t,
}
}
@ -741,7 +803,13 @@ func checkRecvPayload(pf payloadFormat, recvCompress string, haveCompressor bool
type payloadInfo struct {
compressedLength int // The compressed length got from wire.
uncompressedBytes []byte
uncompressedBytes mem.BufferSlice
}
func (p *payloadInfo) free() {
if p != nil && p.uncompressedBytes != nil {
p.uncompressedBytes.Free()
}
}
// recvAndDecompress reads a message from the stream, decompressing it if necessary.
@ -751,96 +819,113 @@ type payloadInfo struct {
// TODO: Refactor this function to reduce the number of arguments.
// See: https://google.github.io/styleguide/go/best-practices.html#function-argument-lists
func recvAndDecompress(p *parser, s *transport.Stream, dc Decompressor, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor, isServer bool,
) (uncompressedBuf []byte, cancel func(), err error) {
pf, compressedBuf, err := p.recvMsg(maxReceiveMessageSize)
) (out mem.BufferSlice, err error) {
pf, compressed, err := p.recvMsg(maxReceiveMessageSize)
if err != nil {
return nil, nil, err
return nil, err
}
compressedLength := compressed.Len()
if st := checkRecvPayload(pf, s.RecvCompress(), compressor != nil || dc != nil, isServer); st != nil {
return nil, nil, st.Err()
compressed.Free()
return nil, st.Err()
}
var size int
if pf == compressionMade {
if pf.isCompressed() {
defer compressed.Free()
// To match legacy behavior, if the decompressor is set by WithDecompressor or RPCDecompressor,
// use this decompressor as the default.
if dc != nil {
uncompressedBuf, err = dc.Do(bytes.NewReader(compressedBuf))
var uncompressedBuf []byte
uncompressedBuf, err = dc.Do(compressed.Reader())
if err == nil {
out = mem.BufferSlice{mem.NewBuffer(&uncompressedBuf, nil)}
}
size = len(uncompressedBuf)
} else {
uncompressedBuf, size, err = decompress(compressor, compressedBuf, maxReceiveMessageSize)
out, size, err = decompress(compressor, compressed, maxReceiveMessageSize, p.bufferPool)
}
if err != nil {
return nil, nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message: %v", err)
return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the received message: %v", err)
}
if size > maxReceiveMessageSize {
out.Free()
// TODO: Revisit the error code. Currently keep it consistent with java
// implementation.
return nil, nil, status.Errorf(codes.ResourceExhausted, "grpc: received message after decompression larger than max (%d vs. %d)", size, maxReceiveMessageSize)
return nil, status.Errorf(codes.ResourceExhausted, "grpc: received message after decompression larger than max (%d vs. %d)", size, maxReceiveMessageSize)
}
} else {
uncompressedBuf = compressedBuf
out = compressed
}
if payInfo != nil {
payInfo.compressedLength = len(compressedBuf)
payInfo.uncompressedBytes = uncompressedBuf
cancel = func() {}
} else {
cancel = func() {
p.recvBufferPool.Put(&compressedBuf)
}
payInfo.compressedLength = compressedLength
out.Ref()
payInfo.uncompressedBytes = out
}
return uncompressedBuf, cancel, nil
return out, nil
}
// Using compressor, decompress d, returning data and size.
// Optionally, if data will be over maxReceiveMessageSize, just return the size.
func decompress(compressor encoding.Compressor, d []byte, maxReceiveMessageSize int) ([]byte, int, error) {
dcReader, err := compressor.Decompress(bytes.NewReader(d))
func decompress(compressor encoding.Compressor, d mem.BufferSlice, maxReceiveMessageSize int, pool mem.BufferPool) (mem.BufferSlice, int, error) {
dcReader, err := compressor.Decompress(d.Reader())
if err != nil {
return nil, 0, err
}
if sizer, ok := compressor.(interface {
DecompressedSize(compressedBytes []byte) int
}); ok {
if size := sizer.DecompressedSize(d); size >= 0 {
if size > maxReceiveMessageSize {
return nil, size, nil
}
// size is used as an estimate to size the buffer, but we
// will read more data if available.
// +MinRead so ReadFrom will not reallocate if size is correct.
//
// TODO: If we ensure that the buffer size is the same as the DecompressedSize,
// we can also utilize the recv buffer pool here.
buf := bytes.NewBuffer(make([]byte, 0, size+bytes.MinRead))
bytesRead, err := buf.ReadFrom(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
return buf.Bytes(), int(bytesRead), err
}
// TODO: Can/should this still be preserved with the new BufferSlice API? Are
// there any actual benefits to allocating a single large buffer instead of
// multiple smaller ones?
//if sizer, ok := compressor.(interface {
// DecompressedSize(compressedBytes []byte) int
//}); ok {
// if size := sizer.DecompressedSize(d); size >= 0 {
// if size > maxReceiveMessageSize {
// return nil, size, nil
// }
// // size is used as an estimate to size the buffer, but we
// // will read more data if available.
// // +MinRead so ReadFrom will not reallocate if size is correct.
// //
// // TODO: If we ensure that the buffer size is the same as the DecompressedSize,
// // we can also utilize the recv buffer pool here.
// buf := bytes.NewBuffer(make([]byte, 0, size+bytes.MinRead))
// bytesRead, err := buf.ReadFrom(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
// return buf.Bytes(), int(bytesRead), err
// }
//}
var out mem.BufferSlice
_, err = io.Copy(mem.NewWriter(&out, pool), io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
if err != nil {
out.Free()
return nil, 0, err
}
// Read from LimitReader with limit max+1. So if the underlying
// reader is over limit, the result will be bigger than max.
d, err = io.ReadAll(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
return d, len(d), err
return out, out.Len(), nil
}
// For the two compressor parameters, both should not be set, but if they are,
// dc takes precedence over compressor.
// TODO(dfawley): wrap the old compressor/decompressor using the new API?
func recv(p *parser, c baseCodec, s *transport.Stream, dc Decompressor, m any, maxReceiveMessageSize int, payInfo *payloadInfo, compressor encoding.Compressor, isServer bool) error {
buf, cancel, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, payInfo, compressor, isServer)
data, err := recvAndDecompress(p, s, dc, maxReceiveMessageSize, payInfo, compressor, isServer)
if err != nil {
return err
}
defer cancel()
if err := c.Unmarshal(buf, m); err != nil {
// If the codec wants its own reference to the data, it can get it. Otherwise, always
// free the buffers.
defer data.Free()
if err := c.Unmarshal(data, m); err != nil {
return status.Errorf(codes.Internal, "grpc: failed to unmarshal the received message: %v", err)
}
return nil
}
@ -943,7 +1028,7 @@ func setCallInfoCodec(c *callInfo) error {
// encoding.Codec (Name vs. String method name). We only support
// setting content subtype from encoding.Codec to avoid a behavior
// change with the deprecated version.
if ec, ok := c.codec.(encoding.Codec); ok {
if ec, ok := c.codec.(encoding.CodecV2); ok {
c.contentSubtype = strings.ToLower(ec.Name())
}
}
@ -952,12 +1037,12 @@ func setCallInfoCodec(c *callInfo) error {
if c.contentSubtype == "" {
// No codec specified in CallOptions; use proto by default.
c.codec = encoding.GetCodec(proto.Name)
c.codec = getCodec(proto.Name)
return nil
}
// c.contentSubtype is already lowercased in CallContentSubtype
c.codec = encoding.GetCodec(c.contentSubtype)
c.codec = getCodec(c.contentSubtype)
if c.codec == nil {
return status.Errorf(codes.Internal, "no codec registered for content-subtype %s", c.contentSubtype)
}

52
rpc_util_test.go
View File

@ -27,21 +27,45 @@ import (
"testing"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/encoding"
protoenc "google.golang.org/grpc/encoding/proto"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/status"
perfpb "google.golang.org/grpc/test/codec_perf"
"google.golang.org/protobuf/proto"
)
type fullReader struct {
reader io.Reader
data []byte
}
func (f fullReader) Read(p []byte) (int, error) {
return io.ReadFull(f.reader, p)
func (f *fullReader) ReadHeader(header []byte) error {
buf, err := f.Read(len(header))
defer buf.Free()
if err != nil {
return err
}
buf.CopyTo(header)
return nil
}
func (f *fullReader) Read(n int) (mem.BufferSlice, error) {
if len(f.data) == 0 {
return nil, io.EOF
}
if len(f.data) < n {
data := f.data
f.data = nil
return mem.BufferSlice{mem.NewBuffer(&data, nil)}, io.ErrUnexpectedEOF
}
buf := f.data[:n]
f.data = f.data[n:]
return mem.BufferSlice{mem.NewBuffer(&buf, nil)}, nil
}
var _ CallOption = EmptyCallOption{} // ensure EmptyCallOption implements the interface
@ -64,10 +88,10 @@ func (s) TestSimpleParsing(t *testing.T) {
// Check that messages with length >= 2^24 are parsed.
{append([]byte{0, 1, 0, 0, 0}, bigMsg...), nil, bigMsg, compressionNone},
} {
buf := fullReader{bytes.NewReader(test.p)}
parser := &parser{r: buf, recvBufferPool: nopBufferPool{}}
buf := &fullReader{test.p}
parser := &parser{r: buf, bufferPool: mem.DefaultBufferPool()}
pt, b, err := parser.recvMsg(math.MaxInt32)
if err != test.err || !bytes.Equal(b, test.b) || pt != test.pt {
if err != test.err || !bytes.Equal(b.Materialize(), test.b) || pt != test.pt {
t.Fatalf("parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, %v", test.p, pt, b, err, test.pt, test.b, test.err)
}
}
@ -76,8 +100,8 @@ func (s) TestSimpleParsing(t *testing.T) {
func (s) TestMultipleParsing(t *testing.T) {
// Set a byte stream consists of 3 messages with their headers.
p := []byte{0, 0, 0, 0, 1, 'a', 0, 0, 0, 0, 2, 'b', 'c', 0, 0, 0, 0, 1, 'd'}
b := fullReader{bytes.NewReader(p)}
parser := &parser{r: b, recvBufferPool: nopBufferPool{}}
b := &fullReader{p}
parser := &parser{r: b, bufferPool: mem.DefaultBufferPool()}
wantRecvs := []struct {
pt payloadFormat
@ -89,7 +113,7 @@ func (s) TestMultipleParsing(t *testing.T) {
}
for i, want := range wantRecvs {
pt, data, err := parser.recvMsg(math.MaxInt32)
if err != nil || pt != want.pt || !reflect.DeepEqual(data, want.data) {
if err != nil || pt != want.pt || !reflect.DeepEqual(data.Materialize(), want.data) {
t.Fatalf("after %d calls, parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, <nil>",
i, p, pt, data, err, want.pt, want.data)
}
@ -113,12 +137,12 @@ func (s) TestEncode(t *testing.T) {
}{
{nil, []byte{0, 0, 0, 0, 0}, []byte{}, nil},
} {
data, err := encode(encoding.GetCodec(protoenc.Name), test.msg)
if err != test.err || !bytes.Equal(data, test.data) {
data, err := encode(getCodec(protoenc.Name), test.msg)
if err != test.err || !bytes.Equal(data.Materialize(), test.data) {
t.Errorf("encode(_, %v) = %v, %v; want %v, %v", test.msg, data, err, test.data, test.err)
continue
}
if hdr, _ := msgHeader(data, nil); !bytes.Equal(hdr, test.hdr) {
if hdr, _ := msgHeader(data, nil, compressionNone); !bytes.Equal(hdr, test.hdr) {
t.Errorf("msgHeader(%v, false) = %v; want %v", data, hdr, test.hdr)
}
}
@ -194,7 +218,7 @@ func (s) TestToRPCErr(t *testing.T) {
// bmEncode benchmarks encoding a Protocol Buffer message containing mSize
// bytes.
func bmEncode(b *testing.B, mSize int) {
cdc := encoding.GetCodec(protoenc.Name)
cdc := getCodec(protoenc.Name)
msg := &perfpb.Buffer{Body: make([]byte, mSize)}
encodeData, _ := encode(cdc, msg)
encodedSz := int64(len(encodeData))

96
server.go
View File

@ -45,6 +45,7 @@ import (
"google.golang.org/grpc/internal/grpcutil"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/stats"
@ -80,7 +81,7 @@ func init() {
}
internal.BinaryLogger = binaryLogger
internal.JoinServerOptions = newJoinServerOption
internal.RecvBufferPool = recvBufferPool
internal.BufferPool = bufferPool
}
var statusOK = status.New(codes.OK, "")
@ -170,7 +171,7 @@ type serverOptions struct {
maxHeaderListSize *uint32
headerTableSize *uint32
numServerWorkers uint32
recvBufferPool SharedBufferPool
bufferPool mem.BufferPool
waitForHandlers bool
}
@ -181,7 +182,7 @@ var defaultServerOptions = serverOptions{
connectionTimeout: 120 * time.Second,
writeBufferSize: defaultWriteBufSize,
readBufferSize: defaultReadBufSize,
recvBufferPool: nopBufferPool{},
bufferPool: mem.DefaultBufferPool(),
}
var globalServerOptions []ServerOption
@ -313,7 +314,7 @@ func KeepaliveEnforcementPolicy(kep keepalive.EnforcementPolicy) ServerOption {
// Will be supported throughout 1.x.
func CustomCodec(codec Codec) ServerOption {
return newFuncServerOption(func(o *serverOptions) {
o.codec = codec
o.codec = newCodecV0Bridge(codec)
})
}
@ -342,7 +343,22 @@ func CustomCodec(codec Codec) ServerOption {
// later release.
func ForceServerCodec(codec encoding.Codec) ServerOption {
return newFuncServerOption(func(o *serverOptions) {
o.codec = codec
o.codec = newCodecV1Bridge(codec)
})
}
// ForceServerCodecV2 is the equivalent of ForceServerCodec, but for the new
// CodecV2 interface.
//
// Will be supported throughout 1.x.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func ForceServerCodecV2(codecV2 encoding.CodecV2) ServerOption {
return newFuncServerOption(func(o *serverOptions) {
o.codec = codecV2
})
}
@ -592,26 +608,9 @@ func WaitForHandlers(w bool) ServerOption {
})
}
// RecvBufferPool returns a ServerOption that configures the server
// to use the provided shared buffer pool for parsing incoming messages. Depending
// on the application's workload, this could result in reduced memory allocation.
//
// If you are unsure about how to implement a memory pool but want to utilize one,
// begin with grpc.NewSharedBufferPool.
//
// Note: The shared buffer pool feature will not be active if any of the following
// options are used: StatsHandler, EnableTracing, or binary logging. In such
// cases, the shared buffer pool will be ignored.
//
// Deprecated: use experimental.WithRecvBufferPool instead. Will be deleted in
// v1.60.0 or later.
func RecvBufferPool(bufferPool SharedBufferPool) ServerOption {
return recvBufferPool(bufferPool)
}
func recvBufferPool(bufferPool SharedBufferPool) ServerOption {
func bufferPool(bufferPool mem.BufferPool) ServerOption {
return newFuncServerOption(func(o *serverOptions) {
o.recvBufferPool = bufferPool
o.bufferPool = bufferPool
})
}
@ -980,6 +979,7 @@ func (s *Server) newHTTP2Transport(c net.Conn) transport.ServerTransport {
ChannelzParent: s.channelz,
MaxHeaderListSize: s.opts.maxHeaderListSize,
HeaderTableSize: s.opts.headerTableSize,
BufferPool: s.opts.bufferPool,
}
st, err := transport.NewServerTransport(c, config)
if err != nil {
@ -1072,7 +1072,7 @@ var _ http.Handler = (*Server)(nil)
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
st, err := transport.NewServerHandlerTransport(w, r, s.opts.statsHandlers)
st, err := transport.NewServerHandlerTransport(w, r, s.opts.statsHandlers, s.opts.bufferPool)
if err != nil {
// Errors returned from transport.NewServerHandlerTransport have
// already been written to w.
@ -1142,20 +1142,35 @@ func (s *Server) sendResponse(ctx context.Context, t transport.ServerTransport,
channelz.Error(logger, s.channelz, "grpc: server failed to encode response: ", err)
return err
}
compData, err := compress(data, cp, comp)
compData, pf, err := compress(data, cp, comp, s.opts.bufferPool)
if err != nil {
data.Free()
channelz.Error(logger, s.channelz, "grpc: server failed to compress response: ", err)
return err
}
hdr, payload := msgHeader(data, compData)
hdr, payload := msgHeader(data, compData, pf)
defer func() {
compData.Free()
data.Free()
// payload does not need to be freed here, it is either data or compData, both of
// which are already freed.
}()
dataLen := data.Len()
payloadLen := payload.Len()
// TODO(dfawley): should we be checking len(data) instead?
if len(payload) > s.opts.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "grpc: trying to send message larger than max (%d vs. %d)", len(payload), s.opts.maxSendMessageSize)
if payloadLen > s.opts.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "grpc: trying to send message larger than max (%d vs. %d)", payloadLen, s.opts.maxSendMessageSize)
}
err = t.Write(stream, hdr, payload, opts)
if err == nil {
for _, sh := range s.opts.statsHandlers {
sh.HandleRPC(ctx, outPayload(false, msg, data, payload, time.Now()))
if len(s.opts.statsHandlers) != 0 {
for _, sh := range s.opts.statsHandlers {
sh.HandleRPC(ctx, outPayload(false, msg, dataLen, payloadLen, time.Now()))
}
}
}
return err
@ -1334,9 +1349,10 @@ func (s *Server) processUnaryRPC(ctx context.Context, t transport.ServerTranspor
var payInfo *payloadInfo
if len(shs) != 0 || len(binlogs) != 0 {
payInfo = &payloadInfo{}
defer payInfo.free()
}
d, cancel, err := recvAndDecompress(&parser{r: stream, recvBufferPool: s.opts.recvBufferPool}, stream, dc, s.opts.maxReceiveMessageSize, payInfo, decomp, true)
d, err := recvAndDecompress(&parser{r: stream, bufferPool: s.opts.bufferPool}, stream, dc, s.opts.maxReceiveMessageSize, payInfo, decomp, true)
if err != nil {
if e := t.WriteStatus(stream, status.Convert(err)); e != nil {
channelz.Warningf(logger, s.channelz, "grpc: Server.processUnaryRPC failed to write status: %v", e)
@ -1347,24 +1363,22 @@ func (s *Server) processUnaryRPC(ctx context.Context, t transport.ServerTranspor
t.IncrMsgRecv()
}
df := func(v any) error {
defer cancel()
if err := s.getCodec(stream.ContentSubtype()).Unmarshal(d, v); err != nil {
return status.Errorf(codes.Internal, "grpc: error unmarshalling request: %v", err)
}
for _, sh := range shs {
sh.HandleRPC(ctx, &stats.InPayload{
RecvTime: time.Now(),
Payload: v,
Length: len(d),
Length: d.Len(),
WireLength: payInfo.compressedLength + headerLen,
CompressedLength: payInfo.compressedLength,
Data: d,
})
}
if len(binlogs) != 0 {
cm := &binarylog.ClientMessage{
Message: d,
Message: d.Materialize(),
}
for _, binlog := range binlogs {
binlog.Log(ctx, cm)
@ -1548,7 +1562,7 @@ func (s *Server) processStreamingRPC(ctx context.Context, t transport.ServerTran
ctx: ctx,
t: t,
s: stream,
p: &parser{r: stream, recvBufferPool: s.opts.recvBufferPool},
p: &parser{r: stream, bufferPool: s.opts.bufferPool},
codec: s.getCodec(stream.ContentSubtype()),
maxReceiveMessageSize: s.opts.maxReceiveMessageSize,
maxSendMessageSize: s.opts.maxSendMessageSize,
@ -1963,12 +1977,12 @@ func (s *Server) getCodec(contentSubtype string) baseCodec {
return s.opts.codec
}
if contentSubtype == "" {
return encoding.GetCodec(proto.Name)
return getCodec(proto.Name)
}
codec := encoding.GetCodec(contentSubtype)
codec := getCodec(contentSubtype)
if codec == nil {
logger.Warningf("Unsupported codec %q. Defaulting to %q for now. This will start to fail in future releases.", contentSubtype, proto.Name)
return encoding.GetCodec(proto.Name)
return getCodec(proto.Name)
}
return codec
}

154
shared_buffer_pool.go
View File

@ -1,154 +0,0 @@
/*
*
* Copyright 2023 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package grpc
import "sync"
// SharedBufferPool is a pool of buffers that can be shared, resulting in
// decreased memory allocation. Currently, in gRPC-go, it is only utilized
// for parsing incoming messages.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
type SharedBufferPool interface {
// Get returns a buffer with specified length from the pool.
//
// The returned byte slice may be not zero initialized.
Get(length int) []byte
// Put returns a buffer to the pool.
Put(*[]byte)
}
// NewSharedBufferPool creates a simple SharedBufferPool with buckets
// of different sizes to optimize memory usage. This prevents the pool from
// wasting large amounts of memory, even when handling messages of varying sizes.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func NewSharedBufferPool() SharedBufferPool {
return &simpleSharedBufferPool{
pools: [poolArraySize]simpleSharedBufferChildPool{
newBytesPool(level0PoolMaxSize),
newBytesPool(level1PoolMaxSize),
newBytesPool(level2PoolMaxSize),
newBytesPool(level3PoolMaxSize),
newBytesPool(level4PoolMaxSize),
newBytesPool(0),
},
}
}
// simpleSharedBufferPool is a simple implementation of SharedBufferPool.
type simpleSharedBufferPool struct {
pools [poolArraySize]simpleSharedBufferChildPool
}
func (p *simpleSharedBufferPool) Get(size int) []byte {
return p.pools[p.poolIdx(size)].Get(size)
}
func (p *simpleSharedBufferPool) Put(bs *[]byte) {
p.pools[p.poolIdx(cap(*bs))].Put(bs)
}
func (p *simpleSharedBufferPool) poolIdx(size int) int {
switch {
case size <= level0PoolMaxSize:
return level0PoolIdx
case size <= level1PoolMaxSize:
return level1PoolIdx
case size <= level2PoolMaxSize:
return level2PoolIdx
case size <= level3PoolMaxSize:
return level3PoolIdx
case size <= level4PoolMaxSize:
return level4PoolIdx
default:
return levelMaxPoolIdx
}
}
const (
level0PoolMaxSize = 16 // 16 B
level1PoolMaxSize = level0PoolMaxSize * 16 // 256 B
level2PoolMaxSize = level1PoolMaxSize * 16 // 4 KB
level3PoolMaxSize = level2PoolMaxSize * 16 // 64 KB
level4PoolMaxSize = level3PoolMaxSize * 16 // 1 MB
)
const (
level0PoolIdx = iota
level1PoolIdx
level2PoolIdx
level3PoolIdx
level4PoolIdx
levelMaxPoolIdx
poolArraySize
)
type simpleSharedBufferChildPool interface {
Get(size int) []byte
Put(any)
}
type bufferPool struct {
sync.Pool
defaultSize int
}
func (p *bufferPool) Get(size int) []byte {
bs := p.Pool.Get().(*[]byte)
if cap(*bs) < size {
p.Pool.Put(bs)
return make([]byte, size)
}
return (*bs)[:size]
}
func newBytesPool(size int) simpleSharedBufferChildPool {
return &bufferPool{
Pool: sync.Pool{
New: func() any {
bs := make([]byte, size)
return &bs
},
},
defaultSize: size,
}
}
// nopBufferPool is a buffer pool just makes new buffer without pooling.
type nopBufferPool struct {
}
func (nopBufferPool) Get(length int) []byte {
return make([]byte, length)
}
func (nopBufferPool) Put(*[]byte) {
}

6
stats/stats.go
View File

@ -77,9 +77,6 @@ type InPayload struct {
// the call to HandleRPC which provides the InPayload returns and must be
// copied if needed later.
Payload any
// Data is the serialized message payload.
// Deprecated: Data will be removed in the next release.
Data []byte
// Length is the size of the uncompressed payload data. Does not include any
// framing (gRPC or HTTP/2).
@ -150,9 +147,6 @@ type OutPayload struct {
// the call to HandleRPC which provides the OutPayload returns and must be
// copied if needed later.
Payload any
// Data is the serialized message payload.
// Deprecated: Data will be removed in the next release.
Data []byte
// Length is the size of the uncompressed payload data. Does not include any
// framing (gRPC or HTTP/2).
Length int

98
stats/stats_test.go
View File

@ -28,6 +28,7 @@ import (
"testing"
"time"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/internal"
@ -38,6 +39,7 @@ import (
"google.golang.org/grpc/stats"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/testing/protocmp"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
@ -538,40 +540,29 @@ func checkInPayload(t *testing.T, d *gotData, e *expectedData) {
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
var idx *int
var payloads []proto.Message
if d.client {
b, err := proto.Marshal(e.responses[e.respIdx])
if err != nil {
t.Fatalf("failed to marshal message: %v", err)
}
if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.responses[e.respIdx]) {
t.Fatalf("st.Payload = %T, want %T", st.Payload, e.responses[e.respIdx])
}
e.respIdx++
if string(st.Data) != string(b) {
t.Fatalf("st.Data = %v, want %v", st.Data, b)
}
if st.Length != len(b) {
t.Fatalf("st.Length = %v, want %v", st.Length, len(b))
}
idx = &e.respIdx
payloads = e.responses
} else {
b, err := proto.Marshal(e.requests[e.reqIdx])
if err != nil {
t.Fatalf("failed to marshal message: %v", err)
}
if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.requests[e.reqIdx]) {
t.Fatalf("st.Payload = %T, want %T", st.Payload, e.requests[e.reqIdx])
}
e.reqIdx++
if string(st.Data) != string(b) {
t.Fatalf("st.Data = %v, want %v", st.Data, b)
}
if st.Length != len(b) {
t.Fatalf("st.Length = %v, want %v", st.Length, len(b))
}
idx = &e.reqIdx
payloads = e.requests
}
wantPayload := payloads[*idx]
if diff := cmp.Diff(wantPayload, st.Payload.(proto.Message), protocmp.Transform()); diff != "" {
t.Fatalf("unexpected difference in st.Payload (-want +got):\n%s", diff)
}
*idx++
if st.Length != proto.Size(wantPayload) {
t.Fatalf("st.Length = %v, want %v", st.Length, proto.Size(wantPayload))
}
// Below are sanity checks that WireLength and RecvTime are populated.
// TODO: check values of WireLength and RecvTime.
if len(st.Data) > 0 && st.CompressedLength == 0 {
if st.Length > 0 && st.CompressedLength == 0 {
t.Fatalf("st.WireLength = %v with non-empty data, want <non-zero>",
st.CompressedLength)
}
@ -657,40 +648,29 @@ func checkOutPayload(t *testing.T, d *gotData, e *expectedData) {
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
var idx *int
var payloads []proto.Message
if d.client {
b, err := proto.Marshal(e.requests[e.reqIdx])
if err != nil {
t.Fatalf("failed to marshal message: %v", err)
}
if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.requests[e.reqIdx]) {
t.Fatalf("st.Payload = %T, want %T", st.Payload, e.requests[e.reqIdx])
}
e.reqIdx++
if string(st.Data) != string(b) {
t.Fatalf("st.Data = %v, want %v", st.Data, b)
}
if st.Length != len(b) {
t.Fatalf("st.Length = %v, want %v", st.Length, len(b))
}
idx = &e.reqIdx
payloads = e.requests
} else {
b, err := proto.Marshal(e.responses[e.respIdx])
if err != nil {
t.Fatalf("failed to marshal message: %v", err)
}
if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.responses[e.respIdx]) {
t.Fatalf("st.Payload = %T, want %T", st.Payload, e.responses[e.respIdx])
}
e.respIdx++
if string(st.Data) != string(b) {
t.Fatalf("st.Data = %v, want %v", st.Data, b)
}
if st.Length != len(b) {
t.Fatalf("st.Length = %v, want %v", st.Length, len(b))
}
idx = &e.respIdx
payloads = e.responses
}
// Below are sanity checks that WireLength and SentTime are populated.
expectedPayload := payloads[*idx]
if !proto.Equal(st.Payload.(proto.Message), expectedPayload) {
t.Fatalf("st.Payload = %v, want %v", st.Payload, expectedPayload)
}
*idx++
if st.Length != proto.Size(expectedPayload) {
t.Fatalf("st.Length = %v, want %v", st.Length, proto.Size(expectedPayload))
}
// Below are sanity checks that Length, CompressedLength and SentTime are populated.
// TODO: check values of WireLength and SentTime.
if len(st.Data) > 0 && st.WireLength == 0 {
if st.Length > 0 && st.WireLength == 0 {
t.Fatalf("st.WireLength = %v with non-empty data, want <non-zero>",
st.WireLength)
}

181
stream.go
View File

@ -41,6 +41,7 @@ import (
"google.golang.org/grpc/internal/serviceconfig"
istatus "google.golang.org/grpc/internal/status"
"google.golang.org/grpc/internal/transport"
"google.golang.org/grpc/mem"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/stats"
@ -359,7 +360,7 @@ func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *Client
cs.attempt = a
return nil
}
if err := cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) }); err != nil {
if err := cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op, nil) }); err != nil {
return nil, err
}
@ -517,7 +518,7 @@ func (a *csAttempt) newStream() error {
}
a.s = s
a.ctx = s.Context()
a.p = &parser{r: s, recvBufferPool: a.cs.cc.dopts.recvBufferPool}
a.p = &parser{r: s, bufferPool: a.cs.cc.dopts.copts.BufferPool}
return nil
}
@ -566,10 +567,15 @@ type clientStream struct {
// place where we need to check if the attempt is nil.
attempt *csAttempt
// TODO(hedging): hedging will have multiple attempts simultaneously.
committed bool // active attempt committed for retry?
onCommit func()
buffer []func(a *csAttempt) error // operations to replay on retry
bufferSize int // current size of buffer
committed bool // active attempt committed for retry?
onCommit func()
replayBuffer []replayOp // operations to replay on retry
replayBufferSize int // current size of replayBuffer
}
type replayOp struct {
op func(a *csAttempt) error
cleanup func()
}
// csAttempt implements a single transport stream attempt within a
@ -607,7 +613,12 @@ func (cs *clientStream) commitAttemptLocked() {
cs.onCommit()
}
cs.committed = true
cs.buffer = nil
for _, op := range cs.replayBuffer {
if op.cleanup != nil {
op.cleanup()
}
}
cs.replayBuffer = nil
}
func (cs *clientStream) commitAttempt() {
@ -732,7 +743,7 @@ func (cs *clientStream) retryLocked(attempt *csAttempt, lastErr error) error {
// the stream is canceled.
return err
}
// Note that the first op in the replay buffer always sets cs.attempt
// Note that the first op in replayBuffer always sets cs.attempt
// if it is able to pick a transport and create a stream.
if lastErr = cs.replayBufferLocked(attempt); lastErr == nil {
return nil
@ -761,7 +772,7 @@ func (cs *clientStream) withRetry(op func(a *csAttempt) error, onSuccess func())
// already be status errors.
return toRPCErr(op(cs.attempt))
}
if len(cs.buffer) == 0 {
if len(cs.replayBuffer) == 0 {
// For the first op, which controls creation of the stream and
// assigns cs.attempt, we need to create a new attempt inline
// before executing the first op. On subsequent ops, the attempt
@ -851,25 +862,26 @@ func (cs *clientStream) Trailer() metadata.MD {
}
func (cs *clientStream) replayBufferLocked(attempt *csAttempt) error {
for _, f := range cs.buffer {
if err := f(attempt); err != nil {
for _, f := range cs.replayBuffer {
if err := f.op(attempt); err != nil {
return err
}
}
return nil
}
func (cs *clientStream) bufferForRetryLocked(sz int, op func(a *csAttempt) error) {
func (cs *clientStream) bufferForRetryLocked(sz int, op func(a *csAttempt) error, cleanup func()) {
// Note: we still will buffer if retry is disabled (for transparent retries).
if cs.committed {
return
}
cs.bufferSize += sz
if cs.bufferSize > cs.callInfo.maxRetryRPCBufferSize {
cs.replayBufferSize += sz
if cs.replayBufferSize > cs.callInfo.maxRetryRPCBufferSize {
cs.commitAttemptLocked()
cleanup()
return
}
cs.buffer = append(cs.buffer, op)
cs.replayBuffer = append(cs.replayBuffer, replayOp{op: op, cleanup: cleanup})
}
func (cs *clientStream) SendMsg(m any) (err error) {
@ -891,23 +903,50 @@ func (cs *clientStream) SendMsg(m any) (err error) {
}
// load hdr, payload, data
hdr, payload, data, err := prepareMsg(m, cs.codec, cs.cp, cs.comp)
hdr, data, payload, pf, err := prepareMsg(m, cs.codec, cs.cp, cs.comp, cs.cc.dopts.copts.BufferPool)
if err != nil {
return err
}
defer func() {
data.Free()
// only free payload if compression was made, and therefore it is a different set
// of buffers from data.
if pf.isCompressed() {
payload.Free()
}
}()
dataLen := data.Len()
payloadLen := payload.Len()
// TODO(dfawley): should we be checking len(data) instead?
if len(payload) > *cs.callInfo.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), *cs.callInfo.maxSendMessageSize)
if payloadLen > *cs.callInfo.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", payloadLen, *cs.callInfo.maxSendMessageSize)
}
// always take an extra ref in case data == payload (i.e. when the data isn't
// compressed). The original ref will always be freed by the deferred free above.
payload.Ref()
op := func(a *csAttempt) error {
return a.sendMsg(m, hdr, payload, data)
return a.sendMsg(m, hdr, payload, dataLen, payloadLen)
}
// onSuccess is invoked when the op is captured for a subsequent retry. If the
// stream was established by a previous message and therefore retries are
// disabled, onSuccess will not be invoked, and payloadRef can be freed
// immediately.
onSuccessCalled := false
err = cs.withRetry(op, func() {
cs.bufferForRetryLocked(len(hdr)+payloadLen, op, payload.Free)
onSuccessCalled = true
})
if !onSuccessCalled {
payload.Free()
}
err = cs.withRetry(op, func() { cs.bufferForRetryLocked(len(hdr)+len(payload), op) })
if len(cs.binlogs) != 0 && err == nil {
cm := &binarylog.ClientMessage{
OnClientSide: true,
Message: data,
Message: data.Materialize(),
}
for _, binlog := range cs.binlogs {
binlog.Log(cs.ctx, cm)
@ -924,6 +963,7 @@ func (cs *clientStream) RecvMsg(m any) error {
var recvInfo *payloadInfo
if len(cs.binlogs) != 0 {
recvInfo = &payloadInfo{}
defer recvInfo.free()
}
err := cs.withRetry(func(a *csAttempt) error {
return a.recvMsg(m, recvInfo)
@ -931,7 +971,7 @@ func (cs *clientStream) RecvMsg(m any) error {
if len(cs.binlogs) != 0 && err == nil {
sm := &binarylog.ServerMessage{
OnClientSide: true,
Message: recvInfo.uncompressedBytes,
Message: recvInfo.uncompressedBytes.Materialize(),
}
for _, binlog := range cs.binlogs {
binlog.Log(cs.ctx, sm)
@ -958,7 +998,7 @@ func (cs *clientStream) CloseSend() error {
// RecvMsg. This also matches historical behavior.
return nil
}
cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op) })
cs.withRetry(op, func() { cs.bufferForRetryLocked(0, op, nil) })
if len(cs.binlogs) != 0 {
chc := &binarylog.ClientHalfClose{
OnClientSide: true,
@ -1034,7 +1074,7 @@ func (cs *clientStream) finish(err error) {
cs.cancel()
}
func (a *csAttempt) sendMsg(m any, hdr, payld, data []byte) error {
func (a *csAttempt) sendMsg(m any, hdr []byte, payld mem.BufferSlice, dataLength, payloadLength int) error {
cs := a.cs
if a.trInfo != nil {
a.mu.Lock()
@ -1052,8 +1092,10 @@ func (a *csAttempt) sendMsg(m any, hdr, payld, data []byte) error {
}
return io.EOF
}
for _, sh := range a.statsHandlers {
sh.HandleRPC(a.ctx, outPayload(true, m, data, payld, time.Now()))
if len(a.statsHandlers) != 0 {
for _, sh := range a.statsHandlers {
sh.HandleRPC(a.ctx, outPayload(true, m, dataLength, payloadLength, time.Now()))
}
}
if channelz.IsOn() {
a.t.IncrMsgSent()
@ -1065,6 +1107,7 @@ func (a *csAttempt) recvMsg(m any, payInfo *payloadInfo) (err error) {
cs := a.cs
if len(a.statsHandlers) != 0 && payInfo == nil {
payInfo = &payloadInfo{}
defer payInfo.free()
}
if !a.decompSet {
@ -1102,14 +1145,12 @@ func (a *csAttempt) recvMsg(m any, payInfo *payloadInfo) (err error) {
}
for _, sh := range a.statsHandlers {
sh.HandleRPC(a.ctx, &stats.InPayload{
Client: true,
RecvTime: time.Now(),
Payload: m,
// TODO truncate large payload.
Data: payInfo.uncompressedBytes,
Client: true,
RecvTime: time.Now(),
Payload: m,
WireLength: payInfo.compressedLength + headerLen,
CompressedLength: payInfo.compressedLength,
Length: len(payInfo.uncompressedBytes),
Length: payInfo.uncompressedBytes.Len(),
})
}
if channelz.IsOn() {
@ -1273,7 +1314,7 @@ func newNonRetryClientStream(ctx context.Context, desc *StreamDesc, method strin
return nil, err
}
as.s = s
as.p = &parser{r: s, recvBufferPool: ac.dopts.recvBufferPool}
as.p = &parser{r: s, bufferPool: ac.dopts.copts.BufferPool}
ac.incrCallsStarted()
if desc != unaryStreamDesc {
// Listen on stream context to cleanup when the stream context is
@ -1370,17 +1411,26 @@ func (as *addrConnStream) SendMsg(m any) (err error) {
}
// load hdr, payload, data
hdr, payld, _, err := prepareMsg(m, as.codec, as.cp, as.comp)
hdr, data, payload, pf, err := prepareMsg(m, as.codec, as.cp, as.comp, as.ac.dopts.copts.BufferPool)
if err != nil {
return err
}
defer func() {
data.Free()
// only free payload if compression was made, and therefore it is a different set
// of buffers from data.
if pf.isCompressed() {
payload.Free()
}
}()
// TODO(dfawley): should we be checking len(data) instead?
if len(payld) > *as.callInfo.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payld), *as.callInfo.maxSendMessageSize)
if payload.Len() > *as.callInfo.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", payload.Len(), *as.callInfo.maxSendMessageSize)
}
if err := as.t.Write(as.s, hdr, payld, &transport.Options{Last: !as.desc.ClientStreams}); err != nil {
if err := as.t.Write(as.s, hdr, payload, &transport.Options{Last: !as.desc.ClientStreams}); err != nil {
if !as.desc.ClientStreams {
// For non-client-streaming RPCs, we return nil instead of EOF on error
// because the generated code requires it. finish is not called; RecvMsg()
@ -1639,18 +1689,31 @@ func (ss *serverStream) SendMsg(m any) (err error) {
}
// load hdr, payload, data
hdr, payload, data, err := prepareMsg(m, ss.codec, ss.cp, ss.comp)
hdr, data, payload, pf, err := prepareMsg(m, ss.codec, ss.cp, ss.comp, ss.p.bufferPool)
if err != nil {
return err
}
defer func() {
data.Free()
// only free payload if compression was made, and therefore it is a different set
// of buffers from data.
if pf.isCompressed() {
payload.Free()
}
}()
dataLen := data.Len()
payloadLen := payload.Len()
// TODO(dfawley): should we be checking len(data) instead?
if len(payload) > ss.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", len(payload), ss.maxSendMessageSize)
if payloadLen > ss.maxSendMessageSize {
return status.Errorf(codes.ResourceExhausted, "trying to send message larger than max (%d vs. %d)", payloadLen, ss.maxSendMessageSize)
}
if err := ss.t.Write(ss.s, hdr, payload, &transport.Options{Last: false}); err != nil {
return toRPCErr(err)
}
if len(ss.binlogs) != 0 {
if !ss.serverHeaderBinlogged {
h, _ := ss.s.Header()
@ -1663,7 +1726,7 @@ func (ss *serverStream) SendMsg(m any) (err error) {
}
}
sm := &binarylog.ServerMessage{
Message: data,
Message: data.Materialize(),
}
for _, binlog := range ss.binlogs {
binlog.Log(ss.ctx, sm)
@ -1671,7 +1734,7 @@ func (ss *serverStream) SendMsg(m any) (err error) {
}
if len(ss.statsHandler) != 0 {
for _, sh := range ss.statsHandler {
sh.HandleRPC(ss.s.Context(), outPayload(false, m, data, payload, time.Now()))
sh.HandleRPC(ss.s.Context(), outPayload(false, m, dataLen, payloadLen, time.Now()))
}
}
return nil
@ -1708,6 +1771,7 @@ func (ss *serverStream) RecvMsg(m any) (err error) {
var payInfo *payloadInfo
if len(ss.statsHandler) != 0 || len(ss.binlogs) != 0 {
payInfo = &payloadInfo{}
defer payInfo.free()
}
if err := recv(ss.p, ss.codec, ss.s, ss.dc, m, ss.maxReceiveMessageSize, payInfo, ss.decomp, true); err != nil {
if err == io.EOF {
@ -1727,11 +1791,9 @@ func (ss *serverStream) RecvMsg(m any) (err error) {
if len(ss.statsHandler) != 0 {
for _, sh := range ss.statsHandler {
sh.HandleRPC(ss.s.Context(), &stats.InPayload{
RecvTime: time.Now(),
Payload: m,
// TODO truncate large payload.
Data: payInfo.uncompressedBytes,
Length: len(payInfo.uncompressedBytes),
RecvTime: time.Now(),
Payload: m,
Length: payInfo.uncompressedBytes.Len(),
WireLength: payInfo.compressedLength + headerLen,
CompressedLength: payInfo.compressedLength,
})
@ -1739,7 +1801,7 @@ func (ss *serverStream) RecvMsg(m any) (err error) {
}
if len(ss.binlogs) != 0 {
cm := &binarylog.ClientMessage{
Message: payInfo.uncompressedBytes,
Message: payInfo.uncompressedBytes.Materialize(),
}
for _, binlog := range ss.binlogs {
binlog.Log(ss.ctx, cm)
@ -1754,23 +1816,26 @@ func MethodFromServerStream(stream ServerStream) (string, bool) {
return Method(stream.Context())
}
// prepareMsg returns the hdr, payload and data
// using the compressors passed or using the
// passed preparedmsg
func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor) (hdr, payload, data []byte, err error) {
// prepareMsg returns the hdr, payload and data using the compressors passed or
// using the passed preparedmsg. The returned boolean indicates whether
// compression was made and therefore whether the payload needs to be freed in
// addition to the returned data. Freeing the payload if the returned boolean is
// false can lead to undefined behavior.
func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor, pool mem.BufferPool) (hdr []byte, data, payload mem.BufferSlice, pf payloadFormat, err error) {
if preparedMsg, ok := m.(*PreparedMsg); ok {
return preparedMsg.hdr, preparedMsg.payload, preparedMsg.encodedData, nil
return preparedMsg.hdr, preparedMsg.encodedData, preparedMsg.payload, preparedMsg.pf, nil
}
// The input interface is not a prepared msg.
// Marshal and Compress the data at this point
data, err = encode(codec, m)
if err != nil {
return nil, nil, nil, err
return nil, nil, nil, 0, err
}
compData, err := compress(data, cp, comp)
compData, pf, err := compress(data, cp, comp, pool)
if err != nil {
return nil, nil, nil, err
data.Free()
return nil, nil, nil, 0, err
}
hdr, payload = msgHeader(data, compData)
return hdr, payload, data, nil
hdr, payload = msgHeader(data, compData, pf)
return hdr, data, payload, pf, nil
}

3
test/context_canceled_test.go
View File

@ -158,4 +158,7 @@ func (s) TestCancelWhileRecvingWithCompression(t *testing.T) {
}
}
}
if err := ss.CC.Close(); err != nil {
t.Fatalf("Close failed with %v, want nil", err)
}
}

6
test/retry_test.go
View File

@ -455,7 +455,7 @@ func (s) TestRetryStreaming(t *testing.T) {
time.Sleep(time.Millisecond)
}
for _, tc := range testCases {
for i, tc := range testCases {
func() {
serverOpIter = 0
serverOps = tc.serverOps
@ -464,9 +464,9 @@ func (s) TestRetryStreaming(t *testing.T) {
if err != nil {
t.Fatalf("%v: Error while creating stream: %v", tc.desc, err)
}
for _, op := range tc.clientOps {
for j, op := range tc.clientOps {
if err := op(stream); err != nil {
t.Errorf("%v: %v", tc.desc, err)
t.Errorf("%d %d %v: %v", i, j, tc.desc, err)
break
}
}