Fix REPACK worker startup and shutdown sequences

A REPACK background worker can fail to start for various reasons, but
the REPACK steering process was not handling them correctly.  For
instance, the steering process waits on a condition variable, but it
would simply hang if the worker doesn't send the signal.  (A signal
would be sent by postmaster, but it would wake up the process only to be
ignored and continue sleeping.)  Fix this by splitting the wait in two:
first wait for the worker to attach to the error queue, and then wait on
the condition variable as before.

The shutdown sequence can also get stuck, if the message queue gets full
while trying to shut down (this can apparently happen if the log level
is set to very noisy).  Previously the worker would send a signal and
then detach; but the steering process can receive that signal, act on
it, then fail to wake up to handle the detach, sleeping indefinitely.
Fix this by having the worker send a Terminate message, which the
steering process can react to.

Author: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Author: Shihao Zhong <zhong950419@gmail.com>
Author: Álvaro Herrera <alvherre@kurilemu.de>
Reviewed-by: Antonin Houska <ah@cybertec.at>
Reviewed-by: Masahiko Sawada <sawada.mshk@gmail.com>
Reviewed-by: Shihao Zhong <zhong950419@gmail.com>
Reported-by: Nathan Bossart <nathandbossart@gmail.com>
Reported-by: Bharath Rupireddy <bharath.rupireddyforpostgres@gmail.com>
Reported-by: Nikolay Samokhvalov <nik@postgres.ai>
Backpatch-through: 19
Discussion: https://postgr.es/m/CALj2ACVAxA9HxvFe8HSspTJ-UO4Aoz%3DkuQdZBeLrod0gqUxH3g%40mail.gmail.com
Discussion: https://postgr.es/m/apBpOVZOyqrakEr_@nathan

3 files changed, 192 insertions, 56 deletions

diff --git a/src/backend/commands/repack.c b/src/backend/commands/repack.c
index 0f7281f0894..74f2dd46aa5 100644
--- a/src/backend/commands/repack.c
+++ b/src/backend/commands/repack.c
@@ -217,6 +217,7 @@ static Oid	determine_clustered_index(Relation rel, bool usingindex,
 									  const char *indexname);
 
 static void start_repack_decoding_worker(Oid relid);
+static void wait_for_repack_decoding_worker(void);
 static void stop_repack_decoding_worker(void);
 static void stop_repack_decoding_worker_cb(int code, Datum arg);
 static Snapshot get_initial_snapshot(DecodingWorker *worker);
