From 3ee389a9ffebf05fb691beb16b0e0a5f94d02f9b Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Mon, 8 Jul 2024 21:30:19 -0500
Subject: [PATCH v5 11/13] parallelize incompatible polymorphics check in
 pg_upgrade

---
 src/bin/pg_upgrade/check.c       | 174 ++++++++++++++++---------------
 src/tools/pgindent/typedefs.list |   1 +
 2 files changed, 92 insertions(+), 83 deletions(-)

diff --git a/src/bin/pg_upgrade/check.c b/src/bin/pg_upgrade/check.c
index 664039b441..614464dfd3 100644
--- a/src/bin/pg_upgrade/check.c
+++ b/src/bin/pg_upgrade/check.c
@@ -1381,6 +1381,81 @@ check_for_user_defined_postfix_ops(ClusterInfo *cluster)
 		check_ok();
 }
 
+typedef struct incompat_polymorphics_state
+{
+	FILE	   *script;
+	PQExpBufferData old_polymorphics;
+	char		output_path[MAXPGPATH];
+} incompat_polymorphics_state;
+
+static char *
+incompat_polymorphics_query(void *arg)
+{
+	incompat_polymorphics_state *state = (incompat_polymorphics_state *) arg;
+
+	/*
+	 * The query below hardcodes FirstNormalObjectId as 16384 rather than
+	 * interpolating that C #define into the query because, if that #define is
+	 * ever changed, the cutoff we want to use is the value used by
+	 * pre-version 14 servers, not that of some future version.
+	 */
+	/* Aggregate transition functions */
+	return psprintf("SELECT 'aggregate' AS objkind, p.oid::regprocedure::text AS objname "
+					"FROM pg_proc AS p "
+					"JOIN pg_aggregate AS a ON a.aggfnoid=p.oid "
+					"JOIN pg_proc AS transfn ON transfn.oid=a.aggtransfn "
+					"WHERE p.oid >= 16384 "
+					"AND a.aggtransfn = ANY(ARRAY[%s]::regprocedure[]) "
+					"AND a.aggtranstype = ANY(ARRAY['anyarray', 'anyelement']::regtype[]) "
+
+	/* Aggregate final functions */
+					"UNION ALL "
+					"SELECT 'aggregate' AS objkind, p.oid::regprocedure::text AS objname "
+					"FROM pg_proc AS p "
+					"JOIN pg_aggregate AS a ON a.aggfnoid=p.oid "
+					"JOIN pg_proc AS finalfn ON finalfn.oid=a.aggfinalfn "
+					"WHERE p.oid >= 16384 "
+					"AND a.aggfinalfn = ANY(ARRAY[%s]::regprocedure[]) "
+					"AND a.aggtranstype = ANY(ARRAY['anyarray', 'anyelement']::regtype[]) "
+
+	/* Operators */
+					"UNION ALL "
+					"SELECT 'operator' AS objkind, op.oid::regoperator::text AS objname "
+					"FROM pg_operator AS op "
+					"WHERE op.oid >= 16384 "
+					"AND oprcode = ANY(ARRAY[%s]::regprocedure[]) "
+					"AND oprleft = ANY(ARRAY['anyarray', 'anyelement']::regtype[]);",
+					state->old_polymorphics.data,
+					state->old_polymorphics.data,
+					state->old_polymorphics.data);
+}
+
+static void
+incompat_polymorphics_process(DbInfo *dbinfo, PGresult *res, void *arg)
+{
+	incompat_polymorphics_state *state = (incompat_polymorphics_state *) arg;
+	bool		db_used = false;
+	int			ntups = PQntuples(res);
+	int			i_objkind = PQfnumber(res, "objkind");
+	int			i_objname = PQfnumber(res, "objname");
+
+	for (int rowno = 0; rowno < ntups; rowno++)
+	{
+		if (state->script == NULL &&
+			(state->script = fopen_priv(state->output_path, "w")) == NULL)
+			pg_fatal("could not open file \"%s\": %m", state->output_path);
+		if (!db_used)
+		{
+			fprintf(state->script, "In database: %s\n", dbinfo->db_name);
+			db_used = true;
+		}
+
+		fprintf(state->script, "  %s: %s\n",
+				PQgetvalue(res, rowno, i_objkind),
+				PQgetvalue(res, rowno, i_objname));
+	}
+}
+
 /*
  *	check_for_incompatible_polymorphics()
  *
@@ -1390,111 +1465,44 @@ check_for_user_defined_postfix_ops(ClusterInfo *cluster)
 static void
 check_for_incompatible_polymorphics(ClusterInfo *cluster)
 {
-	PGresult   *res;
-	FILE	   *script = NULL;
-	char		output_path[MAXPGPATH];
-	PQExpBufferData old_polymorphics;
+	AsyncTask  *task = async_task_create();
+	incompat_polymorphics_state state;
 
 	prep_status("Checking for incompatible polymorphic functions");
 
-	snprintf(output_path, sizeof(output_path), "%s/%s",
+	state.script = NULL;
+	snprintf(state.output_path, sizeof(state.output_path), "%s/%s",
 			 log_opts.basedir,
 			 "incompatible_polymorphics.txt");
 
 	/* The set of problematic functions varies a bit in different versions */
-	initPQExpBuffer(&old_polymorphics);
+	initPQExpBuffer(&state.old_polymorphics);
 
