From 146c195eff2099755f1b7acc13e94d4bc6c1845b Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 26 Jun 2024 10:49:22 +0900
Subject: [PATCH] Extend injection points to work within a postmaster context

The LWLock protecting injection point lookups is gone, as the postmaster
has no PGPROC entry.  The module is able to work in this case with
notices and errors.
---
 src/include/storage/lwlocklist.h              |  3 +-
 src/backend/postmaster/postmaster.c           |  3 +
 .../utils/activity/wait_event_names.txt       |  1 -
 src/backend/utils/misc/injection_point.c      |  9 --
 .../injection_points/injection_points.c       | 96 +++++++++++++++++--
 5 files changed, 93 insertions(+), 19 deletions(-)

diff --git a/src/include/storage/lwlocklist.h b/src/include/storage/lwlocklist.h
index 85f6568b9e..cdce436b45 100644
--- a/src/include/storage/lwlocklist.h
+++ b/src/include/storage/lwlocklist.h
@@ -81,5 +81,4 @@ PG_LWLOCK(47, NotifyQueueTail)
 PG_LWLOCK(48, WaitEventExtension)
 PG_LWLOCK(49, WALSummarizer)
 PG_LWLOCK(50, DSMRegistry)
-PG_LWLOCK(51, InjectionPoint)
-PG_LWLOCK(52, SerialControl)
+PG_LWLOCK(51, SerialControl)
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index bf0241aed0..7da01fa035 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -118,6 +118,7 @@
 #include "tcop/backend_startup.h"
 #include "tcop/tcopprot.h"
 #include "utils/datetime.h"
+#include "utils/injection_point.h"
 #include "utils/memutils.h"
 #include "utils/pidfile.h"
 #include "utils/timestamp.h"
@@ -3574,6 +3575,8 @@ BackendStartup(ClientSocket *client_sock)
 		return STATUS_ERROR;
 	}
 
+	INJECTION_POINT("backend-startup");
+
 	/* Pass down canAcceptConnections state */
 	startup_data.canAcceptConnections = canAcceptConnections(BACKEND_TYPE_NORMAL);
 	bn->dead_end = (startup_data.canAcceptConnections != CAC_OK);
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index 87cbca2811..3f948c5a44 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -343,7 +343,6 @@ NotifyQueueTail	"Waiting to update limit on <command>NOTIFY</command> message st
 WaitEventExtension	"Waiting to read or update custom wait events information for extensions."
 WALSummarizer	"Waiting to read or update WAL summarization state."
 DSMRegistry	"Waiting to read or update the dynamic shared memory registry."
-InjectionPoint	"Waiting to read or update information related to injection points."
 SerialControl	"Waiting to read or update shared <filename>pg_serial</filename> state."
 
 #
diff --git a/src/backend/utils/misc/injection_point.c b/src/backend/utils/misc/injection_point.c
index 5c2a0d2297..03663a42c6 100644
--- a/src/backend/utils/misc/injection_point.c
+++ b/src/backend/utils/misc/injection_point.c
@@ -23,7 +23,6 @@
 #include "miscadmin.h"
 #include "port/pg_bitutils.h"
 #include "storage/fd.h"
-#include "storage/lwlock.h"
 #include "storage/shmem.h"
 #include "utils/hsearch.h"
 #include "utils/injection_point.h"
