Support custom wait events for wait event type "Extension"

Two backend routines are added to allow extension to allocate and define
custom wait events, all of these being allocated in the type
"Extension":
* WaitEventExtensionNew(), that allocates a wait event ID computed from
a counter in shared memory.
* WaitEventExtensionRegisterName(), to associate a custom string to the
wait event ID allocated.

Note that this includes an example of how to use this new facility in
worker_spi with tests in TAP for various scenarios, and some
documentation about how to use them.

Any code in the tree that currently uses WAIT_EVENT_EXTENSION could
switch to this new facility to define custom wait events.  This is left
as work for future patches.

Author: Masahiro Ikeda
Reviewed-by: Andres Freund, Michael Paquier, Tristan Partin, Bharath
Rupireddy
Discussion: https://postgr.es/m/b9f5411acda0cf15c8fbb767702ff43e@oss.nttdata.com

1 files changed, 102 insertions, 1 deletions

diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 903dcddef97..c4317351ced 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,10 +44,28 @@
 
 PG_MODULE_MAGIC;
 
+PG_FUNCTION_INFO_V1(worker_spi_init);
 PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+/* Shared memory state */
+typedef struct worker_spi_state
+{
+	/* the wait event defined during initialization phase */
+	uint32		wait_event;
+} worker_spi_state;
+
+static worker_spi_state *wsstate = NULL;	/* pointer to shared memory */
+
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+static shmem_request_hook_type prev_shmem_startup_hook = NULL;
+
+static void worker_spi_shmem_request(void);
+static void worker_spi_shmem_startup(void);
+static void worker_spi_shmem_init(void);
+static Size worker_spi_memsize(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
@@ -60,6 +78,63 @@ typedef struct worktable
 	const char *name;
 } worktable;
 
+static void
+worker_spi_shmem_request(void)
+{
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
+
+	RequestAddinShmemSpace(worker_spi_memsize());
+}
+
+static void
+worker_spi_shmem_startup(void)
+{
+	if (prev_shmem_startup_hook)
+		prev_shmem_startup_hook();
+
+	worker_spi_shmem_init();
+}
+
+static Size
+worker_spi_memsize(void)
+{
+	return MAXALIGN(sizeof(worker_spi_state));
+}
+
+/*
+ * Initialize the shared memory state of worker_spi.
+ *
+ * This routine allocates a new wait event when called the first time.
+ * On follow-up calls, the name of the wait event associated with the
+ * existing shared memory state is registered.
+ */
+static void
+worker_spi_shmem_init(void)
+{
+	bool		found;
+
+	wsstate = NULL;
+
+	/* Create or attach to the shared memory state */
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	wsstate = ShmemInitStruct("worker_spi State",
+							  sizeof(worker_spi_state),
+							  &found);
+
+	/* Define a new wait event */
+	if (!found)
+		wsstate->wait_event = WaitEventExtensionNew();
+
+	LWLockRelease(AddinShmemInitLock);
+
+	/*
+	 * Register the wait event in the lookup table of the current process.
+	 */
+	WaitEventExtensionRegisterName(wsstate->wait_event, "worker_spi_main");
+	return;
+}
+
 /*
  * Initialize workspace for a worker process: create the schema if it doesn't
  * already exist.
@@ -149,6 +224,9 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +277,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 WAIT_EVENT_EXTENSION);
+						 wsstate->wait_event);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -328,6 +406,11 @@ _PG_init(void)
 
 	MarkGUCPrefixReserved("worker_spi");
 
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = worker_spi_shmem_request;
+	prev_shmem_startup_hook = shmem_startup_hook;
+	shmem_startup_hook = worker_spi_shmem_startup;
+
 	/* set up common data for all our workers */
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -352,6 +435,21 @@ _PG_init(void)
 }
 
 /*
+ * Wrapper to initialize a session with the shared memory state
+ * used by this module.  This is a convenience routine to be able to
+ * see the custom wait event stored in shared memory without loading
+ * through shared_preload_libraries.
+ */
+Datum
+worker_spi_init(PG_FUNCTION_ARGS)
+{
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
+	PG_RETURN_VOID();
+}
+
+/*
  * Dynamically launch an SPI worker.
  */
 Datum
@@ -363,6 +461,9 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	/* Create (if necessary) and attach to our shared memory area. */
+	worker_spi_shmem_init();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;