Introduce the dynamic shared memory registry.

Presently, the most straightforward way for a shared library to use
shared memory is to request it at server startup via a
shmem_request_hook, which requires specifying the library in
shared_preload_libraries.  Alternatively, the library can create a
dynamic shared memory (DSM) segment, but absent a shared location
to store the segment's handle, other backends cannot use it.  This
commit introduces a registry for DSM segments so that these other
backends can look up existing segments with a library-specified
string.  This allows libraries to easily use shared memory without
needing to request it at server startup.

The registry is accessed via the new GetNamedDSMSegment() function.
This function handles allocating the segment and initializing it
via a provided callback.  If another backend already created and
initialized the segment, it simply attaches the segment.
GetNamedDSMSegment() locks the registry appropriately to ensure
that only one backend initializes the segment and that all other
backends just attach it.

The registry itself is comprised of a dshash table that stores the
DSM segment handles keyed by a library-specified string.

Reviewed-by: Michael Paquier, Andrei Lepikhov, Nikita Malakhov, Robert Haas, Bharath Rupireddy, Zhang Mingli, Amul Sul
Discussion: https://postgr.es/m/20231205034647.GA2705267%40nathanxps13

21 files changed, 455 insertions, 3 deletions

diff --git a/doc/src/sgml/xfunc.sgml b/doc/src/sgml/xfunc.sgml
index ede2a5dea66..0ad9f38e902 100644
--- a/doc/src/sgml/xfunc.sgml
+++ b/doc/src/sgml/xfunc.sgml
@@ -3460,6 +3460,45 @@ LWLockRelease(AddinShmemInitLock);
       the <productname>PostgreSQL</productname> source tree.
      </para>
     </sect3>
+
+    <sect3 id="xfunc-shared-addin-after-startup">
+     <title>Requesting Shared Memory After Startup</title>
+
+     <para>
+      There is another, more flexible method of reserving shared memory that
+      can be done after server startup and outside a
+      <literal>shmem_request_hook</literal>.  To do so, each backend that will
+      use the shared memory should obtain a pointer to it by calling:
+<programlisting>
+void *GetNamedDSMSegment(const char *name, size_t size,
+                         void (*init_callback) (void *ptr),
+                         bool *found)
+</programlisting>
+      If a dynamic shared memory segment with the given name does not yet
+      exist, this function will allocate it and initialize it with the provided
+      <function>init_callback</function> callback function.  If the segment has
+      already been allocated and initialized by another backend, this function
+      simply attaches the existing dynamic shared memory segment to the current
+      backend.
+     </para>
+
+     <para>
+      Unlike shared memory reserved at server startup, there is no need to
+      acquire <function>AddinShmemInitLock</function> or otherwise take action
+      to avoid race conditions when reserving shared memory with
+      <function>GetNamedDSMSegment</function>.  This function ensures that only
+      one backend allocates and initializes the segment and that all other
+      backends receive a pointer to the fully allocated and initialized
+      segment.
+     </para>
+
+     <para>
+      A complete usage example of <function>GetNamedDSMSegment</function> can
+      be found in
+      <filename>src/test/modules/test_dsm_registry/test_dsm_registry.c</filename>
+      in the <productname>PostgreSQL</productname> source tree.
+     </para>
+    </sect3>
    </sect2>
 
    <sect2 id="xfunc-addin-lwlocks">
@@ -3469,8 +3508,9 @@ LWLockRelease(AddinShmemInitLock);
      <title>Requesting LWLocks at Startup</title>
 
      <para>
-      Add-ins can reserve LWLocks on server startup.  As with shared memory,
-      the add-in's shared library must be preloaded by specifying it in
+      Add-ins can reserve LWLocks on server startup.  As with shared memory
+      reserved at server startup, the add-in's shared library must be preloaded
+      by specifying it in
       <xref linkend="guc-shared-preload-libraries"/><indexterm><primary>shared_preload_libraries</primary></indexterm>,
       and the shared library should register a
       <literal>shmem_request_hook</literal> in its
