v3-0004-injection_points-Add-option-and-tests-for-WAL-log.patch

From 00f4bc416b214b5d2f7d53ed5eedbd1209f43953 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 10 Jan 2025 13:36:28 +0900
Subject: [PATCH v3 4/4] injection_points: Add option and tests for WAL-logging

This serves as a template for WAL-logging related to pgstats, based on
the variable-numbered stats the module includes.  Some TAP tests are
added to check the backend facility.
---
 .../injection_points/injection_points.c       | 19 ++++++-
 .../injection_points/injection_stats.c        | 54 +++++++++++++++++++
 .../injection_points/injection_stats.h        |  3 +-
 .../modules/injection_points/t/001_stats.pl   | 22 ++++++++
 src/tools/pgindent/typedefs.list              |  1 +
 5 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index ad528d7775..c7f0d6d242 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -104,7 +104,7 @@ extern PGDLLEXPORT void injection_wait(const char *name,
 static bool injection_point_local = false;
 
 /*
- * GUC variable
+ * GUC variables
  *
  * This GUC is useful to control if statistics should be enabled or not
  * during a test with injection points, like for example if a test relies
@@ -112,6 +112,12 @@ static bool injection_point_local = false;
  */
 bool		inj_stats_enabled = false;
 
+/*
+ * This GUC controls if statistics of injection points should be logged
+ * into WAL or not.
+ */
+bool		inj_stats_wal_enabled = false;
+
 /* Shared memory init callbacks */
 static shmem_request_hook_type prev_shmem_request_hook = NULL;
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -534,6 +540,17 @@ _PG_init(void)
 							 NULL,
 							 NULL);
 
+	DefineCustomBoolVariable("injection_points.log_stats",
+							 "Enables WAL-logging for injection points statistics.",
+							 NULL,
+							 &inj_stats_wal_enabled,
+							 false,
+							 PGC_POSTMASTER,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
 	MarkGUCPrefixReserved("injection_points");
 
 	/* Shared memory initialization */
diff --git a/src/test/modules/injection_points/injection_stats.c b/src/test/modules/injection_points/injection_stats.c
index 5db62bca66..df52cb026e 100644
--- a/src/test/modules/injection_points/injection_stats.c
+++ b/src/test/modules/injection_points/injection_stats.c
@@ -21,6 +21,7 @@
 #include "pgstat.h"
 #include "utils/builtins.h"
 #include "utils/pgstat_internal.h"
+#include "utils/pgstat_xlog.h"
 
 /* Structures for statistics of injection points */
 typedef struct PgStat_StatInjEntry
@@ -34,6 +35,14 @@ typedef struct PgStatShared_InjectionPoint
 	PgStat_StatInjEntry stats;
 } PgStatShared_InjectionPoint;
 
