From 74fc0770c68cbf5605264907ca8cd328f675f9ec Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Fri, 3 Feb 2023 11:17:26 +0000
Subject: [PATCH v6 2/2] Introduce smart shutdown for logical walsender

---
 src/backend/postmaster/postmaster.c         |  4 +++
 src/backend/replication/walsender.c         | 37 ++++++++++++++++++---
 src/include/replication/walsender.h         |  9 +++++
 src/include/replication/walsender_private.h |  7 ++++
 4 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index f92dbc2270..8a7d1d4efa 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -3813,6 +3813,10 @@ PostmasterStateMachine(void)
 				if (CheckpointerPID != 0)
 				{
 					signal_child(CheckpointerPID, SIGUSR2);
+					if (Shutdown >= FastShutdown)
+						WalSndChangeShutdownState(WALSNDSHUTDOWNSTATE_FASTSHUTDOWN);
+					else
+						WalSndChangeShutdownState(WALSNDSHUTDOWNSTATE_SMARTSHUTDOWN);
 					pmState = PM_SHUTDOWN;
 				}
 				else
diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c
index 25a052adfc..820d74c5a7 100644
--- a/src/backend/replication/walsender.c
+++ b/src/backend/replication/walsender.c
@@ -260,6 +260,16 @@ static bool TransactionIdInRecentPast(TransactionId xid, uint32 epoch);
 static void WalSndSegmentOpen(XLogReaderState *state, XLogSegNo nextSegNo,
 							  TimeLineID *tli_p);
 
+#define SHUTDOWN_CONDTION_FOR_LOGICAL(send_data)					\
+	((send_data) == XLogSendLogical &&								\
+	 (WalSndCtl->Shutdown == WALSNDSHUTDOWNSTATE_FASTSHUTDOWN ||	\
+	  (WalSndCtl->Shutdown == WALSNDSHUTDOWNSTATE_SMARTSHUTDOWN &&	\
+	   !pq_is_send_pending())))
+
+#define SHUTDOWN_CONDTION_FOR_PHYSICAL(send_data)		\
+	((send_data) == XLogSendPhysical &&					\
+	 sentPtr == replicatedPtr && !pq_is_send_pending())
+
 
 /* Initialize walsender process before entering the main command loop */
 void
@@ -3122,13 +3132,18 @@ WalSndDone(WalSndSendDataCallback send_data)
 	/*
 	 * Exit if we are in the convenient time.
 	 *
-	 * When we are logical replication mode, we don't have to wait that all
-	 * sent data to be flushed on the subscriber because we cannot support
-	 * clean switchover for it.
+	 * When we are logical replication mode, the condition for shutdown is
+	 * changed based on the shutdown mode.
+	 *
+	 * For smart shutdown mode, we confirm that there is no pending data
+	 * in the output buffer.
+	 *
+	 * For fast shutdown mode, we don't have to wait that all sent data to be
+	 * flushed on the subscriber.
 	 */
 	if (WalSndCaughtUp &&
-		(send_data == XLogSendLogical ||
-		 (sentPtr == replicatedPtr && !pq_is_send_pending())))
+		(SHUTDOWN_CONDTION_FOR_LOGICAL(send_data) ||
+		 SHUTDOWN_CONDTION_FOR_PHYSICAL(send_data)))
 	{
 		QueryCompletion qc;
 
@@ -3867,3 +3882,15 @@ LagTrackerRead(int head, XLogRecPtr lsn, TimestampTz now)
 	Assert(time != 0);
 	return now - time;
 }
+
+/*
+ * Change the indicator for shutdown mode.
+ *
+ * This should be called by postmaster
+ */
+void
+WalSndChangeShutdownState(WalSndShutdownState state)
+{
+	Assert(!IsUnderPostmaster);
+	WalSndCtl->Shutdown = state;
+}
diff --git a/src/include/replication/walsender.h b/src/include/replication/walsender.h
index 52bb3e2aae..97a31e1613 100644
--- a/src/include/replication/walsender.h
+++ b/src/include/replication/walsender.h
@@ -24,6 +24,13 @@ typedef enum
 	CRS_USE_SNAPSHOT
 } CRSSnapshotAction;
 
+typedef enum WalSndShutdownState
+{
+	WALSNDSHUTDOWNSTATE_NOSHUTDOWN = 0,
+	WALSNDSHUTDOWNSTATE_SMARTSHUTDOWN,
+	WALSNDSHUTDOWNSTATE_FASTSHUTDOWN,
+} WalSndShutdownState;
+
 /* global state */
 extern PGDLLIMPORT bool am_walsender;
 extern PGDLLIMPORT bool am_cascading_walsender;
@@ -48,6 +55,8 @@ extern void WalSndWaitStopping(void);
 extern void HandleWalSndInitStopping(void);
 extern void WalSndRqstFileReload(void);
 
+extern void WalSndChangeShutdownState(WalSndShutdownState state);
+
 /*
  * Remember that we want to wakeup walsenders later
  *
diff --git a/src/include/replication/walsender_private.h b/src/include/replication/walsender_private.h
index 5310e054c4..6171c16cb2 100644
--- a/src/include/replication/walsender_private.h
+++ b/src/include/replication/walsender_private.h
@@ -105,6 +105,13 @@ typedef struct
 	 */
 	bool		sync_standbys_defined;
 
+	/*
+	 * Indicator for shutdown mode. Basically it is set to zero, which means
+	 * that shutdown has not been requested yet. The value will be changed
+	 * olny once by postmaster.
+	 */
+	int			Shutdown;
+
 	WalSnd		walsnds[FLEXIBLE_ARRAY_MEMBER];
 } WalSndCtlData;
 
-- 
2.27.0