From b537ccf2b5f56eedcb22f066e00a6d61aececa9e Mon Sep 17 00:00:00 2001
From: Rahila Syed <rahilasyed.90@gmail.com>
Date: Thu, 9 Oct 2025 12:07:20 +0530
Subject: [PATCH] Test fixes

---
 src/backend/utils/adt/mcxtfuncs.c             |  2 +
 src/test/modules/Makefile                     |  1 +
 .../test_memcontext_reporting/Makefile        | 32 +++++++++
 .../t/001_memcontext_inj.pl                   | 45 ++++++++++++
 .../test_memcontext_reporting--1.0.sql        |  7 ++
 .../test_memcontext_reporting.c               | 68 +++++++++++++++++++
 .../test_memcontext_reporting.control         |  4 ++
 7 files changed, 159 insertions(+)
 create mode 100644 src/test/modules/test_memcontext_reporting/Makefile
 create mode 100644 src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
 create mode 100644 src/test/modules/test_memcontext_reporting/test_memcontext_reporting--1.0.sql
 create mode 100644 src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c
 create mode 100644 src/test/modules/test_memcontext_reporting/test_memcontext_reporting.control

diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index aa4e0a2e670..8c3d3f321aa 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -27,6 +27,7 @@
 #include "utils/array.h"
 #include "utils/builtins.h"
 #include "utils/hsearch.h"
+#include "utils/injection_point.h"
 #include "utils/memutils.h"
 #include "utils/wait_event_types.h"
 
@@ -541,6 +542,7 @@ pg_get_process_memory_contexts(PG_FUNCTION_ARGS)
 			PG_RETURN_NULL();
 		}
 
+		INJECTION_POINT("memcontext-client-crash", NULL);
 		while (1)
 		{
 			entry = dshash_find_or_insert(MemoryStatsDsHash, key, &found);
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 902a7954101..a31a2578c18 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -31,6 +31,7 @@ SUBDIRS = \
 		  test_json_parser \
 		  test_lfind \
 		  test_lwlock_tranches \
+		  test_memcontext_reporting \
 		  test_misc \
 		  test_oat_hooks \
 		  test_parser \
diff --git a/src/test/modules/test_memcontext_reporting/Makefile b/src/test/modules/test_memcontext_reporting/Makefile
new file mode 100644
index 00000000000..01a7baa0263
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/Makefile
@@ -0,0 +1,32 @@
+# src/test/modules/test_memcontext_reporting/Makefile
+
+EXTRA_INSTALL = src/test/modules/injection_points
+
+export enable_injection_points
+MODULE_big = test_memcontext_reporting
+OBJS = \
+	$(WIN32RES) \
+	test_memcontext_reporting.o
+PGFILEDESC = "test_memcontext_reporting - test code for memory context reporting"
+
+EXTENSION = test_memcontext_reporting
+DATA = test_memcontext_reporting--1.0.sql
+
+REGRESS = test_memcontext_reporting
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_memcontext_reporting
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
+
+check:
+	$(prove_check)
+
+installcheck:
+	$(prove_installcheck)
diff --git a/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl b/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
new file mode 100644
index 00000000000..842f32376fd
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
@@ -0,0 +1,45 @@
+# Copyright (c) 2025, PostgreSQL Global Development Group
+
+# Test suite for testing enabling data checksums in an online cluster,
+# comprising of a primary and a replicated standby, with concurrent activity
+# via pgbench runs
+
+use strict;
+use warnings FATAL => 'all';
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+if ($ENV{enable_injection_points} ne 'yes')
+{
+       plan skip_all => 'Injection points not supported by this build';
+}
+
+# Create and start a cluster with one node
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init(allows_streaming => 1, no_data_checksums => 1);
+# max_connections need to be bumped in order to accommodate for pgbench clients
+# and log_statement is dialled down since it otherwise will generate enormous
+# amounts of logging. Page verification failures are still logged.
+$node->append_conf(
+       'postgresql.conf',
+       qq[
+max_connections = 100
+log_statement = none
+]);
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION test_memcontext_reporting;');
+# Attaching to an injection point that crashes memory context client
+$node->safe_psql('postgres', 'SELECT memcontext_crash_client();');
+
+my $pid = $node->safe_psql('postgres', "SELECT pid from pg_stat_activity where backend_type='checkpointer'");
+
+#Client should have crashed
+$node->safe_psql('postgres', "select pg_get_process_memory_contexts($pid, true, 5);");
+print "PID";
+print $pid;
+#Query the same process for memory context using some other client and it should succeed.
+my $topcontext_name = $node->safe_psql('postgres', "select name from pg_get_process_memory_contexts($pid, true, 5) where path = '{1}';");
+ok($topcontext_name = 'TopMemoryContext');
+done_testing();
diff --git a/src/test/modules/test_memcontext_reporting/test_memcontext_reporting--1.0.sql b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting--1.0.sql
new file mode 100644
index 00000000000..7f628bf24e2
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting--1.0.sql
@@ -0,0 +1,7 @@
+CREATE FUNCTION memcontext_crash_server()
+RETURNS pg_catalog.void
+AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION memcontext_crash_client()
+RETURNS pg_catalog.void
+AS 'MODULE_PATHNAME' LANGUAGE C;  
diff --git a/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c
new file mode 100644
index 00000000000..f77875b437f
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c
@@ -0,0 +1,68 @@
+/*
+ * -------------------------------------------------------------------------
+ *
+ * Copyright (c) 2025, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *             src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+#include "utils/injection_point.h"
+#include "funcapi.h"
+#include "utils/injection_point.h"
+
+PG_MODULE_MAGIC;
+
+void crash(const char *name, const void *private_data, void *arg);
+
+void
+crash(const char *name, const void *private_data, void *arg)
+{
+	abort();
+}
+
+/*
+ * memcontext_crash_client
+ *
+ * Ensure that the client process aborts in between memory context
+ * reporting.
+ */
+PG_FUNCTION_INFO_V1(memcontext_crash_client);
+Datum
+memcontext_crash_client(PG_FUNCTION_ARGS)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointAttach("memcontext-client-crash",
+						"test_memcontext_reporting", "crash", NULL, 0);
+
+#else
+	elog(ERROR,
+		"test is not working as intended when injection points are disabled");
+#endif
+	PG_RETURN_VOID();
+}
+
+/*
+ * memcontext_crash_server
+ *
+ * Ensure that the server process crashes in between memory context
+ * reporting.
+ */
+PG_FUNCTION_INFO_V1(memcontext_crash_server);
+Datum
+memcontext_crash_server(PG_FUNCTION_ARGS)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointAttach("memcontext-server-crash",
+						"test_memcontext_reporting", "crash", NULL, 0);
+
+#else
+	elog(ERROR,
+		"test is not working as intended when injection points are disabled");
+#endif
+	PG_RETURN_VOID();
+}
+
diff --git a/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.control b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.control
new file mode 100644
index 00000000000..48b501682d5
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.control
@@ -0,0 +1,4 @@
+comment = 'Test code for memcontext reporting'
+default_version = '1.0'
+module_pathname = '$libdir/test_memcontext_reporting'
+relocatable = true
-- 
2.34.1

