Skip to content

Commit

Permalink
pgx.Conn: Fix memory leak: Delete items from preparedStatements
Browse files Browse the repository at this point in the history
Previously, items were never removed from the preparedStatements map.
This means workloads that send a large number of unique queries could
run out of memory. Delete items from the map when sending the
deallocate command to Postgres. Add a test to verify this works.

Fixes #1456
  • Loading branch information
evanj authored and jackc committed May 20, 2023
1 parent eab316e commit 0292ede
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
1 change: 1 addition & 0 deletions conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -1326,6 +1326,7 @@ func (c *Conn) deallocateInvalidatedCachedStatements(ctx context.Context) error

for _, sd := range invalidatedStatements {
pipeline.SendDeallocate(sd.Name)
delete(c.preparedStatements, sd.Name)
}

err := pipeline.Sync()
Expand Down
55 changes: 55 additions & 0 deletions conn_internal_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
package pgx

import (
"context"
"fmt"
"os"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func mustParseConfig(t testing.TB, connString string) *ConnConfig {
config, err := ParseConfig(connString)
require.Nil(t, err)
return config
}

func mustConnect(t testing.TB, config *ConnConfig) *Conn {
conn, err := ConnectConfig(context.Background(), config)
if err != nil {
t.Fatalf("Unable to establish connection: %v", err)
}
return conn
}

// Ensures the connection limits the size of its cached objects.
// This test examines the internals of *Conn so must be in the same package.
func TestStmtCacheSizeLimit(t *testing.T) {
const cacheLimit = 16

connConfig := mustParseConfig(t, os.Getenv("PGX_TEST_DATABASE"))
connConfig.StatementCacheCapacity = cacheLimit
conn := mustConnect(t, connConfig)
defer func() {
err := conn.Close(context.Background())
if err != nil {
t.Fatal(err)
}
}()

// run a set of unique queries that should overflow the cache
ctx := context.Background()
for i := 0; i < cacheLimit*2; i++ {
uniqueString := fmt.Sprintf("unique %d", i)
uniqueSQL := fmt.Sprintf("select '%s'", uniqueString)
var output string
err := conn.QueryRow(ctx, uniqueSQL).Scan(&output)
require.NoError(t, err)
require.Equal(t, uniqueString, output)
}
// preparedStatements contains cacheLimit+1 because deallocation happens before the query
assert.Len(t, conn.preparedStatements, cacheLimit+1)
assert.Equal(t, cacheLimit, conn.statementCache.Len())
}

0 comments on commit 0292ede

Please sign in to comment.