From 2fcd773cd4009cffd626640e1553d4efdceb0777 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Thu, 20 Jul 2023 12:48:07 +0900
Subject: [PATCH] Add tweak to allow worker_spi to register wait_event
 dynamically

---
 src/test/modules/worker_spi/worker_spi.c | 35 +++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index b87d6c75d8..4e3da93129 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -48,11 +48,20 @@ PG_FUNCTION_INFO_V1(worker_spi_launch);
 
 PGDLLEXPORT void worker_spi_main(Datum main_arg) pg_attribute_noreturn();
 
+static void worker_spi_init_shmem(void);
+
 /* GUC variables */
 static int	worker_spi_naptime = 10;
 static int	worker_spi_total_workers = 2;
 static char *worker_spi_database;
 
+typedef struct WorkerSpiState
+{
+	uint32 wait_event_info;
+} WorkerSpiState;
+
+/* the pointer to the shared memory */
+static WorkerSpiState *worker_spi_state = NULL;
 
 typedef struct worktable
 {
@@ -149,6 +158,8 @@ worker_spi_main(Datum main_arg)
 	/* We're now ready to receive signals */
 	BackgroundWorkerUnblockSignals();
 
+	worker_spi_init_shmem();
+
 	/* Connect to our database */
 	BackgroundWorkerInitializeConnection(worker_spi_database, NULL, 0);
 
@@ -199,7 +210,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);
+						 worker_spi_state->wait_event_info);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -346,6 +357,26 @@ _PG_init(void)
 	}
 }
 
+static void
+worker_spi_init_shmem(void)
+{
+	bool found = false;
+
+	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
+	worker_spi_state = ShmemInitStruct("worker_spi",
+									   sizeof(WorkerSpiState),
+									   &found);
+	if (!found)
+	{
+		/* First time through ... */
+		worker_spi_state->wait_event_info = WaitEventExtensionNew();
+	}
+	LWLockRelease(AddinShmemInitLock);
+
+	WaitEventExtensionRegisterName(worker_spi_state->wait_event_info,
+								   "worker_spi_custom");
+}
+
 /*
  * Dynamically launch an SPI worker.
  */
@@ -358,6 +389,8 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 
+	worker_spi_init_shmem();
+
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
 		BGWORKER_BACKEND_DATABASE_CONNECTION;
-- 
2.40.1