@@ -3799,6 +3800,19 @@ start_repack_decoding_worker(Oid relid)
 				errhint("You might need to increase \"%s\".", "max_worker_processes"));
 
 	/*
+	 * Now that the worker is registered, connect the error message queue to
+	 * it.
+	 */
+	shm_mq_set_handle(decoding_worker->error_mqh, decoding_worker->handle);
+
+	/*
+	 * Make sure the worker has started before we wait for it to initialize
+	 * decoding below, so that the failure-to-start case does not hang
+	 * forever.
+	 */
+	wait_for_repack_decoding_worker();
+
+	/*
 	 * The decoding setup must be done before the caller can have XID assigned
 	 * for any reason, otherwise the worker might end up in a deadlock,
 	 * waiting for the caller's transaction to end. Therefore wait here until
@@ -3822,6 +3836,74 @@ start_repack_decoding_worker(Oid relid)
 }
 
 /*
+ * Wait for the decoding worker to start up, and throw an error if it fails
+ * to do so.
+ *
+ * This is similar to WaitForParallelWorkersToAttach(). The only reliable way
+ * to tell a worker that failed to start (fork failure, or an exit before it
+ * attached) from one that is merely slow is to check whether it became the
+ * sender on the error message queue. If it stopped without attaching, nothing
+ * was queued and we report the generic failure ourselves. If it attached, any
+ * error it reported is in the queue and is thrown when we process pending
+ * messages, either here or later while we wait for it to initialize decoding.
+ */
+static void
+wait_for_repack_decoding_worker(void)
+{
+	for (;;)
+	{
+		BgwHandleStatus status;
+		shm_mq	   *mq;
+		int			rc;
+		pid_t		pid;
+
+		/*
+		 * This will process any repack messages that are pending and it may
+		 * also throw an error propagated from a worker.
+		 */
+		CHECK_FOR_INTERRUPTS();
+
+		/* If error_mqh is NULL, the worker has exited cleanly */
+		if (decoding_worker->error_mqh == NULL)
+			break;
+
+		status = GetBackgroundWorkerPid(decoding_worker->handle, &pid);
+		if (status == BGWH_STARTED)
+		{
+			/* Has the worker attached to the error message queue? */
+			mq = shm_mq_get_queue(decoding_worker->error_mqh);
+			if (shm_mq_get_sender(mq) != NULL)
+				break;
+		}
+		else if (status == BGWH_STOPPED)
+		{
+			/*
+			 * If the worker stopped without attaching to the error message
+			 * queue, throw an error. Otherwise, assume it attached and
+			 * reported an error before exiting, so mark it attached and let
+			 * the next attempt to process pending messages, here or later
+			 * while the initial snapshot is set up, throw that error.
+			 */
+			mq = shm_mq_get_queue(decoding_worker->error_mqh);
+			if (shm_mq_get_sender(mq) == NULL)
+				ereport(ERROR,
+						errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+						errmsg("REPACK decoding worker failed to start"),
+						errhint("More details may be available in the server log."));
+			break;
+		}
+
+		/* Worker neither started or stopped yet, so wait. */
+		rc = WaitLatch(MyLatch,
+					   WL_LATCH_SET | WL_EXIT_ON_PM_DEATH,
+					   -1, WAIT_EVENT_BGWORKER_STARTUP);
+
+		if (rc & WL_LATCH_SET)
+			ResetLatch(MyLatch);
+	}
+}
+
+/*
  * Stop the decoding worker and cleanup the related resources.
  *
  * The worker stops on its own when it knows there is no more work to do, but
@@ -3834,13 +3916,36 @@ stop_repack_decoding_worker(void)
 	if (decoding_worker == NULL)
 		return;
 
-	/* Terminate the worker process, if one is running. */
+	/* Terminate the decoding worker, if one is running */
+	if (decoding_worker->handle != NULL)
+		TerminateBackgroundWorker(decoding_worker->handle);
+
+	/*
+	 * The error queue should have already been nulled out during worker
+	 * shutdown, but if that didn't happen, do it now.
+	 */
+	if (decoding_worker->error_mqh != NULL)
+	{
+		shm_mq_detach(decoding_worker->error_mqh);
+		decoding_worker->error_mqh = NULL;
+	}
+
+	/*
+	 * Cancel any sleep on the condition variable before detaching the shared
+	 * memory segment, because the CV lives in that segment.  Otherwise later
+	 * cleanup would touch freed memory.
+	 */
+	ConditionVariableCancelSleep();
+
+	/*
+	 * We can't finish the REPACK command until the worker has exited.  This
+	 * means, in particular, that we can't respond to interrupts at this
+	 * stage.
+	 */
 	if (decoding_worker->handle != NULL)
 	{
 		BgwHandleStatus status;
 
-		TerminateBackgroundWorker(decoding_worker->handle);
-		/* The worker should really exit before the REPACK command does. */
 		HOLD_INTERRUPTS();
 		status = WaitForBackgroundWorkerShutdown(decoding_worker->handle);
 		RESUME_INTERRUPTS();
@@ -3852,21 +3957,17 @@ stop_repack_decoding_worker(void)
 	}
 
 	/*
-	 * Now detach from our shared memory segment.  In error cases there might
-	 * still be messages from the worker in the queue, which ProcessInterrupts
-	 * would try to read; this is pointless (and causes an assertion failure),
-	 * so set the global pointer to NULL to have ProcessRepackMessages ignore
-	 * them.
-	 *
-	 * We must also cancel the current sleep, if one is still set up.  This is
-	 * critical because the CV lives in the DSM that we're about to detach, so
-	 * if we omit it, later automatic cleanup tries to clear freed memory.
+	 * Detach from the shared memory segment only now that the worker is gone.
+	 * The worker attaches to the shared file set after it maps the segment,
+	 * so detaching any earlier can destroy the file set under a worker that
+	 * is still starting up.
 	 */
