Show size of DSAs and dshashes in pg_dsm_registry_allocations.

Presently, this view reports NULL for the size of DSAs and dshash
tables because 1) the current backend might not be attached to them
and 2) the registry doesn't save the pointers to the dsa_area or
dshash_table in local memory.  Also, the view doesn't show
partially-initialized entries to avoid ambiguity, since those
entries would report a NULL size as well.

This commit introduces a function that looks up the size of a DSA
given its handle (transiently attaching to the control segment if
needed) and teaches pg_dsm_registry_allocations to use it to show
the size of successfully-initialized DSA and dshash entries.
Furthermore, the view now reports partially-initialized entries
with a NULL size.

Reviewed-by: Rahila Syed <rahilasyed90@gmail.com>
Reviewed-by: Robert Haas <robertmhaas@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/aSeEDeznAsHR1_YF%40nathan

6 files changed, 55 insertions, 26 deletions

diff --git a/doc/src/sgml/system-views.sgml b/doc/src/sgml/system-views.sgml
index 7db8f73eba2..162c76b729a 100644
--- a/doc/src/sgml/system-views.sgml
+++ b/doc/src/sgml/system-views.sgml
@@ -1150,8 +1150,8 @@ AND c1.path[c2.level] = c2.path[c2.level];
        <structfield>size</structfield> <type>int8</type>
       </para>
       <para>
-       Size of the allocation in bytes.  NULL for entries of type
-       <literal>area</literal> and <literal>hash</literal>.
+       Size of the allocation in bytes.  NULL for entries that failed
+       initialization.
       </para></entry>
      </row>
     </tbody>
diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
index ef6533b1100..66240318e83 100644
--- a/src/backend/storage/ipc/dsm_registry.c
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -463,26 +463,19 @@ pg_get_dsm_registry_allocations(PG_FUNCTION_ARGS)
 		Datum		vals[3];
 		bool		nulls[3] = {0};
 
-		/* Do not show partially-initialized entries. */
-		if (entry->type == DSMR_ENTRY_TYPE_DSM &&
-			entry->dsm.handle == DSM_HANDLE_INVALID)
-			continue;
-		if (entry->type == DSMR_ENTRY_TYPE_DSA &&
-			entry->dsa.handle == DSA_HANDLE_INVALID)
-			continue;
-		if (entry->type == DSMR_ENTRY_TYPE_DSH &&
-			entry->dsh.dsa_handle == DSA_HANDLE_INVALID)
-			continue;
-
 		vals[0] = CStringGetTextDatum(entry->name);
 		vals[1] = CStringGetTextDatum(DSMREntryTypeNames[entry->type]);
 
-		/*
-		 * Since we can't know the size of DSA/dshash entries without first
-		 * attaching to them, return NULL for those.
-		 */
-		if (entry->type == DSMR_ENTRY_TYPE_DSM)
+		/* Be careful to only return the sizes of initialized entries. */
+		if (entry->type == DSMR_ENTRY_TYPE_DSM &&
+			entry->dsm.handle != DSM_HANDLE_INVALID)
 			vals[2] = Int64GetDatum(entry->dsm.size);
+		else if (entry->type == DSMR_ENTRY_TYPE_DSA &&
+				 entry->dsa.handle != DSA_HANDLE_INVALID)
+			vals[2] = Int64GetDatum(dsa_get_total_size_from_handle(entry->dsa.handle));
+		else if (entry->type == DSMR_ENTRY_TYPE_DSH &&
+				 entry->dsh.dsa_handle !=DSA_HANDLE_INVALID)
+			vals[2] = Int64GetDatum(dsa_get_total_size_from_handle(entry->dsh.dsa_handle));
 		else
 			nulls[2] = true;
 
diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index be43e9351c3..c8a72686177 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -1051,6 +1051,41 @@ dsa_get_total_size(dsa_area *area)
 }
 
 /*
+ * Same as dsa_get_total_size(), but accepts a DSA handle.  The area must have
+ * been created with dsa_create (not dsa_create_in_place).
+ */
+size_t
+dsa_get_total_size_from_handle(dsa_handle handle)
+{
+	size_t		size;
+	bool		already_attached;
+	dsm_segment *segment;
+	dsa_area_control *control;
+
+	already_attached = dsa_is_attached(handle);
+	if (already_attached)
+		segment = dsm_find_mapping(handle);
+	else
+		segment = dsm_attach(handle);
+
+	if (segment == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("could not attach to dynamic shared area")));
+
+	control = (dsa_area_control *) dsm_segment_address(segment);
+
+	LWLockAcquire(&control->lock, LW_EXCLUSIVE);
+	size = control->total_segment_size;
+	LWLockRelease(&control->lock);
+
+	if (!already_attached)
+		dsm_detach(segment);
+
+	return size;
+}
+
+/*
  * Aggressively free all spare memory in the hope of returning DSM segments to
  * the operating system.
  */
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index f2104dacbfc..42449ff22de 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -161,6 +161,7 @@ extern dsa_pointer dsa_allocate_extended(dsa_area *area, size_t size, int flags)
 extern void dsa_free(dsa_area *area, dsa_pointer dp);
 extern void *dsa_get_address(dsa_area *area, dsa_pointer dp);
 extern size_t dsa_get_total_size(dsa_area *area);
+extern size_t dsa_get_total_size_from_handle(dsa_handle handle);
 extern void dsa_trim(dsa_area *area);
 extern void dsa_dump(dsa_area *area);
 
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 ca8abbb377e..9128e171b1b 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
@@ -1,8 +1,8 @@
-SELECT name, type, size IS DISTINCT FROM 0 AS size
+SELECT name, type, size > 0 AS size_ok
 FROM pg_dsm_registry_allocations
 WHERE name like 'test_dsm_registry%' ORDER BY name;
- name | type | size 
-------+------+------
+ name | type | size_ok 
+------+------+---------
 (0 rows)
 
 CREATE EXTENSION test_dsm_registry;
@@ -32,11 +32,11 @@ SELECT get_val_in_hash('test');
 (1 row)
 
 \c
-SELECT name, type, size IS DISTINCT FROM 0 AS size
+SELECT name, type, size > 0 AS size_ok
 FROM pg_dsm_registry_allocations
 WHERE name like 'test_dsm_registry%' ORDER BY name;
-          name          |  type   | size 
-------------------------+---------+------
+          name          |  type   | size_ok 
+------------------------+---------+---------
  test_dsm_registry_dsa  | area    | t
  test_dsm_registry_dsm  | segment | t
  test_dsm_registry_hash | hash    | t
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 965a3f1ebb6..a606e8872a1 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,4 @@
-SELECT name, type, size IS DISTINCT FROM 0 AS size
+SELECT name, type, size > 0 AS size_ok
 FROM pg_dsm_registry_allocations
 WHERE name like 'test_dsm_registry%' ORDER BY name;
 CREATE EXTENSION test_dsm_registry;
@@ -8,6 +8,6 @@ SELECT set_val_in_hash('test', '1414');
 SELECT get_val_in_shmem();
 SELECT get_val_in_hash('test');
 \c
-SELECT name, type, size IS DISTINCT FROM 0 AS size
+SELECT name, type, size > 0 AS size_ok
 FROM pg_dsm_registry_allocations
 WHERE name like 'test_dsm_registry%' ORDER BY name;