Reset pg_stat_statements when nearly full

input/full.go

@@ -12,7 +12,6 @@ import (
 	"github.com/pganalyze/collector/input/postgres"
 	"github.com/pganalyze/collector/input/system"
 	"github.com/pganalyze/collector/logs"
-	"github.com/pganalyze/collector/scheduler"
 	"github.com/pganalyze/collector/state"
 	"github.com/pganalyze/collector/util"
 )
@@ -116,7 +115,8 @@ func CollectFull(ctx context.Context, server *state.Server, connection *sql.DB,
 		logger.PrintError("Error setting query text timeout: %s", err)
 		return
 	}
-	ts.Statements, ts.StatementTexts, err = postgres.GetStatementTexts(ctx, c, connection)
+	var statementSize int
+	ts.Statements, ts.StatementTexts, statementSize, err = postgres.GetStatementTexts(ctx, c, connection)
 	if err != nil {
 		// Despite query performance data being an essential part of pganalyze, there are
 		// situations where it may not be available (or it timed out), so treat it as a
@@ -151,41 +151,36 @@ func CollectFull(ctx context.Context, server *state.Server, connection *sql.DB,
 		return
 	}
 
-	// Reset query stats and texts if needed (this must run after the query text collection)
-	ps.StatementResetCounter = server.PrevState.StatementResetCounter + 1
-	config := server.Grant.Load().Config
-	if config.Features.StatementResetFrequency != 0 && ps.StatementResetCounter >= int(config.Features.StatementResetFrequency) {
-		// Block concurrent collection of query stats, as that may see the actual Postgres-side
-		// reset before we updated the struct that the collector diffs against.
+	if opts.CollectPostgresSettings {
+		ts.Settings, err = postgres.GetSettings(ctx, connection)
+		if err != nil {
+			logger.PrintError("Error collecting config settings: %s", err)
+			return
+		}
+	}
+
+	shouldReset, err := postgres.ShouldResetStatements(server, &ps, &ts, statementSize)
+	if err != nil {
+		logger.PrintError("Error checking if should reset statements: %s", err)
+		err = nil
+	} else if shouldReset {
 		server.HighFreqStateMutex.Lock()
-		ps.StatementResetCounter = 0
 		err = postgres.ResetStatements(ctx, c, connection)
 		if err != nil {
-			logger.PrintError("Error calling pg_stat_statements_reset() as requested: %s", err)
+			logger.PrintError("Error calling pg_stat_statements_reset(): %s", err)
 			err = nil
-		} else {
-			logger.PrintInfo("Successfully called pg_stat_statements_reset() for all queries, next reset in %d hours", config.Features.StatementResetFrequency/scheduler.FullSnapshotsPerHour)
-
-			// Make sure the next high frequency run has an empty reference point
-			newHighFreqState.LastStatementStatsAt = time.Now()
-			resetStatementStats, err := postgres.GetStatementStats(ctx, c, connection)
-			if err != nil {
-				logger.PrintError("Error collecting pg_stat_statements after reset: %s", err)
-				err = nil
-				newHighFreqState.StatementStats = make(state.PostgresStatementStatsMap)
-			} else {
-				newHighFreqState.StatementStats = resetStatementStats
-			}
 		}
-		server.HighFreqStateMutex.Unlock()
-	}
-
-	if opts.CollectPostgresSettings {
-		ts.Settings, err = postgres.GetSettings(ctx, connection)
+		// Make sure the next high frequency run has an empty reference point
+		newHighFreqState.LastStatementStatsAt = time.Now()
+		resetStatementStats, err := postgres.GetStatementStats(ctx, c, connection)
 		if err != nil {
-			logger.PrintError("Error collecting config settings: %s", err)
-			return
+			logger.PrintError("Error collecting pg_stat_statements after reset: %s", err)
+			err = nil
+			newHighFreqState.StatementStats = make(state.PostgresStatementStatsMap)
+		} else {
+			newHighFreqState.StatementStats = resetStatementStats
 		}
+		server.HighFreqStateMutex.Unlock()
 	}
 
 	// CollectAllSchemas relies on GetBufferCache to access the filenode OIDs before that data is discarded

input/postgres/statements.go

@@ -6,9 +6,12 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/guregu/null"
+	"github.com/pganalyze/collector/scheduler"
 	"github.com/pganalyze/collector/selftest"
 	"github.com/pganalyze/collector/state"
 	"github.com/pganalyze/collector/util"
@@ -59,7 +62,35 @@ func insufficientPrivilege(query string) bool {
 	return query == "<insufficient privilege>"
 }
 
-func ResetStatements(ctx context.Context, c *Collection, db *sql.DB) error {
+func ShouldResetStatements(server *state.Server, ps *state.PersistedState, ts *state.TransientState, size int) (reset bool, err error) {
+	config := server.Grant.Load().Config
+	lastReset := ps.PgStatStatementsStats.Reset
+	resetFreq := config.Features.StatementResetFrequency * scheduler.FullSnapshotMinutes
+	maxSize := int(config.Features.StatementMaxSize)
+	if !lastReset.Valid {
+		return // It's always set on PG14+ with the extension enabled. Older versions aren't supported
+	}
+	if maxSize == 0 {
+		maxSize = 250
+	}
+	count := len(ts.Statements)
+	max := 5_000
+	for _, setting := range ts.Settings {
+		if setting.Name == "pg_stat_statements.max" && setting.CurrentValue.Valid {
+			max, err = strconv.Atoi(setting.CurrentValue.String)
+			if err != nil {
+				return
+			}
+		}
+	}
+	timeElapsed := resetFreq > 0 && time.Since(lastReset.Time).Minutes() >= float64(resetFreq)
+	tooMany := float64(count) >= 0.9*float64(max)
+	tooLarge := size > maxSize*1024*1024
+	reset = timeElapsed && (tooMany || tooLarge)
+	return
+}
+
+func ResetStatements(ctx context.Context, c *Collection, db *sql.DB) (err error) {
 	var method string
 	if c.HelperExists("reset_stat_statements", nil) {
 		c.Logger.PrintVerbose("Found pganalyze.reset_stat_statements() stats helper")
@@ -71,11 +102,9 @@ func ResetStatements(ctx context.Context, c *Collection, db *sql.DB) error {
 		}
 		method = "pg_stat_statements_reset()"
 	}
-	_, err := db.ExecContext(ctx, QueryMarkerSQL+"SELECT "+method)
-	if err != nil {
-		return err
-	}
-	return nil
+	c.Logger.PrintInfo("Resetting pg_stat_statements")
+	_, err = db.ExecContext(ctx, QueryMarkerSQL+"SELECT "+method)
+	return
 }
 
 func GetStatementStats(ctx context.Context, c *Collection, db *sql.DB) (state.PostgresStatementStatsMap, error) {
@@ -152,10 +181,10 @@ func GetStatementStats(ctx context.Context, c *Collection, db *sql.DB) (state.Po
 	return statementStats, nil
 }
 
-func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (state.PostgresStatementMap, state.PostgresStatementTextMap, error) {
+func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (statements state.PostgresStatementMap, statementTextsByFp state.PostgresStatementTextMap, querySize int, err error) {
 	sourceTable, foundExtMinorVersion, err := getStatementSource(ctx, c, db, true)
 	if err != nil {
-		return nil, nil, err
+		return
 	}
 
 	topLevelField := statementSQLTopLevelFieldDefault
@@ -166,28 +195,27 @@ func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (state.Po
 	querySql := QueryMarkerSQL + fmt.Sprintf(statementTextSQL, topLevelField, sourceTable)
 	stmt, err := db.PrepareContext(ctx, querySql)
 	if err != nil {
-		return nil, nil, err
+		return
 	}
 	defer stmt.Close()
 
 	rows, err := stmt.QueryContext(ctx)
 	if err != nil {
-		return nil, nil, err
+		return
 	}
 	defer rows.Close()
 
 	var tmpFile *os.File
 
 	tmpFile, err = os.CreateTemp("", util.TempFilePrefix)
 	if err != nil {
-		return nil, nil, err
+		return
 	}
 	defer tmpFile.Close()
 	defer os.Remove(tmpFile.Name())
 
-	statements := make(state.PostgresStatementMap)
-	statementTextsByFp := make(state.PostgresStatementTextMap)
-
+	statements = make(state.PostgresStatementMap)
+	statementTextsByFp = make(state.PostgresStatementTextMap)
 	queryKeys := make([]state.PostgresStatementKey, 0)
 	queryLengths := make([]int, 0)
 
@@ -198,8 +226,9 @@ func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (state.Po
 
 		err = rows.Scan(&key.DatabaseOid, &key.UserOid, &queryID, &key.TopLevel, &receivedQuery)
 		if err != nil {
-			return nil, nil, err
+			return
 		}
+		querySize += len(receivedQuery.String)
 
 		if queryID.Valid {
 			key.QueryID = queryID.Int64
@@ -214,31 +243,32 @@ func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (state.Po
 	}
 
 	if err = rows.Err(); err != nil {
-		return nil, nil, err
+		return
 	}
 
 	tmpFile.Seek(0, io.SeekStart)
 	for idx, length := range queryLengths {
 		bytes := make([]byte, length)
 		_, err = io.ReadFull(tmpFile, bytes)
 		if err != nil {
-			return nil, nil, err
+			return
 		}
 		query := string(bytes)
 		ignoreIoTiming := ignoreIOTiming(c.PostgresVersion, query)
 		key := queryKeys[idx]
 		select {
 		// Since normalizing can take time, explicitly check for cancellations
 		case <-ctx.Done():
-			return nil, nil, ctx.Err()
+			err = ctx.Err()
+			return
 		default:
 			fingerprintAndNormalize(c, key, key.QueryID, query, statements, statementTextsByFp, ignoreIoTiming)
 		}
 	}
 
 	c.SelfTest.MarkCollectionAspectOk(state.CollectionAspectPgStatStatements)
 
-	return statements, statementTextsByFp, nil
+	return
 }
 
 func getStatementSource(ctx context.Context, c *Collection, db *sql.DB, showtext bool) (string, int16, error) {

output/full.go

@@ -18,6 +18,9 @@ import (
 
 func SendFull(ctx context.Context, server *state.Server, collectionOpts state.CollectionOpts, logger *util.Logger, newState state.PersistedState, diffState state.DiffState, transientState state.TransientState, collectedIntervalSecs uint32) error {
 	s := transform.StateToSnapshot(newState, diffState, transientState, server)
+	if s.ServerStatistic.PgStatStatementsDealloc > 0 {
+		logger.PrintWarning("pg_stat_statements deallocation detected. We recommend enabling automatic resets on the pganalyze server settings page to avoid <query text unavailable>")
+	}
 	s.CollectedIntervalSecs = collectedIntervalSecs
 	err := verifyIntegrity(&s)
 	if err != nil {