Redesign parallel dump/restore's wait-for-workers logic.

The ListenToWorkers/ReapWorkerStatus APIs were messy and hard to use.
Instead, make DispatchJobForTocEntry register a callback function that
will take care of state cleanup, doing whatever had been done by the caller
of ReapWorkerStatus in the old design.  (This callback is essentially just
the old mark_work_done function in the restore case, and a trivial test for
worker failure in the dump case.)  Then we can have ListenToWorkers call
the callback immediately on receipt of a status message, and return the
worker to WRKR_IDLE state; so the WRKR_FINISHED state goes away.

This allows us to design a unified wait-for-worker-messages loop:
WaitForWorkers replaces EnsureIdleWorker and EnsureWorkersFinished as well
as the mess in restore_toc_entries_parallel.  Also, we no longer need the
fragile API spec that the caller of DispatchJobForTocEntry is responsible
for ensuring there's an idle worker, since DispatchJobForTocEntry can just
wait until there is one.

In passing, I got rid of the ParallelArgs struct, which was a net negative
in terms of notational verboseness, and didn't seem to be providing any
noticeable amount of abstraction either.

Tom Lane, reviewed by Kevin Grittner

Discussion: <1188.1464544443@sss.pgh.pa.us>

src/bin/pg_dump/parallel.h

@@ -2,14 +2,11 @@
  *
  * parallel.h
  *
- *	Parallel support header file for the pg_dump archiver
+ *	Parallel support for pg_dump and pg_restore
  *
  * Portions Copyright (c) 1996-2016, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- *	The author is not responsible for loss or damages that may
- *	result from its use.
- *
  * IDENTIFICATION
  *		src/bin/pg_dump/parallel.h
  *
@@ -21,31 +18,53 @@
 
 #include "pg_backup_archiver.h"
 
+/* Function to call in master process on completion of a worker task */
+typedef void (*ParallelCompletionPtr) (ArchiveHandle *AH,
+												   TocEntry *te,
+												   int status,
+												   void *callback_data);
+
+/* Wait options for WaitForWorkers */
+typedef enum
+{
+	WFW_NO_WAIT,
+	WFW_GOT_STATUS,
+	WFW_ONE_IDLE,
+	WFW_ALL_IDLE
+} WFW_WaitOption;
+
+/* Worker process statuses */
 typedef enum
 {
-	WRKR_TERMINATED = 0,
 	WRKR_IDLE,
 	WRKR_WORKING,
-	WRKR_FINISHED
+	WRKR_TERMINATED
 } T_WorkerStatus;
 
-/* Arguments needed for a worker process */
-typedef struct ParallelArgs
-{
-	ArchiveHandle *AH;
-	TocEntry   *te;
-} ParallelArgs;
-
-/* State for each parallel activity slot */
+/*
+ * Per-parallel-worker state of parallel.c.
+ *
+ * Much of this is valid only in the master process (or, on Windows, should
+ * be touched only by the master thread).  But the AH field should be touched
+ * only by workers.  The pipe descriptors are valid everywhere.
+ */
 typedef struct ParallelSlot
 {
-	ParallelArgs *args;
-	T_WorkerStatus workerStatus;
-	int			status;
+	T_WorkerStatus workerStatus;	/* see enum above */
+
+	/* These fields are valid if workerStatus == WRKR_WORKING: */
+	TocEntry   *te;				/* item being worked on */
+	ParallelCompletionPtr callback;		/* function to call on completion */
+	void	   *callback_data;	/* passthru data for it */
+
+	ArchiveHandle *AH;			/* Archive data worker is using */
+
 	int			pipeRead;		/* master's end of the pipes */
 	int			pipeWrite;
 	int			pipeRevRead;	/* child's end of the pipes */
 	int			pipeRevWrite;
+
+	/* Child process/thread identity info: */
 #ifdef WIN32
 	uintptr_t	hThread;
 	unsigned int threadId;
@@ -54,12 +73,11 @@ typedef struct ParallelSlot
 #endif
 } ParallelSlot;
 
-#define NO_SLOT (-1)
-
+/* Overall state for parallel.c */
 typedef struct ParallelState
 {
-	int			numWorkers;
-	ParallelSlot *parallelSlot;
+	int			numWorkers;		/* allowed number of workers */
+	ParallelSlot *parallelSlot; /* array of numWorkers slots */
 } ParallelState;
 
 #ifdef WIN32
