diff --git a/brontide/noise.go b/brontide/noise.go index 3b2191d9..a01d518b 100644 --- a/brontide/noise.go +++ b/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 +}