From 345fa06de44c836c90ce26866a6b7f6b87d1ec6e Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Wed, 4 Jun 2025 14:21:14 -0500
Subject: [PATCH v7 1/1] simplify creating hash table in dsm registry

---
 src/backend/storage/ipc/dsm_registry.c        | 202 +++++++++++++++++-
 src/include/storage/dsm_registry.h            |   7 +-
 .../expected/test_dsm_registry.out            |  26 ++-
 .../sql/test_dsm_registry.sql                 |   6 +-
 .../test_dsm_registry--1.0.sql                |  10 +-
 .../test_dsm_registry/test_dsm_registry.c     | 114 ++++++++--
 src/tools/pgindent/typedefs.list              |   3 +
 7 files changed, 334 insertions(+), 34 deletions(-)

diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index 1d4fd31ffed..962efee5644 100644
--- a/src/backend/storage/ipc/dsm_registry.c
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -15,6 +15,20 @@
  * current backend.  This function guarantees that only one backend
  * initializes the segment and that all other backends just attach it.
  *
+ * A DSA can be created in or retrieved from the registry by calling
+ * GetNamedDSA().  As with GetNamedDSMSegment(), if a DSA with the provided
+ * name does not yet exist, it is created.  Otherwise, GetNamedDSA()
+ * ensures the DSA is attached to the current backend.  This function
+ * guarantees that only one backend initializes the DSA and that all other
+ * backends just attach it.
+ *
+ * A dshash table can be created in or retrieved from the registry by
+ * calling GetNamedDSMHash().  As with GetNamedDSMSegment(), if a hash
+ * table with the provided name does not yet exist, it is created.
+ * Otherwise, GetNamedDSMHash() ensures the hash table is attached to the
+ * current backend.  This function guarantees that only one backend
+ * initializes the table and that all other backends just attach it.
+ *
  * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
@@ -32,6 +46,12 @@
 #include "storage/shmem.h"
 #include "utils/memutils.h"
 