-	if (decoding_worker->error_mqh != NULL)
-		shm_mq_detach(decoding_worker->error_mqh);
-	ConditionVariableCancelSleep();
 	if (decoding_worker->seg != NULL)
+	{
 		dsm_detach(decoding_worker->seg);
+		decoding_worker->seg = NULL;
+	}
+
 	pfree(decoding_worker);
 	decoding_worker = NULL;
 }
@@ -3971,9 +4072,11 @@ ProcessRepackMessages(void)
 
 	/*
 	 * Nothing to do if we haven't launched the worker yet or have already
-	 * terminated it.
+	 * terminated it. Stopping the worker detaches the error message queue
+	 * before clearing decoding_worker, so also bail out once error_mqh is
+	 * gone.
 	 */
-	if (decoding_worker == NULL)
+	if (decoding_worker == NULL || decoding_worker->error_mqh == NULL)
 		return;
 
 	/*
@@ -4003,40 +4106,49 @@ ProcessRepackMessages(void)
 	RepackMessagePending = false;
 
 	/*
-	 * Read as many messages as we can from the worker, but stop when no more
-	 * messages can be read from the worker without blocking.
+	 * Read messages from the worker, but stop if the error queue disappears,
+	 * which happens when a PqRepackMsg_Terminate is received; or as soon as
+	 * no more messages can be read without blocking.  Messages are
+	 * infrequent, so no point optimizing stringinfo allocation.
 	 */
-	while (true)
+	while (decoding_worker->error_mqh != NULL)
 	{
 		shm_mq_result res;
 		Size		nbytes;
 		void	   *data;
+		StringInfoData msg;
 
 		res = shm_mq_receive(decoding_worker->error_mqh, &nbytes,
 							 &data, true);
-		if (res == SHM_MQ_WOULD_BLOCK)
-			break;
-		else if (res == SHM_MQ_SUCCESS)
+		switch (res)
 		{
-			StringInfoData msg;
+			case SHM_MQ_SUCCESS:
+				initStringInfo(&msg);
+				appendBinaryStringInfo(&msg, data, nbytes);
+				ProcessRepackMessage(&msg);
+				pfree(msg.data);
+				break;
 
-			initStringInfo(&msg);
-			appendBinaryStringInfo(&msg, data, nbytes);
-			ProcessRepackMessage(&msg);
-			pfree(msg.data);
-		}
-		else
-		{
-			/*
-			 * The decoding worker is special in that it exits as soon as it
-			 * has its work done. Thus the DETACHED result code is fine.
-			 */
-			Assert(res == SHM_MQ_DETACHED);
+			case SHM_MQ_WOULD_BLOCK:
+				/* no more messages to process for now */
+				goto out;
 
-			break;
+			case SHM_MQ_DETACHED:
+
+				/*
+				 * Normal worker stop sends a Terminate message, causing
+				 * ProcessRepackMessage to set error_mqh to NULL, thus exiting
+				 * this loop; so this case should never happen.  If it does,
+				 * the worker terminated abnormally, so report that.
+				 */
+				ereport(ERROR,
+						errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+						errmsg("lost connection to REPACK decoding worker"),
+						errhint("More details may be available in the server log."));
 		}
 	}
 
+out:
 	MemoryContextSwitchTo(oldcontext);
 
 	/* Might as well clear the context on our way out */
@@ -4046,7 +4158,7 @@ ProcessRepackMessages(void)
 }
 
 /*
- * Process a single protocol message received from a single parallel worker.
+ * Process a single protocol message received from a repack worker.
  */
 static void
 ProcessRepackMessage(StringInfo msg)
@@ -4085,10 +4197,18 @@ ProcessRepackMessage(StringInfo msg)
 				break;
 			}
 
+		case PqRepackMsg_Terminate:
+
+			/*
+			 * The worker has completed its work; stop watching its message
+			 * queue now for orderly shutdown.
+			 */
+			shm_mq_detach(decoding_worker->error_mqh);
+			decoding_worker->error_mqh = NULL;
+			break;
+
 		default:
-			{
-				elog(ERROR, "unrecognized message type received from decoding worker: %c (message length %d bytes)",
-					 msgtype, msg->len);
-			}
+			elog(ERROR, "unrecognized message type received from decoding worker: %c (message length %d bytes)",
+				 msgtype, msg->len);
 	}
 }
