alts: Release read buffer when blocked on socket read

[arjan-bal]

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 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. 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:

// 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 gRPC:

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, 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.

[codecov]

Codecov Report

❌ Patch coverage is 75.53191% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.51%. Comparing base (b6597b3) to head (bc023c9).
⚠️ Report is 8 commits behind head on master.

Files with missing lines	Patch %	Lines
credentials/alts/internal/conn/record.go	69.76%	4 Missing and 9 partials ⚠️
internal/transport/ready_reader.go	75.67%	6 Missing and 3 partials ⚠️
internal/mem/buffer_pool.go	87.50%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #8964      +/-   ##
==========================================
- Coverage   83.42%   82.51%   -0.91%     
==========================================
  Files         410      412       +2     
  Lines       32572    32645      +73     
==========================================
- Hits        27172    26938     -234     
- Misses       4030     4077      +47     
- Partials     1370     1630     +260     

Files with missing lines	Coverage Δ
internal/transport/raw_conn_unix.go	100.00% <100.00%> (ø)
internal/mem/buffer_pool.go	94.78% <87.50%> (-1.65%)	⬇️
internal/transport/ready_reader.go	75.67% <75.67%> (ø)
credentials/alts/internal/conn/record.go	67.36% <69.76%> (-14.85%)	⬇️

... and 33 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.