-	appendPQExpBufferStr(&old_polymorphics,
+	appendPQExpBufferStr(&state.old_polymorphics,
 						 "'array_append(anyarray,anyelement)'"
 						 ", 'array_cat(anyarray,anyarray)'"
 						 ", 'array_prepend(anyelement,anyarray)'");
 
 	if (GET_MAJOR_VERSION(cluster->major_version) >= 903)
-		appendPQExpBufferStr(&old_polymorphics,
+		appendPQExpBufferStr(&state.old_polymorphics,
 							 ", 'array_remove(anyarray,anyelement)'"
 							 ", 'array_replace(anyarray,anyelement,anyelement)'");
 
 	if (GET_MAJOR_VERSION(cluster->major_version) >= 905)
-		appendPQExpBufferStr(&old_polymorphics,
+		appendPQExpBufferStr(&state.old_polymorphics,
 							 ", 'array_position(anyarray,anyelement)'"
 							 ", 'array_position(anyarray,anyelement,integer)'"
 							 ", 'array_positions(anyarray,anyelement)'"
 							 ", 'width_bucket(anyelement,anyarray)'");
 
-	for (int dbnum = 0; dbnum < cluster->dbarr.ndbs; dbnum++)
-	{
-		bool		db_used = false;
-		DbInfo	   *active_db = &cluster->dbarr.dbs[dbnum];
-		PGconn	   *conn = connectToServer(cluster, active_db->db_name);
-		int			ntups;
-		int			i_objkind,
-					i_objname;
-
-		/*
-		 * The query below hardcodes FirstNormalObjectId as 16384 rather than
-		 * interpolating that C #define into the query because, if that
-		 * #define is ever changed, the cutoff we want to use is the value
-		 * used by pre-version 14 servers, not that of some future version.
-		 */
-		res = executeQueryOrDie(conn,
-		/* Aggregate transition functions */
-								"SELECT 'aggregate' AS objkind, p.oid::regprocedure::text AS objname "
-								"FROM pg_proc AS p "
-								"JOIN pg_aggregate AS a ON a.aggfnoid=p.oid "
-								"JOIN pg_proc AS transfn ON transfn.oid=a.aggtransfn "
-								"WHERE p.oid >= 16384 "
-								"AND a.aggtransfn = ANY(ARRAY[%s]::regprocedure[]) "
-								"AND a.aggtranstype = ANY(ARRAY['anyarray', 'anyelement']::regtype[]) "
-
-		/* Aggregate final functions */
-								"UNION ALL "
-								"SELECT 'aggregate' AS objkind, p.oid::regprocedure::text AS objname "
-								"FROM pg_proc AS p "
-								"JOIN pg_aggregate AS a ON a.aggfnoid=p.oid "
-								"JOIN pg_proc AS finalfn ON finalfn.oid=a.aggfinalfn "
-								"WHERE p.oid >= 16384 "
-								"AND a.aggfinalfn = ANY(ARRAY[%s]::regprocedure[]) "
-								"AND a.aggtranstype = ANY(ARRAY['anyarray', 'anyelement']::regtype[]) "
-
-		/* Operators */
-								"UNION ALL "
-								"SELECT 'operator' AS objkind, op.oid::regoperator::text AS objname "
-								"FROM pg_operator AS op "
-								"WHERE op.oid >= 16384 "
-								"AND oprcode = ANY(ARRAY[%s]::regprocedure[]) "
-								"AND oprleft = ANY(ARRAY['anyarray', 'anyelement']::regtype[]);",
-								old_polymorphics.data,
-								old_polymorphics.data,
-								old_polymorphics.data);
-
-		ntups = PQntuples(res);
-
-		i_objkind = PQfnumber(res, "objkind");
-		i_objname = PQfnumber(res, "objname");
-
-		for (int rowno = 0; rowno < ntups; rowno++)
-		{
-			if (script == NULL &&
-				(script = fopen_priv(output_path, "w")) == NULL)
-				pg_fatal("could not open file \"%s\": %m", output_path);
-			if (!db_used)
-			{
-				fprintf(script, "In database: %s\n", active_db->db_name);
-				db_used = true;
-			}
-
-			fprintf(script, "  %s: %s\n",
-					PQgetvalue(res, rowno, i_objkind),
-					PQgetvalue(res, rowno, i_objname));
-		}
-
-		PQclear(res);
-		PQfinish(conn);
-	}
+	async_task_add_step(task, incompat_polymorphics_query,
+						incompat_polymorphics_process, true, &state);
+	async_task_run(task, cluster);
+	async_task_free(task);
 
-	if (script)
+	if (state.script)
 	{
-		fclose(script);
+		fclose(state.script);
 		pg_log(PG_REPORT, "fatal");
 		pg_fatal("Your installation contains user-defined objects that refer to internal\n"
 				 "polymorphic functions with arguments of type \"anyarray\" or \"anyelement\".\n"
@@ -1502,12 +1510,12 @@ check_for_incompatible_polymorphics(ClusterInfo *cluster)
 				 "afterwards, changing them to refer to the new corresponding functions with\n"
 				 "arguments of type \"anycompatiblearray\" and \"anycompatible\".\n"
 				 "A list of the problematic objects is in the file:\n"
-				 "    %s", output_path);
+				 "    %s", state.output_path);
 	}
 	else
 		check_ok();
 
-	termPQExpBuffer(&old_polymorphics);
+	termPQExpBuffer(&state.old_polymorphics);
 }
 
 /*
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index a2330f4346..66db3d3e2c 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3536,6 +3536,7 @@ hstoreUpgrade_t
 hyperLogLogState
 ifState
 import_error_callback_arg
+incompat_polymorphics_state
 indexed_tlist
 inet
 inetKEY
-- 
2.39.3 (Apple Git-146)

