Add pg_query fingerprint cache (#768)

Author: seanlinsley

input/postgres/backends.go

@@ -19,12 +19,14 @@ END
 const activitySQL string = `
 SELECT (extract(epoch from COALESCE(backend_start, pg_catalog.pg_postmaster_start_time()))::int::text || pg_catalog.to_char(pid, 'FM0000000'))::bigint,
 	datid, datname, usesysid, usename, pid, application_name, client_addr::text, client_port,
-	backend_start, xact_start, query_start, state_change, COALESCE(wait_event_type, '') = 'Lock' as waiting, backend_xid, backend_xmin, wait_event_type, wait_event, backend_type, %s, state, query
+	backend_start, xact_start, query_start, state_change, COALESCE(wait_event_type, '') = 'Lock' as waiting,
+	backend_xid, backend_xmin, wait_event_type, wait_event, backend_type, %s, state, query, %s
 FROM %s
 WHERE pid IS NOT NULL`
 
 func GetBackends(ctx context.Context, c *Collection, db *sql.DB) ([]state.PostgresBackend, error) {
 	var blockingPidsField string
+	var queryIdField string
 	var sourceTable string
 
 	if c.GlobalOpts.CollectPostgresLocks {
@@ -33,13 +35,19 @@ func GetBackends(ctx context.Context, c *Collection, db *sql.DB) ([]state.Postgr
 		blockingPidsField = "NULL"
 	}
 
+	if c.PostgresVersion.Numeric >= state.PostgresVersion14 {
+		queryIdField = "coalesce(query_id, 0)"
+	} else {
+		queryIdField = "0"
+	}
+
 	if c.HelperExists("get_stat_activity", nil) {
 		sourceTable = "pganalyze.get_stat_activity()"
 	} else {
 		sourceTable = "pg_catalog.pg_stat_activity"
 	}
 
-	stmt, err := db.PrepareContext(ctx, QueryMarkerSQL+fmt.Sprintf(activitySQL, blockingPidsField, sourceTable))
+	stmt, err := db.PrepareContext(ctx, QueryMarkerSQL+fmt.Sprintf(activitySQL, blockingPidsField, queryIdField, sourceTable))
 	if err != nil {
 		return nil, err
 	}
@@ -63,7 +71,7 @@ func GetBackends(ctx context.Context, c *Collection, db *sql.DB) ([]state.Postgr
 			&row.ClientPort, &row.BackendStart, &row.XactStart, &row.QueryStart,
 			&row.StateChange, &row.Waiting, &row.BackendXid, &row.BackendXmin,
 			&row.WaitEventType, &row.WaitEvent, &row.BackendType, pq.Array(&row.BlockedByPids),
-			&row.State, &row.Query)
+			&row.State, &row.Query, &row.QueryId)
 		if err != nil {
 			return nil, err
 		}

input/postgres/collection.go

@@ -23,6 +23,8 @@ type Collection struct {
 
 	// Information that is specific to the current database we're connected to
 	HelperFunctions map[string][]state.PostgresFunction
+
+	Fingerprints *state.Fingerprints
 }
 
 func helpersFromFunctions(functions []state.PostgresFunction) map[string][]state.PostgresFunction {
@@ -81,6 +83,7 @@ func NewCollection(ctx context.Context, logger *util.Logger, server *state.Serve
 		ConnectedAsSuperUser:      connectedAsSuperUser,
 		ConnectedAsMonitoringRole: connectedAsMonitoringRole,
 		HelperFunctions:           helpersFromFunctions(helperFunctions),
+		Fingerprints:              server.Fingerprints,
 	}, nil
 }
 
@@ -95,6 +98,7 @@ func (c *Collection) ForCurrentDatabase(functions []state.PostgresFunction) *Col
 		ConnectedAsSuperUser:      c.ConnectedAsSuperUser,
 		ConnectedAsMonitoringRole: c.ConnectedAsMonitoringRole,
 		HelperFunctions:           helpersFromFunctions(functions),
+		Fingerprints:              c.Fingerprints,
 	}
 }
 

input/postgres/statements.go

@@ -232,7 +232,7 @@ func GetStatementTexts(ctx context.Context, c *Collection, db *sql.DB) (state.Po
 		case <-ctx.Done():
 			return nil, nil, ctx.Err()
 		default:
-			fingerprintAndNormalize(c, key, query, statements, statementTextsByFp, ignoreIoTiming)
+			fingerprintAndNormalize(c, key, key.QueryID, query, statements, statementTextsByFp, ignoreIoTiming)
 		}
 	}
 
