Add GetNamedDSA() and GetNamedDSHash().

Presently, the dynamic shared memory (DSM) registry only provides
GetNamedDSMSegment(), which allocates a fixed-size segment.  To use
the DSM registry for more sophisticated things like dynamic shared
memory areas (DSAs) or a hash table backed by a DSA (dshash), users
need to create a DSM segment that stores various handles and LWLock
tranche IDs and to write fairly complicated initialization code.
Furthermore, there is likely little variation in this
initialization code between libraries.

This commit introduces functions that simplify allocating a DSA or
dshash within the DSM registry.  These functions are very similar
to GetNamedDSMSegment().  Notable differences include the lack of
an initialization callback parameter and the prohibition of calling
the functions more than once for a given entry in each backend
(which should be trivially avoidable in most circumstances).  While
at it, this commit bumps the maximum DSM registry entry name length
from 63 bytes to 127 bytes.

Also note that even though one could presumably detach/destroy the
DSAs and dshashes created in the registry, such use-cases are not
yet well-supported, if for no other reason than the associated DSM
registry entries cannot be removed.  Adding such support is left as
a future exercise.

The test_dsm_registry test module contains tests for the new
functions and also serves as a complete usage example.

Reviewed-by: Dagfinn Ilmari Mannsåker <ilmari@ilmari.org>
Reviewed-by: Sami Imseih <samimseih@gmail.com>
Reviewed-by: Florents Tselai <florents.tselai@gmail.com>
Reviewed-by: Rahila Syed <rahilasyed90@gmail.com>
Discussion: https://postgr.es/m/aEC8HGy2tRQjZg_8%40nathan

2 files changed, 266 insertions, 14 deletions

diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index 1d4fd31ffed..828c2ff0c7f 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 GetNamedDSHash().  As with GetNamedDSMSegment(), if a hash
+ * table with the provided name does not yet exist, it is created.
+ * Otherwise, GetNamedDSHash() 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;
@@ -40,15 +60,48 @@ typedef struct DSMRegistryCtxStruct
 
 static DSMRegistryCtxStruct *DSMRegistryCtx;
 