@@ -69,17 +87,17 @@ extern DWORD mainThreadId;
 
 extern void init_parallel_dump_utils(void);
 
-extern int	GetIdleWorker(ParallelState *pstate);
 extern bool IsEveryWorkerIdle(ParallelState *pstate);
-extern void ListenToWorkers(ArchiveHandle *AH, ParallelState *pstate, bool do_wait);
-extern int	ReapWorkerStatus(ParallelState *pstate, int *status);
-extern void EnsureIdleWorker(ArchiveHandle *AH, ParallelState *pstate);
-extern void EnsureWorkersFinished(ArchiveHandle *AH, ParallelState *pstate);
+extern void WaitForWorkers(ArchiveHandle *AH, ParallelState *pstate,
+			   WFW_WaitOption mode);
 
 extern ParallelState *ParallelBackupStart(ArchiveHandle *AH);
 extern void DispatchJobForTocEntry(ArchiveHandle *AH,
 					   ParallelState *pstate,
-					   TocEntry *te, T_Action act);
+					   TocEntry *te,
+					   T_Action act,
+					   ParallelCompletionPtr callback,
+					   void *callback_data);
 extern void ParallelBackupEnd(ArchiveHandle *AH, ParallelState *pstate);
 
 extern void set_archive_cancel_info(ArchiveHandle *AH, PGconn *conn);

src/bin/pg_dump/pg_backup_archiver.c

@@ -97,9 +97,14 @@ static void par_list_remove(TocEntry *te);
 static TocEntry *get_next_work_item(ArchiveHandle *AH,
 				   TocEntry *ready_list,
 				   ParallelState *pstate);
-static void mark_work_done(ArchiveHandle *AH, TocEntry *ready_list,
-			   int worker, int status,
-			   ParallelState *pstate);
+static void mark_dump_job_done(ArchiveHandle *AH,
+				   TocEntry *te,
+				   int status,
+				   void *callback_data);
+static void mark_restore_job_done(ArchiveHandle *AH,
+					  TocEntry *te,
+					  int status,
+					  void *callback_data);
 static void fix_dependencies(ArchiveHandle *AH);
 static bool has_lock_conflicts(TocEntry *te1, TocEntry *te2);
 static void repoint_table_dependencies(ArchiveHandle *AH);
@@ -2355,8 +2360,8 @@ WriteDataChunks(ArchiveHandle *AH, ParallelState *pstate)
 			 * If we are in a parallel backup, then we are always the master
 			 * process.  Dispatch each data-transfer job to a worker.
 			 */
-			EnsureIdleWorker(AH, pstate);
-			DispatchJobForTocEntry(AH, pstate, te, ACT_DUMP);
+			DispatchJobForTocEntry(AH, pstate, te, ACT_DUMP,
+								   mark_dump_job_done, NULL);
 		}
 		else
 			WriteDataChunksForTocEntry(AH, te);
@@ -2365,9 +2370,32 @@ WriteDataChunks(ArchiveHandle *AH, ParallelState *pstate)
 	/*
 	 * If parallel, wait for workers to finish.
 	 */
-	EnsureWorkersFinished(AH, pstate);
+	if (pstate && pstate->numWorkers > 1)
+		WaitForWorkers(AH, pstate, WFW_ALL_IDLE);
 }
 
+
+/*
+ * Callback function that's invoked in the master process after a step has
+ * been parallel dumped.
+ *
+ * We don't need to do anything except check for worker failure.
+ */
+static void
+mark_dump_job_done(ArchiveHandle *AH,
+				   TocEntry *te,
+				   int status,
+				   void *callback_data)
+{
+	ahlog(AH, 1, "finished item %d %s %s\n",
+		  te->dumpId, te->desc, te->tag);
+
+	if (status != 0)
+		exit_horribly(modulename, "worker process failed: exit code %d\n",
+					  status);
+}
+
+
 void
 WriteDataChunksForTocEntry(ArchiveHandle *AH, TocEntry *te)
 {
@@ -2751,9 +2779,9 @@ _tocEntryRequired(TocEntry *te, teSection curSection, RestoreOptions *ropt)
 			return 0;
 	}
 
-	if (ropt->schemaExcludeNames.head != NULL
-		&& te->namespace
-		&& simple_string_list_member(&ropt->schemaExcludeNames, te->namespace))
+	if (ropt->schemaExcludeNames.head != NULL &&
+		te->namespace &&
+		simple_string_list_member(&ropt->schemaExcludeNames, te->namespace))
 		return 0;
 
 	if (ropt->selTypes)
