From bbdf4e76ad8d1ea2504a16669f4dfcd6cdcee3a1 Mon Sep 17 00:00:00 2001
From: Zsolt Parragi <zsolt.parragi@percona.com>
Date: Tue, 2 Dec 2025 20:01:33 +0000
Subject: [PATCH] Add a callback data parameter to GetNamedDSMSegment

This function already supports a callback, but there is no way to pass
user data to that callback.

For example, the documentation notes that when using GetNamedDSMSegment,
LWLocks should be initialized in the init callback. However, if the calling
code is generic (for example, when it receives a tranche_name as a
parameter), there is no clean way to pass that parameter into the init
callback.

This change addresses that limitation by adding a generic callback argument
that is forwarded to the callback function.
---
 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     |  3 ++-
 src/test/modules/test_dsa/test_dsa.c                |  8 ++++----
 .../modules/test_dsm_registry/test_dsm_registry.c   |  3 ++-
 7 files changed, 25 insertions(+), 16 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 537ee6fa254..c91defb6a6c 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 *arg),
+                         void *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>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..cb150bac690 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.  callback_arg is passed to
+ * the init_callback as an additional argument.
  */
 void *
 GetNamedDSMSegment(const char *name, size_t size,
-				   void (*init_callback) (void *ptr), bool *found)
+				   void (*init_callback) (void *ptr, void *arg),
+				   void *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), 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..150eba1765c 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 *arg),
+								void *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 b7c1c58ea56..412d2a64ed7 100644
--- a/src/test/modules/injection_points/injection_points.c
+++ b/src/test/modules/injection_points/injection_points.c
@@ -125,7 +125,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;
 
@@ -189,6 +189,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..67c22dc0f5b 100644
--- a/src/test/modules/test_dsa/test_dsa.c
+++ b/src/test/modules/test_dsa/test_dsa.c
@@ -21,11 +21,11 @@
 PG_MODULE_MAGIC;
 
 static void
-init_tranche(void *ptr)
+init_tranche(void *ptr, void *arg)
 {
 	int		   *tranche_id = (int *) ptr;
 
-	*tranche_id = LWLockNewTrancheId("test_dsa");
+	*tranche_id = LWLockNewTrancheId(arg);
 }
 
 /* Test basic DSA functionality */
@@ -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, "test_dsa", &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, "test_dsa", &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..7c5d4ee015b 100644
--- a/src/test/modules/test_dsm_registry/test_dsm_registry.c
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -44,7 +44,7 @@ static const dshash_parameters dsh_params = {
 };
 
 static void
-init_tdr_dsm(void *ptr)
+init_tdr_dsm(void *ptr, void *arg)
 {
 	TestDSMRegistryStruct *dsm = (TestDSMRegistryStruct *) ptr;
 
@@ -60,6 +60,7 @@ tdr_attach_shmem(void)
 	tdr_dsm = GetNamedDSMSegment("test_dsm_registry_dsm",
 								 sizeof(TestDSMRegistryStruct),
 								 init_tdr_dsm,
+								 NULL,
 								 &found);
 
 	if (tdr_dsa == NULL)
-- 
2.43.0