+#define DSMR_NAME_LEN				128
+
+#define DSMR_DSA_TRANCHE_SUFFIX		" DSA"
+#define DSMR_DSA_TRANCHE_SUFFIX_LEN (sizeof(DSMR_DSA_TRANCHE_SUFFIX) - 1)
+#define DSMR_DSA_TRANCHE_NAME_LEN	(DSMR_NAME_LEN + DSMR_DSA_TRANCHE_SUFFIX_LEN)
+
 typedef struct DSMRegistryCtxStruct
 {
 	dsa_handle	dsah;
@@ -42,7 +62,7 @@ static DSMRegistryCtxStruct *DSMRegistryCtx;
 
 typedef struct DSMRegistryEntry
 {
-	char		name[64];
+	char		name[DSMR_NAME_LEN];
 	dsm_handle	handle;
 	size_t		size;
 } DSMRegistryEntry;
@@ -56,6 +76,21 @@ static const dshash_parameters dsh_params = {
 	LWTRANCHE_DSM_REGISTRY_HASH
 };
 
+typedef struct NamedDSAState
+{
+	dsa_handle	dsah;
+	int			dsa_tranche;
+	char		dsa_tranche_name[DSMR_DSA_TRANCHE_NAME_LEN];
+} NamedDSAState;
+
+typedef struct NamedDSMHashState
+{
+	NamedDSAState dsa;
+	dshash_table_handle dshh;
+	int			dsh_tranche;
+	char		dsh_tranche_name[DSMR_NAME_LEN];
+} NamedDSMHashState;
+
 static dsa_area *dsm_registry_dsa;
 static dshash_table *dsm_registry_table;
 
@@ -125,7 +160,8 @@ 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.  Note that should be called
+ * at most once for a given segment in each backend.
  */
 void *
 GetNamedDSMSegment(const char *name, size_t size,
@@ -198,3 +234,165 @@ GetNamedDSMSegment(const char *name, size_t size,
 
 	return ret;
 }
+
+static void
+init_named_dsa_state(void *ptr)
+{
+	NamedDSAState *state = (NamedDSAState *) ptr;
+
+	state->dsah = DSA_HANDLE_INVALID;
+	state->dsa_tranche = -1;
+	state->dsa_tranche_name[0] = '\0';
+}
+
+/*
+ * Initialize or attach a named DSA.
+ *
+ * This routine returns a pointer to the DSA.  A new LWLock tranche ID will be
+ * generated if needed.  Note that the lock tranche will be registered with the
+ * provided name.  Also note that this should be called at most once for a
+ * given DSA in each backend.
+ */
+dsa_area *
+GetNamedDSA(const char *name, bool *found)
+{
+	NamedDSAState *state;
+	MemoryContext oldcontext;
+	dsa_area   *dsa;
+	bool		unused;
+
+	state = GetNamedDSMSegment(name, sizeof(NamedDSAState),
+							   init_named_dsa_state, &unused);
+
+	/* Use a lock to ensure only one process creates the DSA. */
+	LWLockAcquire(DSMRegistryLock, LW_EXCLUSIVE);
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	if (state->dsah == DSA_HANDLE_INVALID)
+	{
+		/* Initialize LWLock tranche for the DSA. */
+		state->dsa_tranche = LWLockNewTrancheId();
+		strcpy(state->dsa_tranche_name, name);
+		LWLockRegisterTranche(state->dsa_tranche, state->dsa_tranche_name);
+
+		/* Initialize DSA. */
+		dsa = dsa_create(state->dsa_tranche);
+		dsa_pin(dsa);
+		dsa_pin_mapping(dsa);
+
+		/* Store handle in the DSM segment for other backends to use. */
+		state->dsah = dsa_get_handle(dsa);
+
+		*found = false;
+	}
+	else
+	{
+		/* Initialize existing LWLock tranche for the DSA. */
+		LWLockRegisterTranche(state->dsa_tranche, state->dsa_tranche_name);
+
+		/* Attach to existing DSA. */
+		dsa = dsa_attach(state->dsah);
+		dsa_pin_mapping(dsa);
+
+		*found = true;
+	}
+
+	MemoryContextSwitchTo(oldcontext);
+	LWLockRelease(DSMRegistryLock);
+
+	return dsa;
+}
+
+static void
+init_named_hash_state(void *ptr)
+{
+	NamedDSMHashState *state = (NamedDSMHashState *) ptr;
+
+	init_named_dsa_state(&state->dsa);
+
+	state->dshh = DSHASH_HANDLE_INVALID;
+	state->dsh_tranche = -1;
+	state->dsh_tranche_name[0] = '\0';
+}
+
+/*
+ * Initialize or attach a named dshash table.
+ *
+ * This routine returns the address of the table.  The tranche_id member of
+ * params is ignored; new tranche IDs will be generated if needed.  Note that
+ * the DSA lock tranche will be registered with the provided name with " DSA"
+ * appended.  The dshash lock tranche will be registered with the provided
+ * name.  Also note that this should be called at most once for a given table
+ * in each backend.
+ */
+dshash_table *
+GetNamedDSMHash(const char *name, const dshash_parameters *params, bool *found)
+{
+	NamedDSMHashState *state;
+	MemoryContext oldcontext;
+	dsa_area   *dsa;
+	dshash_table *dsh;
+	bool		unused;
+
+	state = GetNamedDSMSegment(name, sizeof(NamedDSMHashState),
+							   init_named_hash_state, &unused);
+
+	/* Use a lock to ensure only one process creates the table. */
+	LWLockAcquire(DSMRegistryLock, LW_EXCLUSIVE);
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	if (state->dshh == DSHASH_HANDLE_INVALID)
+	{
+		dshash_parameters params_copy;
+
+		/* Initialize LWLock tranche for the DSA. */
+		state->dsa.dsa_tranche = LWLockNewTrancheId();
+		sprintf(state->dsa.dsa_tranche_name, "%s%s", name, DSMR_DSA_TRANCHE_SUFFIX);
+		LWLockRegisterTranche(state->dsa.dsa_tranche, state->dsa.dsa_tranche_name);
+
+		/* Initialize LWLock tranche for the dshash table. */
+		state->dsh_tranche = LWLockNewTrancheId();
+		strcpy(state->dsh_tranche_name, name);
+		LWLockRegisterTranche(state->dsh_tranche, state->dsh_tranche_name);
+
+		/* Initialize DSA for the hash table. */
+		dsa = dsa_create(state->dsa.dsa_tranche);
+		dsa_pin(dsa);
+		dsa_pin_mapping(dsa);
+
+		/* Initialize the dshash table. */
+		memcpy(&params_copy, params, sizeof(dshash_parameters));
+		params_copy.tranche_id = state->dsh_tranche;
+		dsh = dshash_create(dsa, &params_copy, NULL);
+
+		/* Store handles in the DSM segment for other backends to use. */
+		state->dsa.dsah = dsa_get_handle(dsa);
+		state->dshh = dshash_get_hash_table_handle(dsh);
+
+		*found = false;
+	}
+	else
+	{
+		/* Initialize existing LWLock tranches for the DSA and dshash table. */
+		LWLockRegisterTranche(state->dsa.dsa_tranche, state->dsa.dsa_tranche_name);
+		LWLockRegisterTranche(state->dsh_tranche, state->dsh_tranche_name);
+
+		/* Attach to existing DSA for the hash table. */
+		dsa = dsa_attach(state->dsa.dsah);
+		dsa_pin_mapping(dsa);
+
+		/* Attach to existing dshash table. */
+		dsh = dshash_attach(dsa, params, state->dshh, NULL);
+
+		*found = true;
+	}
+
+	MemoryContextSwitchTo(oldcontext);
+	LWLockRelease(DSMRegistryLock);
+
+	return dsh;
+}
diff --git a/src/include/storage/dsm_registry.h b/src/include/storage/dsm_registry.h
index b381e44bc9d..841bc37af82 100644
--- a/src/include/storage/dsm_registry.h
+++ b/src/include/storage/dsm_registry.h
@@ -13,10 +13,15 @@
 #ifndef DSM_REGISTRY_H
 #define DSM_REGISTRY_H
 
+#include "lib/dshash.h"
+
 extern void *GetNamedDSMSegment(const char *name, size_t size,
 								void (*init_callback) (void *ptr),
 								bool *found);
-
+extern dsa_area *GetNamedDSA(const char *name, bool *found);
+extern dshash_table *GetNamedDSMHash(const char *name,
+									 const dshash_parameters *params,
+									 bool *found);
 extern Size DSMRegistryShmemSize(void);
 extern void DSMRegistryShmemInit(void);
 
diff --git a/src/test/modules/test_dsm_registry/expected/test_dsm_registry.out b/src/test/modules/test_dsm_registry/expected/test_dsm_registry.out
index 8ffbd343a05..7ee02bb51e3 100644
--- a/src/test/modules/test_dsm_registry/expected/test_dsm_registry.out
+++ b/src/test/modules/test_dsm_registry/expected/test_dsm_registry.out
@@ -1,14 +1,26 @@
 CREATE EXTENSION test_dsm_registry;
-SELECT set_val_in_shmem(1236);
- set_val_in_shmem 
-------------------
+SELECT set_val_in_dsm(1236);
+ set_val_in_dsm 
+----------------
+ 
+(1 row)
+
+SELECT set_val_in_hash('test', '1414');
+ set_val_in_hash 
+-----------------
  
 (1 row)
 
 \c
-SELECT get_val_in_shmem();
- get_val_in_shmem 
-------------------
-             1236
+SELECT get_val_in_dsm();
+ get_val_in_dsm 
+----------------
+           1236
+(1 row)
+
+SELECT get_val_in_hash('test');
+ get_val_in_hash 
+-----------------
+ 1414
 (1 row)
 
diff --git a/src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql b/src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql
index b3351be0a16..7076f825260 100644
--- a/src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql
+++ b/src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql
@@ -1,4 +1,6 @@
 CREATE EXTENSION test_dsm_registry;
-SELECT set_val_in_shmem(1236);
+SELECT set_val_in_dsm(1236);
+SELECT set_val_in_hash('test', '1414');
 \c
-SELECT get_val_in_shmem();
+SELECT get_val_in_dsm();
+SELECT get_val_in_hash('test');
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql b/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
index 8c55b0919b1..74ceeccfd3b 100644
--- a/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
@@ -3,8 +3,14 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION test_dsm_registry" to load this file. \quit
 
-CREATE FUNCTION set_val_in_shmem(val INT) RETURNS VOID
+CREATE FUNCTION set_val_in_dsm(val INT) RETURNS VOID
 	AS 'MODULE_PATHNAME' LANGUAGE C;
 
-CREATE FUNCTION get_val_in_shmem() RETURNS INT
+CREATE FUNCTION get_val_in_dsm() RETURNS INT
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION set_val_in_hash(key TEXT, val TEXT) RETURNS VOID
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION get_val_in_hash(key TEXT) RETURNS TEXT
 	AS 'MODULE_PATHNAME' LANGUAGE C;
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 96a890be228..6be142d00d8 100644
--- a/src/test/modules/test_dsm_registry/test_dsm_registry.c
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -15,6 +15,7 @@
 #include "fmgr.h"
 #include "storage/dsm_registry.h"
 #include "storage/lwlock.h"
+#include "utils/builtins.h"
 
 PG_MODULE_MAGIC;
 
@@ -24,15 +25,31 @@ typedef struct TestDSMRegistryStruct
 	LWLock		lck;
 } TestDSMRegistryStruct;
 
-static TestDSMRegistryStruct *tdr_state;
+typedef struct TestDSMRegistryHashEntry
+{
+	char		key[64];
+	dsa_handle	val;
+} TestDSMRegistryHashEntry;
+
+static TestDSMRegistryStruct *tdr_dsm;
+static dsa_area *tdr_dsa;
+static dshash_table *tdr_hash;
+
+static const dshash_parameters dsh_params = {
+	offsetof(TestDSMRegistryHashEntry, val),
+	sizeof(TestDSMRegistryHashEntry),
+	dshash_strcmp,
+	dshash_strhash,
+	dshash_strcpy
+};
 
 static void
-tdr_init_shmem(void *ptr)
+init_tdr_dsm(void *ptr)
 {
-	TestDSMRegistryStruct *state = (TestDSMRegistryStruct *) ptr;
+	TestDSMRegistryStruct *dsm = (TestDSMRegistryStruct *) ptr;
 
-	LWLockInitialize(&state->lck, LWLockNewTrancheId());
-	state->val = 0;
+	LWLockInitialize(&dsm->lck, LWLockNewTrancheId());
+	dsm->val = 0;
 }
 
 static void
@@ -40,37 +57,94 @@ tdr_attach_shmem(void)
 {
 	bool		found;
 
-	tdr_state = GetNamedDSMSegment("test_dsm_registry",
-								   sizeof(TestDSMRegistryStruct),
-								   tdr_init_shmem,
-								   &found);
-	LWLockRegisterTranche(tdr_state->lck.tranche, "test_dsm_registry");
+	if (tdr_dsm == NULL)
+	{
+		tdr_dsm = GetNamedDSMSegment("test_dsm_registry_dsm",
+									 sizeof(TestDSMRegistryStruct),
+									 init_tdr_dsm,
+									 &found);
+		LWLockRegisterTranche(tdr_dsm->lck.tranche, "test_dsm_registry");
+	}
+
+	if (tdr_dsa == NULL)
+		tdr_dsa = GetNamedDSA("test_dsm_registry_dsa", &found);
+
+	if (tdr_hash == NULL)
+		tdr_hash = GetNamedDSMHash("test_dsm_registry_hash", &dsh_params, &found);
 }
 
-PG_FUNCTION_INFO_V1(set_val_in_shmem);
+PG_FUNCTION_INFO_V1(set_val_in_dsm);
 Datum
-set_val_in_shmem(PG_FUNCTION_ARGS)
+set_val_in_dsm(PG_FUNCTION_ARGS)
 {
 	tdr_attach_shmem();
 
-	LWLockAcquire(&tdr_state->lck, LW_EXCLUSIVE);
-	tdr_state->val = PG_GETARG_INT32(0);
-	LWLockRelease(&tdr_state->lck);
+	LWLockAcquire(&tdr_dsm->lck, LW_EXCLUSIVE);
+	tdr_dsm->val = PG_GETARG_INT32(0);
+	LWLockRelease(&tdr_dsm->lck);
 
 	PG_RETURN_VOID();
 }
 
-PG_FUNCTION_INFO_V1(get_val_in_shmem);
+PG_FUNCTION_INFO_V1(get_val_in_dsm);
 Datum
-get_val_in_shmem(PG_FUNCTION_ARGS)
+get_val_in_dsm(PG_FUNCTION_ARGS)
 {
 	int			ret;
 
 	tdr_attach_shmem();
 
-	LWLockAcquire(&tdr_state->lck, LW_SHARED);
-	ret = tdr_state->val;
-	LWLockRelease(&tdr_state->lck);
+	LWLockAcquire(&tdr_dsm->lck, LW_SHARED);
+	ret = tdr_dsm->val;
+	LWLockRelease(&tdr_dsm->lck);
 
 	PG_RETURN_INT32(ret);
 }
+
+PG_FUNCTION_INFO_V1(set_val_in_hash);
+Datum
+set_val_in_hash(PG_FUNCTION_ARGS)
+{
+	TestDSMRegistryHashEntry *entry;
+	char	   *key = TextDatumGetCString(PG_GETARG_DATUM(0));
+	char	   *val = TextDatumGetCString(PG_GETARG_DATUM(1));
+	bool		found;
+
+	if (strlen(key) >= offsetof(TestDSMRegistryHashEntry, val))
+		ereport(ERROR,
+				(errmsg("key too long")));
+
+	tdr_attach_shmem();
+
+	entry = dshash_find_or_insert(tdr_hash, key, &found);
+	if (found)
+		dsa_free(tdr_dsa, entry->val);
+
+	entry->val = dsa_allocate(tdr_dsa, strlen(val) + 1);
+	strcpy(dsa_get_address(tdr_dsa, entry->val), val);
+
+	dshash_release_lock(tdr_hash, entry);
+
+	PG_RETURN_VOID();
+}
+
+PG_FUNCTION_INFO_V1(get_val_in_hash);
+Datum
+get_val_in_hash(PG_FUNCTION_ARGS)
+{
+	TestDSMRegistryHashEntry *entry;
+	char	   *key = TextDatumGetCString(PG_GETARG_DATUM(0));
+	text	   *val = NULL;
+
+	tdr_attach_shmem();
+
+	entry = dshash_find(tdr_hash, key, false);
+	if (entry == NULL)
+		PG_RETURN_NULL();
+
+	val = cstring_to_text(dsa_get_address(tdr_dsa, entry->val));
+
+	dshash_release_lock(tdr_hash, entry);
+
+	PG_RETURN_TEXT_P(val);
+}
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index a8346cda633..245ab978ecb 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1730,6 +1730,8 @@ Name
 NameData
 NameHashEntry
 NamedArgExpr
+NamedDSAState
+NamedDSMHashState
 NamedLWLockTranche
 NamedLWLockTrancheRequest
 NamedTuplestoreScan
@@ -2998,6 +3000,7 @@ Tcl_NotifierProcs
 Tcl_Obj
 Tcl_Time
 TempNamespaceStatus
+TestDSMRegistryHashEntry
 TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
-- 
2.39.5 (Apple Git-154)

