1 files changed, 71 insertions, 6 deletions

diff --git a/src/backend/postmaster/bgworker.c b/src/backend/postmaster/bgworker.c
index 7ef6259eb50..d209b69ec05 100644
--- a/src/backend/postmaster/bgworker.c
+++ b/src/backend/postmaster/bgworker.c
@@ -231,13 +231,15 @@ FindRegisteredWorkerBySlotNumber(int slotno)
 }
 
 /*
- * Notice changes to shared memory made by other backends.  This code
- * runs in the postmaster, so we must be very careful not to assume that
- * shared memory contents are sane.  Otherwise, a rogue backend could take
- * out the postmaster.
+ * Notice changes to shared memory made by other backends.
+ * Accept new worker requests only if allow_new_workers is true.
+ *
+ * This code runs in the postmaster, so we must be very careful not to assume
+ * that shared memory contents are sane.  Otherwise, a rogue backend could
+ * take out the postmaster.
  */
 void
-BackgroundWorkerStateChange(void)
+BackgroundWorkerStateChange(bool allow_new_workers)
 {
 	int			slotno;
 
@@ -298,6 +300,15 @@ BackgroundWorkerStateChange(void)
 		}
 
 		/*
+		 * If we aren't allowing new workers, then immediately mark it for
+		 * termination; the next stanza will take care of cleaning it up.
+		 * Doing this ensures that any process waiting for the worker will get
+		 * awoken, even though the worker will never be allowed to run.
+		 */
+		if (!allow_new_workers)
+			slot->terminate = true;
+
+		/*
 		 * 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
@@ -504,11 +515,54 @@ BackgroundWorkerStopNotifications(pid_t pid)
 }
 
 /*
+ * Cancel any not-yet-started worker requests that have waiting processes.
+ *
+ * This is called during a normal ("smart" or "fast") database shutdown.
+ * After this point, no new background workers will be started, so anything
+ * that might be waiting for them needs to be kicked off its wait.  We do
+ * that by cancelling the bgworker registration entirely, which is perhaps
+ * overkill, but since we're shutting down it does not matter whether the
+ * registration record sticks around.
+ *
+ * This function should only be called from the postmaster.
+ */
+void
+ForgetUnstartedBackgroundWorkers(void)
+{
+	slist_mutable_iter iter;
+
+	slist_foreach_modify(iter, &BackgroundWorkerList)
+	{
+		RegisteredBgWorker *rw;
+		BackgroundWorkerSlot *slot;
+
+		rw = slist_container(RegisteredBgWorker, rw_lnode, iter.cur);
+		Assert(rw->rw_shmem_slot < max_worker_processes);
+		slot = &BackgroundWorkerData->slot[rw->rw_shmem_slot];
+
+		/* If it's not yet started, and there's someone waiting ... */
+		if (slot->pid == InvalidPid &&
+			rw->rw_worker.bgw_notify_pid != 0)
+		{
+			/* ... then zap it, and notify the waiter */
+			int			notify_pid = rw->rw_worker.bgw_notify_pid;
+
+			ForgetBackgroundWorker(&iter);
+			if (notify_pid != 0)
+				kill(notify_pid, SIGUSR1);
+		}
+	}
+}
+
+/*
  * Reset background worker crash state.
  *
  * We assume that, after a crash-and-restart cycle, background workers without
  * the never-restart flag should be restarted immediately, instead of waiting
- * for bgw_restart_time to elapse.
+ * for bgw_restart_time to elapse.  On the other hand, workers with that flag
+ * should be forgotten immediately, since we won't ever restart them.
+ *
+ * This function should only be called from the postmaster.
  */
 void
 ResetBackgroundWorkerCrashTimes(void)
@@ -548,6 +602,11 @@ 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;
 		}
 	}
 }
@@ -1077,6 +1136,9 @@ GetBackgroundWorkerPid(BackgroundWorkerHandle *handle, pid_t *pidp)
  * returned.  However, if the postmaster has died, we give up and return
  * 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.
  */
 BgwHandleStatus
 WaitForBackgroundWorkerStartup(BackgroundWorkerHandle *handle, pid_t *pidp)
@@ -1119,6 +1181,9 @@ WaitForBackgroundWorkerStartup(BackgroundWorkerHandle *handle, pid_t *pidp)
  * and then return BGWH_STOPPED.  However, if the postmaster has died, we give
  * 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.
  */
 BgwHandleStatus
 WaitForBackgroundWorkerShutdown(BackgroundWorkerHandle *handle)