From a7ecd5b36f88af492fb8c2a57c268cabcdfabf98 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Wed, 10 Jul 2024 16:13:18 +1200
Subject: [PATCH v3 1/2] Use latch for bgworker state change notification.

As part of a project to remove the reliance on PIDs to identify
backends, backends that register dynamic background workers will now
have their latch set by the postmaster when the worker state changes.
Previously, the postmaster would send a SIGUSR1 signal, which
relied on the ProcSignal system's handler setting the latch, with the
same end effect.  Now that intermediate step is skipped.

The field bgw_notify_pid still exists for backwards compatibility, but
has no effect and may be removed in a later release.
RegisterBackgroundWorker() now automatically assumes that the caller
would like state change notifications delivered to its proc number,
unless BGW_NO_NOTIFY is set in bgw_flags.

Removing other cases of PIDs from the bgwork API is left for later work;
this patch is concerned only with removing a dependency on ProcSignal
infrastructure that is due to be removed by a later commit.

This represents the first use of SetLatch() in the postmaster.  It might
seem more natural to use condition variables in the
wait-for-state-change functions, so that anyone with a handle can wait,
but condition variables as currently implemented would be a lot less
robust that simple latches.

Reviewed-by: Heikki Linnakangas <hlinnaka@iki.fi>
Reviewed-by:
Discussion: https://postgr.es/m/CA%2BhUKG%2B3MkS21yK4jL4cgZywdnnGKiBg0jatoV6kzaniBmcqbQ%40mail.gmail.com
---
 contrib/pg_prewarm/autoprewarm.c            |   6 -
 doc/src/sgml/bgworker.sgml                  |  27 +++--
 src/backend/access/transam/parallel.c       |   1 -
 src/backend/postmaster/bgworker.c           | 125 +++++++++++---------
 src/backend/postmaster/postmaster.c         |   8 +-
 src/backend/replication/logical/launcher.c  |   2 -
 src/backend/storage/lmgr/proc.c             |  13 ++
 src/include/postmaster/bgworker.h           |  13 +-
 src/include/postmaster/bgworker_internals.h |   1 +
 src/include/storage/proc.h                  |   1 +
 src/test/modules/test_shm_mq/setup.c        |   3 +-
 src/test/modules/test_shm_mq/test_shm_mq.h  |   2 +
 src/test/modules/test_shm_mq/worker.c       |  10 +-
 src/test/modules/worker_spi/worker_spi.c    |   3 -
 14 files changed, 118 insertions(+), 97 deletions(-)

diff --git a/contrib/pg_prewarm/autoprewarm.c b/contrib/pg_prewarm/autoprewarm.c
index d061731706a..bb590f77c54 100644
--- a/contrib/pg_prewarm/autoprewarm.c
+++ b/contrib/pg_prewarm/autoprewarm.c
@@ -824,9 +824,6 @@ apw_start_leader_worker(void)
 		return;
 	}
 
