From 85ddbfeeee841f3ee129fd6f0c167f1dad4c983c Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 4 May 2019 08:55:05 -0500
Subject: [PATCH] Remove 2 mallocs from query path

---
 bench_test.go | 24 ++++++++++++++++++++++++
 conn.go       | 25 +++++++++++++++++++++----
 go.mod        |  1 +
 go.sum        |  2 ++
 4 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/bench_test.go b/bench_test.go
index d212c74f..83c84221 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -13,6 +13,30 @@ import (
 	"github.com/jackc/pgx/v4"
 )
 
+func BenchmarkMinimalPreparedSelect(b *testing.B) {
+	conn := mustConnect(b, mustParseConfig(b, os.Getenv("PGX_TEST_DATABASE")))
+	defer closeConn(b, conn)
+
+	_, err := conn.Prepare(context.Background(), "ps1", "select $1::int8")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	var n int64
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		err = conn.QueryRow(context.Background(), "ps1", int64(i)).Scan(&n)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		if n != int64(i) {
+			b.Fatalf("expected %d, got %d", i, n)
+		}
+	}
+}
+
 func BenchmarkPointerPointerWithNullValues(b *testing.B) {
 	conn := mustConnect(b, mustParseConfig(b, os.Getenv("PGX_TEST_DATABASE")))
 	defer closeConn(b, conn)
diff --git a/conn.go b/conn.go
index 70f608f0..5728c58c 100644
--- a/conn.go
+++ b/conn.go
@@ -34,13 +34,11 @@ type ConnConfig struct {
 // goroutines.
 type Conn struct {
 	pgConn             *pgconn.PgConn
-	wbuf               []byte
 	config             *ConnConfig // config used when establishing this connection
 	preparedStatements map[string]*PreparedStatement
 	logger             Logger
 	logLevel           LogLevel
 	fp                 *fastpath
-	preallocatedRows   []connRows
 
 	causeOfDeath error
 
@@ -48,6 +46,11 @@ type Conn struct {
 	closedChan chan error
 
 	ConnInfo *pgtype.ConnInfo
+
+	wbuf             []byte
+	preallocatedRows []connRows
+	paramFormats     []int16
+	paramValues      [][]byte
 }
 
 // PreparedStatement is a description of a prepared statement
@@ -156,6 +159,8 @@ func connect(ctx context.Context, config *ConnConfig) (c *Conn, err error) {
 	c.doneChan = make(chan struct{})
 	c.closedChan = make(chan error)
 	c.wbuf = make([]byte, 0, 1024)
+	c.paramFormats = make([]int16, 0, 16)
+	c.paramValues = make([][]byte, 0, 16)
 
 	// Replication connections can't execute the queries to
 	// populate the c.PgTypes and c.pgsqlAfInet
@@ -652,8 +657,20 @@ optionLoop:
 		return rows, rows.err
 	}
 
-	paramFormats := make([]int16, len(args))
-	paramValues := make([][]byte, len(args))
+	var paramFormats []int16
+	if len(args) > cap(c.paramFormats) {
+		paramFormats = make([]int16, len(args))
+	} else {
+		paramFormats = c.paramFormats[:len(args)]
+	}
+
+	var paramValues [][]byte
+	if len(args) > cap(c.paramValues) {
+		paramValues = make([][]byte, len(args))
+	} else {
+		paramValues = c.paramValues[:len(args)]
+	}
+
 	for i := range args {
 		paramFormats[i] = chooseParameterFormatCode(c.ConnInfo, ps.ParameterOIDs[i], args[i])
 		paramValues[i], err = newencodePreparedStatementArgument(c.ConnInfo, ps.ParameterOIDs[i], args[i])
diff --git a/go.mod b/go.mod
index 54012172..033beea3 100644
--- a/go.mod
+++ b/go.mod
@@ -15,6 +15,7 @@ require (
 	github.com/shopspring/decimal v0.0.0-20180709203117-cd690d0c9e24
 	github.com/sirupsen/logrus v1.4.1
 	github.com/stretchr/testify v1.3.0
+	go.uber.org/atomic v1.4.0 // indirect
 	go.uber.org/zap v1.9.1
 	golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373
 )
diff --git a/go.sum b/go.sum
index f4bc21aa..1796077e 100644
--- a/go.sum
+++ b/go.sum
@@ -53,6 +53,8 @@ github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 go.uber.org/atomic v1.3.2 h1:2Oa65PReHzfn29GpvgsYwloV9AVFHPDk8tYxt2c2tr4=
 go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/atomic v1.4.0 h1:cxzIVoETapQEqDhQu3QfnvXAV4AlzcvUCxkVUFw3+EU=
+go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/zap v1.9.1 h1:XCJQEf3W6eZaVwhRBof6ImoYGJSITeKWsyeh3HFu/5o=