commit 02513d121ce0b96be5619edfd0317b46b70a44da
Author: Sachin Kumar <ssetiya@amazon.com>
Date:   Tue Jan 2 16:57:42 2024 +0000

    Pass pg_upgrade --jobs parameter to pg_restore
    
    This patch changes pg_upgrade --jobs behaviour for data, post_data
    part of restore. Instead for restoring N databases in parallel,
    they are restored in sequential order but with pg_restore
    --jobs = {original pg_upgrade --jobs value}

diff --git a/src/bin/pg_upgrade/pg_upgrade.c b/src/bin/pg_upgrade/pg_upgrade.c
index 5cfd2282e1..73ffdafd19 100644
--- a/src/bin/pg_upgrade/pg_upgrade.c
+++ b/src/bin/pg_upgrade/pg_upgrade.c
@@ -54,6 +54,7 @@
 static void set_locale_and_encoding(void);
 static void prepare_new_cluster(void);
 static void prepare_new_globals(void);
+static void parallel_pg_restore_dbs(bool);
 static void create_new_objects(void);
 static void copy_xact_xlog_xid(void);
 static void set_frozenxids(bool minmxid_only);
@@ -508,6 +509,96 @@ prepare_new_globals(void)
 	check_ok();
 }
 
+/*
+ * This function will use parallel pg_restore to restore
+ * {data}, {post_data} section of the dump.
+ * Since {pre_data} section cant be parallelized , instead we will
+ * restore {pg_upgrade --jobs} dbs in parallel.
+ */
+static void
+parallel_pg_restore_dbs(bool pre_data)
+{
+	int			dbnum;
+	int         jobs = user_opts.jobs ? user_opts.jobs : 1;
+	/*
+	 * Restore @section of the dump with parallel pg_restore
+	 */
+	for (dbnum = 0; dbnum < old_cluster.dbarr.ndbs; dbnum++)
+	{
+		char		sql_file_name[MAXPGPATH],
+					log_file_name[MAXPGPATH];
+		DbInfo	   *old_db = &old_cluster.dbarr.dbs[dbnum];
+		const char *create_opts;
+
+		/* Skip template1 in this pass */
+		if (strcmp(old_db->db_name, "template1") == 0)
+			continue;
+
+		pg_log(PG_STATUS, "%s", old_db->db_name);
+		snprintf(sql_file_name, sizeof(sql_file_name), DB_DUMP_FILE_MASK, old_db->db_oid);
+		snprintf(log_file_name, sizeof(log_file_name), DB_DUMP_LOG_FILE_MASK, old_db->db_oid);
+
+		/*
+		 * postgres database will already exist in the target installation, so
+		 * tell pg_restore to drop and recreate it; otherwise we would fail to
+		 * propagate its database-level properties.
+		 */
+		if (pre_data)
+		{
+			if (strcmp(old_db->db_name, "postgres") == 0)
+				create_opts = "--clean --create";
+			else if (pre_data)
+				create_opts = "--create";
+		}
+		else
+			create_opts = "";
+
+		/*
+		 * Restore pre-data section of the dump in parallel with single pg_restore job
+		 * This section of dump cant be parallelized with parallel pg_restore
+		 */
+		if (pre_data)
+		{
+			parallel_exec_prog(log_file_name,
+							   NULL,
+							   "\"%s/pg_restore\" %s %s --exit-on-error --verbose "
+							   " --section=pre-data "
+							   "--dbname template1 \"%s/%s\"",
+							   new_cluster.bindir,
+							   cluster_conn_opts(&new_cluster),
+							   create_opts,
+							   log_opts.dumpdir,
+							   sql_file_name);
+
+		}
+		else
+		{
+			exec_prog(log_file_name,
+							   NULL,
+							   true,
+							   true,
+							   "\"%s/pg_restore\" %s %s --exit-on-error --verbose "
+							   "--transaction-size=1000 --jobs %d --section=data --section=post-data "
+							   "--dbname template1 \"%s/%s\"",
+							   new_cluster.bindir,
+							   cluster_conn_opts(&new_cluster),
+							   create_opts,
+							   jobs,
+							   log_opts.dumpdir,
+							   sql_file_name);
+		}
+	}
+	/*
+	 * Wait for child process for pre-data section
+	 */
+	if (pre_data)
+	{
+		/* reap all children */
+		while (reap_child(true) == true)
+			;
+	}
+
+}
 
 static void
 create_new_objects(void)
@@ -559,46 +650,10 @@ create_new_objects(void)
 		break;					/* done once we've processed template1 */
 	}
 
-	for (dbnum = 0; dbnum < old_cluster.dbarr.ndbs; dbnum++)
-	{
-		char		sql_file_name[MAXPGPATH],
-					log_file_name[MAXPGPATH];
-		DbInfo	   *old_db = &old_cluster.dbarr.dbs[dbnum];
-		const char *create_opts;
-
-		/* Skip template1 in this pass */
-		if (strcmp(old_db->db_name, "template1") == 0)
-			continue;
-
-		pg_log(PG_STATUS, "%s", old_db->db_name);
-		snprintf(sql_file_name, sizeof(sql_file_name), DB_DUMP_FILE_MASK, old_db->db_oid);
-		snprintf(log_file_name, sizeof(log_file_name), DB_DUMP_LOG_FILE_MASK, old_db->db_oid);
-
-		/*
-		 * postgres database will already exist in the target installation, so
-		 * tell pg_restore to drop and recreate it; otherwise we would fail to
-		 * propagate its database-level properties.
-		 */
-		if (strcmp(old_db->db_name, "postgres") == 0)
-			create_opts = "--clean --create";
-		else
-			create_opts = "--create";
-
-		parallel_exec_prog(log_file_name,
-						   NULL,
-						   "\"%s/pg_restore\" %s %s --exit-on-error --verbose "
-						   "--transaction-size=1000 "
-						   "--dbname template1 \"%s/%s\"",
-						   new_cluster.bindir,
-						   cluster_conn_opts(&new_cluster),
-						   create_opts,
-						   log_opts.dumpdir,
-						   sql_file_name);
-	}
-
-	/* reap all children */
-	while (reap_child(true) == true)
-		;
+	/* Restore pre_data */
+	parallel_pg_restore_dbs(true);
+	/* Restore data, post_data */
+	parallel_pg_restore_dbs(false);
 
 	end_progress_output();
 	check_ok();