From 3c136f8fb235882c2d26b59f34e5b21c69b8cb6a Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Wed, 8 Feb 2023 09:09:31 +0000
Subject: [PATCH v7 2/2] Extend START_REPLICATION command to accept walsender
 options

This commit extends START_REPLICATION to accept SHUTDOWN_MODE clause. It is
currently implemented only for logical replication.

The shutdown modes are:

1) 'wait_flush' (the default). In this mode, the walsender will wait for all the
sent WALs to be flushed on the subscriber side, before exiting the process.

2) 'immediate'. In this mode, the walsender will exit without confirming the
remote flush. This may break the consistency between publisher and subscriber.
This mode might be useful for a system that has a high-latency network (to
reduce the amount of time for shutdown), or to allow the shutdown of the
publisher even when the worker is stuck.

Author: Hayato Kuroda
Discussion: https://postgr.es/m/TYAPR01MB586668E50FC2447AD7F92491F5E89%40TYAPR01MB5866.jpnprd01.prod.outlook.com
---
 doc/src/sgml/protocol.sgml                    | 16 ++++-
 .../libpqwalreceiver/libpqwalreceiver.c       |  7 +++
 src/backend/replication/logical/worker.c      | 18 +++++-
 src/backend/replication/repl_gram.y           | 11 +++-
 src/backend/replication/repl_scanner.l        |  1 +
 src/backend/replication/walsender.c           | 60 ++++++++++++++++++-
 src/include/nodes/replnodes.h                 |  1 +
 src/include/replication/walreceiver.h         |  2 +
 src/test/subscription/t/001_rep_changes.pl    | 22 ++++++-
 src/tools/pgindent/typedefs.list              |  1 +
 10 files changed, 131 insertions(+), 8 deletions(-)

diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml
index 93fc7167d4..1ba5238060 100644
--- a/doc/src/sgml/protocol.sgml
+++ b/doc/src/sgml/protocol.sgml
@@ -2500,7 +2500,7 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;"
     </varlistentry>
 
     <varlistentry id="protocol-replication-start-replication-slot-logical">
-     <term><literal>START_REPLICATION</literal> <literal>SLOT</literal> <replaceable class="parameter">slot_name</replaceable> <literal>LOGICAL</literal> <replaceable class="parameter">XXX/XXX</replaceable> [ ( <replaceable>option_name</replaceable> [ <replaceable>option_value</replaceable> ] [, ...] ) ]</term>
+     <term><literal>START_REPLICATION</literal> <literal>SLOT</literal> <replaceable class="parameter">slot_name</replaceable> <literal>LOGICAL</literal> <replaceable class="parameter">XXX/XXX</replaceable> [ <literal>SHUTDOWN_MODE</literal> { <literal>'wait_flush'</literal> | <literal>'immediate'</literal> } ] [ ( <replaceable>option_name</replaceable> [ <replaceable>option_value</replaceable> ] [, ...] ) ]</term>
      <listitem>
       <para>
        Instructs server to start streaming WAL for logical replication,
@@ -2555,6 +2555,20 @@ psql "dbname=postgres replication=database" -c "IDENTIFY_SYSTEM;"
         </listitem>
        </varlistentry>
 
+       <varlistentry>
+        <term><literal>shutdown_mode</literal></term>
+        <listitem>
+         <para>
+          Determines the behavior of the walsender process at shutdown. If
+          shutdown_mode is <literal>'wait_flush'</literal>, the walsender waits
+          for all the sent WALs to be flushed on the subscriber side. This is
+          the default when SHUTDOWN_MODE is not specified. If shutdown_mode is
+          <literal>'immediate'</literal>, the walsender exits without
+          confirming the remote flush.
+         </para>
+        </listitem>
+       </varlistentry>
+
        <varlistentry>
         <term><replaceable class="parameter">option_name</replaceable></term>
         <listitem>