@@ -3508,7 +3548,10 @@ void LWLockInitialize(LWLock *lock, int tranche_id)
       process allocates a new <literal>tranche_id</literal> and initializes
       each new LWLock.  One way to do this is to only call these functions in
       your shared memory initialization code with the
-      <function>AddinShmemInitLock</function> held exclusively.
+      <function>AddinShmemInitLock</function> held exclusively.  If using
+      <function>GetNamedDSMSegment</function>, calling these functions in the
+      <function>init_callback</function> callback function is sufficient to
+      avoid race conditions.
      </para>
 
      <para>
diff --git a/src/backend/storage/ipc/Makefile b/src/backend/storage/ipc/Makefile
index 6d5b921038c..d8a1653eb6a 100644
--- a/src/backend/storage/ipc/Makefile
+++ b/src/backend/storage/ipc/Makefile
@@ -12,6 +12,7 @@ OBJS = \
 	barrier.o \
 	dsm.o \
 	dsm_impl.o \
+	dsm_registry.o \
 	ipc.o \
 	ipci.o \
 	latch.o \
diff --git a/src/backend/storage/ipc/dsm_registry.c b/src/backend/storage/ipc/dsm_registry.c
new file mode 100644
index 00000000000..ac11f51375e
--- /dev/null
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -0,0 +1,198 @@
+/*-------------------------------------------------------------------------
+ *
+ * dsm_registry.c
+ *	  Functions for interfacing with the dynamic shared memory registry.
+ *
+ * This provides a way for libraries to use shared memory without needing
+ * to request it at startup time via a shmem_request_hook.  The registry
+ * stores dynamic shared memory (DSM) segment handles keyed by a
+ * library-specified string.
+ *
+ * The registry is accessed by calling GetNamedDSMSegment().  If a segment
+ * with the provided name does not yet exist, it is created and initialized
+ * with the provided init_callback callback function.  Otherwise,
+ * GetNamedDSMSegment() simply ensures that the segment is attached to the
+ * current backend.  This function guarantees that only one backend
+ * initializes the segment and that all other backends just attach it.
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/backend/storage/ipc/dsm_registry.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "lib/dshash.h"
+#include "storage/dsm_registry.h"
+#include "storage/lwlock.h"
+#include "storage/shmem.h"
+#include "utils/memutils.h"
+
+typedef struct DSMRegistryCtxStruct
+{
+	dsa_handle	dsah;
+	dshash_table_handle dshh;
+} DSMRegistryCtxStruct;
+
+static DSMRegistryCtxStruct *DSMRegistryCtx;
+
+typedef struct DSMRegistryEntry
+{
+	char		name[64];
+	dsm_handle	handle;
+	size_t		size;
+} DSMRegistryEntry;
+
+static const dshash_parameters dsh_params = {
+	offsetof(DSMRegistryEntry, handle),
+	sizeof(DSMRegistryEntry),
+	dshash_memcmp,
+	dshash_memhash,
+	LWTRANCHE_DSM_REGISTRY_HASH
+};
+
+static dsa_area *dsm_registry_dsa;
+static dshash_table *dsm_registry_table;
+
+Size
+DSMRegistryShmemSize(void)
+{
+	return MAXALIGN(sizeof(DSMRegistryCtxStruct));
+}
+
+void
+DSMRegistryShmemInit(void)
+{
+	bool		found;
+
+	DSMRegistryCtx = (DSMRegistryCtxStruct *)
+		ShmemInitStruct("DSM Registry Data",
+						DSMRegistryShmemSize(),
+						&found);
+
+	if (!found)
+	{
+		DSMRegistryCtx->dsah = DSA_HANDLE_INVALID;
+		DSMRegistryCtx->dshh = DSHASH_HANDLE_INVALID;
+	}
+}
+
+/*
+ * Initialize or attach to the dynamic shared hash table that stores the DSM
+ * registry entries, if not already done.  This must be called before accessing
+ * the table.
+ */
+static void
+init_dsm_registry(void)
+{
+	/* Quick exit if we already did this. */
+	if (dsm_registry_table)
+		return;
+
+	/* Otherwise, use a lock to ensure only one process creates the table. */
+	LWLockAcquire(DSMRegistryLock, LW_EXCLUSIVE);
+
+	if (DSMRegistryCtx->dshh == DSHASH_HANDLE_INVALID)
+	{
+		/* Initialize dynamic shared hash table for registry. */
+		dsm_registry_dsa = dsa_create(LWTRANCHE_DSM_REGISTRY_DSA);
+		dsa_pin(dsm_registry_dsa);
+		dsa_pin_mapping(dsm_registry_dsa);
+		dsm_registry_table = dshash_create(dsm_registry_dsa, &dsh_params, NULL);
+
+		/* Store handles in shared memory for other backends to use. */
+		DSMRegistryCtx->dsah = dsa_get_handle(dsm_registry_dsa);
+		DSMRegistryCtx->dshh = dshash_get_hash_table_handle(dsm_registry_table);
+	}
+	else
+	{
+		/* Attach to existing dynamic shared hash table. */
+		dsm_registry_dsa = dsa_attach(DSMRegistryCtx->dsah);
+		dsa_pin_mapping(dsm_registry_dsa);
+		dsm_registry_table = dshash_attach(dsm_registry_dsa, &dsh_params,
+										   DSMRegistryCtx->dshh, NULL);
+	}
+
+	LWLockRelease(DSMRegistryLock);
+}
+
+/*
+ * Initialize or attach a named DSM segment.
+ *
+ * This routine returns the address of the segment.  init_callback is called to
+ * initialize the segment when it is first created.
+ */
+void *
+GetNamedDSMSegment(const char *name, size_t size,
+				   void (*init_callback) (void *ptr), bool *found)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+	char		name_padded[offsetof(DSMRegistryEntry, handle)] = {0};
+	void	   *ret;
+
+	Assert(found);
+
+	if (!name || *name == '\0')
+		ereport(ERROR,
+				(errmsg("DSM segment name cannot be empty")));
+
+	if (strlen(name) >= offsetof(DSMRegistryEntry, handle))
+		ereport(ERROR,
+				(errmsg("DSM segment name too long")));
+
+	if (size == 0)
+		ereport(ERROR,
+				(errmsg("DSM segment size must be nonzero")));
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	strcpy(name_padded, name);
+	entry = dshash_find_or_insert(dsm_registry_table, name_padded, found);
+	if (!(*found))
+	{
+		/* Initialize the segment. */
+		dsm_segment *seg = dsm_create(size, 0);
+
+		dsm_pin_segment(seg);
+		dsm_pin_mapping(seg);
+		entry->handle = dsm_segment_handle(seg);
+		entry->size = size;
+		ret = dsm_segment_address(seg);
+
+		if (init_callback)
+			(*init_callback) (ret);
+	}
+	else if (entry->size != size)
+	{
+		ereport(ERROR,
+				(errmsg("requested DSM segment size does not match size of "
+						"existing segment")));
+	}
+	else if (!dsm_find_mapping(entry->handle))
+	{
+		/* Attach to existing segment. */
+		dsm_segment *seg = dsm_attach(entry->handle);
+
+		dsm_pin_mapping(seg);
+		ret = dsm_segment_address(seg);
+	}
+	else
+	{
+		/* Return address of an already-attached segment. */
+		ret = dsm_segment_address(dsm_find_mapping(entry->handle));
+	}
+
+	dshash_release_lock(dsm_registry_table, entry);
+	MemoryContextSwitchTo(oldcontext);
+
+	return ret;
+}
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index e5119ed55dd..fbc62b15637 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -40,6 +40,7 @@
 #include "replication/walsender.h"
 #include "storage/bufmgr.h"
 #include "storage/dsm.h"
