1 files changed, 7 insertions, 102 deletions

diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index c4317351ced..98f8d4194b2 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -44,33 +44,17 @@
 
 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;
 static char *worker_spi_database = NULL;
 
+/* value cached, fetched from shared memory */
+static uint32 worker_spi_wait_event_main = 0;
 
 typedef struct worktable
 {
@@ -78,63 +62,6 @@ 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.
@@ -224,9 +151,6 @@ 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);
 
@@ -268,6 +192,10 @@ worker_spi_main(Datum main_arg)
 	{
 		int			ret;
 
+		/* First time, allocate or get the custom wait event */
+		if (worker_spi_wait_event_main == 0)
+			worker_spi_wait_event_main = WaitEventExtensionNew("worker_spi_main");
+
 		/*
 		 * Background workers mustn't call usleep() or any direct equivalent:
 		 * instead, they may wait on their process latch, which sleeps as
@@ -277,7 +205,7 @@ worker_spi_main(Datum main_arg)
 		(void) WaitLatch(MyLatch,
 						 WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH,
 						 worker_spi_naptime * 1000L,
-						 wsstate->wait_event);
+						 worker_spi_wait_event_main);
 		ResetLatch(MyLatch);
 
 		CHECK_FOR_INTERRUPTS();
@@ -406,11 +334,6 @@ _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 |
@@ -435,21 +358,6 @@ _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
@@ -461,9 +369,6 @@ 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;