From 59ade2d443da696d1a7dc6a25f73ad4a0357e35b Mon Sep 17 00:00:00 2001
From: Florents Tselai <florents.tselai@gmail.com>
Date: Tue, 3 Jun 2025 22:30:33 +0300
Subject: [PATCH v4 2/4] * Rename view to pg_dsm_registry_allocations and
 function to pg_get_dsm_registry_allocations * Simplify implementation for
 pg_get_dsm_registry_allocations * Simplify test

---
 src/backend/catalog/system_views.sql          | 12 +--
 src/backend/storage/ipc/dsm_registry.c        | 78 +++++--------------
 src/include/catalog/pg_proc.dat               |  4 +-
 .../expected/test_dsm_registry.out            |  9 +--
 .../sql/test_dsm_registry.sql                 |  3 +-
 src/test/regress/expected/rules.out           |  4 +-
 6 files changed, 35 insertions(+), 75 deletions(-)

diff --git a/src/backend/catalog/system_views.sql b/src/backend/catalog/system_views.sql
index 7c2246999ec..64e62732f18 100644
--- a/src/backend/catalog/system_views.sql
+++ b/src/backend/catalog/system_views.sql
@@ -666,13 +666,13 @@ GRANT SELECT ON pg_shmem_allocations_numa TO pg_read_all_stats;
 REVOKE EXECUTE ON FUNCTION pg_get_shmem_allocations_numa() FROM PUBLIC;
 GRANT EXECUTE ON FUNCTION pg_get_shmem_allocations_numa() TO pg_read_all_stats;
 
-CREATE VIEW pg_dsm_registry AS
-SELECT * FROM pg_get_dsm_registry();
+CREATE VIEW pg_dsm_registry_allocations AS
+SELECT * FROM pg_get_dsm_registry_allocations();
 
-REVOKE ALL ON pg_dsm_registry FROM PUBLIC;
-GRANT SELECT ON pg_dsm_registry TO pg_read_all_stats;
-REVOKE EXECUTE ON FUNCTION pg_get_dsm_registry() FROM PUBLIC;
-GRANT EXECUTE ON FUNCTION pg_get_dsm_registry() TO pg_read_all_stats;
+REVOKE ALL ON pg_dsm_registry_allocations FROM PUBLIC;
+GRANT SELECT ON pg_dsm_registry_allocations TO pg_read_all_stats;
+REVOKE EXECUTE ON FUNCTION pg_get_dsm_registry_allocations() FROM PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_get_dsm_registry_allocations() TO pg_read_all_stats;
 
 CREATE VIEW pg_backend_memory_contexts AS
     SELECT * FROM pg_get_backend_memory_contexts();
diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index 549cbe07a28..33e3d9ba3f0 100644
--- a/src/backend/storage/ipc/dsm_registry.c
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -203,81 +203,43 @@ GetNamedDSMSegment(const char *name, size_t size,
 	return ret;
 }
 
-void
-iterate_dsm_registry(void (*callback)(DSMRegistryEntry *, void *), void *arg);
-void
-iterate_dsm_registry(void (*callback)(DSMRegistryEntry *, void *), void *arg)
-{
-	DSMRegistryEntry *entry;
-	dshash_seq_status status;
-	/* Ensure DSM registry is initialized */
-	init_dsm_registry();
-
-	/* Use non-exclusive access to avoid blocking other backends */
-	dshash_seq_init(&status, dsm_registry_table, false);
-	while ((entry = dshash_seq_next(&status)) != NULL)
-		callback(entry, arg);
-	dshash_seq_term(&status);
-}
-
-/* SQL SRF showing DSM registry allocated memory */
-PG_FUNCTION_INFO_V1(pg_get_dsm_registry);
-
 typedef struct
 {
 	Tuplestorestate *tupstore;
 	TupleDesc        tupdesc;
 } DSMRegistrySRFContext;
 
-static void
-collect_dsm_registry(DSMRegistryEntry *entry, void *arg)
-{
-	DSMRegistrySRFContext *ctx = (DSMRegistrySRFContext *) arg;
-	Datum values[2];
-	bool nulls[2] = {false, false};
-
-	values[0] = CStringGetTextDatum(entry->name);
-	values[1] = Int64GetDatum(entry->size);
-
-	tuplestore_putvalues(ctx->tupstore, ctx->tupdesc, values, nulls);
-}
-
 Datum
