Optimize ALTS reads

Author: arjan-bal

credentials/alts/internal/conn/common.go

@@ -54,11 +54,10 @@ func SliceForAppend(in []byte, n int) (head, tail []byte) {
 func ParseFramedMsg(b []byte, maxLen uint32) ([]byte, []byte, error) {
 	// If the size field is not complete, return the provided buffer as
 	// remaining buffer.
-	if len(b) < MsgLenFieldSize {
+	length, sufficientBytes := ParseMessageLength(b)
+	if !sufficientBytes {
 		return nil, b, nil
 	}
-	msgLenField := b[:MsgLenFieldSize]
-	length := binary.LittleEndian.Uint32(msgLenField)
 	if length > maxLen {
 		return nil, nil, fmt.Errorf("received the frame length %d larger than the limit %d", length, maxLen)
 	}
@@ -68,3 +67,15 @@ func ParseFramedMsg(b []byte, maxLen uint32) ([]byte, []byte, error) {
 	}
 	return b[:MsgLenFieldSize+length], b[MsgLenFieldSize+length:], nil
 }
+
+// ParseMessageLength returns the message length based on frame header. It also
+// returns a boolean that indicates if the buffer contains sufficient bytes to
+// parse the length header. If there are insufficient bytes, (0, false) is
+// returned.
+func ParseMessageLength(b []byte) (uint32, bool) {
+	if len(b) < MsgLenFieldSize {
+		return 0, false
+	}
+	msgLenField := b[:MsgLenFieldSize]
+	return binary.LittleEndian.Uint32(msgLenField), true
+}

credentials/alts/internal/conn/record.go

@@ -23,10 +23,12 @@ package conn
 import (
 	"encoding/binary"
 	"fmt"
+	"io"
 	"math"
 	"net"
 
 	core "google.golang.org/grpc/credentials/alts/internal"
+	"google.golang.org/grpc/mem"
 )
 
 // ALTSRecordCrypto is the interface for gRPC ALTS record protocol.
@@ -63,6 +65,8 @@ const (
 	// The maximum write buffer size. This *must* be multiple of
 	// altsRecordDefaultLength.
 	altsWriteBufferMaxSize = 512 * 1024 // 512KiB
+	// The initial buffer used to read from the network.
+	altsReadBufferInitialSize = 32 * 1024 // 32KiB
 )
 
 var (
@@ -84,18 +88,27 @@ type conn struct {
 	crypto ALTSRecordCrypto
 	// buf holds data that has been read from the connection and decrypted,
 	// but has not yet been returned by Read.
-	buf                []byte
+	buf []byte
+	// bufPointer holds the entire decrypted record, even bytes that have
+	// been returned by read. It is used to restore buf to it's initial
+	// capacity after each frame is decrypted. It is also used to return the
+	// buffer to the buffer pool.
+	bufPointer         *[]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
+	// protectedPointer holds a pointer to the protected buffer. It is used to
+	// return the protected buffer to the buffer pool.
+	protectedPointer *[]byte
 	// writeBuf is a buffer used to contain encrypted frames before being
 	// written to the network.
 	writeBuf []byte
 	// nextFrame stores the next frame (in protected buffer) info.
 	nextFrame []byte
 	// overhead is the calculated overhead of each frame.
 	overhead int
+	isClosed bool
 }
 
 // NewConn creates a new secure channel instance given the other party role and
@@ -111,39 +124,43 @@ func NewConn(c net.Conn, side core.Side, recordProtocol string, key []byte, prot
 	}
 	overhead := MsgLenFieldSize + msgTypeFieldSize + crypto.EncryptionOverhead()
 	payloadLengthLimit := altsRecordDefaultLength - overhead
-	var protectedBuf []byte
-	if protected == nil {
-		// We pre-allocate protected to be of size
-		// 2*altsRecordDefaultLength-1 during initialization. We only
-		// read from the network into protected when protected does not
-		// contain a complete frame, which is at most
-		// altsRecordDefaultLength-1 (bytes). And we read at most
-		// altsRecordDefaultLength (bytes) data into protected at one
-		// time. Therefore, 2*altsRecordDefaultLength-1 is large enough
-		// to buffer data read from the network.
-		protectedBuf = make([]byte, 0, 2*altsRecordDefaultLength-1)
-	} else {
-		protectedBuf = make([]byte, len(protected))
-		copy(protectedBuf, protected)
-	}
+	// We pre-allocate protected to be of size 32KB during initialization.
+	// We increase the size of the buffer by the required amount if can't hold a
+	// complete encrypted record.
+	protectedPointer := mem.DefaultBufferPool().Get(max(altsReadBufferInitialSize, len(protected)))
+	protectedBuf := (*protectedPointer)[:copy(*protectedPointer, protected)]
 
 	altsConn := &conn{
 		Conn:               c,
 		crypto:             crypto,
 		payloadLengthLimit: payloadLengthLimit,
+		protectedPointer:   protectedPointer,
 		protected:          protectedBuf,
 		writeBuf:           make([]byte, altsWriteBufferInitialSize),
 		nextFrame:          protectedBuf,
 		overhead:           overhead,
+		bufPointer:         mem.DefaultBufferPool().Get(altsReadBufferInitialSize),
 	}
 	return altsConn, nil
 }
 
+func (p *conn) Close() error {
+	if !p.isClosed {
+		p.isClosed = true
+		mem.DefaultBufferPool().Put(p.protectedPointer)
+		mem.DefaultBufferPool().Put(p.bufPointer)
+	}
+	return p.Conn.Close()
+}
+
 // Read reads and decrypts a frame from the underlying connection, and copies the
 // decrypted payload into b. If the size of the payload is greater than len(b),
 // Read retains the remaining bytes in an internal buffer, and subsequent calls
 // to Read will read from this buffer until it is exhausted.
 func (p *conn) Read(b []byte) (n int, err error) {
+	if p.isClosed {
+		return 0, io.EOF
+	}
 	if len(p.buf) == 0 {
 		var framedMsg []byte
 		framedMsg, p.nextFrame, err = ParseFramedMsg(p.nextFrame, altsRecordLengthLimit)
@@ -153,20 +170,30 @@ func (p *conn) Read(b []byte) (n int, err error) {
 		// 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)]
+			p.protected = p.protected[:copy(p.protected, p.nextFrame)]
 			// Always copy next incomplete frame to the beginning of
 			// the protected buffer and reset nextFrame to it.
 			p.nextFrame = p.protected
 		}
 		// Check whether a complete frame has been received yet.
 		for len(framedMsg) == 0 {
 			if len(p.protected) == cap(p.protected) {
-				tmp := make([]byte, len(p.protected), cap(p.protected)+altsRecordDefaultLength)
-				copy(tmp, p.protected)
-				p.protected = tmp
+				// 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)
+				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))
+				}
+				tmp := mem.DefaultBufferPool().Get(int(length))
+				copy(*tmp, p.protected)
+				p.protected = (*tmp)[:len(p.protected)]
+				mem.DefaultBufferPool().Put(p.protectedPointer)
+				p.protectedPointer = tmp
 			}
-			n, err = p.Conn.Read(p.protected[len(p.protected):min(cap(p.protected), len(p.protected)+altsRecordDefaultLength)])
+			n, err = p.Conn.Read(p.protected[len(p.protected):cap(p.protected)])
 			if err != nil {
 				return 0, err
 			}
@@ -185,17 +212,27 @@ func (p *conn) Read(b []byte) (n int, err error) {
 		}
 		ciphertext := msg[msgTypeFieldSize:]
 
-		// Decrypt requires that if the dst and ciphertext alias, they
-		// must alias exactly. Code here used to use msg[:0], but msg
-		// starts MsgLenFieldSize+msgTypeFieldSize bytes earlier than
-		// ciphertext, so they alias inexactly. Using ciphertext[:0]
-		// arranges the appropriate aliasing without needing to copy
-		// ciphertext or use a separate destination buffer. For more info
-		// check: https://golang.org/pkg/crypto/cipher/#AEAD.
-		p.buf, err = p.crypto.Decrypt(ciphertext[:0], ciphertext)
+		// Decrypt directly into the buffer, avoiding a copy to p.buf if
+		// possible.
+		if cap(b) >= len(msg) {
+			dec, err := p.crypto.Decrypt(b[:0], ciphertext)
+			if err != nil {
+				return 0, err
+			}
+			return len(dec), nil
+		}
+		// Resize the read buffer if needed to hold the entire decrypted
+		// frame.
+		if cap(*p.bufPointer) < len(ciphertext) {
+			mem.DefaultBufferPool().Put(p.bufPointer)
+			p.bufPointer = mem.DefaultBufferPool().Get(len(ciphertext))
+		}
+		p.buf = *p.bufPointer
+		dec, err := p.crypto.Decrypt(p.buf[:0], ciphertext)
 		if err != nil {
 			return 0, err
 		}
+		p.buf = p.buf[:len(dec)]
 	}
 
 	n = copy(b, p.buf)

credentials/alts/internal/conn/record_test.go

@@ -188,6 +188,27 @@ func (s) TestLargeMsg(t *testing.T) {
 	}
 }
 
