From 4cf66958ca26a75e4ffdf8527beb5aab3733108a Mon Sep 17 00:00:00 2001
From: Zsolt Parragi <zsolt.parragi@percona.com>
Date: Thu, 11 Dec 2025 22:54:06 +0000
Subject: [PATCH] Add init_callback_arg parameter to GetNamedDSMSegment

Modify GetNamedDSMSegment to accept an optional init_callback_arg
parameter that gets passed to the init_callback function. This allows
callbacks to differentiate initialization logic based on additional
context provided by the caller.

Author: zsolt.parragi@percona.com
Reviewed-by: Nathan Bossart <bossartn@amazon.com>
Reviewed-by: Sami Imseih <samimseih@gmail.com>
Discussion: https://postgr.es/m/CAN4CZFMjh8TrT9ZhWgjVTzBDkYZi2a84BnZ8bM%2BfLPuq7Cirzg%40mail.gmail.com
---
 contrib/pg_prewarm/autoprewarm.c                    |  3 ++-
 doc/src/sgml/xfunc.sgml                             | 13 ++++++++-----
 src/backend/storage/ipc/dsm_registry.c              |  8 +++++---
 src/include/storage/dsm_registry.h                  |  3 ++-
 .../modules/injection_points/injection_points.c     |  5 +++--
 src/test/modules/test_dsa/test_dsa.c                |  6 +++---
 .../modules/test_dsm_registry/test_dsm_registry.c   | 11 +++++++++--
 7 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/contrib/pg_prewarm/autoprewarm.c b/contrib/pg_prewarm/autoprewarm.c
index 5ba1240d51f..9a68e765e54 100644
--- a/contrib/pg_prewarm/autoprewarm.c
+++ b/contrib/pg_prewarm/autoprewarm.c
@@ -858,7 +858,7 @@ autoprewarm_dump_now(PG_FUNCTION_ARGS)
 }
 
 static void