-	/* must set notify PID to wait for startup */
-	worker.bgw_notify_pid = MyProcPid;
-
 	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
 		ereport(ERROR,
 				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
@@ -860,9 +857,6 @@ apw_start_database_worker(void)
 	strcpy(worker.bgw_name, "autoprewarm worker");
 	strcpy(worker.bgw_type, "autoprewarm worker");
 
-	/* must set notify PID to wait for shutdown */
-	worker.bgw_notify_pid = MyProcPid;
-
 	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
 		ereport(ERROR,
 				(errcode(ERRCODE_INSUFFICIENT_RESOURCES),
diff --git a/doc/src/sgml/bgworker.sgml b/doc/src/sgml/bgworker.sgml
index 2c393385a91..c8bff86af09 100644
--- a/doc/src/sgml/bgworker.sgml
+++ b/doc/src/sgml/bgworker.sgml
@@ -63,7 +63,7 @@ typedef struct BackgroundWorker
     char        bgw_function_name[BGW_MAXLEN];
     Datum       bgw_main_arg;
     char        bgw_extra[BGW_EXTRALEN];
-    pid_t       bgw_notify_pid;
+    pid_t       bgw_notify_pid;			/* not used */
 } BackgroundWorker;
 </programlisting>
   </para>
@@ -108,6 +108,18 @@ typedef struct BackgroundWorker
      </listitem>
     </varlistentry>
 
+    <varlistentry>
+     <term><literal>BGWORKER_NO_NOTIFY</literal></term>
+     <listitem>
+      <para>
+       <indexterm><primary>BGWORKER_NO_NOTIFY</primary></indexterm>
+       Normally, the backend that registers a dynamic worker will be notified
+       by having its latch set when the workers state changes, which allows the
+       caller to wait for the worker to start and shut down.  That can be
+       suppressed by setting this flag.
+      </para>
+     </listitem>
+    </varlistentry>
    </variablelist>
 
   </para>
@@ -181,12 +193,8 @@ typedef struct BackgroundWorker
   </para>
 
   <para>
-   <structfield>bgw_notify_pid</structfield> is the PID of a PostgreSQL
-   backend process to which the postmaster should send <literal>SIGUSR1</literal>
-   when the process is started or exits.  It should be 0 for workers registered
-   at postmaster startup time, or when the backend registering the worker does
-   not wish to wait for the worker to start up.  Otherwise, it should be
-   initialized to <literal>MyProcPid</literal>.
+   <structfield>bgw_notify_pid</structfield> is not used and may be removed
+   in a future release.
   </para>
 
   <para>Once running, the process can connect to a database by calling
@@ -258,10 +266,7 @@ typedef struct BackgroundWorker
 
   <para>
    In some cases, a process which registers a background worker may wish to
-   wait for the worker to start up.  This can be accomplished by initializing
-   <structfield>bgw_notify_pid</structfield> to <literal>MyProcPid</literal> and
-   then passing the <type>BackgroundWorkerHandle *</type> obtained at
-   registration time to
+   wait for the worker to start up.  This can be accomplished with the
    <function>WaitForBackgroundWorkerStartup(<parameter>BackgroundWorkerHandle
    *handle</parameter>, <parameter>pid_t *</parameter>)</function> function.
    This function will block until the postmaster has attempted to start the
diff --git a/src/backend/access/transam/parallel.c b/src/backend/access/transam/parallel.c
index 8613fc6fb54..6a14c0057b7 100644
--- a/src/backend/access/transam/parallel.c
+++ b/src/backend/access/transam/parallel.c
@@ -582,7 +582,6 @@ LaunchParallelWorkers(ParallelContext *pcxt)
 	sprintf(worker.bgw_library_name, "postgres");
 	sprintf(worker.bgw_function_name, "ParallelWorkerMain");
 	worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(pcxt->seg));
-	worker.bgw_notify_pid = MyProcPid;
 
 	/*
 	 * Start workers.
diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index 77707bb3847..2e3aeba8b4f 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -78,6 +78,8 @@ typedef struct BackgroundWorkerSlot
 	pid_t		pid;			/* InvalidPid = not started yet; 0 = dead */
 	uint64		generation;		/* incremented when slot is recycled */
 	BackgroundWorker worker;
+	pid_t		notify_pid;
+	ProcNumber	notify_proc_number;
 } BackgroundWorkerSlot;
 
 /*
@@ -192,8 +194,9 @@ BackgroundWorkerShmemInit(void)
 			slot->terminate = false;
 			slot->pid = InvalidPid;
 			slot->generation = 0;
+			slot->notify_pid = InvalidPid;
+			slot->notify_proc_number = INVALID_PROC_NUMBER;
 			rw->rw_shmem_slot = slotno;
-			rw->rw_worker.bgw_notify_pid = 0;	/* might be reinit after crash */
 			memcpy(&slot->worker, &rw->rw_worker, sizeof(BackgroundWorker));
 			++slotno;
 		}
@@ -234,6 +237,14 @@ FindRegisteredWorkerBySlotNumber(int slotno)
 	return NULL;
 }
 