+#include "storage/dsm_registry.h"
 #include "storage/ipc.h"
 #include "storage/pg_shmem.h"
 #include "storage/pmsignal.h"
@@ -115,6 +116,7 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, hash_estimate_size(SHMEM_INDEX_SIZE,
 											 sizeof(ShmemIndexEnt)));
 	size = add_size(size, dsm_estimate_size());
+	size = add_size(size, DSMRegistryShmemSize());
 	size = add_size(size, BufferShmemSize());
 	size = add_size(size, LockShmemSize());
 	size = add_size(size, PredicateLockShmemSize());
@@ -289,6 +291,7 @@ CreateOrAttachShmemStructs(void)
 	InitShmemIndex();
 
 	dsm_shmem_init();
+	DSMRegistryShmemInit();
 
 	/*
 	 * Set up xlog, clog, and buffers
diff --git a/src/backend/storage/ipc/meson.build b/src/backend/storage/ipc/meson.build
index 08bdc718b83..5a936171f73 100644
--- a/src/backend/storage/ipc/meson.build
+++ b/src/backend/storage/ipc/meson.build
@@ -4,6 +4,7 @@ backend_sources += files(
   'barrier.c',
   'dsm.c',
   'dsm_impl.c',
+  'dsm_registry.c',
   'ipc.c',
   'ipci.c',
   'latch.c',
diff --git a/src/backend/storage/lmgr/lwlock.c b/src/backend/storage/lmgr/lwlock.c
index b4b989ac569..2f2de5a5624 100644
--- a/src/backend/storage/lmgr/lwlock.c
+++ b/src/backend/storage/lmgr/lwlock.c
@@ -190,6 +190,10 @@ static const char *const BuiltinTrancheNames[] = {
 	"LogicalRepLauncherDSA",
 	/* LWTRANCHE_LAUNCHER_HASH: */
 	"LogicalRepLauncherHash",
