brontide/noise: take read buffers from pool, return w/ finalizer

brontide/noise.go

@@ -8,12 +8,15 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"runtime"
 	"time"
 
 	"golang.org/x/crypto/chacha20poly1305"
 	"golang.org/x/crypto/hkdf"
 
 	"github.com/btcsuite/btcd/btcec"
+	"github.com/lightningnetwork/lnd/buffer"
+	"github.com/lightningnetwork/lnd/pool"
 )
 
 const (
@@ -57,6 +60,14 @@ var (
 	ephemeralGen = func() (*btcec.PrivateKey, error) {
 		return btcec.NewPrivateKey(btcec.S256())
 	}
+
+	// readBufferPool is a singleton instance of a buffer pool, used to
+	// conserve memory allocations due to read buffers across the entire
+	// brontide package.
+	readBufferPool = pool.NewReadBuffer(
+		pool.DefaultReadBufferGCInterval,
+		pool.DefaultReadBufferExpiryInterval,
+	)
 )
 
 // TODO(roasbeef): free buffer pool?
@@ -375,7 +386,7 @@ type Machine struct {
 	// read the next one. Having a fixed buffer that's re-used also means
 	// that we save on allocations as we don't need to create a new one
 	// each time.
-	nextCipherText [math.MaxUint16 + macSize]byte
+	nextCipherText *buffer.Read
 }
 
 // NewBrontideMachine creates a new instance of the brontide state-machine. If
@@ -739,6 +750,15 @@ func (b *Machine) ReadMessage(r io.Reader) ([]byte, error) {
 		return nil, err
 	}
 
+	// If this is the first message being read, take a read buffer from the
+	// buffer pool. This is delayed until this point to avoid allocating
+	// read buffers until after the peer has successfully completed the
+	// handshake, and is ready to begin sending lnwire messages.
+	if b.nextCipherText == nil {
+		b.nextCipherText = readBufferPool.Take()
+		runtime.SetFinalizer(b, freeReadBuffer)
+	}
+
 	// Next, using the length read from the packet header, read the
 	// encrypted packet itself.
 	pktLen := uint32(binary.BigEndian.Uint16(pktLenBytes)) + macSize
@@ -749,3 +769,12 @@ func (b *Machine) ReadMessage(r io.Reader) ([]byte, error) {
 	// TODO(roasbeef): modify to let pass in slice
 	return b.recvCipher.Decrypt(nil, nil, b.nextCipherText[:pktLen])
 }
+
+// freeReadBuffer returns the Machine's read buffer back to the package wide
+// read buffer pool.
+//
+// NOTE: This method should only be called by a Machine's finalizer.
+func freeReadBuffer(b *Machine) {
+	readBufferPool.Return(b.nextCipherText)
+	b.nextCipherText = nil
+}