From 4f2f9faedd05e7d01b9413ee9679f994a24bd86f Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Mon, 19 May 2025 21:37:00 +0200
Subject: [PATCH] Avoid memory context reporting in aborted transactions

Author:
Reviewed-by:
Discussion:
---
 src/backend/utils/mmgr/mcxt.c                 | 17 +++-
 .../test_misc/t/008_memcxt_reporting.pl       | 78 +++++++++++++++++++
 2 files changed, 91 insertions(+), 4 deletions(-)
 create mode 100644 src/test/modules/test_misc/t/008_memcxt_reporting.pl

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 7d28ca706eb..ee80b552a8b 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -21,6 +21,7 @@
 
 #include "postgres.h"
 
+#include "access/xact.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "nodes/pg_list.h"
@@ -1450,6 +1451,16 @@ ProcessGetMemoryContextInterrupt(void)
 
 	PublishMemoryContextPending = false;
 
+	/*
+	 * Avoid performing any shared memory operations in aborted transaction,
+	 * the caller will get the fallback behaviour of the past known stats.
+	 */
+	if (IsAbortedTransactionBlockState())
+	{
+		ConditionVariableBroadcast(&memCxtState[idx].memcxt_cv);
+		return;
+	}
+
 	/*
 	 * The hash table is used for constructing "path" column of the view,
 	 * similar to its local backend counterpart.
@@ -1569,7 +1580,7 @@ ProcessGetMemoryContextInterrupt(void)
 
 	if (summary)
 	{
-		int			cxt_id = 0;
+		int			cxt_id = 1;
 		List	   *path = NIL;
 
 		/* Copy TopMemoryContext statistics to DSA */
@@ -1577,9 +1588,8 @@ ProcessGetMemoryContextInterrupt(void)
 		(*TopMemoryContext->methods->stats) (TopMemoryContext, NULL, NULL,
 											 &stat, true);
 		path = lcons_int(1, path);
-		PublishMemoryContext(meminfo, cxt_id, TopMemoryContext, path, stat,
+		PublishMemoryContext(meminfo, 0, TopMemoryContext, path, stat,
 							 1, MemoryStatsDsaArea, 100);
-		cxt_id = cxt_id + 1;
 
 		/*
 		 * Copy statistics for each of TopMemoryContexts children.  This
@@ -1609,7 +1619,6 @@ ProcessGetMemoryContextInterrupt(void)
 			PublishMemoryContext(meminfo, cxt_id, c, path,
 								 grand_totals, num_contexts, MemoryStatsDsaArea, 100);
 		}
-		memCxtState[idx].total_stats = cxt_id;
 
 		/* Notify waiting backends and return */
 		end_memorycontext_reporting();
diff --git a/src/test/modules/test_misc/t/008_memcxt_reporting.pl b/src/test/modules/test_misc/t/008_memcxt_reporting.pl
new file mode 100644
index 00000000000..f94c3f6374b
--- /dev/null
+++ b/src/test/modules/test_misc/t/008_memcxt_reporting.pl
@@ -0,0 +1,78 @@
+
+# Copyright (c) 2025, PostgreSQL Global Development Group
+
+# Test process memory context reporting, right now we are testing interacting
+# with aborted transactions.
+
+use strict;
+use warnings FATAL => 'all';
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+my $node = PostgreSQL::Test::Cluster->new('node');
+$node->init;
+$node->start;
+
+# Start up a node which we can keep open in an aborted state, and grab the pid
+# from the backend function. If the pid fails to get registered, the tests will
+# still execute but will fail using pid 0.
+my $bsession = $node->interactive_psql('postgres');
+my $bpid = $bsession->query_safe('SELECT pg_backend_pid();');
+my $pid = 0;
+if ($bpid =~ /([0-9]{3,})/)
+{
+	$pid = $1;
+}
+
+# First get a memory context report from the backend in a good state
+my $pre = $node->safe_psql('postgres',
+	"SELECT pg_get_process_memory_contexts($pid,false,0.5);");
+
+# Then set the backend in an aborted state and retry the process, it should
+# return the previous report due to it being aborted.
+$bsession->query('BEGIN;');
+$bsession->query('SELECT 1/0;');
+my $post = $node->safe_psql('postgres',
+	"SELECT pg_get_process_memory_contexts($pid,false,0.5);");
+ok($pre eq $post, "Check that aborted backend returned previous stats");
+
+# There should be no errors in the log and the backend should still be in an
+# aborted state.
+$node->log_check("Check for backend exit", 0, log_unlike => [qr/FATAL/]);
+ok($bsession->quit, "Check that aborted backend hadn't terminated");
+
+# Make sure the backend has exited and cannot be queried anymore
+my ($stdout, $stderr, $timed_out);
+$post = $node->psql(
+	'postgres', "SELECT pg_get_process_memory_contexts($pid,false,0.5);",
+	stdout => \$stdout,
+	stderr => \$stderr);
+ok(!$post, "Check that we didn't get any results");
+ok($stderr =~ /is not a PostgreSQL server process/,
+	"Check STDERR for WARNING");
+
+# Restart and again set the backend in aborted state, this time without any
+# memory context report from it before aborting.
+$bsession = $node->interactive_psql('postgres');
+$bpid = $bsession->query_safe('SELECT pg_backend_pid();');
+$pid = 0;
+if ($bpid =~ /([0-9]{3,})/)
+{
+	$pid = $1;
+}
+$bsession->query('BEGIN;');
+$bsession->query('SELECT 1/0;');
+$post = $node->safe_psql('postgres',
+	"SELECT pg_get_process_memory_contexts($pid,false,0.5);");
+
+# Since there was no fallback we should have a NULL return and not a report.§
+ok( $post eq '',
+	"Check that an aborted backend returns NULL when no previous stats exist"
+);
+ok($bsession->quit, "Check that aborted backend hadn't terminated");
+
+$node->stop;
+
+done_testing();
-- 
2.39.3 (Apple Git-146)