+	/* LWTRANCHE_DSM_REGISTRY_DSA: */
+	"DSMRegistryDSA",
+	/* LWTRANCHE_DSM_REGISTRY_HASH: */
+	"DSMRegistryHash",
 };
 
 StaticAssertDecl(lengthof(BuiltinTrancheNames) ==
diff --git a/src/backend/storage/lmgr/lwlocknames.txt b/src/backend/storage/lmgr/lwlocknames.txt
index d621f5507f3..ef8542de469 100644
--- a/src/backend/storage/lmgr/lwlocknames.txt
+++ b/src/backend/storage/lmgr/lwlocknames.txt
@@ -55,3 +55,4 @@ WrapLimitsVacuumLock				46
 NotifyQueueTailLock					47
 WaitEventExtensionLock				48
 WALSummarizerLock					49
+DSMRegistryLock						50
diff --git a/src/backend/utils/activity/wait_event_names.txt b/src/backend/utils/activity/wait_event_names.txt
index f625473ad43..6bcb1cca0cb 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -329,6 +329,7 @@ WrapLimitsVacuum	"Waiting to update limits on transaction id and multixact consu
 NotifyQueueTail	"Waiting to update limit on <command>NOTIFY</command> message storage."
 WaitEventExtension	"Waiting to read or update custom wait events information for extensions."
 WALSummarizer	"Waiting to read or update WAL summarization state."
+DSMRegistry	"Waiting to read or update the dynamic shared memory registry."
 
 #
 # END OF PREDEFINED LWLOCKS (DO NOT CHANGE THIS LINE)
@@ -367,6 +368,8 @@ PgStatsHash	"Waiting for stats shared memory hash table access."
 PgStatsData	"Waiting for shared memory stats data access."
 LogicalRepLauncherDSA	"Waiting to access logical replication launcher's dynamic shared memory allocator."
 LogicalRepLauncherHash	"Waiting to access logical replication launcher's shared hash table."
+DSMRegistryDSA	"Waiting to access dynamic shared memory registry's dynamic shared memory allocator."
+DSMRegistryHash	"Waiting to access dynamic shared memory registry's shared hash table."
 
 #
 # Wait Events - Lock
diff --git a/src/include/storage/dsm_registry.h b/src/include/storage/dsm_registry.h
new file mode 100644
index 00000000000..5fa79731d28
--- /dev/null
+++ b/src/include/storage/dsm_registry.h
@@ -0,0 +1,23 @@
+/*-------------------------------------------------------------------------
+ *
+ * dsm_registry.h
+ *	  Functions for interfacing with the dynamic shared memory registry.
+ *
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/storage/dsm_registry.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef DSM_REGISTRY_H
+#define DSM_REGISTRY_H
+
+extern void *GetNamedDSMSegment(const char *name, size_t size,
+								void (*init_callback) (void *ptr),
+								bool *found);
+
+extern Size DSMRegistryShmemSize(void);
+extern void DSMRegistryShmemInit(void);
+
+#endif							/* DSM_REGISTRY_H */
diff --git a/src/include/storage/lwlock.h b/src/include/storage/lwlock.h
index 167ae342088..50a65e046d8 100644
--- a/src/include/storage/lwlock.h
+++ b/src/include/storage/lwlock.h
@@ -207,6 +207,8 @@ typedef enum BuiltinTrancheIds
 	LWTRANCHE_PGSTATS_DATA,
 	LWTRANCHE_LAUNCHER_DSA,
 	LWTRANCHE_LAUNCHER_HASH,
+	LWTRANCHE_DSM_REGISTRY_DSA,
+	LWTRANCHE_DSM_REGISTRY_HASH,
 	LWTRANCHE_FIRST_USER_DEFINED,
 }			BuiltinTrancheIds;
 
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
index 5d33fa6a9a8..f656032589d 100644
--- a/src/test/modules/Makefile
+++ b/src/test/modules/Makefile
@@ -18,6 +18,7 @@ SUBDIRS = \
 		  test_custom_rmgrs \
 		  test_ddl_deparse \
 		  test_dsa \
