From ba5e3f9638ccafb7a7042bc855eccef756dc7e0c Mon Sep 17 00:00:00 2001 From: "Chao Li (Evan)" Date: Mon, 31 Aug 2026 14:53:41 +0800 Subject: [PATCH v1] pg_stat_statements: preserve normalized query text across reset pg_stat_statements_reset() removes entries from the shared hash table. However, a prepared statement can remain cached and be executed again without going through parse analysis. In that case, pgss_ExecutorEnd() recreates the missing entry without a JumbleState, so the representative query text is taken from the original query string instead of being normalized again. This can make squashed query text such as ``` $1 /*, ... */ ``` reappear as ``` $1, $2, $3 ``` after a reset. Instead of removing entries during reset, reset their statistics and leave them as sticky entries so that cached prepared statements can continue to use the existing normalized query text. The entry for the statement performing the reset is still removed so that the current invocation can establish a fresh entry. If a query with the same query ID is parsed again after the reset, update the sticky entry with the newly normalized representative query text. Also track whether the currently executing statement is top-level so that nested reset statements are handled correctly. Add regression coverage for cached custom and generic plans, reparsing the same query ID after reset, and nested reset statements. Reported-by: Chong Peng Author: Chao Li --- .../pg_stat_statements/expected/squashing.out | 96 ++++++++++ .../pg_stat_statements/pg_stat_statements.c | 170 ++++++++++++++---- contrib/pg_stat_statements/sql/squashing.sql | 40 +++++ 3 files changed, 269 insertions(+), 37 deletions(-) diff --git a/contrib/pg_stat_statements/expected/squashing.out b/contrib/pg_stat_statements/expected/squashing.out index 8438235a2ce..4347b069b91 100644 --- a/contrib/pg_stat_statements/expected/squashing.out +++ b/contrib/pg_stat_statements/expected/squashing.out @@ -129,6 +129,102 @@ SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C"; SELECT pg_stat_statements_reset() IS NOT NULL AS t | 1 (3 rows) +-- Cached prepared statements retain normalized text across a reset +SET plan_cache_mode = force_custom_plan; +SELECT * FROM test_squash WHERE id IN ($1, $2, $3) \parse p1 +\bind_named p1 1 2 3 \g + id | data +----+------ +(0 rows) + +SELECT pg_stat_statements_reset() IS NOT NULL AS t; + t +--- + t +(1 row) + +\bind_named p1 1 2 3 \g + id | data +----+------ +(0 rows) + +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'SELECT * FROM test_squash WHERE id IN%'; + query | calls +-------------------------------------------------------+------- + SELECT * FROM test_squash WHERE id IN ($1 /*, ... */) | 1 +(1 row) + +-- Reparse the statement with the same query ID, changing "IN" to "in" to show +-- that its representative query text is updated +SELECT pg_stat_statements_reset() IS NOT NULL AS t; + t +--- + t +(1 row) + +select * from test_squash where id in ($1, $2, $3) \parse p2 +\bind_named p2 1 2 3 \g + id | data +----+------ +(0 rows) + +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'select * from test_squash where id in%'; + query | calls +-------------------------------------------------------+------- + select * from test_squash where id in ($1 /*, ... */) | 1 +(1 row) + +\close_prepared p2 +SET plan_cache_mode = force_generic_plan; +\bind_named p1 1 2 3 \g + id | data +----+------ +(0 rows) + +SELECT pg_stat_statements_reset() IS NOT NULL AS t; + t +--- + t +(1 row) + +\bind_named p1 1 2 3 \g + id | data +----+------ +(0 rows) + +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'select * from test_squash where id in%'; + query | calls +-------------------------------------------------------+------- + select * from test_squash where id in ($1 /*, ... */) | 1 +(1 row) + +\close_prepared p1 +RESET plan_cache_mode; +-- A nested statement performing a reset stores its current query text +SET pg_stat_statements.track = 'all'; +CREATE FUNCTION pgss_nested_reset() RETURNS void AS $$ +BEGIN + PERFORM pg_stat_statements_reset(0, 0, 0); +END; +$$ LANGUAGE plpgsql; +SELECT pgss_nested_reset(); + pgss_nested_reset +------------------- + +(1 row) + +SELECT query, toplevel FROM pg_stat_statements + WHERE query LIKE 'SELECT pg_stat_statements_reset%'; + query | toplevel +------------------------------------------+---------- + SELECT pg_stat_statements_reset(0, 0, 0) | f +(1 row) + +DROP FUNCTION pgss_nested_reset(); +RESET pg_stat_statements.track; -- prepared statements will also be squashed -- the IN and ARRAY forms of this statement will have the same queryId SELECT pg_stat_statements_reset() IS NOT NULL AS t; diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c index 4562222c9fc..bee60ffc262 100644 --- a/contrib/pg_stat_statements/pg_stat_statements.c +++ b/contrib/pg_stat_statements/pg_stat_statements.c @@ -273,6 +273,8 @@ static const ShmemCallbacks pgss_shmem_callbacks = { /* Current nesting depth of planner/ExecutorRun/ProcessUtility calls */ static int nesting_level = 0; +static int64 active_queryid = INT64CONST(0); +static bool active_toplevel = false; /* Saved hook values */ static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL; @@ -379,6 +381,9 @@ static char *qtext_fetch(Size query_offset, int query_len, static bool need_gc_qtexts(void); static void gc_qtexts(void); static TimestampTz entry_reset(Oid userid, Oid dbid, int64 queryid, bool minmax_only); +static void entry_update_sticky(const char *query, int64 queryid, + int query_location, int query_len, + const JumbleState *jstate); static char *generate_normalized_query(const JumbleState *jstate, const char *query, int query_loc, int *query_len_p); @@ -854,6 +859,17 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, const JumbleState *jst } } + /* + * A sticky entry may preserve old query text for a cached statement after + * a reset. Since this query has been parsed, let it update such an entry + * with a new representative text. + */ + entry_update_sticky(pstate->p_sourcetext, + query->queryId, + query->stmt_location, + query->stmt_len, + jstate); + /* * If query jumbling were able to identify any ignorable constants, we * immediately create a hash table entry for the query, so that we can @@ -1015,6 +1031,11 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags) static void pgss_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count) { + int64 saved_queryid = active_queryid; + bool saved_toplevel = active_toplevel; + + active_queryid = queryDesc->plannedstmt->queryId; + active_toplevel = (nesting_level == 0); nesting_level++; PG_TRY(); { @@ -1026,6 +1047,8 @@ pgss_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count) PG_FINALLY(); { nesting_level--; + active_queryid = saved_queryid; + active_toplevel = saved_toplevel; } PG_END_TRY(); } @@ -1036,6 +1059,11 @@ pgss_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count) static void pgss_ExecutorFinish(QueryDesc *queryDesc) { + int64 saved_queryid = active_queryid; + bool saved_toplevel = active_toplevel; + + active_queryid = queryDesc->plannedstmt->queryId; + active_toplevel = (nesting_level == 0); nesting_level++; PG_TRY(); { @@ -1047,6 +1075,8 @@ pgss_ExecutorFinish(QueryDesc *queryDesc) PG_FINALLY(); { nesting_level--; + active_queryid = saved_queryid; + active_toplevel = saved_toplevel; } PG_END_TRY(); } @@ -2665,12 +2695,105 @@ if (e) { \ } \ else \ { \ - /* Remove the key otherwise */ \ - hash_search(pgss_hash, &e->key, HASH_REMOVE, NULL); \ - num_remove++; \ + if (e->key.userid == GetUserId() && \ + e->key.dbid == MyDatabaseId && \ + e->key.queryid == current_queryid && \ + e->key.toplevel == current_toplevel) \ + { \ + /* Let the statement performing the reset store its current text. */ \ + hash_search(pgss_hash, &e->key, HASH_REMOVE, NULL); \ + } \ + else \ + { \ + /* \ + * Keep the representative query text, so that a cached prepared \ + * statement can reuse it without being parsed again. With no calls, \ + * the entry is sticky and therefore hidden from the view. \ + */ \ + memset(&e->counters, 0, sizeof(Counters)); \ + e->counters.usage = pgss->cur_median_usage; \ + e->stats_since = stats_reset; \ + e->minmax_stats_since = stats_reset; \ + } \ + num_reset++; \ } \ } +/* + * Update the representative query text of a sticky entry when its query is + * parsed again. Cached statements do not reach this function, so after a + * reset they can continue to use the representative query text that was + * saved before it. + */ +static void +entry_update_sticky(const char *query, int64 queryid, + int query_location, int query_len, + const JumbleState *jstate) +{ + pgssHashKey key; + pgssEntry *entry; + char *norm_query = NULL; + bool is_sticky = false; + Size query_offset; + int encoding = GetDatabaseEncoding(); + bool stored; + + if (queryid == INT64CONST(0)) + return; + + query = CleanQuerytext(query, &query_location, &query_len); + + memset(&key, 0, sizeof(pgssHashKey)); + key.userid = GetUserId(); + key.dbid = MyDatabaseId; + key.queryid = queryid; + key.toplevel = (nesting_level == 0); + + /* Avoid taking the exclusive lock in the usual case. */ + LWLockAcquire(&pgss->lock.lock, LW_SHARED); + entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_FIND, NULL); + if (entry) + { + SpinLockAcquire(&entry->mutex); + is_sticky = IS_STICKY(entry->counters); + SpinLockRelease(&entry->mutex); + } + LWLockRelease(&pgss->lock.lock); + + if (!is_sticky) + return; + + if (jstate && jstate->clocations_count > 0) + norm_query = generate_normalized_query(jstate, query, + query_location, &query_len); + + /* + * Update the text while holding the exclusive lock, so that an executor + * cannot observe the entry as missing between parse and execution. + */ + LWLockAcquire(&pgss->lock.lock, LW_EXCLUSIVE); + entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_FIND, NULL); + if (!entry || !IS_STICKY(entry->counters)) + goto done; + + stored = qtext_store(norm_query ? norm_query : query, query_len, + &query_offset, NULL); + if (!stored) + goto done; + + entry->query_offset = query_offset; + entry->query_len = query_len; + entry->encoding = encoding; + + if (need_gc_qtexts()) + gc_qtexts(); + +done: + LWLockRelease(&pgss->lock.lock); + if (norm_query) + pfree(norm_query); +} + /* * Reset entries corresponding to parameters passed. */ @@ -2679,9 +2802,10 @@ entry_reset(Oid userid, Oid dbid, int64 queryid, bool minmax_only) { HASH_SEQ_STATUS hash_seq; pgssEntry *entry; - FILE *qfile; int64 num_entries; - int64 num_remove = 0; + int64 num_reset = 0; + int64 current_queryid = active_queryid; + bool current_toplevel = active_toplevel; pgssHashKey key; TimestampTz stats_reset; @@ -2742,47 +2866,19 @@ entry_reset(Oid userid, Oid dbid, int64 queryid, bool minmax_only) } } - /* All entries are removed? */ - if (num_entries != num_remove) + /* Have all entries been reset? */ + if (num_entries != num_reset) goto release_lock; /* - * Reset global statistics for pg_stat_statements since all entries are - * removed. + * Reset global statistics for pg_stat_statements since all entries have + * been reset. */ SpinLockAcquire(&pgss->mutex); pgss->stats.dealloc = 0; pgss->stats.stats_reset = stats_reset; SpinLockRelease(&pgss->mutex); - /* - * Write new empty query file, perhaps even creating a new one to recover - * if the file was missing. - */ - qfile = AllocateFile(PGSS_TEXT_FILE, PG_BINARY_W); - if (qfile == NULL) - { - ereport(LOG, - (errcode_for_file_access(), - errmsg("could not create file \"%s\": %m", - PGSS_TEXT_FILE))); - goto done; - } - - /* If ftruncate fails, log it, but it's not a fatal problem */ - if (ftruncate(fileno(qfile), 0) != 0) - ereport(LOG, - (errcode_for_file_access(), - errmsg("could not truncate file \"%s\": %m", - PGSS_TEXT_FILE))); - - FreeFile(qfile); - -done: - pgss->extent = 0; - /* This counts as a query text garbage collection for our purposes */ - record_gc_qtexts(); - release_lock: LWLockRelease(&pgss->lock.lock); diff --git a/contrib/pg_stat_statements/sql/squashing.sql b/contrib/pg_stat_statements/sql/squashing.sql index fc9e6573873..22b1e8cff28 100644 --- a/contrib/pg_stat_statements/sql/squashing.sql +++ b/contrib/pg_stat_statements/sql/squashing.sql @@ -39,6 +39,46 @@ SELECT * FROM test_squash WHERE id::text = ANY(ARRAY[$1, $2, $3, $4, $5]) \bind ; SELECT query, calls FROM pg_stat_statements ORDER BY query COLLATE "C"; +-- Cached prepared statements retain normalized text across a reset +SET plan_cache_mode = force_custom_plan; +SELECT * FROM test_squash WHERE id IN ($1, $2, $3) \parse p1 +\bind_named p1 1 2 3 \g +SELECT pg_stat_statements_reset() IS NOT NULL AS t; +\bind_named p1 1 2 3 \g +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'SELECT * FROM test_squash WHERE id IN%'; + +-- Reparse the statement with the same query ID, changing "IN" to "in" to show +-- that its representative query text is updated +SELECT pg_stat_statements_reset() IS NOT NULL AS t; +select * from test_squash where id in ($1, $2, $3) \parse p2 +\bind_named p2 1 2 3 \g +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'select * from test_squash where id in%'; +\close_prepared p2 + +SET plan_cache_mode = force_generic_plan; +\bind_named p1 1 2 3 \g +SELECT pg_stat_statements_reset() IS NOT NULL AS t; +\bind_named p1 1 2 3 \g +SELECT query, calls FROM pg_stat_statements + WHERE query LIKE 'select * from test_squash where id in%'; +\close_prepared p1 +RESET plan_cache_mode; + +-- A nested statement performing a reset stores its current query text +SET pg_stat_statements.track = 'all'; +CREATE FUNCTION pgss_nested_reset() RETURNS void AS $$ +BEGIN + PERFORM pg_stat_statements_reset(0, 0, 0); +END; +$$ LANGUAGE plpgsql; +SELECT pgss_nested_reset(); +SELECT query, toplevel FROM pg_stat_statements + WHERE query LIKE 'SELECT pg_stat_statements_reset%'; +DROP FUNCTION pgss_nested_reset(); +RESET pg_stat_statements.track; + -- prepared statements will also be squashed -- the IN and ARRAY forms of this statement will have the same queryId SELECT pg_stat_statements_reset() IS NOT NULL AS t; -- 2.50.1 (Apple Git-155)