@@ -3769,11 +3797,9 @@ static void
 restore_toc_entries_parallel(ArchiveHandle *AH, ParallelState *pstate,
 							 TocEntry *pending_list)
 {
-	int			work_status;
 	bool		skipped_some;
 	TocEntry	ready_list;
 	TocEntry   *next_work_item;
-	int			ret_child;
 
 	ahlog(AH, 2, "entering restore_toc_entries_parallel\n");
 
@@ -3850,54 +3876,29 @@ restore_toc_entries_parallel(ArchiveHandle *AH, ParallelState *pstate,
 
 			par_list_remove(next_work_item);
 
-			DispatchJobForTocEntry(AH, pstate, next_work_item, ACT_RESTORE);
+			DispatchJobForTocEntry(AH, pstate, next_work_item, ACT_RESTORE,
+								   mark_restore_job_done, &ready_list);
 		}
 		else
 		{
 			/* at least one child is working and we have nothing ready. */
 		}
 
-		for (;;)
-		{
-			int			nTerm = 0;
-
-			/*
-			 * In order to reduce dependencies as soon as possible and
-			 * especially to reap the status of workers who are working on
-			 * items that pending items depend on, we do a non-blocking check
-			 * for ended workers first.
-			 *
-			 * However, if we do not have any other work items currently that
-			 * workers can work on, we do not busy-loop here but instead
-			 * really wait for at least one worker to terminate. Hence we call
-			 * ListenToWorkers(..., ..., do_wait = true) in this case.
-			 */
-			ListenToWorkers(AH, pstate, !next_work_item);
-
-			while ((ret_child = ReapWorkerStatus(pstate, &work_status)) != NO_SLOT)
-			{
-				nTerm++;
-				mark_work_done(AH, &ready_list, ret_child, work_status, pstate);
-			}
-
-			/*
-			 * We need to make sure that we have an idle worker before
-			 * re-running the loop. If nTerm > 0 we already have that (quick
-			 * check).
-			 */
-			if (nTerm > 0)
-				break;
-
-			/* if nobody terminated, explicitly check for an idle worker */
-			if (GetIdleWorker(pstate) != NO_SLOT)
-				break;
-
-			/*
-			 * If we have no idle worker, read the result of one or more
-			 * workers and loop the loop to call ReapWorkerStatus() on them.
-			 */
-			ListenToWorkers(AH, pstate, true);
-		}
+		/*
+		 * Before dispatching another job, check to see if anything has
+		 * finished.  We should check every time through the loop so as to
+		 * reduce dependencies as soon as possible.  If we were unable to
+		 * dispatch any job this time through, wait until some worker finishes
+		 * (and, hopefully, unblocks some pending item).  If we did dispatch
+		 * something, continue as soon as there's at least one idle worker.
+		 * Note that in either case, there's guaranteed to be at least one
+		 * idle worker when we return to the top of the loop.  This ensures we
+		 * won't block inside DispatchJobForTocEntry, which would be
+		 * undesirable: we'd rather postpone dispatching until we see what's
+		 * been unblocked by finished jobs.
+		 */
+		WaitForWorkers(AH, pstate,
+					   next_work_item ? WFW_ONE_IDLE : WFW_GOT_STATUS);
 	}
 
 	ahlog(AH, 1, "finished main parallel loop\n");
@@ -4025,9 +4026,11 @@ get_next_work_item(ArchiveHandle *AH, TocEntry *ready_list,
 		int			count = 0;
 
 		for (k = 0; k < pstate->numWorkers; k++)
-			if (pstate->parallelSlot[k].args->te != NULL &&
-				pstate->parallelSlot[k].args->te->section == SECTION_DATA)
+		{
+			if (pstate->parallelSlot[k].workerStatus == WRKR_WORKING &&
+				pstate->parallelSlot[k].te->section == SECTION_DATA)
 				count++;
+		}
 		if (pstate->numWorkers == 0 || count * 4 < pstate->numWorkers)
 			pref_non_data = false;
 	}