+/* Structure for data of injection points logged in WAL */
+typedef struct PgStat_StatInjRecord
+{
+	uint64		objid;			/* hash of the point name */
+	PgStat_StatInjEntry entry;	/* stats data */
+} PgStat_StatInjRecord;
+
+static void injection_stats_redo_cb(void *data, size_t data_size);
 static bool injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait);
 
 static const PgStat_KindInfo injection_stats = {
@@ -48,6 +57,7 @@ static const PgStat_KindInfo injection_stats = {
 	.shared_data_off = offsetof(PgStatShared_InjectionPoint, stats),
 	.shared_data_len = sizeof(((PgStatShared_InjectionPoint *) 0)->stats),
 	.pending_size = sizeof(PgStat_StatInjEntry),
+	.redo_cb = injection_stats_redo_cb,
 	.flush_pending_cb = injection_stats_flush_cb,
 };
 
@@ -64,6 +74,31 @@ static const PgStat_KindInfo injection_stats = {
 /* Track if stats are loaded */
 static bool inj_stats_loaded = false;
 
+/*
+ * REDO callback for injection point stats.
+ */
+static void
+injection_stats_redo_cb(void *data, size_t data_size)
+{
+	PgStat_StatInjRecord *record_data = (PgStat_StatInjRecord *) data;
+	PgStat_StatInjEntry record_entry = record_data->entry;
+	PgStat_StatInjEntry *stat_entry;
+	PgStatShared_InjectionPoint *shstatent;
+	PgStat_EntryRef *entry_ref;
+
+	Assert(data_size == sizeof(PgStat_StatInjRecord));
+
+	/* create or fetch existing entry */
+	entry_ref = pgstat_prep_pending_entry(PGSTAT_KIND_INJECTION, InvalidOid,
+										  record_data->objid, NULL);
+
+	shstatent = (PgStatShared_InjectionPoint *) entry_ref->shared_stats;
+	stat_entry = &shstatent->stats;
+
+	/* Update the injection point statistics, overwriting any existing data */
+	*stat_entry = record_entry;
+}
+
 /*
  * Callback for stats handling
  */
@@ -72,6 +107,9 @@ injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 {
 	PgStat_StatInjEntry *localent;
 	PgStatShared_InjectionPoint *shfuncent;
+	PgStat_StatInjRecord record_data;
+
+	memset(&record_data, 0, sizeof(PgStat_StatInjRecord));
 
 	localent = (PgStat_StatInjEntry *) entry_ref->pending;
 	shfuncent = (PgStatShared_InjectionPoint *) entry_ref->shared_stats;
@@ -81,8 +119,24 @@ injection_stats_flush_cb(PgStat_EntryRef *entry_ref, bool nowait)
 
 	shfuncent->stats.numcalls += localent->numcalls;
 
+	record_data.objid = entry_ref->shared_entry->key.objid;
+	record_data.entry = shfuncent->stats;
+
 	pgstat_unlock_entry(entry_ref);
 
+	/*
+	 * Store the data in WAL, if not in recovery and if the option is enabled.
+	 */
+	if (!RecoveryInProgress() && inj_stats_wal_enabled)
+	{
+		XLogRecPtr	lsn;
+
+		lsn = pgstat_xlog_data(PGSTAT_KIND_INJECTION, &record_data,
+							   sizeof(PgStat_StatInjRecord));
+
+		/* Force a flush, to ensure persistency */
+		XLogFlush(lsn);
+	}
 	return true;
 }
 
diff --git a/src/test/modules/injection_points/injection_stats.h b/src/test/modules/injection_points/injection_stats.h
index ba310c52c7..baa8bb506c 100644
--- a/src/test/modules/injection_points/injection_stats.h
+++ b/src/test/modules/injection_points/injection_stats.h
@@ -15,8 +15,9 @@
 #ifndef INJECTION_STATS
 #define INJECTION_STATS
 
-/* GUC variable */
+/* GUC variables */
 extern bool inj_stats_enabled;
+extern bool inj_stats_wal_enabled;
 
 /* injection_stats.c */
 extern void pgstat_register_inj(void);
diff --git a/src/test/modules/injection_points/t/001_stats.pl b/src/test/modules/injection_points/t/001_stats.pl
index d4539fe872..4b34dd0a5a 100644
--- a/src/test/modules/injection_points/t/001_stats.pl
+++ b/src/test/modules/injection_points/t/001_stats.pl
@@ -75,4 +75,26 @@ $node->stop;
 $node->adjust_conf('postgresql.conf', 'shared_preload_libraries', "''");
 $node->start;
 
+# Stop the server and enable WAL, the stats are preserved after recovery.
+$node->stop;
+$node->append_conf(
+	'postgresql.conf', qq(
+shared_preload_libraries = 'injection_points'
+injection_points.log_stats = true
+));
+$node->start;
+
+# Two calls, both WAL-logged.
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('stats-wal-notice', 'notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_run('stats-wal-notice');");
+$node->safe_psql('postgres',
+	"SELECT injection_points_run('stats-wal-notice');");
+$node->stop('immediate');
+$node->start;
+$numcalls = $node->safe_psql('postgres',
+	"SELECT injection_points_stats_numcalls('stats-wal-notice');");
+is($numcalls, '2', 'number of stats after crash with WAL-logging enabled');
+
 done_testing();
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 4db1e7cf4b..303a06b3fb 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2167,6 +2167,7 @@ PgStat_StatDBEntry
 PgStat_StatFuncEntry
 PgStat_StatInjEntry
 PgStat_StatInjFixedEntry
+PgStat_StatInjRecord
 PgStat_StatReplSlotEntry
 PgStat_StatSubEntry
 PgStat_StatTabEntry
-- 
2.47.1