On Wed, Dec 27, 2023 at 01:53:27PM -0600, Nathan Bossart wrote:
> Here is a new version of the patch.  In addition to various small changes,
> I've rewritten the test suite to use an integer and a lock, added a
> dsm_registry_find() function, and adjusted autoprewarm to use the registry.

Here's a v3 that fixes a silly mistake in the test.

-- 
Nathan Bossart
Amazon Web Services: https://aws.amazon.com
>From b9b1638c9eb55eba51356690a2b3da1fe0b2b14e Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nat...@postgresql.org>
Date: Wed, 11 Oct 2023 22:07:26 -0500
Subject: [PATCH v3 1/2] add dsm registry

---
 src/backend/storage/ipc/Makefile              |   1 +
 src/backend/storage/ipc/dsm_registry.c        | 209 ++++++++++++++++++
 src/backend/storage/ipc/ipci.c                |   3 +
 src/backend/storage/ipc/meson.build           |   1 +
 src/backend/storage/lmgr/lwlock.c             |   4 +
 src/backend/storage/lmgr/lwlocknames.txt      |   1 +
 .../utils/activity/wait_event_names.txt       |   3 +
 src/include/storage/dsm_registry.h            |  24 ++
 src/include/storage/lwlock.h                  |   2 +
 src/test/modules/Makefile                     |   1 +
 src/test/modules/meson.build                  |   1 +
 src/test/modules/test_dsm_registry/.gitignore |   4 +
 src/test/modules/test_dsm_registry/Makefile   |  23 ++
 .../expected/test_dsm_registry.out            |  14 ++
 .../modules/test_dsm_registry/meson.build     |  33 +++
 .../sql/test_dsm_registry.sql                 |   4 +
 .../test_dsm_registry--1.0.sql                |  10 +
 .../test_dsm_registry/test_dsm_registry.c     |  75 +++++++
 .../test_dsm_registry.control                 |   4 +
 src/tools/pgindent/typedefs.list              |   3 +
 20 files changed, 420 insertions(+)
 create mode 100644 src/backend/storage/ipc/dsm_registry.c
 create mode 100644 src/include/storage/dsm_registry.h
 create mode 100644 src/test/modules/test_dsm_registry/.gitignore
 create mode 100644 src/test/modules/test_dsm_registry/Makefile
 create mode 100644 src/test/modules/test_dsm_registry/expected/test_dsm_registry.out
 create mode 100644 src/test/modules/test_dsm_registry/meson.build
 create mode 100644 src/test/modules/test_dsm_registry/sql/test_dsm_registry.sql
 create mode 100644 src/test/modules/test_dsm_registry/test_dsm_registry--1.0.sql
 create mode 100644 src/test/modules/test_dsm_registry/test_dsm_registry.c
 create mode 100644 src/test/modules/test_dsm_registry/test_dsm_registry.control

