From 54dc1690bb2ec185c91ef08a02914df1e24cc6a0 Mon Sep 17 00:00:00 2001 From: Artem Gavrilov Date: Thu, 19 Mar 2026 17:54:24 +0200 Subject: [PATCH 1/2] Explicitly reset memory context Attaching memory reset callback to MessageContext is wrong as for background workers that use SPI for query execution such context doesn't exist. There is no other context which lifetime matches our needs, so instead reset our memory context explicitly when we no longer need its data. --- pg_stat_monitor.c | 334 ++++++++++++++++++++++++---------------------- 1 file changed, 174 insertions(+), 160 deletions(-) diff --git a/pg_stat_monitor.c b/pg_stat_monitor.c index dcdee2fb..8db9bc72 100644 --- a/pg_stat_monitor.c +++ b/pg_stat_monitor.c @@ -202,18 +202,9 @@ static void pgsm_add_to_list(pgsmEntry *entry, char *query_text, int query_len); static void pgsm_delete_entry(uint64 queryid); static pgsmEntry *pgsm_get_entry_for_query(int64 queryid, PlanInfo *plan_info, const char *query_text, int query_len, bool create, CmdType cmd_type); static int64 get_pgsm_query_id_hash(const char *norm_query, int len); - -static void pgsm_cleanup_callback(void *arg); +static void pgsm_cleanup_memory(void); static void pgsm_store_error(const char *query, ErrorData *edata); -/*---- Local variables ----*/ -static MemoryContextCallback mem_cxt_reset_callback = -{ - .func = pgsm_cleanup_callback, - .arg = NULL -}; -static volatile bool callback_setup = false; - static void pgsm_update_entry(pgsmEntry *entry, const char *query, char *comments, @@ -389,19 +380,6 @@ pgsm_post_parse_analyze_internal(ParseState *pstate, Query *query, JumbleState * if (!IsSystemInitialized()) return; - if (callback_setup == false) - { - /* - * If MessageContext is valid setup a callback to cleanup our local - * stats list when the MessagContext gets reset - */ - if (MemoryContextIsValid(MessageContext)) - { - MemoryContextRegisterResetCallback(MessageContext, &mem_cxt_reset_callback); - callback_setup = true; - } - } - if (!pgsm_enabled(nesting_level)) return; @@ -757,6 +735,15 @@ pgsm_ExecutorEnd(QueryDesc *queryDesc) pgsm_delete_entry(queryDesc->plannedstmt->queryId); num_relations = 0; + +#if PG_VERSION_NUM >= 170000 + if (nesting_level == 0) +#else + if ((nesting_level + plan_nested_level) == 0) +#endif + { + pgsm_cleanup_memory(); + } } static bool @@ -1013,175 +1000,203 @@ pgsm_ProcessUtility(PlannedStmt *pstmt, const char *queryString, * * Likewise, we don't track execution of DEALLOCATE. */ - if (enabled && - !IsA(parsetree, ExecuteStmt) && - !IsA(parsetree, PrepareStmt) && - !IsA(parsetree, DeallocateStmt)) + PG_TRY(); { - pgsmEntry *entry; - char *query_text; - int location; - int query_len; - instr_time start; - instr_time duration; - uint64 rows; - SysInfo sys_info; - BufferUsage bufusage; - BufferUsage bufusage_start = pgBufferUsage; - WalUsage walusage; - WalUsage walusage_start = pgWalUsage; - - if (getrusage(RUSAGE_SELF, &rusage_start) != 0) - elog(DEBUG1, "[pg_stat_monitor] pgsm_ProcessUtility: Failed to execute getrusage."); - - INSTR_TIME_SET_CURRENT(start); - nesting_level++; - - PG_TRY(); + if (enabled && + !IsA(parsetree, ExecuteStmt) && + !IsA(parsetree, PrepareStmt) && + !IsA(parsetree, DeallocateStmt)) { - if (prev_ProcessUtility) - prev_ProcessUtility(pstmt, queryString, - readOnlyTree, - context, params, queryEnv, - dest, - qc); - else - standard_ProcessUtility(pstmt, queryString, + pgsmEntry *entry; + char *query_text; + int location; + int query_len; + instr_time start; + instr_time duration; + uint64 rows; + SysInfo sys_info; + BufferUsage bufusage; + BufferUsage bufusage_start = pgBufferUsage; + WalUsage walusage; + WalUsage walusage_start = pgWalUsage; + + if (getrusage(RUSAGE_SELF, &rusage_start) != 0) + elog(DEBUG1, "[pg_stat_monitor] pgsm_ProcessUtility: Failed to execute getrusage."); + + INSTR_TIME_SET_CURRENT(start); + nesting_level++; + +#if PG_VERSION_NUM < 160000 + PG_TRY(); +#else + PG_TRY(1); +#endif + { + if (prev_ProcessUtility) + prev_ProcessUtility(pstmt, queryString, readOnlyTree, context, params, queryEnv, dest, qc); - nesting_level--; - } - PG_CATCH(); - { - nesting_level--; - PG_RE_THROW(); - } - - sys_info.utime = 0; - sys_info.stime = 0; + else + standard_ProcessUtility(pstmt, queryString, + readOnlyTree, + context, params, queryEnv, + dest, + qc); + nesting_level--; + } +#if PG_VERSION_NUM < 160000 + PG_CATCH(); +#else + PG_CATCH(1); +#endif + { + nesting_level--; + PG_RE_THROW(); + } - PG_END_TRY(); + sys_info.utime = 0; + sys_info.stime = 0; +#if PG_VERSION_NUM < 160000 + PG_END_TRY(); +#else + PG_END_TRY(1); +#endif - if (getrusage(RUSAGE_SELF, &rusage_end) != 0) - elog(DEBUG1, "[pg_stat_monitor] pgsm_ProcessUtility: Failed to execute getrusage."); - else - { - sys_info.utime = time_diff(rusage_end.ru_utime, rusage_start.ru_utime); - sys_info.stime = time_diff(rusage_end.ru_stime, rusage_start.ru_stime); - } + if (getrusage(RUSAGE_SELF, &rusage_end) != 0) + elog(DEBUG1, "[pg_stat_monitor] pgsm_ProcessUtility: Failed to execute getrusage."); + else + { + sys_info.utime = time_diff(rusage_end.ru_utime, rusage_start.ru_utime); + sys_info.stime = time_diff(rusage_end.ru_stime, rusage_start.ru_stime); + } - INSTR_TIME_SET_CURRENT(duration); - INSTR_TIME_SUBTRACT(duration, start); + INSTR_TIME_SET_CURRENT(duration); + INSTR_TIME_SUBTRACT(duration, start); - rows = (qc && (qc->commandTag == CMDTAG_COPY || - qc->commandTag == CMDTAG_FETCH || - qc->commandTag == CMDTAG_SELECT || - qc->commandTag == CMDTAG_REFRESH_MATERIALIZED_VIEW)) - ? qc->nprocessed - : 0; + rows = (qc && (qc->commandTag == CMDTAG_COPY || + qc->commandTag == CMDTAG_FETCH || + qc->commandTag == CMDTAG_SELECT || + qc->commandTag == CMDTAG_REFRESH_MATERIALIZED_VIEW)) + ? qc->nprocessed + : 0; - /* calc differences of WAL counters. */ - memset(&walusage, 0, sizeof(WalUsage)); - WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start); + /* calc differences of WAL counters. */ + memset(&walusage, 0, sizeof(WalUsage)); + WalUsageAccumDiff(&walusage, &pgWalUsage, &walusage_start); - /* calc differences of buffer counters. */ - memset(&bufusage, 0, sizeof(BufferUsage)); - BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start); + /* calc differences of buffer counters. */ + memset(&bufusage, 0, sizeof(BufferUsage)); + BufferUsageAccumDiff(&bufusage, &pgBufferUsage, &bufusage_start); - /* Create an entry for this query */ - entry = pgsm_create_hash_entry(0, queryId, NULL); + /* Create an entry for this query */ + entry = pgsm_create_hash_entry(0, queryId, NULL); - location = pstmt->stmt_location; - query_len = pstmt->stmt_len; - query_text = (char *) CleanQuerytext(queryString, &location, &query_len); + location = pstmt->stmt_location; + query_len = pstmt->stmt_len; + query_text = (char *) CleanQuerytext(queryString, &location, &query_len); - entry->pgsm_query_id = get_pgsm_query_id_hash(query_text, query_len); - entry->counters.info.cmd_type = pstmt->commandType; + entry->pgsm_query_id = get_pgsm_query_id_hash(query_text, query_len); + entry->counters.info.cmd_type = pstmt->commandType; - pgsm_add_to_list(entry, query_text, query_len); + pgsm_add_to_list(entry, query_text, query_len); - /* Check that we've not exceeded max_stack_depth */ - Assert(list_length(lentries) <= max_stack_depth); + /* Check that we've not exceeded max_stack_depth */ + Assert(list_length(lentries) <= max_stack_depth); - /* The plan details are captured when the query finishes */ - pgsm_update_entry(entry, /* entry */ - (char *) query_text, /* query */ - NULL, /* comments */ - 0, /* comments length */ - NULL, /* PlanInfo */ - &sys_info, /* SysInfo */ - NULL, /* ErrorInfo */ - 0, /* plan_total_time */ - INSTR_TIME_GET_MILLISEC(duration), /* exec_total_time */ - rows, /* rows */ - &bufusage, /* bufusage */ - &walusage, /* walusage */ - NULL, /* jitusage */ - 0, /* parallel_workers_to_launch */ - 0, /* parallel_workers_launched */ - false, /* reset */ - PGSM_EXEC); /* kind */ + /* The plan details are captured when the query finishes */ + pgsm_update_entry(entry, /* entry */ + (char *) query_text, /* query */ + NULL, /* comments */ + 0, /* comments length */ + NULL, /* PlanInfo */ + &sys_info, /* SysInfo */ + NULL, /* ErrorInfo */ + 0, /* plan_total_time */ + INSTR_TIME_GET_MILLISEC(duration), /* exec_total_time */ + rows, /* rows */ + &bufusage, /* bufusage */ + &walusage, /* walusage */ + NULL, /* jitusage */ + 0, /* parallel_workers_to_launch */ + 0, /* parallel_workers_launched */ + false, /* reset */ + PGSM_EXEC); /* kind */ - pgsm_store(entry); - } - else - { - /* - * Even though we're not tracking execution time for this statement, - * we must still increment the nesting level, to ensure that functions - * evaluated within it are not seen as top-level calls. But don't do - * so for EXECUTE; that way, when control reaches pgss_planner or - * pgss_ExecutorStart, we will treat the costs as top-level if - * appropriate. Likewise, don't bump for PREPARE, so that parse - * analysis will treat the statement as top-level if appropriate. - * - * Likewise, we don't track execution of DEALLOCATE. - * - * To be absolutely certain we don't mess up the nesting level, - * evaluate the bump_level condition just once. - */ + pgsm_store(entry); + } + else + { + /* + * Even though we're not tracking execution time for this + * statement, we must still increment the nesting level, to ensure + * that functions evaluated within it are not seen as top-level + * calls. But don't do so for EXECUTE; that way, when control + * reaches pgss_planner or pgss_ExecutorStart, we will treat the + * costs as top-level if appropriate. Likewise, don't bump for + * PREPARE, so that parse analysis will treat the statement as + * top-level if appropriate. + * + * + * Likewise, we don't track execution of DEALLOCATE. + * + * To be absolutely certain we don't mess up the nesting level, + * evaluate the bump_level condition just once. + */ #if PG_VERSION_NUM >= 170000 - bool bump_level = - !IsA(parsetree, ExecuteStmt) && - !IsA(parsetree, PrepareStmt) && - !IsA(parsetree, DeallocateStmt); + bool bump_level = + !IsA(parsetree, ExecuteStmt) && + !IsA(parsetree, PrepareStmt) && + !IsA(parsetree, DeallocateStmt); - if (bump_level) - nesting_level++; + if (bump_level) + nesting_level++; - PG_TRY(); - { + PG_TRY(2); + { #endif - if (prev_ProcessUtility) - prev_ProcessUtility(pstmt, queryString, - readOnlyTree, - context, params, queryEnv, - dest, - qc); - else - standard_ProcessUtility(pstmt, queryString, + if (prev_ProcessUtility) + prev_ProcessUtility(pstmt, queryString, readOnlyTree, context, params, queryEnv, dest, qc); + else + standard_ProcessUtility(pstmt, queryString, + readOnlyTree, + context, params, queryEnv, + dest, + qc); #if PG_VERSION_NUM >= 170000 - if (bump_level) - nesting_level--; + if (bump_level) + nesting_level--; + } + PG_CATCH(2); + { + if (bump_level) + nesting_level--; + PG_RE_THROW(); + } + PG_END_TRY(2); +#endif } - PG_CATCH(); + pgsm_delete_entry(pstmt->queryId); + } + PG_FINALLY(); + { +#if PG_VERSION_NUM >= 170000 + if (nesting_level == 0) +#else + if ((nesting_level + plan_nested_level) == 0) +#endif { - if (bump_level) - nesting_level--; - PG_RE_THROW(); + pgsm_cleanup_memory(); } - PG_END_TRY(); -#endif } - pgsm_delete_entry(pstmt->queryId); + PG_END_TRY(); } /* @@ -1678,13 +1693,12 @@ pgsm_get_entry_for_query(int64 queryid, PlanInfo *plan_info, const char *query_t } static void -pgsm_cleanup_callback(void *arg) +pgsm_cleanup_memory() { /* Reset the memory context holding the list */ MemoryContextReset(GetPgsmMemoryContext()); lentries = NIL; - callback_setup = false; } /* From 34aa683c5283e0a57468b1679e23e40df68d7b5d Mon Sep 17 00:00:00 2001 From: Artem Gavrilov Date: Thu, 19 Mar 2026 17:58:15 +0200 Subject: [PATCH 2/2] Free list entry memory explicitly on delete As we delete entry from the list we should also free its memory, even if we later reset whole memory context. --- pg_stat_monitor.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pg_stat_monitor.c b/pg_stat_monitor.c index 8db9bc72..950cedd3 100644 --- a/pg_stat_monitor.c +++ b/pg_stat_monitor.c @@ -1625,6 +1625,7 @@ pgsm_delete_entry(uint64 queryid) pfree(entry->query_text.query_pointer); entry->query_text.query_pointer = NULL; lentries = list_delete_last(lentries); + pfree(entry); return; } @@ -1642,6 +1643,7 @@ pgsm_delete_entry(uint64 queryid) pfree(entry->query_text.query_pointer); entry->query_text.query_pointer = NULL; lentries = list_delete_cell(lentries, lc); + pfree(entry); return; } }