@@ -4044,13 +4047,13 @@ get_next_work_item(ArchiveHandle *AH, TocEntry *ready_list,
 		 * that a currently running item also needs lock on, or vice versa. If
 		 * so, we don't want to schedule them together.
 		 */
-		for (i = 0; i < pstate->numWorkers && !conflicts; i++)
+		for (i = 0; i < pstate->numWorkers; i++)
 		{
 			TocEntry   *running_te;
 
 			if (pstate->parallelSlot[i].workerStatus != WRKR_WORKING)
 				continue;
-			running_te = pstate->parallelSlot[i].args->te;
+			running_te = pstate->parallelSlot[i].te;
 
 			if (has_lock_conflicts(te, running_te) ||
 				has_lock_conflicts(running_te, te))
@@ -4091,10 +4094,8 @@ get_next_work_item(ArchiveHandle *AH, TocEntry *ready_list,
  * our work is finished, the master process will assign us a new work item.
  */
 int
-parallel_restore(ParallelArgs *args)
+parallel_restore(ArchiveHandle *AH, TocEntry *te)
 {
-	ArchiveHandle *AH = args->AH;
-	TocEntry   *te = args->te;
 	int			status;
 
 	Assert(AH->connection != NULL);
@@ -4110,22 +4111,18 @@ parallel_restore(ParallelArgs *args)
 
 
 /*
- * Housekeeping to be done after a step has been parallel restored.
+ * Callback function that's invoked in the master process after a step has
+ * been parallel restored.
  *
- * Clear the appropriate slot, free all the extra memory we allocated,
- * update status, and reduce the dependency count of any dependent items.
+ * Update status and reduce the dependency count of any dependent items.
  */
 static void
-mark_work_done(ArchiveHandle *AH, TocEntry *ready_list,
-			   int worker, int status,
-			   ParallelState *pstate)
+mark_restore_job_done(ArchiveHandle *AH,
+					  TocEntry *te,
+					  int status,
+					  void *callback_data)
 {
-	TocEntry   *te = NULL;
-
-	te = pstate->parallelSlot[worker].args->te;
-
-	if (te == NULL)
-		exit_horribly(modulename, "could not find slot of finished worker\n");
+	TocEntry   *ready_list = (TocEntry *) callback_data;
 
 	ahlog(AH, 1, "finished item %d %s %s\n",
 		  te->dumpId, te->desc, te->tag);

src/bin/pg_dump/pg_backup_archiver.h

@@ -111,7 +111,6 @@ typedef z_stream *z_streamp;
 
 typedef struct _archiveHandle ArchiveHandle;
 typedef struct _tocEntry TocEntry;
-struct ParallelArgs;
 struct ParallelState;
 
 #define READ_ERROR_EXIT(fd) \
@@ -375,7 +374,7 @@ struct _tocEntry
 	int			nLockDeps;		/* number of such dependencies */
 };
 
-extern int	parallel_restore(struct ParallelArgs *args);
+extern int	parallel_restore(ArchiveHandle *AH, TocEntry *te);
 extern void on_exit_close_archive(Archive *AHX);
 
 extern void warn_or_exit_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt,...) pg_attribute_printf(3, 4);

src/bin/pg_dump/pg_backup_custom.c

@@ -820,13 +820,9 @@ _WorkerJobRestoreCustom(ArchiveHandle *AH, TocEntry *te)
 	 */
 	const int	buflen = 64;
 	char	   *buf = (char *) pg_malloc(buflen);
-	ParallelArgs pargs;
 	int			status;
 
-	pargs.AH = AH;
-	pargs.te = te;
-
-	status = parallel_restore(&pargs);
+	status = parallel_restore(AH, te);
 
 	snprintf(buf, buflen, "OK RESTORE %d %d %d", te->dumpId, status,
 			 status == WORKER_IGNORED_ERRORS ? AH->public.n_errors : 0);

src/bin/pg_dump/pg_backup_directory.c

@@ -826,13 +826,9 @@ _WorkerJobRestoreDirectory(ArchiveHandle *AH, TocEntry *te)
 	 */
 	const int	buflen = 64;
 	char	   *buf = (char *) pg_malloc(buflen);
-	ParallelArgs pargs;
 	int			status;
 
-	pargs.AH = AH;
-	pargs.te = te;
-
-	status = parallel_restore(&pargs);
+	status = parallel_restore(AH, te);
 
 	snprintf(buf, buflen, "OK RESTORE %d %d %d", te->dumpId, status,
 			 status == WORKER_IGNORED_ERRORS ? AH->public.n_errors : 0);