alts: Release read buffer when blocked on socket read (#8964)
## Problem
Normally, since the Go `net.Conn` interface provides the abstraction of
a blocking read to hide the complexity of non-blocking I/O (epoll,
kqueue), users need to pass a read buffer to the `net.Conn.Read` call.
Even when the TCP socket doesn't have data, the application needs to
hold onto the read buffer. This results in the ALTS conn and the gRPC
HTTP/2 stack holding on to 32KB read buffers each, even for non-readable
transports.
## Solution
On Unix platforms, there is a
[RawConn](https://pkg.go.dev/syscall#RawConn#Read) interface that
exposes a non-blocking mechanism. The idea is that the Go runtime will
call a user-registered callback when the socket is readable. gRPC can
use this callback to allocate a buffer from the pool and return it once
it has passed the plaintext to the HTTP/2 layer. The same RawConn
interface is also available in other OSs, but with slightly different
method signatures for the blocking `Read` method. In the future, we can
add the same optimization for them and have CI runners to catch
regressions.
The main abstraction that allows these non-memory-pinning reads is the
following interface:
```go
// ReadyReader is an optional interface that can be implemented by net.Conn
// implementations to enable gRPC to perform non-memory-pinning reads.
type ReadyReader interface {
// ReadOnReady waits for data to arrive, fetches a buffer, and performs a
// read. It returns a pointer to the buffer so you can return it to the pool
// later.
ReadOnReady(bufSize int, pool mem.BufferPool) (*[]byte, int, error)
}
```
In this PR, an implementation is provided that wraps a `RawConn`. This
allows the ALTS conn to perform efficient reads.
In a future PR, the following changes will enable getting rid of the
bufio.Reader in the HTTP/2 layer:
1. `ReadyReader` will be implemented by the ALTS conn.
2. gRPC will implement its own buffered reader that releases the buffer
when it's empty.
3. The buffered reader will call `ReadOnReady()` instead of `Read()` on
the underlying conn, if supported, to delay the re-allocation of the
buffer.
## Benchmarks
The following micro-benchmarks show no regression in QPS (LargeMessage
test) and a significant reduction in memory usage while performing reads
(ReadMemoryUsage). There is an increase in 2 allocs in conn construction
due to the use of pointer fields for the `ReadyReader` and read buffer
handle, but these happen only when creating a subchannel, not per-RPC.
There is an increase in sec/op for the `WriteMemoryUsage` test, but
these tests are not meant to measure conn construction time, only the
memory effeciency.
```
goos: linux
goarch: amd64
pkg: google.golang.org/grpc/credentials/alts/internal/conn
cpu: Intel(R) Xeon(R) CPU @ 2.60GHz
│ old.txt │ new.txt │
│ sec/op │ sec/op vs base │
LargeMessage-48 77.54m ± 2% 76.51m ± 1% ~ (p=0.202 n=15)
WriteMemoryUsage-48 6.816µ ± 1% 7.519µ ± 1% +10.31% (p=0.000 n=15)
ReadMemoryUsage-48 9.754µ ± 1% 7.021µ ± 1% -28.02% (p=0.000 n=15)
geomean 172.7µ 159.3µ -7.81%
│ old.txt │ new.txt │
│ B/op │ B/op vs base │
LargeMessage-48 4.578Mi ± 0% 4.579Mi ± 7% ~ (p=0.373 n=15)
WriteMemoryUsage-48 41.60Ki ± 0% 41.77Ki ± 0% +0.39% (p=0.000 n=15)
ReadMemoryUsage-48 83.06Ki ± 0% 43.25Ki ± 0% -47.94% (p=0.000 n=15)
geomean 253.0Ki 203.8Ki -19.44%
│ old.txt │ new.txt │
│ allocs/op │ allocs/op vs base │
LargeMessage-48 2.000 ± 0% 2.000 ± 0% ~ (p=1.000 n=15) ¹
WriteMemoryUsage-48 5.000 ± 0% 7.000 ± 0% +40.00% (p=0.000 n=15)
ReadMemoryUsage-48 16.00 ± 0% 18.00 ± 0% +12.50% (p=0.000 n=15)
geomean 5.429 6.316 +16.35%
```
In a [real-world
benchmark](https://github.com/arjan-bal/custom-go-client-benchmark/tree/retry-dp),
where a GCS directpath client downloads a file in a loop, the average
"in use" memory falls from ~35MB to ~20MB.
RELEASE NOTES:
* alts: pool read buffers to lower memory utilization when sockets are
unreadable.diff --git a/credentials/alts/internal/conn/record.go b/credentials/alts/internal/conn/record.go
index c7fcf5e..7fdf349 100644
--- a/credentials/alts/internal/conn/record.go
+++ b/credentials/alts/internal/conn/record.go
@@ -28,6 +28,7 @@
core "google.golang.org/grpc/credentials/alts/internal"
"google.golang.org/grpc/internal/mem"
+ "google.golang.org/grpc/internal/transport"
)
// ALTSRecordCrypto is the interface for gRPC ALTS record protocol.
@@ -75,6 +76,11 @@
var (
protocols = make(map[string]ALTSRecordFunc)
writeBufPool *mem.BinaryTieredBufferPool
+ // readBufPool pools buffers of at least `altsReadBufferInitialSize` size.
+ // Since the read buffer size is slightly larger than 32KB, using a regular
+ // BinaryTieredBufferPool results in allocating buffers of almost double the
+ // required length.
+ readBufPool = mem.NewDirtySimplePool()
)
func init() {
@@ -105,14 +111,18 @@
// conn represents a secured connection. It implements the net.Conn interface.
type conn struct {
net.Conn
+ reader transport.ReadyReader
crypto ALTSRecordCrypto
// buf holds data that has been read from the connection and decrypted,
// but has not yet been returned by Read. It is a sub-slice of protected.
buf []byte
payloadLengthLimit int
- // protected holds data read from the network but have not yet been
- // decrypted. This data might not compose a complete frame.
- protected []byte
+ // protectedHandle buffer holds data read from the network but have not yet
+ // been decrypted. This data might not compose a complete frame.
+ //
+ // The buffer pointer to points to a buffer from the readBufPool. The handle
+ // should only be returned to the pool once nextFrame and buf are empty.
+ protectedHandle *[]byte
// nextFrame stores the next frame (in protected buffer) info.
nextFrame []byte
// overhead is the calculated overhead of each frame.
@@ -135,16 +145,18 @@
// We pre-allocate protected to be of size 32KB during initialization.
// We increase the size of the buffer by the required amount if it can't
// hold a complete encrypted record.
- protectedBuf := make([]byte, max(altsReadBufferInitialSize, len(protected)))
+ protectedHandle := readBufPool.Get(max(altsReadBufferInitialSize, len(protected)))
+ protectedBuf := *protectedHandle
// Copy additional data from hanshaker service.
copy(protectedBuf, protected)
protectedBuf = protectedBuf[:len(protected)]
altsConn := &conn{
Conn: c,
+ reader: transport.NewReadyReader(c),
crypto: crypto,
payloadLengthLimit: payloadLengthLimit,
- protected: protectedBuf,
+ protectedHandle: protectedHandle,
nextFrame: protectedBuf,
overhead: overhead,
}
@@ -158,47 +170,65 @@
func (p *conn) Read(b []byte) (n int, err error) {
if len(p.buf) == 0 {
var framedMsg []byte
+ var protected []byte
+ if p.protectedHandle != nil {
+ protected = *p.protectedHandle
+ protected = protected[:cap(protected)]
+ }
framedMsg, p.nextFrame, err = ParseFramedMsg(p.nextFrame, altsRecordLengthLimit)
if err != nil {
- return n, err
+ return 0, err
}
// Check whether the next frame to be decrypted has been
// completely received yet.
if len(framedMsg) == 0 {
- copy(p.protected, p.nextFrame)
- p.protected = p.protected[:len(p.nextFrame)]
+ copy(protected, p.nextFrame)
+ protected = protected[:len(p.nextFrame)]
// Always copy next incomplete frame to the beginning of
// the protected buffer and reset nextFrame to it.
- p.nextFrame = p.protected
+ p.nextFrame = protected
}
// Check whether a complete frame has been received yet.
for len(framedMsg) == 0 {
- if len(p.protected) == cap(p.protected) {
+ if p.protectedHandle != nil && len(protected) == cap(protected) {
// We can parse the length header to know exactly how large
// the buffer needs to be to hold the entire frame.
- length, didParse := parseMessageLength(p.protected)
+ length, didParse := parseMessageLength(protected)
if !didParse {
// The protected buffer is initialized with a capacity of
// larger than 4B. It should always hold the message length
// header.
- panic(fmt.Sprintf("protected buffer length shorter than expected: %d vs %d", len(p.protected), MsgLenFieldSize))
+ panic(fmt.Sprintf("protected buffer length shorter than expected: %d vs %d", len(protected), MsgLenFieldSize))
}
- oldProtectedBuf := p.protected
+ oldProtectedBuf := protected
+ oldBufHandle := p.protectedHandle
// The new buffer must be able to hold the message length header
// and the entire message.
requiredCapacity := int(length) + MsgLenFieldSize
- p.protected = make([]byte, requiredCapacity)
+ p.protectedHandle = readBufPool.Get(requiredCapacity)
+ protected = *p.protectedHandle
// Copy the contents of the old buffer and set the length of the
// new buffer to the number of bytes already read.
- copy(p.protected, oldProtectedBuf)
- p.protected = p.protected[:len(oldProtectedBuf)]
+ copy(protected, oldProtectedBuf)
+ protected = protected[:len(oldProtectedBuf)]
+ readBufPool.Put(oldBufHandle)
}
- n, err = p.Conn.Read(p.protected[len(p.protected):cap(p.protected)])
- if err != nil {
- return 0, err
+ if p.protectedHandle == nil {
+ // Connection was idle, need to re-allocate the read buffer.
+ newBuf, nRead, err := p.reader.ReadOnReady(altsReadBufferInitialSize, readBufPool)
+ if err != nil {
+ return 0, err
+ }
+ p.protectedHandle = newBuf
+ protected = (*newBuf)[:nRead]
+ } else {
+ nRead, err := p.Conn.Read(protected[len(protected):cap(protected)])
+ if err != nil {
+ return 0, err
+ }
+ protected = protected[:len(protected)+nRead]
}
- p.protected = p.protected[:len(p.protected)+n]
- framedMsg, p.nextFrame, err = ParseFramedMsg(p.protected, altsRecordLengthLimit)
+ framedMsg, p.nextFrame, err = ParseFramedMsg(protected, altsRecordLengthLimit)
if err != nil {
return 0, err
}
@@ -219,6 +249,7 @@
if err != nil {
return 0, err
}
+ p.dropProtectedIfEmtpy()
return len(dec), nil
}
// Decrypt requires that if the dst and ciphertext alias, they
@@ -236,9 +267,23 @@
n = copy(b, p.buf)
p.buf = p.buf[n:]
+ p.dropProtectedIfEmtpy()
return n, nil
}
+func (p *conn) dropProtectedIfEmtpy() {
+ if len(p.buf) > 0 || len(p.nextFrame) > 0 {
+ return
+ }
+ // Potentially idle connection, release the read buffer.
+ p.nextFrame = nil
+ p.buf = nil
+ if p.protectedHandle != nil {
+ readBufPool.Put(p.protectedHandle)
+ p.protectedHandle = nil
+ }
+}
+
// Write encrypts, frames, and writes bytes from b to the underlying connection.
func (p *conn) Write(b []byte) (n int, err error) {
n = len(b)
@@ -257,17 +302,11 @@
writeBuf := *bufHandle
for partialBStart := 0; partialBStart < len(b); partialBStart += partialBSize {
- partialBEnd := partialBStart + partialBSize
- if partialBEnd > len(b) {
- partialBEnd = len(b)
- }
+ partialBEnd := min(partialBStart+partialBSize, len(b))
partialB := b[partialBStart:partialBEnd]
writeBufIndex := 0
for len(partialB) > 0 {
- payloadLen := len(partialB)
- if payloadLen > p.payloadLengthLimit {
- payloadLen = p.payloadLengthLimit
- }
+ payloadLen := min(len(partialB), p.payloadLengthLimit)
buf := partialB[:payloadLen]
partialB = partialB[payloadLen:]
diff --git a/credentials/alts/internal/conn/record_test.go b/credentials/alts/internal/conn/record_test.go
index d9968da..ae5642d 100644
--- a/credentials/alts/internal/conn/record_test.go
+++ b/credentials/alts/internal/conn/record_test.go
@@ -370,16 +370,54 @@
}
}
-// BenchmarkMemoryUsage measures the allocations per ALTS connection.
-// Run this with: go test -bench=BenchmarkMemoryUsage -benchmem
-func BenchmarkMemoryUsage(b *testing.B) {
+// BenchmarkWriteMemoryUsage measures the allocations per ALTS connection.
+// Run this with: go test -bench=BenchmarkWriteMemoryUsage -benchmem
+func BenchmarkWriteMemoryUsage(b *testing.B) {
b.ReportAllocs()
+ data := []byte(strings.Repeat("d", 256))
+ key := []byte{
+ // 16 arbitrary bytes.
+ 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xe2, 0xd2, 0x4c, 0xce, 0x4f, 0x49}
+ conn := &noopConn{}
- for i := 0; i < b.N; i++ {
- c, _ := newConnPair(rekeyRecordProtocol, nil, nil)
+ for b.Loop() {
+ c, err := NewConn(conn, core.ClientSide, rekeyRecordProtocol, key, nil)
+ if err != nil {
+ b.Fatal(err)
+ }
- if _, err := c.Write([]byte("d")); err != nil {
+ if _, err := c.Write(data); err != nil {
b.Fatalf("Write failed: %v", err)
}
}
}
+
+// BenchmarkReadMemoryUsage measures the allocations per ALTS connection.
+// Run this with: go test -bench=BenchmarkReadMemoryUsage -benchmem
+func BenchmarkReadMemoryUsage(b *testing.B) {
+ b.ReportAllocs()
+ data := []byte(strings.Repeat("d", 257))
+ readBuf := make([]byte, 1024)
+
+ for b.Loop() {
+ clientConn, serverConn := newConnPair(rekeyRecordProtocol, nil, nil)
+ b.StopTimer()
+ clientConn.Write(data)
+ b.StartTimer()
+ if _, err := serverConn.Read(readBuf); err != nil {
+ b.Fatalf("Write failed: %v", err)
+ }
+ }
+}
+
+type noopConn struct {
+ net.Conn
+}
+
+func (c *noopConn) Write(b []byte) (n int, err error) {
+ return len(b), nil
+}
+
+func (c *noopConn) Close() error {
+ return nil
+}
diff --git a/internal/mem/buffer_pool.go b/internal/mem/buffer_pool.go
index c2348a8..2d83b2e 100644
--- a/internal/mem/buffer_pool.go
+++ b/internal/mem/buffer_pool.go
@@ -73,7 +73,7 @@
func NewBinaryTieredBufferPool(powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
return newBinaryTiered(func(size int) bufferPool {
return newSizedBufferPool(size, true)
- }, &simpleBufferPool{shouldZero: true}, powerOfTwoExponents...)
+ }, &SimpleBufferPool{shouldZero: true}, powerOfTwoExponents...)
}
// NewDirtyBinaryTieredBufferPool returns a BufferPool backed by multiple
@@ -82,7 +82,7 @@
func NewDirtyBinaryTieredBufferPool(powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
return newBinaryTiered(func(size int) bufferPool {
return newSizedBufferPool(size, false)
- }, &simpleBufferPool{shouldZero: false}, powerOfTwoExponents...)
+ }, NewDirtySimplePool(), powerOfTwoExponents...)
}
func newBinaryTiered(sizedPoolFactory func(int) bufferPool, fallbackPool bufferPool, powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
@@ -258,7 +258,7 @@
// buffer pools for different sizes of buffers.
type TieredBufferPool struct {
sizedPools []*sizedBufferPool
- fallbackPool simpleBufferPool
+ fallbackPool SimpleBufferPool
}
// NewTieredBufferPool returns a BufferPool implementation that uses multiple
@@ -271,7 +271,7 @@
}
return &TieredBufferPool{
sizedPools: pools,
- fallbackPool: simpleBufferPool{shouldZero: true},
+ fallbackPool: SimpleBufferPool{shouldZero: true},
}
}
@@ -297,16 +297,26 @@
return p.sizedPools[poolIdx]
}
-// simpleBufferPool is an implementation of the BufferPool interface that
+// SimpleBufferPool is an implementation of the mem.BufferPool interface that
// attempts to pool buffers with a sync.Pool. When Get is invoked, it tries to
// acquire a buffer from the pool but if that buffer is too small, it returns it
// to the pool and creates a new one.
-type simpleBufferPool struct {
+type SimpleBufferPool struct {
pool sync.Pool
shouldZero bool
}
-func (p *simpleBufferPool) Get(size int) *[]byte {
+// NewDirtySimplePool constructs a [SimpleBufferPool]. It does not initialize
+// the buffers before returning them. Callers must ensure they don't read the
+// buffers before writing data to them.
+func NewDirtySimplePool() *SimpleBufferPool {
+ return &SimpleBufferPool{
+ shouldZero: false,
+ }
+}
+
+// Get returns a buffer with specified length from the pool.
+func (p *SimpleBufferPool) Get(size int) *[]byte {
bs, ok := p.pool.Get().(*[]byte)
if ok && cap(*bs) >= size {
if p.shouldZero {
@@ -333,6 +343,7 @@
return &b
}
-func (p *simpleBufferPool) Put(buf *[]byte) {
+// Put returns a buffer to the pool.
+func (p *SimpleBufferPool) Put(buf *[]byte) {
p.pool.Put(buf)
}
diff --git a/internal/transport/raw_conn_linux.go b/internal/transport/raw_conn_linux.go
new file mode 100644
index 0000000..b897862
--- /dev/null
+++ b/internal/transport/raw_conn_linux.go
@@ -0,0 +1,37 @@
+/*
+ *
+ * Copyright 2026 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 transport
+
+import (
+ "golang.org/x/sys/unix"
+)
+
+func isRawConnSupported() bool {
+ return true
+}
+
+// sysRead uses the modern unix package for Unix-like systems.
+func sysRead(fd uintptr, p []byte) (int, error) {
+ return unix.Read(int(fd), p)
+}
+
+// wouldBlock checks standard Unix non-blocking errors.
+func wouldBlock(err error) bool {
+ return err == unix.EAGAIN || err == unix.EWOULDBLOCK
+}
diff --git a/internal/transport/raw_conn_nonlinux.go b/internal/transport/raw_conn_nonlinux.go
new file mode 100644
index 0000000..4edbb9e
--- /dev/null
+++ b/internal/transport/raw_conn_nonlinux.go
@@ -0,0 +1,35 @@
+//go:build !linux
+
+/*
+ *
+ * Copyright 2026 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 transport
+
+func isRawConnSupported() bool {
+ return false
+}
+
+// sysRead is not implemented. Support can be added in the future if necessary.
+func sysRead(fd uintptr, p []byte) (int, error) {
+ panic("RawConn functionality is not implemented for non-unix platforms.")
+}
+
+// wouldBlock is not implemented. Support can be added in the future if necessary.
+func wouldBlock(err error) bool {
+ panic("RawConn functionality is not implemented for non-unix platforms.")
+}
diff --git a/internal/transport/ready_reader.go b/internal/transport/ready_reader.go
new file mode 100644
index 0000000..5810bbb
--- /dev/null
+++ b/internal/transport/ready_reader.go
@@ -0,0 +1,105 @@
+/*
+ *
+ * Copyright 2026 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 transport
+
+import (
+ "net"
+ "syscall"
+
+ "google.golang.org/grpc/mem"
+)
+
+// ReadyReader is an optional interface that can be implemented by [net.Conn]
+// implementations to enable gRPC to perform non-memory-pinning reads.
+type ReadyReader interface {
+ // ReadOnReady waits for data to arrive, fetches a buffer, and performs a
+ // read. When the underlying IO is readable, it allocates a buffer of size
+ // bufSize from the pool and reads up to bufSize bytes into the buffer.
+ //
+ // It returns a pointer to the buffer so it can be returned to the pool
+ // later, the number of bytes read, and an error.
+ //
+ // Callers should always process the n > 0 bytes returned before considering
+ // the error. Doing so correctly handles I/O errors that happen after
+ // reading some bytes, as well as both of the allowed EOF behaviors.
+ ReadOnReady(bufSize int, pool mem.BufferPool) (b *[]byte, n int, err error)
+}
+
+// nonBlockingReader is optimized for non-memory-pinning reads using the RawConn
+// interface.
+type nonBlockingReader struct {
+ raw syscall.RawConn
+}
+
+func (c *nonBlockingReader) ReadOnReady(bufSize int, pool mem.BufferPool) (buf *[]byte, n int, err error) {
+ var readErr error
+ err = c.raw.Read(func(fd uintptr) bool {
+ buf = pool.Get(bufSize)
+ n, readErr = sysRead(fd, *buf)
+ if readErr != nil {
+ pool.Put(buf)
+ buf = nil
+ }
+
+ if wouldBlock(readErr) {
+ return false // Wait for readiness
+ }
+ return true // Done
+ })
+
+ if err != nil {
+ if buf != nil {
+ pool.Put(buf)
+ }
+ return nil, 0, err
+ }
+ if readErr != nil {
+ // buffer is already released in the callback.
+ return nil, 0, readErr
+ }
+ return buf, n, nil
+}
+
+type blockingReader struct {
+ conn net.Conn
+}
+
+func (c *blockingReader) ReadOnReady(bufSize int, pool mem.BufferPool) (*[]byte, int, error) {
+ buf := pool.Get(bufSize)
+ n, err := c.conn.Read(*buf)
+ if err != nil {
+ pool.Put(buf)
+ return nil, 0, err
+ }
+ return buf, n, nil
+}
+
+// NewReadyReader detects if [syscall.RawConn] is available for
+// non-memory-pinning reads. If [syscall.RawConn] is unavailable, it falls back
+// to using the simpler [net.Conn] interface for reads.
+func NewReadyReader(conn net.Conn) ReadyReader {
+ sysConn, ok := conn.(syscall.Conn)
+ if !ok || !isRawConnSupported() {
+ return &blockingReader{conn: conn}
+ }
+ if raw, err := sysConn.SyscallConn(); err == nil {
+ return &nonBlockingReader{raw: raw}
+ }
+ return &blockingReader{conn: conn}
+}
diff --git a/internal/transport/ready_reader_test.go b/internal/transport/ready_reader_test.go
new file mode 100644
index 0000000..92efc35
--- /dev/null
+++ b/internal/transport/ready_reader_test.go
@@ -0,0 +1,159 @@
+/*
+ *
+ * Copyright 2026 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 transport
+
+import (
+ "bytes"
+ "context"
+ "net"
+ "testing"
+
+ "google.golang.org/grpc/internal/testutils"
+ "google.golang.org/grpc/mem"
+)
+
+// TestReadyReader_NonRawConn verifies that ReadOnReady correctly reads data
+// from a net.Conn that doesn't support non-memory-pinning reads.
+func (s) TestReadyReader_NonRawConn(t *testing.T) {
+ data := []byte("hello world")
+ reader, writer := net.Pipe()
+ go writer.Write(data)
+
+ pool := mem.DefaultBufferPool()
+ readyReader := NewReadyReader(reader)
+
+ bufHandle, n, err := readyReader.ReadOnReady(1024, pool)
+ if err != nil {
+ t.Fatalf("ReadOnReady() failed: %v", err)
+ }
+ defer pool.Put(bufHandle)
+
+ if n != len(data) {
+ t.Errorf("n = %d; want %d", n, len(data))
+ }
+ if !bytes.Equal((*bufHandle)[:n], data) {
+ t.Errorf("Read data = %s; want %s", string((*bufHandle)[:n]), string(data))
+ }
+}
+
+func (s) TestReadyReader_TCP_Blocking(t *testing.T) {
+ ctx, cancel := context.WithTimeout(t.Context(), defaultTestTimeout)
+ defer cancel()
+ ln, err := net.Listen("tcp", "localhost:0")
+ if err != nil {
+ t.Fatalf("net.Listen failed: %v", err)
+ }
+ defer ln.Close()
+
+ data := []byte("hello tcp delayed")
+ connCh := make(chan net.Conn)
+ go func() {
+ conn, err := ln.Accept()
+ if err != nil {
+ return
+ }
+ connCh <- conn
+ }()
+
+ conn, err := net.Dial("tcp", ln.Addr().String())
+ if err != nil {
+ t.Fatalf("net.Dial failed: %v", err)
+ }
+ defer conn.Close()
+
+ serverConn := <-connCh
+ defer serverConn.Close()
+
+ pool := newTrackingPool(mem.DefaultBufferPool())
+ ac := NewReadyReader(conn)
+
+ resCh := make(chan []byte)
+ const readBufSize = 1024
+
+ go func() {
+ bufHandle, n, err := ac.ReadOnReady(readBufSize, pool)
+ if err != nil {
+ t.Errorf("Failed to read: %v", err)
+ return
+ }
+ defer pool.Put(bufHandle)
+ resCh <- (*bufHandle)[:n]
+ }()
+
+ // Verify that no read buffer is allocated for a short while. If it is
+ // allocated (e.g. for a probe), it must be returned immediately.
+ sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+ defer sCancel()
+ if n, err := pool.requestChan.Receive(sCtx); err == nil {
+ if n != readBufSize {
+ t.Fatalf("Unexpected request buffer size: got %d, want %d", n, readBufSize)
+ }
+ if n, err := pool.putChan.Receive(ctx); err != nil {
+ t.Fatal("Read buffer allocated and NOT returned while idle.")
+ } else if n != readBufSize {
+ t.Fatalf("Unexpected returned buffer size: got %d, want %d", n, readBufSize)
+ }
+ }
+
+ // Write data to unblock.
+ serverConn.Write(data)
+
+ if n, err := pool.requestChan.Receive(ctx); err != nil {
+ t.Fatal("Context timed out while waiting for a read buffer to be allocated.")
+ } else if n != readBufSize {
+ t.Fatalf("Unexpected request buffer size: got %d, want %d", n, readBufSize)
+ }
+
+ var res []byte
+ select {
+ case res = <-resCh:
+ case <-ctx.Done():
+ t.Fatal("Context timed out waiting for read to complete.")
+ }
+ if !bytes.Equal(res, data) {
+ t.Errorf("Read data = %s; want %s", string(res), string(data))
+ }
+}
+
+// trackingBufferPool wraps a mem.BufferPool and provides channels to track
+// when buffers are requested and returned, useful for verifying allocation
+// behavior in tests.
+type trackingBufferPool struct {
+ mem.BufferPool
+ requestChan testutils.Channel
+ putChan testutils.Channel
+}
+
+func newTrackingPool(pool mem.BufferPool) *trackingBufferPool {
+ return &trackingBufferPool{
+ BufferPool: pool,
+ requestChan: *testutils.NewChannelWithSize(1),
+ putChan: *testutils.NewChannelWithSize(1),
+ }
+}
+
+func (p *trackingBufferPool) Get(size int) *[]byte {
+ p.requestChan.Replace(size)
+ return p.BufferPool.Get(size)
+}
+
+func (p *trackingBufferPool) Put(b *[]byte) {
+ p.putChan.Replace(len(*b))
+ p.BufferPool.Put(b)
+}
