1 files changed, 251 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;
+}