From a9de0367fc504de403b93be179f36615a3884fc7 Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Tue, 31 Mar 2026 23:30:11 -0700
Subject: [PATCH vnocfbot-3] Use GetNamedDSA / GetNamedDSHash for shared memory
 allocations

This allows utilizing pg_get_dsm_registry_allocations to see the amount
of memory used by the internal storage.
---
 contrib/pg_stash_advice/pg_stash_advice.c | 108 +++-------------------
 1 file changed, 11 insertions(+), 97 deletions(-)

diff --git a/contrib/pg_stash_advice/pg_stash_advice.c b/contrib/pg_stash_advice/pg_stash_advice.c
index 22122236694..c74d0eebf45 100644
--- a/contrib/pg_stash_advice/pg_stash_advice.c
+++ b/contrib/pg_stash_advice/pg_stash_advice.c
@@ -43,13 +43,7 @@ PG_FUNCTION_INFO_V1(pg_set_stashed_advice);
 typedef struct pgsa_shared_state
 {
 	LWLock		lock;
-	int			dsa_tranche;
-	int			stash_tranche;
-	int			entry_tranche;
 	uint64		next_stash_id;
-	dsa_handle	area;
-	dshash_table_handle stash_hash;
-	dshash_table_handle entry_hash;
 } pgsa_shared_state;
 
 typedef struct pgsa_stash
@@ -113,29 +107,27 @@ static dshash_table *pgsa_stash_dshash;
 static dshash_table *pgsa_entry_dshash;
 
 /* Shared memory hash table parameters */