+		  test_dsm_registry \
 		  test_extensions \
 		  test_ginpostinglist \
 		  test_integerset \
diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build
index 00ff1d77d19..2c3b8d73bce 100644
--- a/src/test/modules/meson.build
+++ b/src/test/modules/meson.build
@@ -15,6 +15,7 @@ subdir('test_copy_callbacks')
 subdir('test_custom_rmgrs')
 subdir('test_ddl_deparse')
 subdir('test_dsa')
+subdir('test_dsm_registry')
 subdir('test_extensions')
 subdir('test_ginpostinglist')
 subdir('test_integerset')
diff --git a/src/test/modules/test_dsm_registry/.gitignore b/src/test/modules/test_dsm_registry/.gitignore
new file mode 100644
index 00000000000..5dcb3ff9723
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_dsm_registry/Makefile b/src/test/modules/test_dsm_registry/Makefile
new file mode 100644
index 00000000000..b13e99a354f
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/Makefile
@@ -0,0 +1,23 @@
+# src/test/modules/test_dsm_registry/Makefile
+
+MODULE_big = test_dsm_registry
+OBJS = \
+	$(WIN32RES) \
+	test_dsm_registry.o
+PGFILEDESC = "test_dsm_registry - test code for the DSM registry"
+
+EXTENSION = test_dsm_registry
+DATA = test_dsm_registry--1.0.sql
+
+REGRESS = test_dsm_registry
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_dsm_registry
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
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
new file mode 100644
index 00000000000..8ffbd343a05
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/expected/test_dsm_registry.out
@@ -0,0 +1,14 @@
+CREATE EXTENSION test_dsm_registry;
+SELECT set_val_in_shmem(1236);
+ set_val_in_shmem 
+------------------
+ 
+(1 row)
+
+\c
+SELECT get_val_in_shmem();
+ get_val_in_shmem 
+------------------
+             1236
+(1 row)
+
diff --git a/src/test/modules/test_dsm_registry/meson.build b/src/test/modules/test_dsm_registry/meson.build
new file mode 100644
index 00000000000..a4045fea37e
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2024, PostgreSQL Global Development Group
+
+test_dsm_registry_sources = files(
+  'test_dsm_registry.c',
+)
+
+if host_system == 'windows'
+  test_dsm_registry_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'test_dsm_registry',
+    '--FILEDESC', 'test_dsm_registry - test code for the DSM registry',])
+endif
+
+test_dsm_registry = shared_module('test_dsm_registry',
+  test_dsm_registry_sources,
+  kwargs: pg_test_mod_args,
+)
+test_install_libs += test_dsm_registry
+
+test_install_data += files(
+  'test_dsm_registry.control',
+  'test_dsm_registry--1.0.sql',
+)
+
+tests += {
+  'name': 'test_dsm_registry',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'regress': {
+    'sql': [
+      'test_dsm_registry',
+    ],
+  },
+}
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
new file mode 100644
index 00000000000..b3351be0a16
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql
@@ -0,0 +1,4 @@
+CREATE EXTENSION test_dsm_registry;
+SELECT set_val_in_shmem(1236);
+\c
+SELECT get_val_in_shmem();
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql b/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
new file mode 100644
index 00000000000..8c55b0919b1
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
@@ -0,0 +1,10 @@
+/* src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_dsm_registry" to load this file. \quit
+
+CREATE FUNCTION set_val_in_shmem(val INT) RETURNS VOID
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION get_val_in_shmem() RETURNS INT
+	AS 'MODULE_PATHNAME' LANGUAGE C;
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry.c b/src/test/modules/test_dsm_registry/test_dsm_registry.c
new file mode 100644
index 00000000000..96eaa850bfa
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -0,0 +1,76 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_dsm_registry.c
+ *	  Test the dynamic shared memory registry.
+ *
+ * Copyright (c) 2024, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_dsm_registry/test_dsm_registry.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "storage/dsm_registry.h"
+#include "storage/lwlock.h"
+
+PG_MODULE_MAGIC;
+
+typedef struct TestDSMRegistryStruct
+{
+	int			val;
+	LWLock		lck;
+} TestDSMRegistryStruct;
+
+static TestDSMRegistryStruct *tdr_state;
+
+static void
+tdr_init_shmem(void *ptr)
+{
+	TestDSMRegistryStruct *state = (TestDSMRegistryStruct *) ptr;
+
+	LWLockInitialize(&state->lck, LWLockNewTrancheId());
+	state->val = 0;
+}
+
+static void
+tdr_attach_shmem(void)
+{
+	bool		found;
+
+	tdr_state = GetNamedDSMSegment("test_dsm_registry",
+								   sizeof(TestDSMRegistryStruct),
+								   tdr_init_shmem,
+								   &found);
+	LWLockRegisterTranche(tdr_state->lck.tranche, "test_dsm_registry");
+}
+
+PG_FUNCTION_INFO_V1(set_val_in_shmem);
+Datum
+set_val_in_shmem(PG_FUNCTION_ARGS)
+{
+	tdr_attach_shmem();
+
+	LWLockAcquire(&tdr_state->lck, LW_EXCLUSIVE);
+	tdr_state->val = PG_GETARG_UINT32(0);
+	LWLockRelease(&tdr_state->lck);
+
+	PG_RETURN_VOID();
+}
+
+PG_FUNCTION_INFO_V1(get_val_in_shmem);
+Datum
+get_val_in_shmem(PG_FUNCTION_ARGS)
+{
+	int			ret;
+
+	tdr_attach_shmem();
+
+	LWLockAcquire(&tdr_state->lck, LW_SHARED);
+	ret = tdr_state->val;
+	LWLockRelease(&tdr_state->lck);
+
+	PG_RETURN_UINT32(ret);
+}
diff --git a/src/test/modules/test_dsm_registry/test_dsm_registry.control b/src/test/modules/test_dsm_registry/test_dsm_registry.control
new file mode 100644
index 00000000000..813f099889b
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.control
@@ -0,0 +1,4 @@
+comment = 'Test code for the DSM registry'
+default_version = '1.0'
+module_pathname = '$libdir/test_dsm_registry'
+relocatable = true
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 16421f034cb..7e228c5aea3 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -611,6 +611,8 @@ DropSubscriptionStmt
 DropTableSpaceStmt
 DropUserMappingStmt
 DropdbStmt
+DSMRegistryCtxStruct
+DSMRegistryEntry
 DumpComponents
 DumpId
 DumpOptions
@@ -2799,6 +2801,7 @@ Tcl_NotifierProcs
 Tcl_Obj
 Tcl_Time
 TempNamespaceStatus
+TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
 TestSpec