-pg_get_dsm_registry(PG_FUNCTION_ARGS)
+pg_get_dsm_registry_allocations(PG_FUNCTION_ARGS)
 {
 	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
-	MemoryContext per_query_ctx;
-	MemoryContext oldcontext;
-	TupleDesc tupdesc;
-	Tuplestorestate *tupstore;
-	DSMRegistrySRFContext ctx;
+	DSMRegistryEntry *entry;
+	dshash_seq_status status;
 
 	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
-		ereport(ERROR, (errmsg("pg_get_dsm_registry must be used in a SRF context")));
+		ereport(ERROR, (errmsg("pg_get_dsm_registry_allocations must be used in a SRF context")));
+
+	InitMaterializedSRF(fcinfo, MAT_SRF_USE_EXPECTED_DESC);
 
-	/* Set up tuple descriptor */
-	tupdesc = CreateTemplateTupleDesc(2);
-	TupleDescInitEntry(tupdesc, (AttrNumber) 1, "name", TEXTOID, -1, 0);
-	TupleDescInitEntry(tupdesc, (AttrNumber) 2, "size", INT8OID, -1, 0);
+	/* Ensure DSM registry initialized */
+	init_dsm_registry();
 
-	/* Switch to per-query memory context */
-	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
-	oldcontext = MemoryContextSwitchTo(per_query_ctx);
+	/* Use non-exclusive access to avoid blocking other backends */
+	dshash_seq_init(&status, dsm_registry_table, false);
 
-	/* Initialize tuplestore */
-	tupstore = tuplestore_begin_heap(false, false, work_mem);
+	while ((entry = dshash_seq_next(&status)) != NULL)
+	{
+		Datum values[2];
+		bool nulls[2] = {false, false};
 
-	ctx.tupstore = tupstore;
-	ctx.tupdesc = tupdesc;
+		values[0] = CStringGetTextDatum(entry->name);
+		values[1] = Int64GetDatum(entry->size);
 
-	/* Collect registry data */
-	iterate_dsm_registry(collect_dsm_registry, &ctx);
+		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc,
+							 values, nulls);
+	}
 
-	/* Switch back and return results */
-	MemoryContextSwitchTo(oldcontext);
-	rsinfo->returnMode = SFRM_Materialize;
-	rsinfo->setResult = tupstore;
-	rsinfo->setDesc = tupdesc;
+	dshash_seq_term(&status);
 
 	return (Datum) 0;
 }
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index aa3fac2afb8..bef92be8250 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8542,11 +8542,11 @@
   proargnames => '{name,off,size,allocated_size}',
   prosrc => 'pg_get_shmem_allocations' },
 { oid => '6062', descr => 'dsm registry allocations',
-  proname => 'pg_get_dsm_registry', prorows => '50', proretset => 't',
+  proname => 'pg_get_dsm_registry_allocations', prorows => '50', proretset => 't',
   provolatile => 'v', prorettype => 'record', proargtypes => '',
   proallargtypes => '{text,int8}', proargmodes => '{o,o}',
   proargnames => '{name,size}',
-  prosrc => 'pg_get_dsm_registry' },
+  prosrc => 'pg_get_dsm_registry_allocations' },
 
 { oid => '4099', descr => 'Is NUMA support available?',
   proname => 'pg_numa_available', provolatile => 's', prorettype => 'bool',
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 229abf926b6..3c279378b5b 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
@@ -12,10 +12,9 @@ SELECT get_val_in_shmem();
              1236
 (1 row)
 
--- 20 bytes = int (4 bytes) + LWLock (16bytes)
-SELECT * FROM pg_dsm_registry;
-       name        | size 
--------------------+------
- test_dsm_registry |   20
+SELECT size > 0 FROM pg_dsm_registry_allocations WHERE name = 'test_dsm_registry';
+ ?column? 
+----------
+ t
 (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 aad402b5e64..d27ee2e6a1a 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
@@ -3,5 +3,4 @@ SELECT set_val_in_shmem(1236);
 \c
 SELECT get_val_in_shmem();
 
--- 20 bytes = int (4 bytes) + LWLock (16bytes)
-SELECT * FROM pg_dsm_registry;
+SELECT size > 0 FROM pg_dsm_registry_allocations WHERE name = 'test_dsm_registry';
diff --git a/src/test/regress/expected/rules.out b/src/test/regress/expected/rules.out
index ec257f93398..893c5b2b933 100644
--- a/src/test/regress/expected/rules.out
+++ b/src/test/regress/expected/rules.out
@@ -1340,9 +1340,9 @@ pg_cursors| SELECT name,
     is_scrollable,
     creation_time
    FROM pg_cursor() c(name, statement, is_holdable, is_binary, is_scrollable, creation_time);
-pg_dsm_registry| SELECT name,
+pg_dsm_registry_allocations| SELECT name,
     size
-   FROM pg_get_dsm_registry() pg_get_dsm_registry(name, size);
+   FROM pg_get_dsm_registry_allocations() pg_get_dsm_registry_allocations(name, size);
 pg_file_settings| SELECT sourcefile,
     sourceline,
     seqno,
-- 
2.49.0

