From d7683a095d4d2c1574005eb41504a5be256d6480 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Fri, 28 Jun 2024 11:02:44 -0500
Subject: [PATCH v2 1/6] introduce framework for parallelizing pg_upgrade tasks

---
 src/bin/pg_upgrade/Makefile      |   1 +
 src/bin/pg_upgrade/async.c       | 323 +++++++++++++++++++++++++++++++
 src/bin/pg_upgrade/meson.build   |   1 +
 src/bin/pg_upgrade/pg_upgrade.h  |  16 ++
 src/tools/pgindent/typedefs.list |   4 +
 5 files changed, 345 insertions(+)
 create mode 100644 src/bin/pg_upgrade/async.c

diff --git a/src/bin/pg_upgrade/Makefile b/src/bin/pg_upgrade/Makefile
index bde91e2beb..3bc4f5d740 100644
--- a/src/bin/pg_upgrade/Makefile
+++ b/src/bin/pg_upgrade/Makefile
@@ -12,6 +12,7 @@ include $(top_builddir)/src/Makefile.global
 
 OBJS = \
 	$(WIN32RES) \
+	async.o \
 	check.o \
 	controldata.o \
 	dump.o \
diff --git a/src/bin/pg_upgrade/async.c b/src/bin/pg_upgrade/async.c
new file mode 100644
index 0000000000..7df1e7712d
--- /dev/null
+++ b/src/bin/pg_upgrade/async.c
@@ -0,0 +1,323 @@
+/*
+ * async.c
+ *
+ * parallelization via libpq's async APIs
+ *
+ * Copyright (c) 2024, PostgreSQL Global Development Group
+ * src/bin/pg_upgrade/async.c
+ */
+
+#include "postgres_fe.h"
+
+#include "common/connect.h"
+#include "fe_utils/string_utils.h"
+#include "pg_upgrade.h"
+
+static int	dbs_complete;
+static int	dbs_processing;
+
+typedef struct AsyncTaskCallbacks
+{
+	AsyncTaskGetQueryCB query_cb;
+	AsyncTaskProcessCB process_cb;
+	bool		free_result;
+	void	   *arg;
+} AsyncTaskCallbacks;
+
+typedef struct AsyncTask
+{
+	AsyncTaskCallbacks *cbs;
+	int			num_cb_sets;
+} AsyncTask;
+
+typedef enum
+{
+	FREE,
+	CONNECTING,
+	SETTING_SEARCH_PATH,
+	RUNNING_QUERY,
+} AsyncSlotState;
+
+typedef struct
+{
+	AsyncSlotState state;
+	int			db;
+	int			query;
+	PGconn	   *conn;
+} AsyncSlot;
+
+AsyncTask *
+async_task_create(void)
+{
+	return pg_malloc0(sizeof(AsyncTask));
+}
+
+void
+async_task_free(AsyncTask *task)
+{
+	if (task->cbs)
+		pg_free(task->cbs);
+
+	pg_free(task);
+}
+
+void
+async_task_add_step(AsyncTask *task,
+					AsyncTaskGetQueryCB query_cb,
+					AsyncTaskProcessCB process_cb, bool free_result,
+					void *arg)
+{
+	AsyncTaskCallbacks *new_cbs;
+
+	task->cbs = pg_realloc(task->cbs,
+						   ++task->num_cb_sets * sizeof(AsyncTaskCallbacks));
+
+	new_cbs = &task->cbs[task->num_cb_sets - 1];
+	new_cbs->query_cb = query_cb;
+	new_cbs->process_cb = process_cb;
+	new_cbs->free_result = free_result;
+	new_cbs->arg = arg;
+}
+
+static void
+conn_failure(PGconn *conn)
+{
+	pg_log(PG_REPORT, "%s", PQerrorMessage(conn));
+	printf(_("Failure, exiting\n"));
+	exit(1);
+}
+
+static void
+start_conn(const ClusterInfo *cluster, AsyncSlot *slot)
+{
+	PQExpBufferData conn_opts;
+
+	/* Build connection string with proper quoting */
+	initPQExpBuffer(&conn_opts);
+	appendPQExpBufferStr(&conn_opts, "dbname=");
+	appendConnStrVal(&conn_opts, cluster->dbarr.dbs[slot->db].db_name);
+	appendPQExpBufferStr(&conn_opts, " user=");
+	appendConnStrVal(&conn_opts, os_info.user);
+	appendPQExpBuffer(&conn_opts, " port=%d", cluster->port);
+	if (cluster->sockdir)
+	{
+		appendPQExpBufferStr(&conn_opts, " host=");
+		appendConnStrVal(&conn_opts, cluster->sockdir);
+	}
+
+	slot->conn = PQconnectStart(conn_opts.data);
+	termPQExpBuffer(&conn_opts);
+
+	if (!slot->conn)
+		conn_failure(slot->conn);
+}
+
+static void
+dispatch_query(const ClusterInfo *cluster, AsyncSlot *slot,
+			   const AsyncTask *task)
+{
+	AsyncTaskCallbacks *cbs = &task->cbs[slot->query];
+	AsyncTaskGetQueryCB get_query = cbs->query_cb;
+	DbInfo	   *dbinfo = &cluster->dbarr.dbs[slot->db];
+	char	   *query = (*get_query) (dbinfo, cbs->arg);
+
+	if (!PQsendQuery(slot->conn, query))
+		conn_failure(slot->conn);
+
+	pg_free(query);
+}
+
+static PGresult *
+get_last_result(PGconn *conn)
+{
+	PGresult   *tmp;
+	PGresult   *res = NULL;
+
+	while ((tmp = PQgetResult(conn)) != NULL)
+	{
+		PQclear(res);
+		res = tmp;
+		if (PQstatus(conn) == CONNECTION_BAD)
+			conn_failure(conn);
+	}
+
+	if (PQresultStatus(res) != PGRES_COMMAND_OK &&
+		PQresultStatus(res) != PGRES_TUPLES_OK)
+		conn_failure(conn);
+
+	return res;
+}
+
+static void
+process_query_result(const ClusterInfo *cluster, AsyncSlot *slot,
+					 const AsyncTask *task)
+{
+	AsyncTaskCallbacks *cbs = &task->cbs[slot->query];
+	AsyncTaskProcessCB process_cb = cbs->process_cb;
+	DbInfo	   *dbinfo = &cluster->dbarr.dbs[slot->db];
+	PGresult   *res = get_last_result(slot->conn);
+
+	(*process_cb) (dbinfo, res, cbs->arg);
+
+	if (cbs->free_result)
+		PQclear(res);
+}
+
+static void
+process_slot(const ClusterInfo *cluster, AsyncSlot *slot, const AsyncTask *task)
+{
+	switch (slot->state)
+	{
+		case FREE:
+			if (dbs_processing >= cluster->dbarr.ndbs)
+				return;
+			slot->db = dbs_processing++;
+			slot->state = CONNECTING;
+			start_conn(cluster, slot);
+			return;
+
+		case CONNECTING:
+			if (PQconnectPoll(slot->conn) == PGRES_POLLING_FAILED)
+				conn_failure(slot->conn);
+			if (PQconnectPoll(slot->conn) != PGRES_POLLING_OK)
+				return;
+			slot->state = SETTING_SEARCH_PATH;
+			if (!PQsendQuery(slot->conn, ALWAYS_SECURE_SEARCH_PATH_SQL))
+				conn_failure(slot->conn);
+			return;
+
+		case SETTING_SEARCH_PATH:
+			if (!PQconsumeInput(slot->conn))
+				conn_failure(slot->conn);
+			if (PQisBusy(slot->conn))
+				return;
+			PQclear(get_last_result(slot->conn));
+			slot->state = RUNNING_QUERY;
+			dispatch_query(cluster, slot, task);
+			return;
+
+		case RUNNING_QUERY:
+			if (!PQconsumeInput(slot->conn))
+				conn_failure(slot->conn);
+			if (PQisBusy(slot->conn))
+				return;
+			process_query_result(cluster, slot, task);
+			if (++slot->query >= task->num_cb_sets)
+			{
+				dbs_complete++;
+				PQfinish(slot->conn);
+				memset(slot, 0, sizeof(AsyncSlot));
+				return;
+			}
+			dispatch_query(cluster, slot, task);
+			return;
+	}
+}
+
+/*
+ * Wait on the slots to either finish connecting or to receive query results if
+ * possible.  This avoids a tight loop in async_task_run().
+ */
+static void
+wait_on_slots(AsyncSlot *slots, int numslots)
+{
+	fd_set		input_mask;
+	fd_set		output_mask;
+	fd_set		except_mask;
+	int			maxFd = 0;
+
+	FD_ZERO(&input_mask);
+	FD_ZERO(&output_mask);
+	FD_ZERO(&except_mask);
+
+	for (int i = 0; i < numslots; i++)
+	{
+		int			sock;
+		bool		read = false;
+
+		switch (slots[i].state)
+		{
+			case FREE:
+
+				/*
+				 * If we see a free slot, return right away so that it can be
+				 * reused immediately for the next database.  This might cause
+				 * us to spin more than necessary as we finish processing the
+				 * last few databases, but that shouldn't cause too much harm.
+				 */
+				return;
+
+			case CONNECTING:
+
+				/*
+				 * If we are waiting for the connection to establish, choose
+				 * whether to wait for reading or for writing on the socket as
+				 * appropriate.  If neither apply, just return immediately so
+				 * that we can handle the slot.
+				 */
+				{
+					PostgresPollingStatusType status;
+
+					status = PQconnectPoll(slots[i].conn);
+					if (status == PGRES_POLLING_READING)
+						read = true;
+					else if (status != PGRES_POLLING_WRITING)
+						return;
+				}
+				break;
+
+			case SETTING_SEARCH_PATH:
+			case RUNNING_QUERY:
+
+				/*
+				 * If we've sent a query, we must wait for the socket to be
+				 * read-ready.  Note that process_slot() handles calling
+				 * PQconsumeInput() as required.
+				 */
+				read = true;
+				break;
+		}
+
+		/*
+		 * If there's some problem retrieving the socket, just pretend this
+		 * slot doesn't exist.  We don't expect this to happen regularly in
+		 * practice, so it seems unlikely to cause too much harm.
+		 */
+		sock = PQsocket(slots[i].conn);
+		if (sock < 0)
+			continue;
+
+		/*
+		 * Add the socket to the set.
+		 */
+		FD_SET(sock, read ? &input_mask : &output_mask);
+		FD_SET(sock, &except_mask);
+		maxFd = Max(maxFd, sock);
+	}
+
+	/*
+	 * If we found socket(s) to wait on, wait.
+	 */
+	if (maxFd != 0)
+		(void) select(maxFd + 1, &input_mask, &output_mask, &except_mask, NULL);
+}
+
+void
+async_task_run(const AsyncTask *task, const ClusterInfo *cluster)
+{
+	int			jobs = Max(1, user_opts.jobs);
+	AsyncSlot  *slots = pg_malloc0(sizeof(AsyncSlot) * jobs);
+
+	dbs_complete = 0;
+	dbs_processing = 0;
+
+	while (dbs_complete < cluster->dbarr.ndbs)
+	{
+		for (int i = 0; i < jobs; i++)
+			process_slot(cluster, &slots[i], task);
+
+		wait_on_slots(slots, jobs);
+	}
+
+	pg_free(slots);
+}
diff --git a/src/bin/pg_upgrade/meson.build b/src/bin/pg_upgrade/meson.build
index 9825fa3305..9eb48e176c 100644
--- a/src/bin/pg_upgrade/meson.build
+++ b/src/bin/pg_upgrade/meson.build
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2024, PostgreSQL Global Development Group
 
 pg_upgrade_sources = files(
+  'async.c',
   'check.c',
   'controldata.c',
   'dump.c',
diff --git a/src/bin/pg_upgrade/pg_upgrade.h b/src/bin/pg_upgrade/pg_upgrade.h
index 8afe240bdf..1ebad3bd74 100644
--- a/src/bin/pg_upgrade/pg_upgrade.h
+++ b/src/bin/pg_upgrade/pg_upgrade.h
@@ -494,3 +494,19 @@ void		parallel_transfer_all_new_dbs(DbInfoArr *old_db_arr, DbInfoArr *new_db_arr
 										  char *old_pgdata, char *new_pgdata,
 										  char *old_tablespace);
 bool		reap_child(bool wait_for_child);
+
+/* async.c */
+
+typedef char *(*AsyncTaskGetQueryCB) (DbInfo *dbinfo, void *arg);
+typedef void (*AsyncTaskProcessCB) (DbInfo *dbinfo, PGresult *res, void *arg);
+
+/* struct definition is private to async.c */
+typedef struct AsyncTask AsyncTask;
+
+AsyncTask  *async_task_create(void);
+void		async_task_add_step(AsyncTask *task,
+								AsyncTaskGetQueryCB query_cb,
+								AsyncTaskProcessCB process_cb, bool free_result,
+								void *arg);
+void		async_task_run(const AsyncTask *task, const ClusterInfo *cluster);
+void		async_task_free(AsyncTask *task);
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index e6c1caf649..3d219cbfe2 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -153,6 +153,10 @@ ArrayMetaState
 ArraySubWorkspace
 ArrayToken
 ArrayType
+AsyncSlot
+AsyncSlotState
+AsyncTask
+AsyncTaskCallbacks
 AsyncQueueControl
 AsyncQueueEntry
 AsyncRequest
-- 
2.39.3 (Apple Git-146)