diff --git a/src/backend/storage/ipc/Makefile b/src/backend/storage/ipc/Makefile
index 6d5b921038..d8a1653eb6 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 0000000000..5fc970001e
--- /dev/null
+++ b/src/backend/storage/ipc/dsm_registry.c
@@ -0,0 +1,209 @@
+/*-------------------------------------------------------------------------
+ *
+ * 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.
+ *
+ * Portions Copyright (c) 1996-2023, 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		key[256];
+	dsm_handle	handle;
+} 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;
+
+static void init_dsm_registry(void);
+
+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, 0);
+
+		/* 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, 0);
+	}
+
+	LWLockRelease(DSMRegistryLock);
+}
+
+/*
+ * Find an existing DSM entry.
+ *
+ * This routine returns whether a segment is registered with the provided key.
+ */
+bool
+dsm_registry_find(const char *key)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+	char		key_padded[offsetof(DSMRegistryEntry, handle)] = {0};
+
+	Assert(key);
+
+	if (strlen(key) >= offsetof(DSMRegistryEntry, handle))
+		elog(ERROR, "DSM registry key too long");
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	strcpy(key_padded, key);
+	entry = dshash_find(dsm_registry_table, key_padded, false);
+	if (entry)
+		dshash_release_lock(dsm_registry_table, entry);
+
+	MemoryContextSwitchTo(oldcontext);
+
+	return entry != NULL;
+}
+
+/*
+ * Initialize or attach a DSM entry.
+ *
+ * This routine returns the address of the segment.  init_callback is called to
+ * initialize the segment when it is first created.
+ */
+void *
+dsm_registry_init_or_attach(const char *key, size_t size,
+							void (*init_callback) (void *ptr), bool *found)
+{
+	DSMRegistryEntry *entry;
+	MemoryContext oldcontext;
+	char		key_padded[offsetof(DSMRegistryEntry, handle)] = {0};
+	void	   *ret;
+
+	Assert(key);
+	Assert(size);
+	Assert(found);
+
+	if (strlen(key) >= offsetof(DSMRegistryEntry, handle))
+		elog(ERROR, "DSM registry key too long");
+
+	/* Be sure any local memory allocated by DSM/DSA routines is persistent. */
+	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
+
+	/* Connect to the registry. */
+	init_dsm_registry();
+
+	strcpy(key_padded, key);
+	entry = dshash_find_or_insert(dsm_registry_table, key_padded, found);
+	if (!(*found))
+	{
+		/* Initialize DSM registry entry. */
+		dsm_segment *seg = dsm_create(size, 0);
+
+		dsm_pin_segment(seg);
+		dsm_pin_mapping(seg);
+		entry->handle = dsm_segment_handle(seg);
+		ret = dsm_segment_address(seg);
+
+		if (init_callback)
+			(*init_callback) (ret);
+	}
+	else if (!dsm_find_mapping(entry->handle))
+	{
+		/* Attach to existing DSM registry entry. */
+		dsm_segment *seg = dsm_attach(entry->handle);
+
+		dsm_pin_mapping(seg);
+		ret = dsm_segment_address(seg);
+	}
+	else
+	{
+		/* Return address of already-attached DSM registry entry. */
+		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 706140eb9f..44d6a243e5 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 79a16d077f..88fef448be 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 315a78cda9..f3faa991d1 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 d621f5507f..ef8542de46 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 7e79163466..de1bc986e1 100644
--- a/src/backend/utils/activity/wait_event_names.txt
+++ b/src/backend/utils/activity/wait_event_names.txt
@@ -324,6 +324,7 @@ XactTruncation	"Waiting to execute <function>pg_xact_status</function> or update
 WrapLimitsVacuum	"Waiting to update limits on transaction id and multixact consumption."
 NotifyQueueTail	"Waiting to update limit on <command>NOTIFY</command> message storage."
 WaitEventExtension	"Waiting to read or update custom wait events information for extensions."
+DSMRegistry	"Waiting to read or update the dynamic shared memory registry."
 
 XactBuffer	"Waiting for I/O on a transaction status SLRU buffer."
 CommitTsBuffer	"Waiting for I/O on a commit timestamp SLRU buffer."
@@ -354,6 +355,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 0000000000..5e9c1000bc
--- /dev/null
+++ b/src/include/storage/dsm_registry.h
@@ -0,0 +1,24 @@
+/*-------------------------------------------------------------------------
+ *
+ * dsm_registry.h
+ *
+ *
+ * Portions Copyright (c) 1996-2023, 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 bool dsm_registry_find(const char *key);
+extern void *dsm_registry_init_or_attach(const char *key, 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 b038e599c0..e0e7b3b765 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 5d33fa6a9a..f656032589 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 b76f588559..bd53d52a3f 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 0000000000..5dcb3ff972
--- /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 0000000000..b13e99a354
--- /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 0000000000..8ffbd343a0
--- /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 0000000000..4a7992109b
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/meson.build
@@ -0,0 +1,33 @@
+# Copyright (c) 2023, 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 0000000000..b3351be0a1
--- /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 0000000000..8c55b0919b
--- /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 0000000000..068f47ed02
--- /dev/null
+++ b/src/test/modules/test_dsm_registry/test_dsm_registry.c
@@ -0,0 +1,75 @@
+/*--------------------------------------------------------------------------
+ *
+ * test_dsm_registry.c
+ *		Test the DSM registry
+ *
+ * Copyright (c) 2023, 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
+init_state(void *ptr)
+{
+	TestDSMRegistryStruct *state = (TestDSMRegistryStruct *) ptr;
+
+	LWLockInitialize(&state->lck, LWLockNewTrancheId());
+	state->val = 0;
+}
+
+static void
+dsm_registry_attach(void)
+{
+	bool		found;
+
+	tdr_state = dsm_registry_init_or_attach("test_dsm_registry",
+											sizeof(TestDSMRegistryStruct),
+											init_state, &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)
+{
+	dsm_registry_attach();
+
+	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;
+
+	dsm_registry_attach();
+
+	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 0000000000..813f099889
--- /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 e37ef9aa76..4270109fb8 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -610,6 +610,8 @@ DropSubscriptionStmt
 DropTableSpaceStmt
 DropUserMappingStmt
 DropdbStmt
+DSMRegistryCtxStruct
+DSMRegistryEntry
 DumpComponents
 DumpId
 DumpOptions
@@ -2796,6 +2798,7 @@ Tcl_NotifierProcs
 Tcl_Obj
 Tcl_Time
 TempNamespaceStatus
+TestDSMRegistryStruct
 TestDecodingData
 TestDecodingTxnData
 TestSpec
-- 
2.25.1

>From 4f7a4672aee298ee86bdb54a2cf1438fa80ced35 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nat...@postgresql.org>
Date: Tue, 26 Dec 2023 22:25:45 -0600
Subject: [PATCH v3 2/2] use dsm registry for pg_prewarm

---
 contrib/pg_prewarm/autoprewarm.c | 44 ++++++++++----------------------
 1 file changed, 14 insertions(+), 30 deletions(-)

diff --git a/contrib/pg_prewarm/autoprewarm.c b/contrib/pg_prewarm/autoprewarm.c
index 0993bd2453..ed12142d04 100644
--- a/contrib/pg_prewarm/autoprewarm.c
+++ b/contrib/pg_prewarm/autoprewarm.c
@@ -38,6 +38,7 @@
 #include "postmaster/interrupt.h"
 #include "storage/buf_internals.h"
 #include "storage/dsm.h"
+#include "storage/dsm_registry.h"
 #include "storage/fd.h"
 #include "storage/ipc.h"
 #include "storage/latch.h"
@@ -95,8 +96,6 @@ static void apw_start_database_worker(void);
 static bool apw_init_shmem(void);
 static void apw_detach_shmem(int code, Datum arg);
 static int	apw_compare_blockinfo(const void *p, const void *q);
-static void autoprewarm_shmem_request(void);
-static shmem_request_hook_type prev_shmem_request_hook = NULL;
 
 /* Pointer to shared-memory state. */
 static AutoPrewarmSharedState *apw_state = NULL;
@@ -140,26 +139,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("pg_prewarm");
 
-	prev_shmem_request_hook = shmem_request_hook;
-	shmem_request_hook = autoprewarm_shmem_request;
-
 	/* Register autoprewarm worker, if enabled. */
 	if (autoprewarm)
 		apw_start_leader_worker();
 }
 
-/*
- * Requests any additional shared memory required for autoprewarm.
- */
-static void
-autoprewarm_shmem_request(void)
-{
-	if (prev_shmem_request_hook)
-		prev_shmem_request_hook();
-
-	RequestAddinShmemSpace(MAXALIGN(sizeof(AutoPrewarmSharedState)));
-}
-
 /*
  * Main entry point for the leader autoprewarm process.  Per-database workers
  * have a separate entry point.
@@ -767,6 +751,16 @@ autoprewarm_dump_now(PG_FUNCTION_ARGS)
 	PG_RETURN_INT64((int64) num_blocks);
 }
 
+static void
+init_state(void *ptr)
+{
+	AutoPrewarmSharedState *state = (AutoPrewarmSharedState *) ptr;
+
+	LWLockInitialize(&state->lock, LWLockNewTrancheId());
+	state->bgworker_pid = InvalidPid;
+	state->pid_using_dumpfile = InvalidPid;
+}
+
 /*
  * Allocate and initialize autoprewarm related shared memory, if not already
  * done, and set up backend-local pointer to that state.  Returns true if an
@@ -777,19 +771,9 @@ apw_init_shmem(void)
 {
 	bool		found;
 
-	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-	apw_state = ShmemInitStruct("autoprewarm",
-								sizeof(AutoPrewarmSharedState),
-								&found);
-	if (!found)
-	{
-		/* First time through ... */
-		LWLockInitialize(&apw_state->lock, LWLockNewTrancheId());
-		apw_state->bgworker_pid = InvalidPid;
-		apw_state->pid_using_dumpfile = InvalidPid;
-	}
-	LWLockRelease(AddinShmemInitLock);
-
+	apw_state = dsm_registry_init_or_attach("autoprewarm",
+											sizeof(AutoPrewarmSharedState),
+											init_state, &found);
 	LWLockRegisterTranche(apw_state->lock.tranche, "autoprewarm");
 
 	return found;
-- 
2.25.1

Reply via email to