-typedef struct DSMRegistryEntry
+typedef struct NamedDSMState
 {
-	char		name[64];
 	dsm_handle	handle;
 	size_t		size;
+} NamedDSMState;
+
+typedef struct NamedDSAState
+{
+	dsa_handle	handle;
+	int			tranche;
+	char		tranche_name[DSMR_DSA_TRANCHE_NAME_LEN];
+} NamedDSAState;
+
+typedef struct NamedDSHState
+{
+	NamedDSAState dsa;
+	dshash_table_handle handle;
+	int			tranche;
+	char		tranche_name[DSMR_NAME_LEN];
+} NamedDSHState;
+
+typedef enum DSMREntryType
+{
+	DSMR_ENTRY_TYPE_DSM,
+	DSMR_ENTRY_TYPE_DSA,
+	DSMR_ENTRY_TYPE_DSH,
+} DSMREntryType;
+
+typedef struct DSMRegistryEntry
+{
+	char		name[DSMR_NAME_LEN];
+	DSMREntryType type;
+	union
+	{
+		NamedDSMState dsm;
+		NamedDSAState dsa;
+		NamedDSHState dsh;
+	}			data;
 } DSMRegistryEntry;
 
 static const dshash_parameters dsh_params = {
-	offsetof(DSMRegistryEntry, handle),
+	offsetof(DSMRegistryEntry, type),
 	sizeof(DSMRegistryEntry),
 	dshash_strcmp,
 	dshash_strhash,
@@ -141,7 +194,7 @@ GetNamedDSMSegment(const char *name, size_t size,
 		ereport(ERROR,
 				(errmsg("DSM segment name cannot be empty")));
 
-	if (strlen(name) >= offsetof(DSMRegistryEntry, handle))
+	if (strlen(name) >= offsetof(DSMRegistryEntry, type))
 		ereport(ERROR,
 				(errmsg("DSM segment name too long")));
 
@@ -158,32 +211,39 @@ GetNamedDSMSegment(const char *name, size_t size,
 	entry = dshash_find_or_insert(dsm_registry_table, name, found);
 	if (!(*found))
 	{
+		NamedDSMState *state = &entry->data.dsm;
+		dsm_segment *seg;
+
+		entry->type = DSMR_ENTRY_TYPE_DSM;
+
 		/* Initialize the segment. */
-		dsm_segment *seg = dsm_create(size, 0);
+		seg = dsm_create(size, 0);
 
 		dsm_pin_segment(seg);
 		dsm_pin_mapping(seg);
-		entry->handle = dsm_segment_handle(seg);
-		entry->size = size;
+		state->handle = dsm_segment_handle(seg);
+		state->size = size;
 		ret = dsm_segment_address(seg);
 
 		if (init_callback)
 			(*init_callback) (ret);
 	}
-	else if (entry->size != size)
-	{
+	else if (entry->type != DSMR_ENTRY_TYPE_DSM)
 		ereport(ERROR,
-				(errmsg("requested DSM segment size does not match size of "
-						"existing segment")));
-	}
+				(errmsg("requested DSM segment does not match type of existing entry")));
+	else if (entry->data.dsm.size != size)
+		ereport(ERROR,
+				(errmsg("requested DSM segment size does not match size of existing segment")));
 	else
 	{
-		dsm_segment *seg = dsm_find_mapping(entry->handle);
+		NamedDSMState *state = &entry->data.dsm;
+		dsm_segment *seg;
 
 		/* If the existing segment is not already attached, attach it now. */
+		seg = dsm_find_mapping(state->handle);
 		if (seg == NULL)
 		{
-			seg = dsm_attach(entry->handle);
+			seg = dsm_attach(state->handle);
 			if (seg == NULL)
 				elog(ERROR, "could not map dynamic shared memory segment");
 
@@ -198,3 +258,180 @@ GetNamedDSMSegment(const char *name, size_t size,
 
 	return ret;
 }
+
+/*
+ * 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)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+	dsa_area   *ret;
+
+	Assert(found);
+
+	if (!name || *name == '\0')
+		ereport(ERROR,
+				(errmsg("DSA name cannot be empty")));
+
+	if (strlen(name) >= offsetof(DSMRegistryEntry, type))
+		ereport(ERROR,
+				(errmsg("DSA name too long")));
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	entry = dshash_find_or_insert(dsm_registry_table, name, found);
+	if (!(*found))
+	{
+		NamedDSAState *state = &entry->data.dsa;
+
+		entry->type = DSMR_ENTRY_TYPE_DSA;
+
+		/* Initialize the LWLock tranche for the DSA. */
+		state->tranche = LWLockNewTrancheId();
+		strcpy(state->tranche_name, name);
+		LWLockRegisterTranche(state->tranche, state->tranche_name);
+
+		/* Initialize the DSA. */
+		ret = dsa_create(state->tranche);
+		dsa_pin(ret);
+		dsa_pin_mapping(ret);
+
+		/* Store handle for other backends to use. */
+		state->handle = dsa_get_handle(ret);
+	}
+	else if (entry->type != DSMR_ENTRY_TYPE_DSA)
+		ereport(ERROR,
+				(errmsg("requested DSA does not match type of existing entry")));
+	else
+	{
+		NamedDSAState *state = &entry->data.dsa;
+
+		if (dsa_is_attached(state->handle))
+			ereport(ERROR,
+					(errmsg("requested DSA already attached to current process")));
+
+		/* Initialize existing LWLock tranche for the DSA. */
+		LWLockRegisterTranche(state->tranche, state->tranche_name);
+
+		/* Attach to existing DSA. */
+		ret = dsa_attach(state->handle);
+		dsa_pin_mapping(ret);
+	}
+
+	dshash_release_lock(dsm_registry_table, entry);
+	MemoryContextSwitchTo(oldcontext);
+
+	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; 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 *
+GetNamedDSHash(const char *name, const dshash_parameters *params, bool *found)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+	dshash_table *ret;
+
+	Assert(params);
+	Assert(found);
+
+	if (!name || *name == '\0')
+		ereport(ERROR,
+				(errmsg("DSHash name cannot be empty")));
+
+	if (strlen(name) >= offsetof(DSMRegistryEntry, type))
+		ereport(ERROR,
+				(errmsg("DSHash name too long")));
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	entry = dshash_find_or_insert(dsm_registry_table, name, found);
+	if (!(*found))
+	{
+		NamedDSAState *dsa_state = &entry->data.dsh.dsa;
+		NamedDSHState *dsh_state = &entry->data.dsh;
+		dshash_parameters params_copy;
+		dsa_area   *dsa;
+
+		entry->type = DSMR_ENTRY_TYPE_DSH;
+
+		/* Initialize the LWLock tranche for the DSA. */
+		dsa_state->tranche = LWLockNewTrancheId();
+		sprintf(dsa_state->tranche_name, "%s%s", name, DSMR_DSA_TRANCHE_SUFFIX);
+		LWLockRegisterTranche(dsa_state->tranche, dsa_state->tranche_name);
+
+		/* Initialize the LWLock tranche for the dshash table. */
+		dsh_state->tranche = LWLockNewTrancheId();
+		strcpy(dsh_state->tranche_name, name);
+		LWLockRegisterTranche(dsh_state->tranche, dsh_state->tranche_name);
+
+		/* Initialize the DSA for the hash table. */
+		dsa = dsa_create(dsa_state->tranche);
+		dsa_pin(dsa);
+		dsa_pin_mapping(dsa);
+
+		/* Initialize the dshash table. */
+		memcpy(&params_copy, params, sizeof(dshash_parameters));
+		params_copy.tranche_id = dsh_state->tranche;
+		ret = dshash_create(dsa, &params_copy, NULL);
+
+		/* Store handles for other backends to use. */
+		dsa_state->handle = dsa_get_handle(dsa);
+		dsh_state->handle = dshash_get_hash_table_handle(ret);
+	}
+	else if (entry->type != DSMR_ENTRY_TYPE_DSH)
+		ereport(ERROR,
+				(errmsg("requested DSHash does not match type of existing entry")));
+	else
+	{
+		NamedDSAState *dsa_state = &entry->data.dsh.dsa;
+		NamedDSHState *dsh_state = &entry->data.dsh;
+		dsa_area   *dsa;
+
+		/* XXX: Should we verify params matches what table was created with? */
+
+		if (dsa_is_attached(dsa_state->handle))
+			ereport(ERROR,
+					(errmsg("requested DSHash already attached to current process")));
+
+		/* Initialize existing LWLock tranches for the DSA and dshash table. */
+		LWLockRegisterTranche(dsa_state->tranche, dsa_state->tranche_name);
+		LWLockRegisterTranche(dsh_state->tranche, dsh_state->tranche_name);
+
+		/* Attach to existing DSA for the hash table. */
+		dsa = dsa_attach(dsa_state->handle);
+		dsa_pin_mapping(dsa);
+
+		/* Attach to existing dshash table. */
+		ret = dshash_attach(dsa, params, dsh_state->handle, NULL);
+	}
+
+	dshash_release_lock(dsm_registry_table, entry);
+	MemoryContextSwitchTo(oldcontext);
+
+	return ret;
+}
diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 17d4f7a7a06..be43e9351c3 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -532,6 +532,21 @@ dsa_attach(dsa_handle handle)
 }
 
 /*
+ * Returns whether the area with the given handle was already attached by the
+ * current process.  The area must have been created with dsa_create (not
+ * dsa_create_in_place).
+ */
+bool
+dsa_is_attached(dsa_handle handle)
+{
+	/*
+	 * An area handle is really a DSM segment handle for the first segment, so
+	 * we can just search for that.
+	 */
+	return dsm_find_mapping(handle) != NULL;
+}
+
+/*
  * Attach to an area that was created with dsa_create_in_place.  The caller
  * must somehow know the location in memory that was used when the area was
  * created, though it may be mapped at a different virtual address in this