From 440625011c4d01c9fc29b59f09c5c1006189cfac Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 30 Apr 2025 12:27:33 +0900
Subject: [PATCH v2 2/3] injection_points: Add function to flush injection
 point data to disk

The function injection_points_flush() can be called in a session to
persist to disk all the injection points currently attached to the
cluster.  These are saved in file called injection_points.data at the
root of PGDATA, and loaded by the cluster at startup with all the points
registered automatically attached.

This will be useful to test scenarios where injection points need to be
for example attached at very early stages of startup, before it is
possible to attach anything via SQL.
---
 .../injection_points--1.0.sql                 |  10 ++
 .../injection_points/injection_points.c       | 163 ++++++++++++++++++
 src/test/modules/injection_points/meson.build |   1 +
 .../injection_points/t/002_data_persist.pl    |  53 ++++++
 4 files changed, 227 insertions(+)
 create mode 100644 src/test/modules/injection_points/t/002_data_persist.pl

diff --git a/src/test/modules/injection_points/injection_points--1.0.sql b/src/test/modules/injection_points/injection_points--1.0.sql
index cc76b1bf99ae..966e1342e4ae 100644
--- a/src/test/modules/injection_points/injection_points--1.0.sql
+++ b/src/test/modules/injection_points/injection_points--1.0.sql
@@ -3,6 +3,16 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION injection_points" to load this file. \quit
 
+--
+-- injection_points_flush()
+--
+-- Flush to disk all the data of the injection points attached.
+--
+CREATE FUNCTION injection_points_flush()
+RETURNS void
+AS 'MODULE_PATHNAME', 'injection_points_flush'
+LANGUAGE C STRICT;
+
 --
 -- injection_points_attach()
 --
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 3da0cbc10e08..c6c541113680 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -24,6 +24,7 @@
 #include "nodes/value.h"
 #include "storage/condition_variable.h"
 #include "storage/dsm_registry.h"
+#include "storage/fd.h"
 #include "storage/ipc.h"
 #include "storage/lwlock.h"
 #include "storage/shmem.h"
@@ -39,6 +40,14 @@ PG_MODULE_MAGIC;
 #define INJ_MAX_WAIT	8
 #define INJ_NAME_MAXLEN	64
 
