From e21832c610574f48496c961a7318070b720568de Mon Sep 17 00:00:00 2001
From: Rahila Syed <rahilasyed.90@gmail.com>
Date: Thu, 23 Oct 2025 18:01:36 +0530
Subject: [PATCH 2/2] Test module to test memory context reporting with
 injection points

---
 src/test/modules/Makefile                     |   1 +
 .../test_memcontext_reporting/Makefile        |  32 +++++
 .../t/001_memcontext_inj.pl                   |  58 +++++++++
 .../test_memcontext_reporting--1.0.sql        |  11 ++
 .../test_memcontext_reporting.c               | 123 ++++++++++++++++++
 .../test_memcontext_reporting.control         |   4 +
 6 files changed, 229 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/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..69d8489eb37
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/t/001_memcontext_inj.pl
@@ -0,0 +1,58 @@
+# Copyright (c) 2025, PostgreSQL Global Development Group
+
+# Test suite for testing memory context statistics reporting
+
+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';
+}
+my $psql_err;
+# Create and start a cluster with one node
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init(allows_streaming => 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;');
+$node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
+# Attaching to a client process injection point that throws an error
+$node->safe_psql('postgres', "select injection_points_attach('memcontext-client-crash', 'error');");
+
+my $pid = $node->safe_psql('postgres', "SELECT pid from pg_stat_activity where backend_type='checkpointer'");
+print "PID";
+print $pid;
+
+#Client should have thrown error
+$node->psql('postgres', qq(select pg_get_process_memory_contexts($pid, true);), stderr => \$psql_err);
+like ( $psql_err, qr/error triggered for injection point memcontext-client-crash/);
+
+#Query the same process after detaching the injection point, using some other client and it should succeed.
+$node->safe_psql('postgres', "select injection_points_detach('memcontext-client-crash');");
+my $topcontext_name = $node->safe_psql('postgres', "select name from pg_get_process_memory_contexts($pid, true) where path = '{1}';");
+ok($topcontext_name = 'TopMemoryContext');
+
+# Attaching to a target process injection point that throws an error
+$node->safe_psql('postgres', "select injection_points_attach('memcontext-server-crash', 'error');");
+
+#Server should have thrown error
+$node->psql('postgres', qq(select pg_get_process_memory_contexts($pid, true);), stderr => \$psql_err);
+
+#Query the same process after detaching the injection point, using some other client and it should succeed.
+$node->safe_psql('postgres', "select injection_points_detach('memcontext-server-crash');");
+$topcontext_name = $node->safe_psql('postgres', "select name from pg_get_process_memory_contexts($pid, true) 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..181daf429d0
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting--1.0.sql
@@ -0,0 +1,11 @@
+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;
+
+CREATE FUNCTION dsa_dump_sql()
+RETURNS bigint
+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..955155524c2
--- /dev/null
+++ b/src/test/modules/test_memcontext_reporting/test_memcontext_reporting.c
@@ -0,0 +1,123 @@
+/*
+ * -------------------------------------------------------------------------
+ *
+ * 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"
+#include "storage/dsm_registry.h"
+
+PG_MODULE_MAGIC;
+
+extern PGDLLEXPORT 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();
+}
+
+PG_FUNCTION_INFO_V1(memcontext_detach_client);
+Datum
+memcontext_detach_client(PG_FUNCTION_ARGS)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointDetach("memcontext-client-crash");
+
+#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();
+}
+
+/*
+ * memcontext_detach_server
+ *
+ * Detach the injection point which crashes the server
+ * reporting.
+ */
+PG_FUNCTION_INFO_V1(memcontext_detach_server);
+Datum
+memcontext_detach_server(PG_FUNCTION_ARGS)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointDetach("memcontext-server-crash");
+
+#else
+	elog(ERROR,
+		 "test is not working as intended when injection points are disabled");
+#endif
+	PG_RETURN_VOID();
+}
+
+/*
+ * dsa_dump_sql
+ */
+PG_FUNCTION_INFO_V1(dsa_dump_sql);
+Datum
+dsa_dump_sql(PG_FUNCTION_ARGS)
+{
+	bool		found;
+	size_t		tot_size;
+	dsa_area   *memstats_dsa_area;
+
+	memstats_dsa_area = pg_get_memstats_dsa_area();
+
+	if (memstats_dsa_area == NULL)
+		memstats_dsa_area = GetNamedDSA("memory_context_statistics_dsa", &found);
+
+	tot_size = dsa_get_total_size(memstats_dsa_area);
+	dsa_detach(memstats_dsa_area);
+	PG_RETURN_INT64(tot_size);
+}
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