@@ -229,13 +228,11 @@ InjectionPointAttach(const char *name,
 	 * Allocate and register a new injection point.  A new point should not
 	 * exist.  For testing purposes this should be fine.
 	 */
-	LWLockAcquire(InjectionPointLock, LW_EXCLUSIVE);
 	entry_by_name = (InjectionPointEntry *)
 		hash_search(InjectionPointHash, name,
 					HASH_ENTER, &found);
 	if (found)
 	{
-		LWLockRelease(InjectionPointLock);
 		elog(ERROR, "injection point \"%s\" already defined", name);
 	}
 
@@ -249,8 +246,6 @@ InjectionPointAttach(const char *name,
 	if (private_data != NULL)
 		memcpy(entry_by_name->private_data, private_data, private_data_size);
 
-	LWLockRelease(InjectionPointLock);
-
 #else
 	elog(ERROR, "injection points are not supported by this build");
 #endif
@@ -267,9 +262,7 @@ InjectionPointDetach(const char *name)
 #ifdef USE_INJECTION_POINTS
 	bool		found;
 
-	LWLockAcquire(InjectionPointLock, LW_EXCLUSIVE);
 	hash_search(InjectionPointHash, name, HASH_REMOVE, &found);
-	LWLockRelease(InjectionPointLock);
 
 	if (!found)
 		return false;
@@ -296,11 +289,9 @@ InjectionPointRun(const char *name)
 	InjectionPointCallback injection_callback;
 	const void *private_data;
 
-	LWLockAcquire(InjectionPointLock, LW_SHARED);
 	entry_by_name = (InjectionPointEntry *)
 		hash_search(InjectionPointHash, name,
 					HASH_FIND, &found);
-	LWLockRelease(InjectionPointLock);
 
 	/*
 	 * If not found, do nothing and remove it from the local cache if it
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 5c44625d1d..42b13483eb 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -37,6 +37,9 @@ PG_MODULE_MAGIC;
 #define INJ_MAX_WAIT	8
 #define INJ_NAME_MAXLEN	64
 
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
+
 /*
  * Conditions related to injection points.  This tracks in shared memory the
  * runtime conditions under which an injection point is allowed to run,
@@ -67,7 +70,12 @@ typedef struct InjectionPointCondition
  */
 static List *inj_list_local = NIL;
 
-/* Shared state information for injection points. */
+/*
+ * Shared state information for injection points.
+ *
+ * Can be initialized with shared_preload_libraries, or dynamically with the
+ * DSM registry.
+ */
 typedef struct InjectionPointSharedState
 {
 	/* Protects access to other fields */
@@ -97,7 +105,8 @@ extern PGDLLEXPORT void injection_wait(const char *name,
 static bool injection_point_local = false;
 
 /*
- * Callback for shared memory area initialization.
+ * Callback for shared memory area initialization, used for DSM registry
+ * or shared memory initialization.
  */
 static void
 injection_point_init_state(void *ptr)
@@ -111,10 +120,62 @@ injection_point_init_state(void *ptr)
 }
 
 /*
- * Initialize shared memory area for this module.
+ * Shared memory size to request.
+ */
+static Size
+inj_shmem_size(void)
+{
+	Size		size;
+
+	size = MAXALIGN(sizeof(InjectionPointSharedState));
+	return size;
+}
+
+/*
+ * Shared memory request hook, when loaded with shared_preload_libraries.
  */
 static void
-injection_init_shmem(void)
+inj_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(inj_shmem_size());
+}
+
+/*
+ * Shared memory startup hook, when loaded with shared_preload_libraries.
+ */
+static void
+inj_shmem_startup(void)
+{
+	bool		found;
+
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	/*
+	 * Create or attach to the shared memory state.
+	 */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+
+	inj_state = ShmemInitStruct("injection_points",
+								sizeof(InjectionPointSharedState),
+								&found);
+	if (!found)
+	{
+		/* First time through */
+		injection_point_init_state(inj_state);
+	}
+
+	LWLockRelease(AddinShmemInitLock);
+}
+
+/*
+ * Initialize shared memory area for this module, with a DSM.
+ */
+static void
+injection_init_shmem_dsm(void)
 {
 	bool		found;
 
@@ -206,7 +267,7 @@ injection_wait(const char *name, const void *private_data)
 	InjectionPointCondition *condition = (InjectionPointCondition *) private_data;
 
 	if (inj_state == NULL)
-		injection_init_shmem();
+		injection_init_shmem_dsm();
 
 	if (!injection_point_allowed(condition))
 		return;
@@ -327,7 +388,7 @@ injection_points_wakeup(PG_FUNCTION_ARGS)
 	int			index = -1;
 
 	if (inj_state == NULL)
-		injection_init_shmem();
+		injection_init_shmem_dsm();
 
 	/* First bump the wait counter for the injection point to wake up */
 	SpinLockAcquire(&inj_state->lock);
@@ -367,7 +428,7 @@ injection_points_set_local(PG_FUNCTION_ARGS)
 	injection_point_local = true;
 
 	if (inj_state == NULL)
-		injection_init_shmem();
+		injection_init_shmem_dsm();
 
 	/*
 	 * Register a before_shmem_exit callback to remove any injection points
@@ -402,3 +463,24 @@ injection_points_detach(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+/*
+ * Module load callback.
+ */
+void
+_PG_init(void)
+{
+	/*
+	 * If not loading through shared_preload_libraries, just leave.  This
+	 * module's shared state will initialize with the DSM registry in this
+	 * case.  Allowing the shared_preload_libraries case is useful for
+	 * injection points in the postmaster.
+	 */
+	if (!process_shared_preload_libraries_in_progress)
+		return;
+
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = inj_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = inj_shmem_startup;
+}
-- 
2.45.2

