From b2fc61d5267e184951aa2740ff7de2bf7130ed7a Mon Sep 17 00:00:00 2001
From: Masahiko Sawada <sawada.mshk@gmail.com>
Date: Fri, 10 Dec 2021 18:01:36 +0900
Subject: [PATCH v2] Move replication slot release to before_shmem_exit().

Previously, the wal sender releases the replication slot in ProcKill()
on error, resulting in sending the replication slot drop message for
the ephemeral slot to the stats collector after it shut down.

To fix this problem, move replication slot release to a
before_shmem_exit() hook that is called before the stats collector
shuts down.
---
 src/backend/replication/slot.c    | 26 ++++++++++++++++++++++++++
 src/backend/storage/lmgr/proc.c   |  7 -------
 src/backend/tcop/postgres.c       |  5 +++--
 src/backend/utils/init/postinit.c |  7 +++++++
 src/include/replication/slot.h    |  1 +
 5 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/src/backend/replication/slot.c b/src/backend/replication/slot.c
index 90ba9b417d..d686f7c0ad 100644
--- a/src/backend/replication/slot.c
+++ b/src/backend/replication/slot.c
@@ -46,6 +46,7 @@
 #include "pgstat.h"
 #include "replication/slot.h"
 #include "storage/fd.h"
+#include "storage/ipc.h"
 #include "storage/proc.h"
 #include "storage/procarray.h"
 #include "utils/builtins.h"
@@ -107,6 +108,8 @@ static void RestoreSlotFromDisk(const char *name);
 static void CreateSlotOnDisk(ReplicationSlot *slot);
 static void SaveSlotToPath(ReplicationSlot *slot, const char *path, int elevel);
 
+static void ReplicationSlotBeforeShmemExit(int code, Datum arg);
+
 /*
  * Report shared-memory space needed by ReplicationSlotsShmemInit.
  */
@@ -160,6 +163,29 @@ ReplicationSlotsShmemInit(void)
 	}
 }
 
+/*
+ * Register the callback for replication slot cleanup and releasing.
+ */
+void
+ReplicationSlotInitialize(void)
+{
+	before_shmem_exit(ReplicationSlotBeforeShmemExit, 0);
+}
+
+/*
+ * Release and cleanup replication slots.
+ */
+static void
+ReplicationSlotBeforeShmemExit(int code, Datum arg)
+{
+	/* Make sure active replication slots are released */
+	if (MyReplicationSlot != NULL)
+		ReplicationSlotRelease();
+
+	/* Also cleanup all the temporary slots. */
+	ReplicationSlotCleanup();
+}
+
 /*
  * Check whether the passed slot name is valid and report errors at elevel.
  *
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index b7d9da0aa9..031a73104e 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -847,13 +847,6 @@ ProcKill(int code, Datum arg)
 	/* Cancel any pending condition variable sleep, too */
 	ConditionVariableCancelSleep();
 
-	/* Make sure active replication slots are released */
-	if (MyReplicationSlot != NULL)
-		ReplicationSlotRelease();
-
-	/* Also cleanup all the temporary slots. */
-	ReplicationSlotCleanup();
-
 	/*
 	 * Detach from any lock group of which we are a member.  If the leader
 	 * exist before all other group members, its PGPROC will remain allocated
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 82de01cdc6..dc9d836297 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -4261,8 +4261,9 @@ PostgresMain(const char *dbname, const char *username)
 		 * We can't release replication slots inside AbortTransaction() as we
 		 * need to be able to start and abort transactions while having a slot
 		 * acquired. But we never need to hold them across top level errors,
-		 * so releasing here is fine. There's another cleanup in ProcKill()
-		 * ensuring we'll correctly cleanup on FATAL errors as well.
+		 * so releasing here is fine. There's another cleanup in
+		 * ReplicationSlotBeforeShmemExit() callback ensuring we'll correctly
+		 * cleanup on FATAL errors as well.
 		 */
 		if (MyReplicationSlot != NULL)
 			ReplicationSlotRelease();
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index 646126edee..6cf3b3c961 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -40,6 +40,7 @@
 #include "pgstat.h"
 #include "postmaster/autovacuum.h"
 #include "postmaster/postmaster.h"
+#include "replication/slot.h"
 #include "replication/walsender.h"
 #include "storage/bufmgr.h"
 #include "storage/fd.h"
@@ -547,6 +548,12 @@ BaseInit(void)
 	 * ever try to insert XLOG.
 	 */
 	InitXLogInsert();
+
+	/*
+	 * Register callback to make sure cleanup and releasing the replication
+	 * slot on exit.
+	 */
+	ReplicationSlotInitialize();
 }
 
 
diff --git a/src/include/replication/slot.h b/src/include/replication/slot.h
index 53d773ccff..2f4b123c87 100644
--- a/src/include/replication/slot.h
+++ b/src/include/replication/slot.h
@@ -206,6 +206,7 @@ extern void ReplicationSlotSave(void);
 extern void ReplicationSlotMarkDirty(void);
 
 /* misc stuff */
+extern void ReplicationSlotInitialize(void);
 extern bool ReplicationSlotValidateName(const char *name, int elevel);
 extern void ReplicationSlotReserveWal(void);
 extern void ReplicationSlotsComputeRequiredXmin(bool already_locked);
-- 
2.24.3 (Apple Git-128)