diff --git a/src/backend/commands/repack_worker.c b/src/backend/commands/repack_worker.c
index bf2bc2dca13..f51a896ebd9 100644
--- a/src/backend/commands/repack_worker.c
+++ b/src/backend/commands/repack_worker.c
@@ -21,6 +21,7 @@
 #include "access/xlogwait.h"
 #include "commands/repack.h"
 #include "commands/repack_internal.h"
+#include "libpq/libpq.h"
 #include "libpq/pqmq.h"
 #include "replication/snapbuild.h"
 #include "storage/ipc.h"
@@ -44,9 +45,6 @@ static bool am_repack_worker = false;
 /* The WAL segment being decoded. */
 static XLogSegNo repack_current_segment = 0;
 
-/* Our DSM segment, for shutting down */
-static dsm_segment *worker_dsm_segment = NULL;
-
 /*
  * Keep track of the table we're processing, to skip logical decoding of data
  * from other relations.
@@ -76,12 +74,11 @@ RepackWorkerMain(Datum main_arg)
 		ereport(ERROR,
 				errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
 				errmsg("could not map dynamic shared memory segment"));
-	worker_dsm_segment = seg;
 
 	shared = (DecodingWorkerShared *) dsm_segment_address(seg);
 
-	/* Arrange to signal the leader if we exit. */
-	before_shmem_exit(RepackWorkerShutdown, PointerGetDatum(shared));
+	/* Arrange to signal the steering process if we exit. */
+	before_shmem_exit(RepackWorkerShutdown, PointerGetDatum(seg));
 
 	/*
 	 * Join locking group - see the comments around the call of
@@ -161,27 +158,43 @@ RepackWorkerMain(Datum main_arg)
 
 		if (stop)
 			break;
-
 	}
 
-	/* Cleanup. */
+	/* Clean up and report termination to our steering process */
 	repack_cleanup_logical_decoding(decoding_ctx);
 	CommitTransactionCommand();
+	pq_putmessage(PqRepackMsg_Terminate, NULL, 0);
 }
 
 /*
- * See ParallelWorkerShutdown for details.
+ * Make sure the repack steering process tries to read from our error queue one
+ * more time.  This guards against the case where we exit uncleanly without
+ * sending an ErrorResponse to the leader, for example because some code calls
+ * proc_exit directly.
  */
 static void
 RepackWorkerShutdown(int code, Datum arg)
 {
-	DecodingWorkerShared *shared = (DecodingWorkerShared *) DatumGetPointer(arg);
+	dsm_segment *seg;
+	DecodingWorkerShared *shared;
+	pid_t		pid;
+	ProcNumber	procno;
 
-	SendProcSignal(shared->backend_pid,
-				   PROCSIG_REPACK_MESSAGE,
-				   shared->backend_proc_number);
+	seg = (dsm_segment *) DatumGetPointer(arg);
+	shared = (DecodingWorkerShared *) dsm_segment_address(seg);
+	pid = shared->backend_pid;
+	procno = shared->backend_proc_number;
+
+	/*
+	 * Detach from the shared memory segment before we signal the backend.
+	 * Detaching also detaches the error message queue, and the backend learns
+	 * that we are gone by reading that queue when it handles our signal. If
+	 * we signaled first, the backend could read the queue while it still
+	 * looks attached, and nothing would make it read again.
+	 */
+	dsm_detach(seg);
 
-	dsm_detach(worker_dsm_segment);
+	SendProcSignal(pid, PROCSIG_REPACK_MESSAGE, procno);
 }
 
 bool
diff --git a/src/include/libpq/protocol.h b/src/include/libpq/protocol.h
index eae8f0e7238..1879b0cebb4 100644
--- a/src/include/libpq/protocol.h
+++ b/src/include/libpq/protocol.h
@@ -69,6 +69,9 @@
 
 #define PqMsg_Progress              'P'
 
+/* Additional codes sent by repack workers to their steering process */
+
+#define PqRepackMsg_Terminate		'X'
 
 /* Replication codes sent by the primary (wrapped in CopyData messages). */