+/* Location of injection point data files, if flush has been requested */
+#define INJ_DUMP_FILE	"injection_points.data"
+#define INJ_DUMP_FILE_TMP	INJ_DUMP_FILE ".tmp"
+
+/* Magic number identifying the injection file */
+static const uint32 INJ_FILE_HEADER = 0xFF345678;
+
+
 /*
  * Conditions related to injection points.  This tracks in shared memory the
  * runtime conditions under which an injection point is allowed to run,
@@ -151,6 +160,9 @@ static void
 injection_shmem_startup(void)
 {
 	bool		found;
+	int32		num_inj_points;
+	uint32		header;
+	FILE	   *file;
 
 	if (prev_shmem_startup_hook)
 		prev_shmem_startup_hook();
@@ -172,6 +184,84 @@ injection_shmem_startup(void)
 	}
 
 	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Done if some other process already completed the initialization.
+	 */
+	if (found)
+		return;
+
+	/*
+	 * Note: there should be no need to bother with locks here, because there
+	 * should be no other processes running when this code is reached.
+	 */
+
+	/* Load injection point data, if any has been found while starting up */
+	file = AllocateFile(INJ_DUMP_FILE, PG_BINARY_R);
+
+	if (file == NULL)
+	{
+		if (errno != ENOENT)
+			goto error;
+
+		/* No file?  We are done. */
+		return;
+	}
+
+	if (fread(&header, sizeof(uint32), 1, file) != 1 ||
+		fread(&num_inj_points, sizeof(int32), 1, file) != 1)
+		goto error;
+
+	if (header != INJ_FILE_HEADER)
+		goto error;
+
+	for (int i = 0; i < num_inj_points; i++)
+	{
+		const char *name;
+		const char *library;
+		const char *function;
+		uint32		len;
+		char		buf[1024];
+
+		if (fread(&len, sizeof(uint32), 1, file) != 1)
+			goto error;
+		if (fread(buf, 1, len + 1, file) != len + 1)
+			goto error;
+		name = pstrdup(buf);
+
+		if (fread(&len, sizeof(uint32), 1, file) != 1)
+			goto error;
+		if (fread(buf, 1, len + 1, file) != len + 1)
+			goto error;
+		library = pstrdup(buf);
+
+		if (fread(&len, sizeof(uint32), 1, file) != 1)
+			goto error;
+		if (fread(buf, 1, len + 1, file) != len + 1)
+			goto error;
+		function = pstrdup(buf);
+
+		/* No private data handled here */
+		InjectionPointAttach(name, library, function, NULL, 0);
+	}
+
+	/*
+	 * Remove the persisted injection point file, we do not need it anymore.
+	 */
+	unlink(INJ_DUMP_FILE);
+	FreeFile(file);
+
+	return;
+
+error:
+	ereport(LOG,
+			(errcode_for_file_access(),
+			 errmsg("could not read file \"%s\": %m",
+					INJ_DUMP_FILE)));
+	if (file)
+		FreeFile(file);
+
+	unlink(INJ_DUMP_FILE);
 }
 
 /*
@@ -343,6 +433,79 @@ injection_wait(const char *name, const void *private_data, void *arg)
 	SpinLockRelease(&inj_state->lock);
 }
 
+/*
+ * SQL function for flushing injection point data to disk.
+ */
+PG_FUNCTION_INFO_V1(injection_points_flush);
+Datum
+injection_points_flush(PG_FUNCTION_ARGS)
+{
+	FILE	   *file = NULL;
+	List	   *inj_points = NIL;
+	ListCell   *lc;
+	int32		num_inj_points;
+
+	inj_points = InjectionPointList();
+	if (inj_points == NIL)
+		PG_RETURN_VOID();
+
+	num_inj_points = list_length(inj_points);
+
+	file = AllocateFile(INJ_DUMP_FILE ".tmp", PG_BINARY_W);
+	if (file == NULL)
+		goto error;
+
+	if (fwrite(&INJ_FILE_HEADER, sizeof(uint32), 1, file) != 1)
+		goto error;
+
+	if (fwrite(&num_inj_points, sizeof(int32), 1, file) != 1)
+		goto error;
+
+	foreach(lc, inj_points)
+	{
+		InjectionPointData *inj_point = lfirst(lc);
+		uint32		len;
+
+		len = strlen(inj_point->name);
+		if (fwrite(&len, sizeof(uint32), 1, file) != 1 ||
+			fwrite(inj_point->name, 1, len + 1, file) != len + 1)
+			goto error;
+
+		len = strlen(inj_point->library);
+		if (fwrite(&len, sizeof(uint32), 1, file) != 1 ||
+			fwrite(inj_point->library, 1, len + 1, file) != len + 1)
+			goto error;
+
+		len = strlen(inj_point->function);
+		if (fwrite(&len, sizeof(uint32), 1, file) != 1 ||
+			fwrite(inj_point->function, 1, len + 1, file) != len + 1)
+			goto error;
+	}
+
+	if (FreeFile(file))
+	{
+		file = NULL;
+		goto error;
+	}
+
+	/*
+	 * Rename file into place, so we atomically replace any old one.
+	 */
+	(void) durable_rename(INJ_DUMP_FILE_TMP, INJ_DUMP_FILE, LOG);
+
+	PG_RETURN_VOID();
+
+error:
+	ereport(LOG,
+			(errcode_for_file_access(),
+			 errmsg("could not write file \"%s\": %m",
+					INJ_DUMP_FILE_TMP)));
+	if (file)
+		FreeFile(file);
+	unlink(INJ_DUMP_FILE_TMP);
+	PG_RETURN_VOID();
+}
+
 /*
  * SQL function for creating an injection point.
  */
diff --git a/src/test/modules/injection_points/meson.build b/src/test/modules/injection_points/meson.build
index d61149712fd7..b994f8d413d3 100644
--- a/src/test/modules/injection_points/meson.build
+++ b/src/test/modules/injection_points/meson.build
@@ -56,6 +56,7 @@ tests += {
     },
     'tests': [
       't/001_stats.pl',
+      't/002_data_persist.pl',
     ],
   },
 }
diff --git a/src/test/modules/injection_points/t/002_data_persist.pl b/src/test/modules/injection_points/t/002_data_persist.pl
new file mode 100644
index 000000000000..9ecb05230931
--- /dev/null
+++ b/src/test/modules/injection_points/t/002_data_persist.pl
@@ -0,0 +1,53 @@
+
+# Copyright (c) 2024-2025, PostgreSQL Global Development Group
+
+# Tests for persistence of injection point data.
+
+use strict;
+use warnings FATAL => 'all';
+use locale;
+
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+# Test persistency of statistics generated for injection points.
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
+# Node initialization
+my $node = PostgreSQL::Test::Cluster->new('master');
+$node->init;
+$node->append_conf(
+	'postgresql.conf', qq(
+shared_preload_libraries = 'injection_points'
+));
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
+
+# Attach a couple of points, which are going to be made persistent.
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('persist-notice', 'notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('persist-error', 'error');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('persist-notice-2', 'notice');");
+
+# Flush and restart, the injection points still exist.
+$node->safe_psql('postgres', "SELECT injection_points_flush();");
+$node->restart;
+
+my ($result, $stdout, $stderr) =
+  $node->psql('postgres', "SELECT injection_points_run('persist-notice-2')");
+ok( $stderr =~
+	  /NOTICE:  notice triggered for injection point persist-notice-2/,
+	"injection point triggering NOTICE exists");
+
+($result, $stdout, $stderr) =
+  $node->psql('postgres', "SELECT injection_points_run('persist-error')");
+ok($stderr =~ /ERROR:  error triggered for injection point persist-error/,
+	"injection point triggering ERROR exists");
+
+done_testing();
-- 
2.49.0

