From b2b023aa3706ec6b3978708545301f7436205c6d Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 23 Aug 2024 10:12:58 +0900
Subject: [PATCH] injection_points: Add initialization of shmem state when
 loading module
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commits adds callbacks to initialize the shared memory state of the
module when loaded with shared_preload_libraries.  This is necessary to
be able to update the test introduced in 768a9fd5535f to use the macros
INJECTION_POINT_{LOAD,CACHED}() rather than a SQL function in the module
injection_points forcing a load, as this test runs a callback in a
critical section where no memory allocation should happen.

Initializing the shared memory state of the module while loading
provides a strict control on the timing of its allocation.  If the
module is not loaded at startup, it will use a GetNamedDSMSegment()
instead to initialize its shmem state on-the-fly.

Per discussion with Álvaro Herrera.

Author: Michael Paquier
Discussion: https://postgr.es/m/ZsUnJUlSOBNAzwW1@paquier.xyz
---
 .../injection_points/injection_points.c       | 62 ++++++++++++++++++-
 1 file changed, 59 insertions(+), 3 deletions(-)

diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 4e775c7ec60..abb1516e124 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -68,7 +68,12 @@ typedef struct InjectionPointCondition
  */
 static List *inj_list_local = NIL;
 
-/* Shared state information for injection points. */
+/*
+ * Shared state information for injection points.
+ *
+ * This state data can be initialized in two ways: dynamically with a DSM
+ * or when loading the module.
+ */
 typedef struct InjectionPointSharedState
 {
 	/* Protects access to other fields */
@@ -97,8 +102,13 @@ extern PGDLLEXPORT void injection_wait(const char *name,
 /* track if injection points attached in this process are linked to it */
 static bool injection_point_local = 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;
+
 /*
- * Callback for shared memory area initialization.
+ * Routine for shared memory area initialization, used as a callback
+ * when initializing dynamically with a DSM or when loading the module.
  */
 static void
 injection_point_init_state(void *ptr)
@@ -111,8 +121,48 @@ injection_point_init_state(void *ptr)
 	ConditionVariableInit(&state->wait_point);
 }
 
+/* Shared memory initialization when loading module */
+static void
+injection_shmem_request(void)
+{
+	Size		size;
+
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	size = MAXALIGN(sizeof(InjectionPointSharedState));
+	RequestAddinShmemSpace(size);
+}
+
+static void
+injection_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, so initialize.  This is shared with the dynamic
+		 * initialization using a DSM.
+		 */
+		injection_point_init_state(inj_state);
+	}
+
+	LWLockRelease(AddinShmemInitLock);
+}
+
 /*
- * Initialize shared memory area for this module.
+ * Initialize shared memory area for this module through DSM.
  */
 static void
 injection_init_shmem(void)
@@ -463,6 +513,12 @@ _PG_init(void)
 	if (!process_shared_preload_libraries_in_progress)
 		return;
 
+	/* Shared memory initialization */
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = injection_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = injection_shmem_startup;
+
 	pgstat_register_inj();
 	pgstat_register_inj_fixed();
 }