From 27803f93ffb3411dac58c3940b9f1a69dd28134c Mon Sep 17 00:00:00 2001
From: Ian Barwick <ian@2ndquadrant.com>
Date: Fri, 12 Apr 2019 12:26:39 +0900
Subject: [PATCH] repmgrd: always unset upstream node ID when monitoring a
 primary

---
 dbutils.c            | 30 ++++++++++++++++++++++++++++++
 dbutils.h            |  1 +
 repmgr--4.3--4.4.sql |  5 +++++
 repmgr--4.4.sql      |  4 ++++
 repmgr.c             | 34 +++++++++++++++++++++++++++++++++-
 repmgrd-physical.c   |  1 +
 6 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/dbutils.c b/dbutils.c
index c26bce5c..04ae684b 100644
--- a/dbutils.c
+++ b/dbutils.c
@@ -1965,6 +1965,36 @@ repmgrd_get_upstream_node_id(PGconn *conn)
 	return upstream_node_id;
 }
 
+
+bool
+repmgrd_set_upstream_node_id(PGconn *conn, int node_id)
+{
+	PQExpBufferData query;
+	PGresult   *res = NULL;
+	bool		success = true;
+
+	initPQExpBuffer(&query);
+	appendPQExpBuffer(&query,
+					  " SELECT repmgr.set_upstream_node_id(%i) ",
+					  node_id);
+
+	log_verbose(LOG_DEBUG, "repmgrd_set_upstream_node_id():\n  %s", query.data);
+
+	res = PQexec(conn, query.data);
+
+	if (PQresultStatus(res) != PGRES_TUPLES_OK)
+	{
+		log_db_error(conn, query.data,
+					 _("repmgrd_set_upstream_node_id(): unable to set upstream node ID (provided value: %i)"), node_id);
+		success = false;
+	}
+
+	termPQExpBuffer(&query);
+	PQclear(res);
+
+	return success;
+}
+
 /* ================ */
 /* result functions */
 /* ================ */
diff --git a/dbutils.h b/dbutils.h
index bddb8c21..e9bbd475 100644
--- a/dbutils.h
+++ b/dbutils.h
@@ -442,6 +442,7 @@ bool		repmgrd_is_paused(PGconn *conn);
 bool		repmgrd_pause(PGconn *conn, bool pause);
 pid_t		get_wal_receiver_pid(PGconn *conn);
 int			repmgrd_get_upstream_node_id(PGconn *conn);
+bool		repmgrd_set_upstream_node_id(PGconn *conn, int node_id);
 
 /* extension functions */
 ExtensionStatus get_repmgr_extension_status(PGconn *conn, t_extension_versions *extversions);
diff --git a/repmgr--4.3--4.4.sql b/repmgr--4.3--4.4.sql
index b5666bec..4b4b8f97 100644
--- a/repmgr--4.3--4.4.sql
+++ b/repmgr--4.3--4.4.sql
@@ -12,3 +12,8 @@ CREATE FUNCTION get_upstream_node_id()
   RETURNS INT
   AS 'MODULE_PATHNAME', 'get_upstream_node_id'
   LANGUAGE C STRICT;
+
+CREATE FUNCTION set_upstream_node_id(INT)
+  RETURNS VOID
+  AS 'MODULE_PATHNAME', 'set_upstream_node_id'
+  LANGUAGE C STRICT;
diff --git a/repmgr--4.4.sql b/repmgr--4.4.sql
index a4ff77d6..39eed17f 100644
--- a/repmgr--4.4.sql
+++ b/repmgr--4.4.sql
@@ -133,6 +133,10 @@ CREATE FUNCTION get_upstream_node_id()
   AS 'MODULE_PATHNAME', 'get_upstream_node_id'
   LANGUAGE C STRICT;
 
+CREATE FUNCTION set_upstream_node_id(INT)
+  RETURNS VOID
+  AS 'MODULE_PATHNAME', 'set_upstream_node_id'
+  LANGUAGE C STRICT;
 
 /* failover functions */
 
diff --git a/repmgr.c b/repmgr.c
index ccb3b21a..39e3616c 100644
--- a/repmgr.c
+++ b/repmgr.c
@@ -119,6 +119,9 @@ PG_FUNCTION_INFO_V1(get_upstream_last_seen);
 Datum		get_upstream_node_id(PG_FUNCTION_ARGS);
 PG_FUNCTION_INFO_V1(get_upstream_node_id);
 
+Datum		set_upstream_node_id(PG_FUNCTION_ARGS);
+PG_FUNCTION_INFO_V1(set_upstream_node_id);
+
 Datum		notify_follow_primary(PG_FUNCTION_ARGS);
 PG_FUNCTION_INFO_V1(notify_follow_primary);
 
@@ -427,7 +430,6 @@ get_upstream_last_seen(PG_FUNCTION_ARGS)
 }
 
 
-
 Datum
 get_upstream_node_id(PG_FUNCTION_ARGS)
 {
@@ -443,6 +445,36 @@ get_upstream_node_id(PG_FUNCTION_ARGS)
 	PG_RETURN_INT32(upstream_node_id);
 }
 
+Datum
+set_upstream_node_id(PG_FUNCTION_ARGS)
+{
+	int			upstream_node_id = UNKNOWN_NODE_ID;
+	int			local_node_id = UNKNOWN_NODE_ID;
+
+	if (!shared_state)
+		PG_RETURN_NULL();
+
+	if (PG_ARGISNULL(0))
+		PG_RETURN_NULL();
+
+	upstream_node_id = PG_GETARG_INT32(0);
+
+	LWLockAcquire(shared_state->lock, LW_SHARED);
+	local_node_id = shared_state->local_node_id;
+	LWLockRelease(shared_state->lock);
+
+	if (local_node_id == upstream_node_id)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 (errmsg("upstream node id cannot be the same as the local node id"))));
+
+	LWLockAcquire(shared_state->lock, LW_EXCLUSIVE);
+	shared_state->upstream_node_id = upstream_node_id;
+	LWLockRelease(shared_state->lock);
+
+	PG_RETURN_VOID();
+}
+
 
 /* ===================*/
 /* failover functions */
diff --git a/repmgrd-physical.c b/repmgrd-physical.c
index c2f575d2..23b04a37 100644
--- a/repmgrd-physical.c
+++ b/repmgrd-physical.c
@@ -231,6 +231,7 @@ monitor_streaming_primary(void)
 	instr_time	log_status_interval_start;
 
 	reset_node_voting_status();
+	repmgrd_set_upstream_node_id(local_conn, NO_UPSTREAM_NODE);
 
 	{
 		PQExpBufferData event_details;