-apw_init_state(void *ptr)
+apw_init_state(void *ptr, void *arg)
 {
 	AutoPrewarmSharedState *state = (AutoPrewarmSharedState *) ptr;
 
@@ -880,6 +880,7 @@ apw_init_shmem(void)
 	apw_state = GetNamedDSMSegment("autoprewarm",
 								   sizeof(AutoPrewarmSharedState),
 								   apw_init_state,
+								   NULL,
 								   &found);
 
 	return found;
diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index 7c76ab8c349..91284abb679 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3696,15 +3696,18 @@ LWLockRelease(AddinShmemInitLock);
       use the shared memory should obtain a pointer to it by calling:
 <programlisting>
 void *GetNamedDSMSegment(const char *name, size_t size,
-                         void (*init_callback) (void *ptr),
+                         void (*init_callback) (void *ptr, void *init_callback_arg),
+                         void *init_callback_arg,
                          bool *found)
 </programlisting>
       If a dynamic shared memory segment with the given name does not yet
       exist, this function will allocate it and initialize it with the provided
-      <function>init_callback</function> callback function.  If the segment has
-      already been allocated and initialized by another backend, this function
-      simply attaches the existing dynamic shared memory segment to the current
-      backend.
+      <function>init_callback</function> callback function.  The
+      <parameter>init_callback_arg</parameter> parameter is passed to the callback
+      function, allowing additional context to be provided during initialization.
+      If the segment has already been allocated and initialized by another backend,
+      this function simply attaches the existing dynamic shared memory segment to
+      the current backend.
      </para>
 
      <para>
diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index 66240318e83..c68fad52c54 100644
--- a/src/backend/storage/ipc/dsm_registry.c
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -180,11 +180,13 @@ init_dsm_registry(void)
  * Initialize or attach a named DSM segment.
  *
  * This routine returns the address of the segment.  init_callback is called to
- * initialize the segment when it is first created.
+ * initialize the segment when it is first created. An optional init_callback_arg is
+ * passed to init_callback to provide additional context during initialization.
  */
 void *
 GetNamedDSMSegment(const char *name, size_t size,
-				   void (*init_callback) (void *ptr), bool *found)
+				   void (*init_callback) (void *ptr, void *init_callback_arg),
+				   void *init_callback_arg, bool *found)
 {
 	DSMRegistryEntry *entry;
 	MemoryContext oldcontext;
@@ -235,7 +237,7 @@ GetNamedDSMSegment(const char *name, size_t size,
 		seg = dsm_create(size, 0);
 
 		if (init_callback)
-			(*init_callback) (dsm_segment_address(seg));
+			(*init_callback) (dsm_segment_address(seg), init_callback_arg);
 
 		dsm_pin_segment(seg);
 		dsm_pin_mapping(seg);
diff --git a/src/include/storage/dsm_registry.h b/src/include/storage/dsm_registry.h
index 4871ed509eb..b2670ebefc7 100644
--- a/src/include/storage/dsm_registry.h
+++ b/src/include/storage/dsm_registry.h
@@ -16,7 +16,8 @@
 #include "lib/dshash.h"
 
 extern void *GetNamedDSMSegment(const char *name, size_t size,
-								void (*init_callback) (void *ptr),
+								void (*init_callback) (void *ptr, void *init_callback_arg),
+								void *init_callback_arg,
 								bool *found);
 extern dsa_area *GetNamedDSA(const char *name, bool *found);
 extern dshash_table *GetNamedDSHash(const char *name,
diff --git a/src/test/modules/injection_points/injection_points.c b/src/test/modules/injection_points/injection_points.c
index 417b61f31c5..36f07295818 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -115,7 +115,7 @@ static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
  * when initializing dynamically with a DSM or when loading the module.
  */
 static void
-injection_point_init_state(void *ptr)
+injection_point_init_state(void *ptr, void *arg)
 {
 	InjectionPointSharedState *state = (InjectionPointSharedState *) ptr;
 
@@ -159,7 +159,7 @@ injection_shmem_startup(void)
 		 * First time through, so initialize.  This is shared with the dynamic
 		 * initialization using a DSM.
 		 */
-		injection_point_init_state(inj_state);
+		injection_point_init_state(inj_state, NULL);
 	}
 
 	LWLockRelease(AddinShmemInitLock);
@@ -179,6 +179,7 @@ injection_init_shmem(void)
 	inj_state = GetNamedDSMSegment("injection_points",
 								   sizeof(InjectionPointSharedState),
 								   injection_point_init_state,
+								   NULL,
 								   &found);
 }
 
diff --git a/src/test/modules/test_dsa/test_dsa.c b/src/test/modules/test_dsa/test_dsa.c
index 21e4ce6a745..1b4fe963998 100644
--- a/src/test/modules/test_dsa/test_dsa.c
+++ b/src/test/modules/test_dsa/test_dsa.c
@@ -21,7 +21,7 @@
 PG_MODULE_MAGIC;
 
 static void
-init_tranche(void *ptr)
+init_tranche(void *ptr, void *arg)
 {
 	int		   *tranche_id = (int *) ptr;
 
@@ -39,7 +39,7 @@ test_dsa_basic(PG_FUNCTION_ARGS)
 	dsa_pointer p[100];
 
 	tranche_id = GetNamedDSMSegment("test_dsa", sizeof(int),
-									init_tranche, &found);
+									init_tranche, NULL, &found);
 
 	a = dsa_create(*tranche_id);
 	for (int i = 0; i < 100; i++)
@@ -80,7 +80,7 @@ test_dsa_resowners(PG_FUNCTION_ARGS)
 	ResourceOwner childowner;
 
 	tranche_id = GetNamedDSMSegment("test_dsa", sizeof(int),
-									init_tranche, &found);
+									init_tranche, NULL, &found);
 
 	/* Create DSA in parent resource owner */
 	a = dsa_create(*tranche_id);
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry.c b/src/test/modules/test_dsm_registry/test_dsm_registry.c
index 4cc2ccdac3f..cc1cdad0d60 100644
--- a/src/test/modules/test_dsm_registry/test_dsm_registry.c
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -43,12 +43,18 @@ static const dshash_parameters dsh_params = {
 	dshash_strcpy
 };
 
+void *sample_arg_ptr = &tdr_dsm;
+
 static void
-init_tdr_dsm(void *ptr)
+init_tdr_dsm(void *ptr, void *arg)
 {
 	TestDSMRegistryStruct *dsm = (TestDSMRegistryStruct *) ptr;
 
-	LWLockInitialize(&dsm->lck, LWLockNewTrancheId("test_dsm_registry"));
+	if (arg != sample_arg_ptr)
+		ereport(ERROR,
+				(errmsg("Unexpected arg pointer address")));
+
+	LWLockInitialize(&dsm->lck, LWLockNewTrancheId("test_dsm_registry_dsm"));
 	dsm->val = 0;
 }
 
@@ -60,6 +66,7 @@ tdr_attach_shmem(void)
 	tdr_dsm = GetNamedDSMSegment("test_dsm_registry_dsm",
 								 sizeof(TestDSMRegistryStruct),
 								 init_tdr_dsm,
+								 sample_arg_ptr,
 								 &found);
 
 	if (tdr_dsa == NULL)
-- 
2.43.0