+ProcNumber
+GetNotifyProcNumberForRegisteredWorker(RegisteredBgWorker *rw)
+{
+	BackgroundWorkerSlot *slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
+
+	return slot->notify_proc_number;
+}
+
 /*
  * Notice changes to shared memory made by other backends.
  * Accept new worker requests only if allow_new_workers is true.
@@ -315,20 +326,20 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 		/*
 		 * If the worker is marked for termination, we don't need to add it to
 		 * the registered workers list; we can just free the slot. However, if
-		 * bgw_notify_pid is set, the process that registered the worker may
-		 * need to know that we've processed the terminate request, so be sure
-		 * to signal it.
+		 * bgw_notify_proc_number is set, the process that registered the
+		 * worker may need to know that we've processed the terminate request,
+		 * so be sure to signal it.
 		 */
 		if (slot->terminate)
 		{
-			int			notify_pid;
+			int			notify_proc_number;
 
 			/*
 			 * We need a memory barrier here to make sure that the load of
-			 * bgw_notify_pid and the update of parallel_terminate_count
-			 * complete before the store to in_use.
+			 * bgw_notify_proc_number and the update of
+			 * parallel_terminate_count complete before the store to in_use.
 			 */
-			notify_pid = slot->worker.bgw_notify_pid;
+			notify_proc_number = slot->notify_proc_number;
 			if ((slot->worker.bgw_flags & BGWORKER_CLASS_PARALLEL) != 0)
 				BackgroundWorkerData->parallel_terminate_count++;
 			slot->pid = 0;
@@ -336,8 +347,8 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 			pg_memory_barrier();
 			slot->in_use = false;
 
-			if (notify_pid != 0)
-				kill(notify_pid, SIGUSR1);
+			if (notify_proc_number != INVALID_PROC_NUMBER)
+				ProcSetLatch(notify_proc_number);
 
 			continue;
 		}
@@ -383,23 +394,6 @@ BackgroundWorkerStateChange(bool allow_new_workers)
 		rw->rw_worker.bgw_main_arg = slot->worker.bgw_main_arg;
 		memcpy(rw->rw_worker.bgw_extra, slot->worker.bgw_extra, BGW_EXTRALEN);
 
-		/*
-		 * Copy the PID to be notified about state changes, but only if the
-		 * postmaster knows about a backend with that PID.  It isn't an error
-		 * if the postmaster doesn't know about the PID, because the backend
-		 * that requested the worker could have died (or been killed) just
-		 * after doing so.  Nonetheless, at least until we get some experience
-		 * with how this plays out in the wild, log a message at a relative
-		 * high debug level.
-		 */
-		rw->rw_worker.bgw_notify_pid = slot->worker.bgw_notify_pid;
-		if (!PostmasterMarkPIDForWorkerNotify(rw->rw_worker.bgw_notify_pid))
-		{
-			elog(DEBUG1, "worker notification PID %d is not valid",
-				 (int) rw->rw_worker.bgw_notify_pid);
-			rw->rw_worker.bgw_notify_pid = 0;
-		}
-
 		/* Initialize postmaster bookkeeping. */
 		rw->rw_backend = NULL;
 		rw->rw_pid = 0;
@@ -424,7 +418,7 @@ BackgroundWorkerStateChange(bool allow_new_workers)
  * points to it.  This convention allows deletion of workers during
  * searches of the worker list, and saves having to search the list again.
  *
- * Caller is responsible for notifying bgw_notify_pid, if appropriate.
+ * Caller is responsible for notifying bgw_notify_proc_number, if appropriate.
  *
  * This function must be invoked only in the postmaster.
  */
@@ -472,8 +466,8 @@ ReportBackgroundWorkerPID(RegisteredBgWorker *rw)
 	slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 	slot->pid = rw->rw_pid;
 
