v37-0002-Parameterizing-parallel-slot-result-handling.patch

From de650441202f54710d5e84692e096200f8d053bb Mon Sep 17 00:00:00 2001
From: Mark Dilger <mark.dilger@enterprisedb.com>
Date: Tue, 2 Feb 2021 12:35:56 -0800
Subject: [PATCH v37 2/4] Parameterizing parallel slot result handling

The function consumeQueryResult was being used to handle all results
returned by queries executed through the parallel slot interface,
but this hardcodes knowledge about the expectations of reindexdb and
vacuumdb such as the expected result status being PGRES_COMMAND_OK
(as opposed to, say, PGRES_TUPLES_OK).

Reworking the slot interface to optionally include a
ParallelSlotResultHandler and context variable per slot.  The idea
is that a caller who executes a command or query through the slot
can set the handler to be called when the query completes with
necessary context information stored for by the handler when
processing the result.

The old logic of consumeQueryResults is moved into a new callback
function, TableCommandResultHandler(), which gets registered as the
slot handler explicitly from vacuumdb and reindexdb.  This is
defined in fe_utils/parallel_slot.c rather than somewhere in
src/bin/scripts where its only callers reside, partly to keep it
close to the rest of the shared parallel slot handling code and
partly in anticipation that other utility programs will eventually
want to use it also.

The expectation of this commit is that pg_amcheck will have handlers
for table and index checks which will process the PGresults of calls
to the amcheck functions.  This commit sets up the infrastructure
necessary to support those handlers being different from the one
used by vacuumdb and reindexdb.
---
 src/bin/scripts/reindexdb.c          |  1 +
 src/bin/scripts/vacuumdb.c           |  1 +
 src/fe_utils/parallel_slot.c         | 92 +++++++++++++++++++---------
 src/include/fe_utils/parallel_slot.h | 29 +++++++++
 4 files changed, 95 insertions(+), 28 deletions(-)

diff --git a/src/bin/scripts/reindexdb.c b/src/bin/scripts/reindexdb.c
index 7781fb1151..9f072ac49a 100644
--- a/src/bin/scripts/reindexdb.c
+++ b/src/bin/scripts/reindexdb.c
@@ -466,6 +466,7 @@ reindex_one_database(const ConnParams *cparams, ReindexType type,
 			goto finish;
 		}
 
+		ParallelSlotSetHandler(free_slot, TableCommandResultHandler, NULL);
 		run_reindex_command(free_slot->connection, process_type, objname,
 							echo, verbose, concurrently, true);
 
diff --git a/src/bin/scripts/vacuumdb.c b/src/bin/scripts/vacuumdb.c
index ed320817bc..9dc8aca29f 100644
--- a/src/bin/scripts/vacuumdb.c
+++ b/src/bin/scripts/vacuumdb.c
@@ -713,6 +713,7 @@ vacuum_one_database(const ConnParams *cparams,
 		 * Execute the vacuum.  All errors are handled in processQueryResult
 		 * through ParallelSlotsGetIdle.
 		 */
+		ParallelSlotSetHandler(free_slot, TableCommandResultHandler, NULL);
 		run_vacuum_command(free_slot->connection, sql.data,
 						   echo, tabname);
 
diff --git a/src/fe_utils/parallel_slot.c b/src/fe_utils/parallel_slot.c
index 3987a4702b..b75dc26a49 100644
--- a/src/fe_utils/parallel_slot.c
+++ b/src/fe_utils/parallel_slot.c
@@ -30,7 +30,7 @@
 
 static void init_slot(ParallelSlot *slot, PGconn *conn);
 static int	select_loop(int maxFd, fd_set *workerset);
-static bool processQueryResult(PGconn *conn, PGresult *result);
+static bool processQueryResult(ParallelSlot *slot, PGresult *result);
 
 static void
 init_slot(ParallelSlot *slot, PGconn *conn)
@@ -38,53 +38,45 @@ init_slot(ParallelSlot *slot, PGconn *conn)
 	slot->connection = conn;
 	/* Initially assume connection is idle */
 	slot->isFree = true;
+	ParallelSlotClearHandler(slot);
 }
 
 /*
- * Process (and delete) a query result.  Returns true if there's no error,
- * false otherwise -- but errors about trying to work on a missing relation
- * are reported and subsequently ignored.
+ * Invoke the slot's handler for a single query result, or fall back to the
+ * default handler if none is defined for the slot.  Returns true if the
+ * handler reports that there's no error, false otherwise.
  */
 static bool
-processQueryResult(PGconn *conn, PGresult *result)
+processQueryResult(ParallelSlot *slot, PGresult *result)
 {
-	/*
-	 * If it's an error, report it.  Errors about a missing table are harmless
-	 * so we continue processing; but die for other errors.
-	 */
-	if (PQresultStatus(result) != PGRES_COMMAND_OK)
-	{
-		char	   *sqlState = PQresultErrorField(result, PG_DIAG_SQLSTATE);
+	ParallelSlotResultHandler handler = slot->handler;
 
-		pg_log_error("processing of database \"%s\" failed: %s",
-					 PQdb(conn), PQerrorMessage(conn));
+	Assert(PointerIsValid(handler));
 
-		if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) != 0)
-		{
-			PQclear(result);
-			return false;
-		}
-	}
+	/* On failure, the handler should return NULL after freeing the result */
+	if (!handler(result, slot->connection, slot->handler_context))
+		return false;
 