@@ -351,7 +351,7 @@ func ignoreIOTiming(postgresVersion state.PostgresVersion, receivedQuery string)
 var collectorQueryFingerprint = util.FingerprintText(util.QueryTextCollector)
 var insufficientPrivsQueryFingerprint = util.FingerprintText(util.QueryTextInsufficientPrivs)
 
-func fingerprintAndNormalize(c *Collection, key state.PostgresStatementKey, text string, statements state.PostgresStatementMap, statementTextsByFp state.PostgresStatementTextMap, ignoreIoTiming bool) {
+func fingerprintAndNormalize(c *Collection, key state.PostgresStatementKey, queryID int64, text string, statements state.PostgresStatementMap, statementTextsByFp state.PostgresStatementTextMap, ignoreIoTiming bool) {
 	if insufficientPrivilege(text) {
 		statements[key] = state.PostgresStatement{
 			InsufficientPrivilege: true,
@@ -365,7 +365,7 @@ func fingerprintAndNormalize(c *Collection, key state.PostgresStatementKey, text
 			IgnoreIoTiming: ignoreIoTiming,
 		}
 	} else {
-		fp := util.FingerprintQuery(text, c.Config.FilterQueryText, -1)
+		fp := c.Fingerprints.LoadOrStore(queryID, text, c.Config.FilterQueryText, -1)
 		statements[key] = state.PostgresStatement{Fingerprint: fp, IgnoreIoTiming: ignoreIoTiming}
 		_, ok := statementTextsByFp[fp]
 		if !ok {

output/transform/activity.go

@@ -35,6 +35,7 @@ func ActivityStateToCompactActivitySnapshot(server *state.Server, activityState
 				b.RoleIdx,
 				b.DatabaseIdx,
 				backend.Query.String,
+				backend.QueryId,
 				activityState.TrackActivityQuerySize,
 			)
 			b.HasQueryIdx = true

output/transform/logs.go

@@ -46,6 +46,7 @@ func transformPostgresQuerySamples(server *state.Server, s snapshot.CompactLogSn
 			roleIdx,
 			databaseIdx,
 			sampleIn.Query,
+			0,
 			-1,
 		)
 
@@ -184,6 +185,7 @@ func transformSystemLogLine(server *state.Server, r *snapshot.CompactSnapshot_Ba
 			logLine.RoleIdx,
 			logLine.DatabaseIdx,
 			logLineIn.Query,
+			0,
 			-1,
 		)
 		logLine.HasQueryIdx = true

output/transform/util.go

@@ -61,8 +61,15 @@ func upsertQueryReferenceAndInformation(s *snapshot.FullSnapshot, statementTexts
 	return idx
 }
 
-func upsertQueryReferenceAndInformationSimple(server *state.Server, refs []*snapshot.QueryReference, infos []*snapshot.QueryInformation, roleIdx int32, databaseIdx int32, originalQuery string, trackActivityQuerySize int) (int32, []*snapshot.QueryReference, []*snapshot.QueryInformation) {
-	fingerprint := util.FingerprintQuery(originalQuery, server.Config.FilterQueryText, trackActivityQuerySize)
+func upsertQueryReferenceAndInformationSimple(server *state.Server, refs []*snapshot.QueryReference, infos []*snapshot.QueryInformation, roleIdx int32, databaseIdx int32, originalQuery string, queryID int64, trackActivityQuerySize int) (int32, []*snapshot.QueryReference, []*snapshot.QueryInformation) {
+	var fingerprint uint64
+	// When the query ID is missing, always fingerprint the query instead of trying to use the fingerprint cache.
+	// It's always zero for query samples from logs, but can also be zero from pg_stat_activity in old Postgres versions.
+	if queryID == 0 {
+		fingerprint = util.FingerprintQuery(originalQuery, server.Config.FilterQueryText, trackActivityQuerySize)
+	} else {
+		fingerprint = server.Fingerprints.LoadOrStore(queryID, originalQuery, server.Config.FilterQueryText, trackActivityQuerySize)
+	}
 
 	fpBuf := make([]byte, 8)
 	binary.BigEndian.PutUint64(fpBuf, fingerprint)

state/fingerprints.go

@@ -0,0 +1,67 @@
+package state
+
+import (
+	"github.com/pganalyze/collector/util"
+	"sync"
+	"sync/atomic"
+)
+
+// This uses around 9 MB per server, as measured by MemStats
+const MAX_SIZE = 450_000
+
+type Fingerprints struct {
+	cache map[int64]uint64
+	lock  sync.RWMutex
+	size  atomic.Int32
+}
+
+func NewFingerprints() *Fingerprints {
+	return &Fingerprints{
+		cache: make(map[int64]uint64, MAX_SIZE),
+		lock:  sync.RWMutex{},
+	}
+}
+
+func (c *Fingerprints) Load(queryID int64) (uint64, bool) {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	fingerprint, exists := c.cache[queryID]
+	return fingerprint, exists
+}
+
+func (c *Fingerprints) LoadOrStore(queryID int64, text string, filterQueryText string, trackActivityQuerySize int) uint64 {
+	fingerprint, exists := c.Load(queryID)
+	if exists {
+		return fingerprint
+	}
+	fingerprint, virtual := util.TryFingerprintQuery(text, filterQueryText, trackActivityQuerySize)
+	if virtual {
+		// Don't store virtual fingerprints so we can cache real fingerprints later
+		return fingerprint
+	}
+	c.cleanup()
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	c.cache[queryID] = fingerprint
+	c.size.Add(1)
+	return fingerprint
+}
+
+// Retains a random 50% sample of entries if the cache grows too large
+func (c *Fingerprints) cleanup() {
+	if c.size.Load() < MAX_SIZE {
+		return
+	}
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	cache := make(map[int64]uint64, MAX_SIZE)
+	index := 0
+	for key, value := range c.cache {
+		if index%2 == 0 {
+			cache[key] = value
+		}
+		index += 1
+	}
+	c.size.Store(int32(len(cache)))
+	c.cache = cache
+}

state/postgres_backend.go

@@ -29,7 +29,8 @@ type PostgresBackend struct {
 
 	BackendType null.String // 10+ The process type of this backend
 
-	Query null.String // Text of this backend's most recent query
+	Query   null.String // Text of this backend's most recent query
+	QueryId int64
 
 	// Current overall state of this backend. Possible values are:
 	// - active: The backend is executing a query.

state/state.go

@@ -331,6 +331,9 @@ type Server struct {
 	// differences (see https://groups.google.com/g/golang-nuts/c/eIqkhXh9PLg),
 	// as we access this in high frequency log-related code paths.
 	LogIgnoreFlags uint32
+
+	// Cache of Postgres query_id -> pg_query fingerprint mappings
+	Fingerprints *Fingerprints
 }
 
 func MakeServer(config config.ServerConfig, testRun bool) *Server {
@@ -347,6 +350,7 @@ func MakeServer(config config.ServerConfig, testRun bool) *Server {
 		QueryRuns:             make(map[int64]*QueryRun),
 		QueryRunsMutex:        &sync.Mutex{},
 		LogParseMutex:         &sync.RWMutex{},
+		Fingerprints:          NewFingerprints(),
 	}
 	server.Grant.Store(&Grant{Config: pganalyze_collector.ServerMessage_Config{Features: &pganalyze_collector.ServerMessage_Features{}}})
 	server.Pause.Store(false)

util/fingerprint.go

@@ -4,10 +4,16 @@ import (
 	pg_query "github.com/pganalyze/pg_query_go/v6"
 )
 
-// FingerprintQuery - Generates a unique fingerprint for the given query
-func FingerprintQuery(query string, filterQueryText string, trackActivityQuerySize int) (fp uint64) {
+// TryFingerprintQuery - Generates a unique fingerprint for the given query,
+// and whether the query text had to be modified to generate a fingerprint.
+//
+// This matters for some callers to decide what to do with the fingerprint, for
+// example if we see a truncated version of a query in pg_stat_activity, and then
+// we see the full query text later via pg_stat_statements.
+func TryFingerprintQuery(query string, filterQueryText string, trackActivityQuerySize int) (fp uint64, virtual bool) {
 	fp, err := pg_query.FingerprintToUInt64(query)
 	if err != nil {
+		virtual = true
 		fixedQuery := fixTruncatedQuery(query)
 
 		fp, err = pg_query.FingerprintToUInt64(fixedQuery)
@@ -20,6 +26,12 @@ func FingerprintQuery(query string, filterQueryText string, trackActivityQuerySi
 	return
 }
 
+// FingerprintQuery - Generates a unique fingerprint for the given query
+func FingerprintQuery(query string, filterQueryText string, trackActivityQuerySize int) (fp uint64) {
+	fp, _ = TryFingerprintQuery(query, filterQueryText, trackActivityQuerySize)
+	return
+}
+
 // FingerprintText - Generates a fingerprint for static texts (used for error scenarios)
 func FingerprintText(query string) (fp uint64) {
 	return pg_query.HashXXH3_64([]byte(query), 0xee)