-	if (rw->rw_worker.bgw_notify_pid != 0)
-		kill(rw->rw_worker.bgw_notify_pid, SIGUSR1);
+	if (slot->notify_proc_number != INVALID_PROC_NUMBER)
+		ProcSetLatch(slot->notify_proc_number);
 }
 
 /*
@@ -487,14 +481,14 @@ ReportBackgroundWorkerExit(slist_mutable_iter *cur)
 {
 	RegisteredBgWorker *rw;
 	BackgroundWorkerSlot *slot;
-	int			notify_pid;
+	int			notify_proc_number;
 
 	rw = slist_container(RegisteredBgWorker, rw_lnode, cur->cur);
 
 	Assert(rw->rw_shmem_slot < max_worker_processes);
 	slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
 	slot->pid = rw->rw_pid;
-	notify_pid = rw->rw_worker.bgw_notify_pid;
+	notify_proc_number = slot->notify_proc_number;
 
 	/*
 	 * If this worker is slated for deregistration, do that before notifying
@@ -507,12 +501,12 @@ ReportBackgroundWorkerExit(slist_mutable_iter *cur)
 		rw->rw_worker.bgw_restart_time == BGW_NEVER_RESTART)
 		ForgetBackgroundWorker(cur);
 
-	if (notify_pid != 0)
-		kill(notify_pid, SIGUSR1);
+	if (notify_proc_number != INVALID_PROC_NUMBER)
+		ProcSetLatch(notify_proc_number);
 }
 
 /*
- * Cancel SIGUSR1 notifications for a PID belonging to an exiting backend.
+ * Cancel latch notifications for a PID belonging to an exiting backend.
  *
  * This function should only be called from the postmaster.
  */
@@ -524,10 +518,17 @@ BackgroundWorkerStopNotifications(pid_t pid)
 	slist_foreach(siter, &BackgroundWorkerList)
 	{
 		RegisteredBgWorker *rw;
+		BackgroundWorkerSlot *slot;
 
 		rw = slist_container(RegisteredBgWorker, rw_lnode, siter.cur);
-		if (rw->rw_worker.bgw_notify_pid == pid)
-			rw->rw_worker.bgw_notify_pid = 0;
+		Assert(rw->rw_shmem_slot < max_worker_processes);
+		slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
+
+		if (slot->notify_pid == pid)
+		{
+			slot->notify_pid = InvalidPid;
+			slot->notify_proc_number = INVALID_PROC_NUMBER;
+		}
 	}
 }
 
@@ -559,14 +560,14 @@ ForgetUnstartedBackgroundWorkers(void)
 
 		/* If it's not yet started, and there's someone waiting ... */
 		if (slot->pid == InvalidPid &&
-			rw->rw_worker.bgw_notify_pid != 0)
+			slot->notify_proc_number != INVALID_PROC_NUMBER)
 		{
 			/* ... then zap it, and notify the waiter */
-			int			notify_pid = rw->rw_worker.bgw_notify_pid;
+			int			notify_proc_number = slot->notify_proc_number;
 
 			ForgetBackgroundWorker(&iter);
-			if (notify_pid != 0)
-				kill(notify_pid, SIGUSR1);
+			if (notify_proc_number != INVALID_PROC_NUMBER)
+				ProcSetLatch(notify_proc_number);
 		}
 	}
 }
@@ -619,11 +620,6 @@ ResetBackgroundWorkerCrashTimes(void)
 			 * resetting.
 			 */
 			rw->rw_crashed_at = 0;
-
-			/*
-			 * If there was anyone waiting for it, they're history.
-			 */
-			rw->rw_worker.bgw_notify_pid = 0;
 		}
 	}
 }
@@ -987,15 +983,6 @@ RegisterBackgroundWorker(BackgroundWorker *worker)
 	if (!SanityCheckBackgroundWorker(worker, LOG))
 		return;
 