diff --git a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
index 560ec974fa..7aad751a86 100644
--- a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
+++ b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
@@ -403,6 +403,12 @@ libpqrcv_startstreaming(WalReceiverConn *conn,
 		List	   *pubnames;
 		char	   *pubnames_literal;
 
+		/* Add SHUTDOWN_MODE clause if not using the default shutdown_mode */
+		if (options->proto.logical.shutdown_mode_str &&
+			PQserverVersion(conn->streamConn) >= 160000)
+			appendStringInfo(&cmd, " SHUTDOWN_MODE '%s'",
+							 options->proto.logical.shutdown_mode_str);
+
 		appendStringInfoString(&cmd, " (");
 
 		appendStringInfo(&cmd, "proto_version '%u'",
@@ -449,6 +455,7 @@ libpqrcv_startstreaming(WalReceiverConn *conn,
 		appendStringInfo(&cmd, " TIMELINE %u",
 						 options->proto.physical.startpointTLI);
 
+
 	/* Start streaming. */
 	res = libpqrcv_PQexec(conn->streamConn, cmd.data);
 	pfree(cmd.data);
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 6b86723b60..be8be077c5 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -4043,10 +4043,16 @@ maybe_reread_subscription(void)
 
 	/*
 	 * Exit if any parameter that affects the remote connection was changed.
+	 *
 	 * The launcher will start a new worker but note that the parallel apply
 	 * worker won't restart if the streaming option's value is changed from
 	 * 'parallel' to any other value or the server decides not to stream the
 	 * in-progress transaction.
+	 *
+	 * Time-delayed logical replication affects the SHUTDOWN_MODE clause. The
+	 * 'immediate' shutdown mode will be specified if min_apply_delay is
+	 * non-zero, otherwise the default shutdown mode will be used. See
+	 * ApplyWorkerMain.
 	 */
 	if (strcmp(newsub->conninfo, MySubscription->conninfo) != 0 ||
 		strcmp(newsub->name, MySubscription->name) != 0 ||
@@ -4055,7 +4061,8 @@ maybe_reread_subscription(void)
 		newsub->stream != MySubscription->stream ||
 		strcmp(newsub->origin, MySubscription->origin) != 0 ||
 		newsub->owner != MySubscription->owner ||
-		!equal(newsub->publications, MySubscription->publications))
+		!equal(newsub->publications, MySubscription->publications) ||
+		(newsub->minapplydelay == 0) != (MySubscription->minapplydelay == 0))
 	{
 		if (am_parallel_apply_worker())
 			ereport(LOG,
@@ -4774,9 +4781,18 @@ ApplyWorkerMain(Datum main_arg)
 
 	options.proto.logical.twophase = false;
 	options.proto.logical.origin = pstrdup(MySubscription->origin);
+	options.proto.logical.shutdown_mode_str = NULL;
 
 	if (!am_tablesync_worker())
 	{
+		/*
+		 * Time-delayed logical replication does not support tablesync
+		 * workers, so only the leader apply worker can request walsenders to
+		 * exit before confirming remote flush.
+		 */
+		if (server_version >= 160000 && MySubscription->minapplydelay > 0)
+			options.proto.logical.shutdown_mode_str = pstrdup("immediate");
+
 		/*
 		 * Even when the two_phase mode is requested by the user, it remains
 		 * as the tri-state PENDING until all tablesyncs have reached READY
diff --git a/src/backend/replication/repl_gram.y b/src/backend/replication/repl_gram.y
index 0c874e33cf..11a01e6b60 100644
--- a/src/backend/replication/repl_gram.y
+++ b/src/backend/replication/repl_gram.y
@@ -76,6 +76,7 @@ Node *replication_parse_result;
 %token K_EXPORT_SNAPSHOT
 %token K_NOEXPORT_SNAPSHOT
 %token K_USE_SNAPSHOT
+%token K_SHUTDOWN_MODE
 
 %type <node>	command
 %type <node>	base_backup start_replication start_logical_replication
@@ -91,6 +92,7 @@ Node *replication_parse_result;
 %type <boolval>	opt_temporary
 %type <list>	create_slot_options create_slot_legacy_opt_list
 %type <defelt>	create_slot_legacy_opt
+%type <str>		opt_shutdown_mode
 
 %%
 
@@ -276,14 +278,15 @@ start_replication:
 
 /* START_REPLICATION SLOT slot LOGICAL %X/%X options */
 start_logical_replication:
-			K_START_REPLICATION K_SLOT IDENT K_LOGICAL RECPTR plugin_options
+			K_START_REPLICATION K_SLOT IDENT K_LOGICAL RECPTR opt_shutdown_mode plugin_options
 				{
 					StartReplicationCmd *cmd;
 					cmd = makeNode(StartReplicationCmd);
 					cmd->kind = REPLICATION_KIND_LOGICAL;
 					cmd->slotname = $3;
 					cmd->startpoint = $5;
-					cmd->options = $6;
+					cmd->shutdownmode = $6;
+					cmd->options = $7;
 					$$ = (Node *) cmd;
 				}
 			;
@@ -336,6 +339,10 @@ opt_timeline:
 				| /* EMPTY */			{ $$ = 0; }
 			;
 
+opt_shutdown_mode:
+			K_SHUTDOWN_MODE SCONST			{ $$ = $2; }
+			| /* EMPTY */					{ $$ = NULL; }
+		;
 
 plugin_options:
 			'(' plugin_opt_list ')'			{ $$ = $2; }
diff --git a/src/backend/replication/repl_scanner.l b/src/backend/replication/repl_scanner.l
index cb467ca46f..fcc6f6feda 100644
--- a/src/backend/replication/repl_scanner.l
+++ b/src/backend/replication/repl_scanner.l
@@ -136,6 +136,7 @@ EXPORT_SNAPSHOT		{ return K_EXPORT_SNAPSHOT; }
 NOEXPORT_SNAPSHOT	{ return K_NOEXPORT_SNAPSHOT; }
 USE_SNAPSHOT		{ return K_USE_SNAPSHOT; }
 WAIT				{ return K_WAIT; }
+SHUTDOWN_MODE		{ return K_SHUTDOWN_MODE; }
 
 {space}+		{ /* do nothing */ }
 
diff --git a/src/backend/replication/walsender.c b/src/backend/replication/walsender.c
index 75e8363e24..b0f8a2d320 100644
--- a/src/backend/replication/walsender.c
+++ b/src/backend/replication/walsender.c
@@ -219,6 +219,15 @@ typedef struct
 
 static LagTracker *lag_tracker;
 
+/* Shutdown modes */
+typedef enum
+{
+	WALSND_SHUTDOWN_MODE_WAIT_FLUSH,
+	WALSND_SHUTDOWN_MODE_IMMEDIATE
+} WalSndShutdownMode;
+
+static WalSndShutdownMode shutdown_mode = WALSND_SHUTDOWN_MODE_WAIT_FLUSH;
+
 /* Signal handlers */
 static void WalSndLastCycleHandler(SIGNAL_ARGS);
 
@@ -260,6 +269,7 @@ static bool TransactionIdInRecentPast(TransactionId xid, uint32 epoch);
 static void WalSndSegmentOpen(XLogReaderState *state, XLogSegNo nextSegNo,
 							  TimeLineID *tli_p);
 
+static void CheckWalSndOptions(const StartReplicationCmd *cmd);
 
 /* Initialize walsender process before entering the main command loop */
 void
@@ -1272,6 +1282,9 @@ StartLogicalReplication(StartReplicationCmd *cmd)
 		got_STOPPING = true;
 	}
 
+	/* Check given options and set flags accordingly */
+	CheckWalSndOptions(cmd);
+
 	/*
 	 * Create our decoding context, making it start at the previously ack'ed
 	 * position.
@@ -1450,6 +1463,16 @@ ProcessPendingWrites(void)
 		/* Try to flush pending output to the client */
 		if (pq_flush_if_writable() != 0)
 			WalSndShutdown();
+
+		/*
+		 * In this function, there is a possibility that the walsender is
+		 * stuck. This can happen when the receiving worker is stuck, and then
+		 * the send-buffer of the walsender becomes full. Therefore, we must
+		 * add an additional path for shutdown for immediate shutdown mode.
+		 */
+		if (got_STOPPING &&
+			shutdown_mode == WALSND_SHUTDOWN_MODE_IMMEDIATE)
+			WalSndDone(XLogSendLogical);
 	}
 
 	/* reactivate latch so WalSndLoop knows to continue */
@@ -3114,19 +3137,25 @@ WalSndDone(WalSndSendDataCallback send_data)
 	 * To figure out whether all WAL has successfully been replicated, check
 	 * flush location if valid, write otherwise. Tools like pg_receivewal will
 	 * usually (unless in synchronous mode) return an invalid flush location.
+	 *
+	 * If we are in the 'immediate' shutdown mode, flush location and output
+	 * buffer is not checked. This may break the consistency between nodes,
+	 * but it may be useful for the system that has high-latency network to
+	 * reduce the amount of time for shutdown.
 	 */
 	replicatedPtr = XLogRecPtrIsInvalid(MyWalSnd->flush) ?
 		MyWalSnd->write : MyWalSnd->flush;
 
-	if (WalSndCaughtUp && sentPtr == replicatedPtr &&
-		!pq_is_send_pending())
+	if (WalSndCaughtUp &&
+		(shutdown_mode == WALSND_SHUTDOWN_MODE_IMMEDIATE ||
+		 (sentPtr == replicatedPtr && !pq_is_send_pending())))
 	{
 		QueryCompletion qc;
 
 		/* Inform the standby that XLOG streaming is done */
 		SetQueryCompletion(&qc, CMDTAG_COPY, 0);
 		EndCommand(&qc, DestRemote, false);
-		pq_flush();
+		pq_flush_if_writable();
 
 		proc_exit(0);
 	}
@@ -3849,3 +3878,28 @@ LagTrackerRead(int head, XLogRecPtr lsn, TimestampTz now)
 	Assert(time != 0);
 	return now - time;
 }
+
+/*
+ * Check options for walsender itself and set flags accordingly.
+ *
+ * Currently only one option is accepted.
+ */
+static void
+CheckWalSndOptions(const StartReplicationCmd *cmd)
+{
+	if (cmd->shutdownmode)
+	{
+		char	   *mode = cmd->shutdownmode;
+
+		if (pg_strcasecmp(mode, "wait_flush") == 0)
+			shutdown_mode = WALSND_SHUTDOWN_MODE_WAIT_FLUSH;
+		else if (pg_strcasecmp(mode, "immediate") == 0)
+			shutdown_mode = WALSND_SHUTDOWN_MODE_IMMEDIATE;
+		else
+			ereport(ERROR,
+					errcode(ERRCODE_SYNTAX_ERROR),
+					errmsg("invalid value for shutdown mode: \"%s\"", mode),
+					errhint("Available values: wait_flush, immediate."));
+	}
+
+}
diff --git a/src/include/nodes/replnodes.h b/src/include/nodes/replnodes.h
index 4321ba8f86..c96e85e859 100644
--- a/src/include/nodes/replnodes.h
+++ b/src/include/nodes/replnodes.h
@@ -83,6 +83,7 @@ typedef struct StartReplicationCmd
 	char	   *slotname;
 	TimeLineID	timeline;
 	XLogRecPtr	startpoint;
+	char	   *shutdownmode;
 	List	   *options;
 } StartReplicationCmd;
 
diff --git a/src/include/replication/walreceiver.h b/src/include/replication/walreceiver.h
index decffe352d..91d9f7b3b1 100644
--- a/src/include/replication/walreceiver.h
+++ b/src/include/replication/walreceiver.h
@@ -187,6 +187,8 @@ typedef struct
 									 * prepare time */
 			char	   *origin; /* Only publish data originating from the
 								 * specified origin */
+			char	   *shutdown_mode_str;	/* The specified shutdown mode
+											 * string, or NULL. */
 		}			logical;
 	}			proto;
 } WalRcvStreamOptions;
