perf(p2p/secretconn): Buffer secret connection writes cometbft#3346 (c…

…ometbft#115) * Buffer secret connection writes * Add changelog * Add changelog v2
berachain · Jul 4, 2024 · 5f565c9 · 5f565c9
1 parent 1fd6499
commit 5f565c9
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 21 deletions.
diff --git a/.changelog/unreleased/improvements/3346-buffer-secret-conn-writs.md b/.changelog/unreleased/improvements/3346-buffer-secret-conn-writs.md
@@ -0,0 +1,2 @@
+- `[p2p/secretconn]` Speedup secretconnection large writes, by buffering the write to the underlying connection.
+  ([\#3346](https://github.com/cometbft/cometbft/pull/3346))
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,15 @@
 # CHANGELOG
 
+<<<<<<< HEAD
 ## Unreleased
+=======
+## v0.37.4-v25-osmo-10
+
+* [#115](https://github.com/osmosis-labs/cometbft/pull/115) perf(p2p/secretconn): Buffer secret connection writes (#3346)
+
+
+## v0.37.4-v25-osmo-9
+>>>>>>> 9f773defd ( perf(p2p/secretconn): Buffer secret connection writes #3346  (#115))
 
 *July 1, 2024*
 

diff --git a/p2p/conn/evil_secret_connection_test.go b/p2p/conn/evil_secret_connection_test.go
@@ -1,6 +1,7 @@
 package conn
 
 import (
+	"bufio"
 	"bytes"
 	"errors"
 	"io"
@@ -222,12 +223,14 @@ func (c *evilConn) signChallenge() []byte {
 
 	b := &buffer{}
 	c.secretConn = &SecretConnection{
-		conn:       b,
-		recvBuffer: nil,
-		recvNonce:  new([aeadNonceSize]byte),
-		sendNonce:  new([aeadNonceSize]byte),
-		recvAead:   recvAead,
-		sendAead:   sendAead,
+		underlyingConn: b,
+		connReader:     b,
+		connWriter:     bufio.NewWriterSize(b, 65536),
+		recvBuffer:     nil,
+		recvNonce:      new([aeadNonceSize]byte),
+		sendNonce:      new([aeadNonceSize]byte),
+		recvAead:       recvAead,
+		sendAead:       sendAead,
 	}
 	c.buffer = b
 

diff --git a/p2p/conn/secret_connection.go b/p2p/conn/secret_connection.go
@@ -1,6 +1,7 @@
 package conn
 
 import (
+	"bufio"
 	"bytes"
 	"crypto/cipher"
 	crand "crypto/rand"
@@ -43,6 +44,9 @@ const (
 	labelEphemeralUpperPublicKey = "EPHEMERAL_UPPER_PUBLIC_KEY"
 	labelDHSecret                = "DH_SECRET"
 	labelSecretConnectionMac     = "SECRET_CONNECTION_MAC"
+
+	defaultWriteBufferSize = 1024 * 1024
+	defaultReadBufferSize  = 65536
 )
 
 var (
@@ -66,7 +70,10 @@ type SecretConnection struct {
 	sendAead cipher.AEAD
 
 	remPubKey crypto.PubKey
-	conn      io.ReadWriteCloser
+
+	underlyingConn io.ReadWriteCloser
+	connWriter     *bufio.Writer
+	connReader     io.Reader
 
 	// net.Conn must be thread safe:
 	// https://golang.org/pkg/net/#Conn.
@@ -141,12 +148,14 @@ func MakeSecretConnection(conn io.ReadWriteCloser, locPrivKey crypto.PrivKey) (*
 	}
 
 	sc := &SecretConnection{
-		conn:       conn,
-		recvBuffer: nil,
-		recvNonce:  new([aeadNonceSize]byte),
-		sendNonce:  new([aeadNonceSize]byte),
-		recvAead:   recvAead,
-		sendAead:   sendAead,
+		underlyingConn: conn,
+		connWriter:     bufio.NewWriterSize(conn, defaultWriteBufferSize),
+		connReader:     conn,
+		recvBuffer:     nil,
+		recvNonce:      new([aeadNonceSize]byte),
+		sendNonce:      new([aeadNonceSize]byte),
+		recvAead:       recvAead,
+		sendAead:       sendAead,
 	}
 
 	// Sign the challenge bytes for authentication.
@@ -210,7 +219,7 @@ func (sc *SecretConnection) Write(data []byte) (n int, err error) {
 			incrNonce(sc.sendNonce)
 			// end encryption
 
-			_, err = sc.conn.Write(sealedFrame)
+			_, err = sc.connWriter.Write(sealedFrame)
 			if err != nil {
 				return err
 			}
@@ -220,6 +229,7 @@ func (sc *SecretConnection) Write(data []byte) (n int, err error) {
 			return n, err
 		}
 	}
+	sc.connWriter.Flush()
 	return n, err
 }
 
@@ -238,7 +248,7 @@ func (sc *SecretConnection) Read(data []byte) (n int, err error) {
 	// read off the conn
 	sealedFrame := pool.Get(aeadSizeOverhead + totalFrameSize)
 	defer pool.Put(sealedFrame)
-	_, err = io.ReadFull(sc.conn, sealedFrame)
+	_, err = io.ReadFull(sc.connReader, sealedFrame)
 	if err != nil {
 		return n, err
 	}
@@ -270,16 +280,19 @@ func (sc *SecretConnection) Read(data []byte) (n int, err error) {
 }
 
 // Implements net.Conn.
-func (sc *SecretConnection) Close() error                  { return sc.conn.Close() }
-func (sc *SecretConnection) LocalAddr() net.Addr           { return sc.conn.(net.Conn).LocalAddr() }
-func (sc *SecretConnection) RemoteAddr() net.Addr          { return sc.conn.(net.Conn).RemoteAddr() }
-func (sc *SecretConnection) SetDeadline(t time.Time) error { return sc.conn.(net.Conn).SetDeadline(t) }
+func (sc *SecretConnection) Close() error         { return sc.underlyingConn.Close() }
+func (sc *SecretConnection) LocalAddr() net.Addr  { return sc.underlyingConn.(net.Conn).LocalAddr() }
+func (sc *SecretConnection) RemoteAddr() net.Addr { return sc.underlyingConn.(net.Conn).RemoteAddr() }
+func (sc *SecretConnection) SetDeadline(t time.Time) error {
+	return sc.underlyingConn.(net.Conn).SetDeadline(t)
+}
+
 func (sc *SecretConnection) SetReadDeadline(t time.Time) error {
-	return sc.conn.(net.Conn).SetReadDeadline(t)
+	return sc.underlyingConn.(net.Conn).SetReadDeadline(t)
 }
 
 func (sc *SecretConnection) SetWriteDeadline(t time.Time) error {
-	return sc.conn.(net.Conn).SetWriteDeadline(t)
+	return sc.underlyingConn.(net.Conn).SetWriteDeadline(t)
 }
 
 func genEphKeys() (ephPub, ephPriv *[32]byte) {