-static dshash_parameters pgsa_stash_dshash_parameters = {
+static const dshash_parameters pgsa_stash_dshash_parameters = {
 	NAMEDATALEN,
 	sizeof(pgsa_stash),
 	dshash_strcmp,
 	dshash_strhash,
 	dshash_strcpy,
-	LWTRANCHE_INVALID			/* gets set at runtime */
+	LWTRANCHE_INVALID			/* assigned by GetNamedDSHash */
 };
 
-static dshash_parameters pgsa_entry_dshash_parameters = {
+static const dshash_parameters pgsa_entry_dshash_parameters = {
 	sizeof(pgsa_entry_key),
 	sizeof(pgsa_entry),
 	dshash_memcmp,
 	dshash_memhash,
 	dshash_memcpy,
-	LWTRANCHE_INVALID			/* gets set at runtime */
+	LWTRANCHE_INVALID			/* assigned by GetNamedDSHash */
 };
 
 /* GUC variable */
 static char *pg_stash_advice_stash_name = "";
 
-/* Other global variables */
-static MemoryContext pg_stash_advice_mcxt;
 
 /* Function prototypes */
 static char *pgsa_advisor(PlannerGlobal *glob,
@@ -519,17 +511,6 @@ static void
 pgsa_attach(void)
 {
 	bool		found;
-	MemoryContext oldcontext;
-
-	/*
-	 * Create a memory context to make sure that any control structures
-	 * allocated in local memory are sufficiently persistent.
-	 */
-	if (pg_stash_advice_mcxt == NULL)
-		pg_stash_advice_mcxt = AllocSetContextCreate(TopMemoryContext,
-													 "pg_stash_advice",
-													 ALLOCSET_DEFAULT_SIZES);
-	oldcontext = MemoryContextSwitchTo(pg_stash_advice_mcxt);
 
 	/* Attach to the fixed-size state object if not already done. */
 	if (pgsa_state == NULL)
@@ -540,80 +521,19 @@ pgsa_attach(void)
 
 	/* Attach to the DSA area if not already done. */
 	if (pgsa_dsa_area == NULL)
-	{
-		dsa_handle	area_handle;
-
-		LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
-		area_handle = pgsa_state->area;
-		if (area_handle == DSA_HANDLE_INVALID)
-		{
-			pgsa_dsa_area = dsa_create(pgsa_state->dsa_tranche);
-			dsa_pin(pgsa_dsa_area);
-			pgsa_state->area = dsa_get_handle(pgsa_dsa_area);
-			LWLockRelease(&pgsa_state->lock);
-		}
-		else
-		{
-			LWLockRelease(&pgsa_state->lock);
-			pgsa_dsa_area = dsa_attach(area_handle);
-		}
-		dsa_pin_mapping(pgsa_dsa_area);
-	}
+		pgsa_dsa_area = GetNamedDSA("pg_stash_advice_dsa", &found);
 
 	/* Attach to the stash_name->stash_id hash table if not already done. */
 	if (pgsa_stash_dshash == NULL)
-	{
-		dshash_table_handle stash_handle;
-
-		LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
-		pgsa_stash_dshash_parameters.tranche_id = pgsa_state->stash_tranche;
-		stash_handle = pgsa_state->stash_hash;
-		if (stash_handle == DSHASH_HANDLE_INVALID)
-		{
-			pgsa_stash_dshash = dshash_create(pgsa_dsa_area,
-											  &pgsa_stash_dshash_parameters,
-											  NULL);
-			pgsa_state->stash_hash =
-				dshash_get_hash_table_handle(pgsa_stash_dshash);
-			LWLockRelease(&pgsa_state->lock);
-		}
-		else
-		{
-			LWLockRelease(&pgsa_state->lock);
-			pgsa_stash_dshash = dshash_attach(pgsa_dsa_area,
-											  &pgsa_stash_dshash_parameters,
-											  stash_handle, NULL);
-		}
-	}
+		pgsa_stash_dshash = GetNamedDSHash("pg_stash_advice_stash",
+										   &pgsa_stash_dshash_parameters,
+										   &found);
 
 	/* Attach to the entry hash table if not already done. */
 	if (pgsa_entry_dshash == NULL)
-	{
-		dshash_table_handle entry_handle;
-
-		LWLockAcquire(&pgsa_state->lock, LW_EXCLUSIVE);
-		pgsa_entry_dshash_parameters.tranche_id = pgsa_state->entry_tranche;
-		entry_handle = pgsa_state->entry_hash;
-		if (entry_handle == DSHASH_HANDLE_INVALID)
-		{
-			pgsa_entry_dshash = dshash_create(pgsa_dsa_area,
-											  &pgsa_entry_dshash_parameters,
-											  NULL);
-			pgsa_state->entry_hash =
-				dshash_get_hash_table_handle(pgsa_entry_dshash);
-			LWLockRelease(&pgsa_state->lock);
-		}
-		else
-		{
-			LWLockRelease(&pgsa_state->lock);
-			pgsa_entry_dshash = dshash_attach(pgsa_dsa_area,
-											  &pgsa_entry_dshash_parameters,
-											  entry_handle, NULL);
-		}
-	}
-
-	/* Restore previous memory context. */
-	MemoryContextSwitchTo(oldcontext);
+		pgsa_entry_dshash = GetNamedDSHash("pg_stash_advice_entry",
+										   &pgsa_entry_dshash_parameters,
+										   &found);
 }
 
 /*
@@ -789,13 +709,7 @@ pgsa_init_shared_state(void *ptr, void *arg)
 
 	LWLockInitialize(&state->lock,
 					 LWLockNewTrancheId("pg_stash_advice_lock"));
-	state->dsa_tranche = LWLockNewTrancheId("pg_stash_advice_dsa");
-	state->stash_tranche = LWLockNewTrancheId("pg_stash_advice_stash");
-	state->entry_tranche = LWLockNewTrancheId("pg_stash_advice_entry");
 	state->next_stash_id = UINT64CONST(1);
-	state->area = DSA_HANDLE_INVALID;
-	state->stash_hash = DSHASH_HANDLE_INVALID;
-	state->entry_hash = DSHASH_HANDLE_INVALID;
 }
 
 /*
-- 
2.47.1

