v2-0003-reuse-connection-when-tablesync-workers-change-th.patch

From 7e9731104d17056509df919b6b6122e6a3f1fce8 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Tue, 27 Jun 2023 07:10:45 +0000
Subject: [PATCH v2 3/6] reuse connection when tablesync workers change the
 target

---
 src/backend/replication/logical/tablesync.c | 33 ++++++++++++++++-----
 src/backend/replication/logical/worker.c    | 33 +++++++++++++--------
 src/backend/replication/walsender.c         |  6 ++++
 src/include/replication/worker_internal.h   |  2 +-
 4 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/src/backend/replication/logical/tablesync.c b/src/backend/replication/logical/tablesync.c
index 37f073b968..8d44ed207f 100644
--- a/src/backend/replication/logical/tablesync.c
+++ b/src/backend/replication/logical/tablesync.c
@@ -1268,7 +1268,7 @@ ReplicationSlotNameForTablesync(Oid suboid, Oid relid,
  * The returned slot name is palloc'ed in current memory context.
  */
 char *
-LogicalRepSyncTableStart(XLogRecPtr *origin_startpos)
+LogicalRepSyncTableStart(XLogRecPtr *origin_startpos, int worker_slot)
 {
 	char	   *slotname;
 	char	   *err;
@@ -1321,14 +1321,31 @@ LogicalRepSyncTableStart(XLogRecPtr *origin_startpos)
 									NAMEDATALEN);
 
 	/*
-	 * Here we use the slot name instead of the subscription name as the
-	 * application_name, so that it is different from the leader apply worker,
-	 * so that synchronous replication can distinguish them.
+	 * Connect to publisher if not yet. The application_name must be also
+	 * different from the leader apply worker because synchronous replication
+	 * must distinguish them.
 	 */
-	LogRepWorkerWalRcvConn =
-		walrcv_connect(MySubscription->conninfo, true,
-					   must_use_password,
-					   slotname, &err);
+	if (LogRepWorkerWalRcvConn == NULL)
+	{
+		char application_name[NAMEDATALEN];
+
+		/*
+		 * FIXME: set appropriate application_name. Previously, the slot name
+		 * was used becasue the lifetime of the tablesync worker was same as
+		 * that, but now the tablesync worker handles many slots during the
+		 * synchronization so that it is not suitable. So what should be?
+		 * Note that if the tablesync worker starts to reuse the replication
+		 * slot during synchronization, we should use the slot name as
+		 * application_name again.
+		 */
+		snprintf(application_name, NAMEDATALEN, "pg_%u_sync_%i",
+				 MySubscription->oid, worker_slot);
+		LogRepWorkerWalRcvConn =
+			walrcv_connect(MySubscription->conninfo, true,
+						   must_use_password,
+						   application_name, &err);
+	}
+
 	if (LogRepWorkerWalRcvConn == NULL)
 		ereport(ERROR,
 				(errcode(ERRCODE_CONNECTION_FAILURE),
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 4a0f402ad4..5420675ce8 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -3498,19 +3498,21 @@ LogicalRepApplyLoop(XLogRecPtr last_received)
 
 	/*
 	 * Init the ApplyMessageContext which we clean up after each replication
-	 * protocol message.
+	 * protocol message, if needed.
 	 */
-	ApplyMessageContext = AllocSetContextCreate(ApplyContext,
-												"ApplyMessageContext",
-												ALLOCSET_DEFAULT_SIZES);
+	if (!ApplyMessageContext)
+		ApplyMessageContext = AllocSetContextCreate(ApplyContext,
+													"ApplyMessageContext",
+													ALLOCSET_DEFAULT_SIZES);
 
 	/*
 	 * This memory context is used for per-stream data when the streaming mode
 	 * is enabled. This context is reset on each stream stop.
 	 */
-	LogicalStreamingContext = AllocSetContextCreate(ApplyContext,
-													"LogicalStreamingContext",
-													ALLOCSET_DEFAULT_SIZES);
+	if (!LogicalStreamingContext)
+		LogicalStreamingContext = AllocSetContextCreate(ApplyContext,
+														"LogicalStreamingContext",
+														ALLOCSET_DEFAULT_SIZES);
 
 	/* mark as idle, before starting to loop */
 	pgstat_report_activity(STATE_IDLE, NULL);
@@ -4462,7 +4464,7 @@ TwoPhaseTransactionGid(Oid subid, TransactionId xid, char *gid, int szgid)
  * are not repeatable.
  */
 static void
-start_table_sync(XLogRecPtr *origin_startpos, char **myslotname)
+start_table_sync(XLogRecPtr *origin_startpos, char **myslotname, int worker_slot)
 {
 	char	   *syncslotname = NULL;
 
@@ -4471,7 +4473,7 @@ start_table_sync(XLogRecPtr *origin_startpos, char **myslotname)
 	PG_TRY();
 	{
 		/* Call initial sync. */
-		syncslotname = LogicalRepSyncTableStart(origin_startpos);
+		syncslotname = LogicalRepSyncTableStart(origin_startpos, worker_slot);
 	}
 	PG_CATCH();
 	{
@@ -4541,12 +4543,13 @@ run_tablesync_worker(WalRcvStreamOptions *options,
 					 char *slotname,
 					 char *originname,
 					 int originname_size,
-					 XLogRecPtr *origin_startpos)
+					 XLogRecPtr *origin_startpos,
+					 int worker_slot)
 {
 	MyLogicalRepWorker->is_sync_completed = false;
 
 	/* Start table synchronization. */
-	start_table_sync(origin_startpos, &slotname);
+	start_table_sync(origin_startpos, &slotname, worker_slot);
 
 	ReplicationOriginNameForLogicalRep(MySubscription->oid,
 									   MyLogicalRepWorker->relid,
@@ -4850,7 +4853,9 @@ TablesyncWorkerMain(Datum main_arg)
 		ListCell   *lc;
 		bool 	is_table_found = false;
 
-		run_tablesync_worker(&options, myslotname, originname, sizeof(originname), &origin_startpos);
+		run_tablesync_worker(&options, myslotname, originname,
+							 sizeof(originname), &origin_startpos,
+							 worker_slot);
 
 		if (IsTransactionState())
 			CommitTransactionCommand();
@@ -4891,7 +4896,9 @@ TablesyncWorkerMain(Datum main_arg)
 
 					/* found a table for next iteration */
 					is_table_found = true;
-					clean_sync_worker();
+
+					CommitTransactionCommand();
+					pgstat_report_stat(true);
 
 					StartTransactionCommand();
 					ereport(LOG,
diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c
index d3a136b6f5..429d00f2f0 100644
--- a/src/backend/replication/walsender.c
+++ b/src/backend/replication/walsender.c
@@ -1828,6 +1828,12 @@ exec_replication_command(const char *cmd_string)
 				set_ps_display(cmdtag);
 				PreventInTransactionBlock(true, cmdtag);
 
+				/*
+				 * Initialize the flag again because this streaming may be
+				 * second time.
+				 */
+				streamingDoneSending = streamingDoneReceiving = false;
+
 				if (cmd->kind == REPLICATION_KIND_PHYSICAL)
 					StartReplication(cmd);
 				else
diff --git a/src/include/replication/worker_internal.h b/src/include/replication/worker_internal.h
index 1e9f8e6e72..d3cabf4033 100644
--- a/src/include/replication/worker_internal.h
+++ b/src/include/replication/worker_internal.h
@@ -249,7 +249,7 @@ extern int	logicalrep_sync_worker_count(Oid subid);
 
 extern void ReplicationOriginNameForLogicalRep(Oid suboid, Oid relid,
 											   char *originname, Size szoriginname);
-extern char *LogicalRepSyncTableStart(XLogRecPtr *origin_startpos);
+extern char *LogicalRepSyncTableStart(XLogRecPtr *origin_startpos, int worker_slot);
 
 extern bool AllTablesyncsReady(void);
 extern void UpdateTwoPhaseState(Oid suboid, char new_state);
-- 
2.27.0