+	/* Ok, we have to free it ourself */
 	PQclear(result);
 	return true;
 }
 
 /*
- * Consume all the results generated for the given connection until
+ * Handle all the results generated for the given connection until
  * nothing remains.  If at least one error is encountered, return false.
  * Note that this will block if the connection is busy.
  */
 static bool
-consumeQueryResult(PGconn *conn)
+consumeQueryResult(ParallelSlot *slot)
 {
 	bool		ok = true;
 	PGresult   *result;
 
-	SetCancelConn(conn);
-	while ((result = PQgetResult(conn)) != NULL)
+	SetCancelConn(slot->connection);
+	while ((result = PQgetResult(slot->connection)) != NULL)
 	{
-		if (!processQueryResult(conn, result))
+		if (!processQueryResult(slot, result))
 			ok = false;
 	}
 	ResetCancelConn();
@@ -227,14 +219,15 @@ ParallelSlotsGetIdle(ParallelSlot *slots, int numslots)
 
 				if (result != NULL)
 				{
-					/* Check and discard the command result */
-					if (!processQueryResult(slots[i].connection, result))
+					/* Handle and discard the command result */
+					if (!processQueryResult(slots + i, result))
 						return NULL;
 				}
 				else
 				{
 					/* This connection has become idle */
 					slots[i].isFree = true;
+					ParallelSlotClearHandler(slots + i);
 					if (firstFree < 0)
 						firstFree = i;
 					break;
@@ -329,9 +322,52 @@ ParallelSlotsWaitCompletion(ParallelSlot *slots, int numslots)
 
 	for (i = 0; i < numslots; i++)
 	{
-		if (!consumeQueryResult((slots + i)->connection))
+		if (!consumeQueryResult(slots + i))
 			return false;
 	}
 
 	return true;
 }
+
+/*
+ * TableCommandResultHandler
+ *
+ * ParallelSlotResultHandler for results of commands (not queries) against tables.
+ *
+ * Requires that the result status is either PGRES_COMMAND_OK or an error about
+ * a missing table.  This is useful for utilities that compile a list of tables
+ * to process and then run commands (vacuum, reindex, or whatever) against
+ * those tables, as there is a race condition between the time the list is
+ * compiled and the time the command attempts to open the table.
+ *
+ * For missing tables, logs an error but allows processing to continue.
+ *
+ * For all other errors, logs an error and terminates further processing.
+ *
+ * res: PGresult from the query executed on the slot's connection
+ * conn: connection belonging to the slot
+ * context: unused
+ */
+PGresult *
+TableCommandResultHandler(PGresult *res, PGconn *conn, void *context)
+{
+	/*
+	 * If it's an error, report it.  Errors about a missing table are harmless
+	 * so we continue processing; but die for other errors.
+	 */
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
+	{
+		char	   *sqlState = PQresultErrorField(res, PG_DIAG_SQLSTATE);
+
+		pg_log_error("processing of database \"%s\" failed: %s",
+					 PQdb(conn), PQerrorMessage(conn));
+
+		if (sqlState && strcmp(sqlState, ERRCODE_UNDEFINED_TABLE) != 0)
+		{
+			PQclear(res);
+			return NULL;
+		}
+	}
+
+	return res;
+}
diff --git a/src/include/fe_utils/parallel_slot.h b/src/include/fe_utils/parallel_slot.h
index 99eeb3328d..6fe58d2a26 100644
--- a/src/include/fe_utils/parallel_slot.h
+++ b/src/include/fe_utils/parallel_slot.h
@@ -15,12 +15,39 @@
 #include "fe_utils/connect_utils.h"
 #include "libpq-fe.h"
 
+typedef PGresult *(*ParallelSlotResultHandler) (PGresult *res, PGconn *conn,
+												void *context);
+
 typedef struct ParallelSlot
 {
 	PGconn	   *connection;		/* One connection */
 	bool		isFree;			/* Is it known to be idle? */
+
+	/*
+	 * Prior to issuing a command or query on 'connection', a handler callback
+	 * function may optionally be registered to be invoked to process the
+	 * results, and context information may optionally be registered for use
+	 * by the handler.  If unset, these fields should be NULL.
+	 */
+	ParallelSlotResultHandler handler;
+	void	   *handler_context;
 } ParallelSlot;
 
+static inline void
+ParallelSlotSetHandler(ParallelSlot *slot, ParallelSlotResultHandler handler,
+					   void *context)
+{
+	slot->handler = handler;
+	slot->handler_context = context;
+}
+
+static inline void
+ParallelSlotClearHandler(ParallelSlot *slot)
+{
+	slot->handler = NULL;
+	slot->handler_context = NULL;
+}
+
 extern ParallelSlot *ParallelSlotsGetIdle(ParallelSlot *slots, int numslots);
 
 extern ParallelSlot *ParallelSlotsSetup(const ConnParams *cparams,
@@ -31,5 +58,7 @@ extern void ParallelSlotsTerminate(ParallelSlot *slots, int numslots);
 
 extern bool ParallelSlotsWaitCompletion(ParallelSlot *slots, int numslots);
 
+extern PGresult *TableCommandResultHandler(PGresult *res, PGconn *conn,
+										   void *context);
 
 #endif							/* PARALLEL_SLOT_H */
-- 
2.21.1 (Apple Git-122.3)