From 2c7d86a54338e69dc33c358b093c8567e1e68a5d Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Fri, 27 Jan 2023 20:34:21 -0600
Subject: [PATCH] Only create RawConn.Write callback once

This saves an allocation on every call.

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

diff --git a/internal/nbconn/nbconn.go b/internal/nbconn/nbconn.go
index eb751509..f4419f0f 100644
--- a/internal/nbconn/nbconn.go
+++ b/internal/nbconn/nbconn.go
@@ -73,10 +73,11 @@ type NetConn struct {
 
 	readFlushLock sync.Mutex
 	// non-blocking writes 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.
-	nonblockWriteBuf []byte
-	nonblockWriteErr error
-	nonblockWriteN   int
+	// callback functions closure to pass the buf argument and receive the n and err results we avoid some allocations.
+	nonblockWriteFunc func(fd uintptr) (done bool)
+	nonblockWriteBuf  []byte
+	nonblockWriteErr  error
+	nonblockWriteN    int
 
 	readDeadlineLock sync.Mutex
 	readDeadline     time.Time
diff --git a/internal/nbconn/nbconn_real_non_block.go b/internal/nbconn/nbconn_real_non_block.go
index ee48d129..00c13adb 100644
--- a/internal/nbconn/nbconn_real_non_block.go
+++ b/internal/nbconn/nbconn_real_non_block.go
@@ -12,13 +12,17 @@ import (
 
 // realNonblockingWrite does a non-blocking write. readFlushLock must already be held.
 func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
+	if c.nonblockWriteFunc == nil {
+		c.nonblockWriteFunc = func(fd uintptr) (done bool) {
+			c.nonblockWriteN, c.nonblockWriteErr = syscall.Write(int(fd), c.nonblockWriteBuf)
+			return true
+		}
+	}
 	c.nonblockWriteBuf = b
 	c.nonblockWriteN = 0
 	c.nonblockWriteErr = nil
-	err = c.rawConn.Write(func(fd uintptr) (done bool) {
-		c.nonblockWriteN, c.nonblockWriteErr = syscall.Write(int(fd), c.nonblockWriteBuf)
-		return true
-	})
+
+	err = c.rawConn.Write(c.nonblockWriteFunc)
 	n = c.nonblockWriteN
 	if err == nil && c.nonblockWriteErr != nil {
 		if errors.Is(c.nonblockWriteErr, syscall.EWOULDBLOCK) {