+// BenchmarkLargeMessage measures the performance of ALTS conns for sending and
+// receiving a large message.
+func BenchmarkLargeMessage(b *testing.B) {
+	msgLen := 20 * 1024 * 1024 // 20 MiB
+	msg := make([]byte, msgLen)
+	rcvMsg := make([]byte, len(msg))
+	b.ResetTimer()
+	clientConn, serverConn := newConnPair(rekeyRecordProtocol, nil, nil)
+	for range b.N {
+		// Write 20 MiB 5 times to transfer a total of 100 MiB.
+		for range 5 {
+			if n, err := clientConn.Write(msg); n != len(msg) || err != nil {
+				b.Fatalf("Write() = %v, %v; want %v, <nil>", n, err, len(msg))
+			}
+			if n, err := io.ReadFull(serverConn, rcvMsg); n != len(rcvMsg) || err != nil {
+				b.Fatalf("Read() = %v, %v; want %v, <nil>", n, err, len(rcvMsg))
+			}
+		}
+	}
+}
+
 func testIncorrectMsgType(t *testing.T, rp string) {
 	// framedMsg is an empty ciphertext with correct framing but wrong
 	// message type.

credentials/alts/internal/handshaker/handshaker.go

@@ -37,6 +37,7 @@ import (
 	altsgrpc "google.golang.org/grpc/credentials/alts/internal/proto/grpc_gcp"
 	altspb "google.golang.org/grpc/credentials/alts/internal/proto/grpc_gcp"
 	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/mem"
 )
 
 const (
@@ -308,6 +309,8 @@ func (h *altsHandshaker) accessHandshakerService(req *altspb.HandshakerReq) (*al
 // whatever received from the network and send it to the handshaker service.
 func (h *altsHandshaker) processUntilDone(resp *altspb.HandshakerResp, extra []byte) (*altspb.HandshakerResult, []byte, error) {
 	var lastWriteTime time.Time
+	buf := mem.DefaultBufferPool().Get(frameLimit)
+	defer mem.DefaultBufferPool().Put(buf)
 	for {
 		if len(resp.OutFrames) > 0 {
 			lastWriteTime = time.Now()
@@ -318,8 +321,7 @@ func (h *altsHandshaker) processUntilDone(resp *altspb.HandshakerResp, extra []b
 		if resp.Result != nil {
 			return resp.Result, extra, nil
 		}
-		buf := make([]byte, frameLimit)
-		n, err := h.conn.Read(buf)
+		n, err := h.conn.Read(*buf)
 		if err != nil && err != io.EOF {
 			return nil, nil, err
 		}
@@ -333,7 +335,7 @@ func (h *altsHandshaker) processUntilDone(resp *altspb.HandshakerResp, extra []b
 		}
 		// Append extra bytes from the previous interaction with the
 		// handshaker service with the current buffer read from conn.
-		p := append(extra, buf[:n]...)
+		p := append(extra, (*buf)[:n]...)
 		// Compute the time elapsed since the last write to the peer.
 		timeElapsed := time.Since(lastWriteTime)
 		timeElapsedMs := uint32(timeElapsed.Milliseconds())