From bc754291c135f91559124ad38e0475f172a6a3d2 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Fri, 27 Jan 2023 20:53:30 -0600
Subject: [PATCH] Save memory on non blocking read path

Only create RawConn.Read callback once and have it use NetConn fields.
Avoids the closure and some allocations.

https://github.com/jackc/pgx/issues/1481
---
 internal/nbconn/nbconn.go                |  7 +++++++
 internal/nbconn/nbconn_real_non_block.go | 25 ++++++++++++++++--------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/internal/nbconn/nbconn.go b/internal/nbconn/nbconn.go
index f4419f0f..44daebfb 100644
--- a/internal/nbconn/nbconn.go
+++ b/internal/nbconn/nbconn.go
@@ -79,6 +79,13 @@ type NetConn struct {
 	nonblockWriteErr  error
 	nonblockWriteN    int
 
+	// non-blocking reads with syscall.RawConn are done with a callback function. By using these fields instead of the
+	// callback functions closure to pass the buf argument and receive the n and err results we avoid some allocations.
+	nonblockReadFunc func(fd uintptr) (done bool)
+	nonblockReadBuf  []byte
+	nonblockReadErr  error
+	nonblockReadN    int
+
 	readDeadlineLock sync.Mutex
 	readDeadline     time.Time
 	readNonblocking  bool
diff --git a/internal/nbconn/nbconn_real_non_block.go b/internal/nbconn/nbconn_real_non_block.go
index 00c13adb..6d554769 100644
--- a/internal/nbconn/nbconn_real_non_block.go
+++ b/internal/nbconn/nbconn_real_non_block.go
@@ -24,6 +24,7 @@ func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
 
 	err = c.rawConn.Write(c.nonblockWriteFunc)
 	n = c.nonblockWriteN
+	c.nonblockWriteBuf = nil // ensure that no reference to b is kept.
 	if err == nil && c.nonblockWriteErr != nil {
 		if errors.Is(c.nonblockWriteErr, syscall.EWOULDBLOCK) {
 			err = ErrWouldBlock
@@ -44,16 +45,24 @@ func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
 }
 
 func (c *NetConn) realNonblockingRead(b []byte) (n int, err error) {
-	var funcErr error
-	err = c.rawConn.Read(func(fd uintptr) (done bool) {
-		n, funcErr = syscall.Read(int(fd), b)
-		return true
-	})
-	if err == nil && funcErr != nil {
-		if errors.Is(funcErr, syscall.EWOULDBLOCK) {
+	if c.nonblockReadFunc == nil {
+		c.nonblockReadFunc = func(fd uintptr) (done bool) {
+			c.nonblockReadN, c.nonblockReadErr = syscall.Read(int(fd), c.nonblockReadBuf)
+			return true
+		}
+	}
+	c.nonblockReadBuf = b
+	c.nonblockReadN = 0
+	c.nonblockReadErr = nil
+
+	err = c.rawConn.Read(c.nonblockReadFunc)
+	n = c.nonblockReadN
+	c.nonblockReadBuf = nil // ensure that no reference to b is kept.
+	if err == nil && c.nonblockReadErr != nil {
+		if errors.Is(c.nonblockReadErr, syscall.EWOULDBLOCK) {
 			err = ErrWouldBlock
 		} else {
-			err = funcErr
+			err = c.nonblockReadErr
 		}
 	}
 	if err != nil {