diff --git a/src/test/subscription/t/001_rep_changes.pl b/src/test/subscription/t/001_rep_changes.pl
index 75fd77b891..d649c78de7 100644
--- a/src/test/subscription/t/001_rep_changes.pl
+++ b/src/test/subscription/t/001_rep_changes.pl
@@ -523,10 +523,18 @@ $node_publisher->poll_query_until('postgres',
 # inserted on the publisher, and b) when those changes are replicated on the
 # subscriber. Even on slow machines, this strategy will give predictable behavior.
 
-# Set min_apply_delay parameter to 3 seconds
+# Check restart on changing min_apply_delay to 3 seconds
 my $delay = 3;
+$oldpid = $node_publisher->safe_psql('postgres',
+	"SELECT pid FROM pg_stat_replication WHERE application_name = 'tap_sub_renamed' AND state = 'streaming';"
+);
 $node_subscriber->safe_psql('postgres',
 	"ALTER SUBSCRIPTION tap_sub_renamed SET (min_apply_delay = '${delay}s')");
+$node_publisher->poll_query_until('postgres',
+	"SELECT pid != $oldpid FROM pg_stat_replication WHERE application_name = 'tap_sub_renamed' AND state = 'streaming';"
+  )
+  or die
+  "Timed out while waiting for the walsender to restart after changing min_apply_delay to non-zero value";
 
 # Before doing the insertion, get the current timestamp that will be
 # used as a comparison base.
@@ -543,6 +551,18 @@ ok( time() - $publisher_insert_time >= $delay,
 	"subscriber applies WAL only after replication delay for non-streaming transaction"
 );
 
+# Check restart on changing min_apply_delay to zero
+$oldpid = $node_publisher->safe_psql('postgres',
+	"SELECT pid FROM pg_stat_replication WHERE application_name = 'tap_sub_renamed' AND state = 'streaming';"
+);
+$node_subscriber->safe_psql('postgres',
+	"ALTER SUBSCRIPTION tap_sub_renamed SET (min_apply_delay = '0')");
+$node_publisher->poll_query_until('postgres',
+	"SELECT pid != $oldpid FROM pg_stat_replication WHERE application_name = 'tap_sub_renamed' AND state = 'streaming';"
+  )
+  or die
+  "Timed out while waiting for the walsender to restart after changing min_apply_delay to zero";
+
 # check all the cleanup
 $node_subscriber->safe_psql('postgres', "DROP SUBSCRIPTION tap_sub_renamed");
 
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 36d1dc0117..d06a7868ca 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2976,6 +2976,7 @@ WalReceiverFunctionsType
 WalSnd
 WalSndCtlData
 WalSndSendDataCallback
+WalSndShutdownMode
 WalSndState
 WalTimeSample
 WalUsage
-- 
2.27.0