-	if (worker->bgw_notify_pid != 0)
-	{
-		ereport(LOG,
-				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-				 errmsg("background worker \"%s\": only dynamic background workers can request notification",
-						worker->bgw_name)));
-		return;
-	}
-
 	/*
 	 * Enforce maximum number of workers.  Note this is overly restrictive: we
 	 * could allow more non-shmem-connected workers, because these don't count
@@ -1113,6 +1100,24 @@ RegisterDynamicBackgroundWorker(BackgroundWorker *worker,
 			if (parallel)
 				BackgroundWorkerData->parallel_register_count++;
 
+			if ((slot->worker.bgw_flags & BGWORKER_SHMEM_ACCESS) != 0 &&
+				(slot->worker.bgw_flags & BGWORKER_NO_NOTIFY) == 0)
+			{
+				/*
+				 * We'll set the latch of the backend that created this
+				 * worker, but we also remember the PID so we can stop doing
+				 * that if it exits.
+				 */
+				slot->notify_pid = MyProcPid;
+				slot->notify_proc_number = MyProcNumber;
+			}
+			else
+			{
+				/* No notifications. */
+				slot->notify_pid = InvalidPid;
+				slot->notify_proc_number = INVALID_PROC_NUMBER;
+			}
+
 			/*
 			 * Make sure postmaster doesn't see the slot as in use before it
 			 * sees the new contents.
@@ -1213,8 +1218,9 @@ GetBackgroundWorkerPid(BackgroundWorkerHandle *handle, pid_t *pidp)
  * BGWH_POSTMASTER_DIED, since it that case we know that startup will not
  * take place.
  *
- * The caller *must* have set our PID as the worker's bgw_notify_pid,
- * else we will not be awoken promptly when the worker's state changes.
+ * This works only if the worker was registered with BGWORKER_SHMEM_ACCESS and
+ * without BGWORKER_NO_NOTIFY, else we will not be awoken promptly when the
+ * worker's state changes.
  */
 BgwHandleStatus
 WaitForBackgroundWorkerStartup(BackgroundWorkerHandle *handle, pid_t *pidp)
@@ -1258,8 +1264,9 @@ WaitForBackgroundWorkerStartup(BackgroundWorkerHandle *handle, pid_t *pidp)
  * up and return BGWH_POSTMASTER_DIED, because it's the postmaster that
  * notifies us when a worker's state changes.
  *
- * The caller *must* have set our PID as the worker's bgw_notify_pid,
- * else we will not be awoken promptly when the worker's state changes.
+ * This works only if the worker was registered with BGWORKER_SHMEM_ACCESS and
+ * without BGWORKER_NO_NOTIFY, else we will not be awoken promptly when the
+ * worker's state changes.
  */
 BgwHandleStatus
 WaitForBackgroundWorkerShutdown(BackgroundWorkerHandle *handle)
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index 6f974a8d21a..3b5aa4ec160 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -4373,15 +4373,15 @@ maybe_start_bgworkers(void)
 		{
 			if (rw->rw_worker.bgw_restart_time == BGW_NEVER_RESTART)
 			{
-				int			notify_pid;
+				ProcNumber	notify_proc_number;
 
-				notify_pid = rw->rw_worker.bgw_notify_pid;
+				notify_proc_number = GetNotifyProcNumberForRegisteredWorker(rw);
 
 				ForgetBackgroundWorker(&iter);
 
 				/* Report worker is gone now. */
-				if (notify_pid != 0)
-					kill(notify_pid, SIGUSR1);
+				if (notify_proc_number != INVALID_PROC_NUMBER)
+					ProcSetLatch(notify_proc_number);
 
 				continue;
 			}
diff --git a/src/backend/replication/logical/launcher.c b/src/backend/replication/logical/launcher.c
index 27c3a91fb75..9894415cd9b 100644
--- a/src/backend/replication/logical/launcher.c
+++ b/src/backend/replication/logical/launcher.c
@@ -497,7 +497,6 @@ retry:
 	}
 
 	bgw.bgw_restart_time = BGW_NEVER_RESTART;
-	bgw.bgw_notify_pid = MyProcPid;
 	bgw.bgw_main_arg = Int32GetDatum(slot);
 
 	if (!RegisterDynamicBackgroundWorker(&bgw, &bgw_handle))
@@ -940,7 +939,6 @@ ApplyLauncherRegister(void)
 	snprintf(bgw.bgw_type, BGW_MAXLEN,
 			 "logical replication launcher");
 	bgw.bgw_restart_time = 5;
-	bgw.bgw_notify_pid = 0;
 	bgw.bgw_main_arg = (Datum) 0;
 
 	RegisterBackgroundWorker(&bgw);
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index 1b23efb26f3..08bfa485b50 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -1883,6 +1883,19 @@ ProcSendSignal(ProcNumber procNumber)
 	SetLatch(&ProcGlobal->allProcs[procNumber].procLatch);
 }
 
