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

---
 src/backend/storage/ipc/dsm_registry.c        | 103 +++++++++++++++++-
 src/include/storage/dsm_registry.h            |   6 +-
 .../expected/test_dsm_registry.out            |  12 ++
 .../sql/test_dsm_registry.sql                 |   2 +
 .../test_dsm_registry--1.0.sql                |   6 +
 .../test_dsm_registry/test_dsm_registry.c     |  65 +++++++++++
 src/tools/pgindent/typedefs.list              |   2 +
 7 files changed, 194 insertions(+), 2 deletions(-)

diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index 1d4fd31ffed..f44848ce847 100644
--- a/src/backend/storage/ipc/dsm_registry.c
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -15,6 +15,13 @@
  * current backend.  This function guarantees that only one backend
  * initializes the segment 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 gurantees 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 +39,16 @@
 #include "storage/shmem.h"
 #include "utils/memutils.h"
 
+#define DSMR_NAME_LEN				64
+
+#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)
+
+#define DSMR_DSH_TRANCHE_SUFFIX		"_dsh"
+#define DSMR_DSH_TRANCHE_SUFFIX_LEN (sizeof(DSMR_DSH_TRANCHE_SUFFIX) - 1)
+#define DSMR_DSH_TRANCHE_NAME_LEN	(DSMR_NAME_LEN + DSMR_DSH_TRANCHE_SUFFIX_LEN)
+
 typedef struct DSMRegistryCtxStruct
 {
 	dsa_handle	dsah;
@@ -42,7 +59,7 @@ static DSMRegistryCtxStruct *DSMRegistryCtx;
 
 typedef struct DSMRegistryEntry
 {
-	char		name[64];
+	char		name[DSMR_NAME_LEN];
 	dsm_handle	handle;
 	size_t		size;
 } DSMRegistryEntry;
@@ -56,6 +73,16 @@ static const dshash_parameters dsh_params = {
 	LWTRANCHE_DSM_REGISTRY_HASH
 };
 
+typedef struct NamedDSMHashState
+{
+	dsa_handle	dsah;
+	dshash_table_handle dshh;
+	int			dsa_tranche;
+	char		dsa_tranche_name[DSMR_DSA_TRANCHE_NAME_LEN];
+	int			dsh_tranche;
+	char		dsh_tranche_name[DSMR_DSH_TRANCHE_NAME_LEN];
+} NamedDSMHashState;
+
 static dsa_area *dsm_registry_dsa;
 static dshash_table *dsm_registry_table;
 
@@ -198,3 +225,77 @@ GetNamedDSMSegment(const char *name, size_t size,
 
 	return ret;
 }
+
+/*
+ * Initialize or attach a named dshash table.
+ *
+ * This routine returns the address of the table.  The tranche_id member of
+ * params is ignored; a new tranche ID will be generated if needed.
+ */
+dshash_table *
+GetNamedDSMHash(const char *name, const dshash_parameters *params, bool *found)
+{
+	NamedDSMHashState *state;
+	MemoryContext oldcontext;
+	dsa_area   *dsa;
+	dshash_table *dsh;
+
+	state = GetNamedDSMSegment(name, sizeof(NamedDSMHashState), NULL, found);
+
+	/* 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_tranche = LWLockNewTrancheId();
+		sprintf(state->dsa_tranche_name, "%s%s", name, DSMR_DSA_TRANCHE_SUFFIX);
+		LWLockRegisterTranche(state->dsa_tranche, state->dsa_tranche_name);
+
+		/* Initialize LWLock tranche for the dshash table. */
+		state->dsh_tranche = LWLockNewTrancheId();
+		sprintf(state->dsh_tranche_name, "%s%s", name, DSMR_DSH_TRANCHE_SUFFIX);
+		LWLockRegisterTranche(state->dsh_tranche, state->dsh_tranche_name);
+
+		/* Initialize DSA for the hash table. */
+		dsa = dsa_create(state->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->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_tranche, state->dsa_tranche_name);
+		LWLockRegisterTranche(state->dsh_tranche, state->dsh_tranche_name);
+
+		/* Attach to existing DSA for the hash table. */
+		dsa = dsa_attach(state->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..be9f2338417 100644
--- a/src/include/storage/dsm_registry.h
+++ b/src/include/storage/dsm_registry.h
@@ -13,10 +13,14 @@
 #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 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..aab1b2a90da 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
@@ -5,6 +5,12 @@ SELECT set_val_in_shmem(1236);
  
 (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 
@@ -12,3 +18,9 @@ SELECT get_val_in_shmem();
              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..741c77098e0 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_hash('test', 1414);
 \c
 SELECT get_val_in_shmem();
+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..95fd6586bf5 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
@@ -8,3 +8,9 @@ CREATE FUNCTION set_val_in_shmem(val INT) RETURNS VOID
 
 CREATE FUNCTION get_val_in_shmem() RETURNS INT
 	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION set_val_in_hash(key TEXT, val INT) RETURNS VOID
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION get_val_in_hash(key TEXT) RETURNS INT
+	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 462a80f8790..5449dd8c22c 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,8 +25,24 @@ typedef struct TestDSMRegistryStruct
 	LWLock		lck;
 } TestDSMRegistryStruct;
 
+typedef struct TestDSMRegistryHashEntry
+{
+	char		key[64];
+	int			val;
+} TestDSMRegistryHashEntry;
+
 static TestDSMRegistryStruct *tdr_state;
 
+static const dshash_parameters dsh_params = {
+	offsetof(TestDSMRegistryHashEntry, val),
+	sizeof(TestDSMRegistryHashEntry),
+	dshash_strcmp,
+	dshash_strhash,
+	dshash_strcpy
+};
+
+static dshash_table *tdr_hash;
+
 static void
 tdr_init_shmem(void *ptr)
 {
@@ -74,3 +91,51 @@ get_val_in_shmem(PG_FUNCTION_ARGS)
 
 	PG_RETURN_UINT32(ret);
 }
+
+static void
+tdr_attach_hash(void)
+{
+	bool		found;
+
+	if (tdr_hash)
+		return;
+
+	tdr_hash = GetNamedDSMHash("test_dsm_registry_hash", &dsh_params, &found);
+}
+
+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));
+	int			val = PG_GETARG_INT32(1);
+	bool		found;
+
+	tdr_attach_hash();
+
+	entry = dshash_find_or_insert(tdr_hash, key, &found);
+	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));
+	int			val = -1;
+
+	tdr_attach_hash();
+
+	entry = dshash_find(tdr_hash, key, false);
+	val = entry->val;
+
+	dshash_release_lock(tdr_hash, entry);
+
+	PG_RETURN_INT32(val);
+}
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index a8346cda633..fc7c04f2dac 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1730,6 +1730,7 @@ Name
 NameData
 NameHashEntry
 NamedArgExpr
+NamedDSMHashState
 NamedLWLockTranche
 NamedLWLockTrancheRequest
 NamedTuplestoreScan
@@ -2998,6 +2999,7 @@ Tcl_NotifierProcs
 Tcl_Obj
 Tcl_Time
 TempNamespaceStatus
+TestDSMRegistryHashEntry
 TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
-- 
2.39.5 (Apple Git-154)