+/*
+ * ProcSetLatch - set the latch of a backend identified by ProcNumber, but
+ * ignore out of range values.
+ */
+void
+ProcSetLatch(ProcNumber procNumber)
+{
+	if (procNumber < 0 || procNumber >= ProcGlobal->allProcCount)
+		return;
+
+	SetLatch(&ProcGlobal->allProcs[procNumber].procLatch);
+}
+
 /*
  * BecomeLockGroupLeader - designate process as lock group leader
  *
diff --git a/src/include/postmaster/bgworker.h b/src/include/postmaster/bgworker.h
index 22fc49ec27f..bb7dca7ff0e 100644
--- a/src/include/postmaster/bgworker.h
+++ b/src/include/postmaster/bgworker.h
@@ -41,6 +41,8 @@
 #ifndef BGWORKER_H
 #define BGWORKER_H
 
+#include "storage/procnumber.h"
+
 /*---------------------------------------------------------------------
  * External module API.
  *---------------------------------------------------------------------
@@ -59,6 +61,14 @@
  */
 #define BGWORKER_BACKEND_DATABASE_CONNECTION		0x0002
 
+/*
+ * Dynamic workers created with shared memory access usually set the latch of
+ * the creating backend when they start and stop, allowing
+ * WaitForBackgroundWorker{Startup,Shutdown}() to work.  Such notificaitons
+ * can be suppressed with this flag.
+ */
+#define BGWORKER_NO_NOTIFY							0x0004
+
 /*
  * This class is used internally for parallel queries, to keep track of the
  * number of active parallel workers and make sure we never launch more than
@@ -97,7 +107,7 @@ typedef struct BackgroundWorker
 	char		bgw_function_name[BGW_MAXLEN];
 	Datum		bgw_main_arg;
 	char		bgw_extra[BGW_EXTRALEN];
-	pid_t		bgw_notify_pid; /* SIGUSR1 this backend on start/stop */
+	pid_t		bgw_notify_pid; /* not used */
 } BackgroundWorker;
 
 typedef enum BgwHandleStatus
@@ -126,6 +136,7 @@ extern BgwHandleStatus
 			WaitForBackgroundWorkerShutdown(BackgroundWorkerHandle *);
 extern const char *GetBackgroundWorkerTypeByPid(pid_t pid);
 
+
 /* Terminate a bgworker */
 extern void TerminateBackgroundWorker(BackgroundWorkerHandle *handle);
 
diff --git a/src/include/postmaster/bgworker_internals.h b/src/include/postmaster/bgworker_internals.h
index 9106a0ef3f0..b609b2fea2c 100644
--- a/src/include/postmaster/bgworker_internals.h
+++ b/src/include/postmaster/bgworker_internals.h
@@ -53,6 +53,7 @@ extern void ReportBackgroundWorkerExit(slist_mutable_iter *cur);
 extern void BackgroundWorkerStopNotifications(pid_t pid);
 extern void ForgetUnstartedBackgroundWorkers(void);
 extern void ResetBackgroundWorkerCrashTimes(void);
+extern ProcNumber GetNotifyProcNumberForRegisteredWorker(RegisteredBgWorker *rw);
 
 /* Entry point for background worker processes */
 extern void BackgroundWorkerMain(char *startup_data, size_t startup_data_len) pg_attribute_noreturn();
diff --git a/src/include/storage/proc.h b/src/include/storage/proc.h
index 7d3fc2bfa60..ae483928ba7 100644
--- a/src/include/storage/proc.h
+++ b/src/include/storage/proc.h
@@ -481,6 +481,7 @@ extern void LockErrorCleanup(void);
 
 extern void ProcWaitForSignal(uint32 wait_event_info);
 extern void ProcSendSignal(ProcNumber procNumber);
+extern void ProcSetLatch(ProcNumber procNumber);
 
 extern PGPROC *AuxiliaryPidGetProc(int pid);
 
diff --git a/src/test/modules/test_shm_mq/setup.c b/src/test/modules/test_shm_mq/setup.c
index b3dac44d97a..43a25dd0b12 100644
--- a/src/test/modules/test_shm_mq/setup.c
+++ b/src/test/modules/test_shm_mq/setup.c
@@ -134,6 +134,7 @@ setup_dynamic_shared_memory(int64 queue_size, int nworkers,
 
 	/* Set up the header region. */
 	hdr = shm_toc_allocate(toc, sizeof(test_shm_mq_header));
+	hdr->leader_proc_number = MyProcNumber;
 	SpinLockInit(&hdr->mutex);
 	hdr->workers_total = nworkers;
 	hdr->workers_attached = 0;
@@ -223,8 +224,6 @@ setup_background_workers(int nworkers, dsm_segment *seg)
 	sprintf(worker.bgw_function_name, "test_shm_mq_main");
 	snprintf(worker.bgw_type, BGW_MAXLEN, "test_shm_mq");
 	worker.bgw_main_arg = UInt32GetDatum(dsm_segment_handle(seg));
-	/* set bgw_notify_pid, so we can detect if the worker stops */
-	worker.bgw_notify_pid = MyProcPid;
 
 	/* Register the workers. */
 	for (i = 0; i < nworkers; ++i)
diff --git a/src/test/modules/test_shm_mq/test_shm_mq.h b/src/test/modules/test_shm_mq/test_shm_mq.h
index 0ae7bd64cd0..7f27a61aff8 100644
--- a/src/test/modules/test_shm_mq/test_shm_mq.h
+++ b/src/test/modules/test_shm_mq/test_shm_mq.h
@@ -28,6 +28,8 @@
  */
 typedef struct
 {
+	ProcNumber	leader_proc_number;
+
 	slock_t		mutex;
 	int			workers_total;
 	int			workers_attached;
diff --git a/src/test/modules/test_shm_mq/worker.c b/src/test/modules/test_shm_mq/worker.c
index 6c4fbc78274..e4f187f4a81 100644
--- a/src/test/modules/test_shm_mq/worker.c
+++ b/src/test/modules/test_shm_mq/worker.c
@@ -52,7 +52,6 @@ test_shm_mq_main(Datum main_arg)
 	shm_mq_handle *outqh;
 	volatile test_shm_mq_header *hdr;
 	int			myworkernumber;
-	PGPROC	   *registrant;
 
 	/*
 	 * Establish signal handlers.
@@ -122,13 +121,8 @@ test_shm_mq_main(Datum main_arg)
 	SpinLockAcquire(&hdr->mutex);
 	++hdr->workers_ready;
 	SpinLockRelease(&hdr->mutex);
-	registrant = BackendPidGetProc(MyBgworkerEntry->bgw_notify_pid);
-	if (registrant == NULL)
-	{
-		elog(DEBUG1, "registrant backend has exited prematurely");
-		proc_exit(1);
-	}
-	SetLatch(&registrant->procLatch);
+
+	ProcSetLatch(hdr->leader_proc_number);
 
 	/* Do the work. */
 	copy_messages(inqh, outqh);
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index 7e1042f4ab6..fe51b7db73f 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -379,7 +379,6 @@ _PG_init(void)
 	worker.bgw_restart_time = BGW_NEVER_RESTART;
 	sprintf(worker.bgw_library_name, "worker_spi");
 	sprintf(worker.bgw_function_name, "worker_spi_main");
-	worker.bgw_notify_pid = 0;
 
 	/*
 	 * Now fill in worker-specific data, and do the actual registrations.
@@ -428,8 +427,6 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	snprintf(worker.bgw_name, BGW_MAXLEN, "worker_spi dynamic worker %d", i);
 	snprintf(worker.bgw_type, BGW_MAXLEN, "worker_spi dynamic");
 	worker.bgw_main_arg = Int32GetDatum(i);
-	/* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
-	worker.bgw_notify_pid = MyProcPid;
 
 	/* extract flags, if any */
 	ndim = ARR_NDIM(arr);
-- 
2.45.2

