From 7e30e717eff342cb7c4d231a31c53d2be8255c13 Mon Sep 17 00:00:00 2001
From: naisila <nicypp@gmail.com>
Date: Fri, 11 Dec 2020 22:35:45 +0300
Subject: [PATCH] Add AppendDistributedTablePlacementNodeList function

---
 .../distributed/operations/stage_protocol.c   |  2 +-
 .../operations/worker_node_manager.c          | 29 +++++++++++++++++++
 src/include/distributed/worker_manager.h      |  2 ++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/operations/stage_protocol.c b/src/backend/distributed/operations/stage_protocol.c
index 71c178d16..ac62eaedf 100644
--- a/src/backend/distributed/operations/stage_protocol.c
+++ b/src/backend/distributed/operations/stage_protocol.c
@@ -166,7 +166,7 @@ master_create_empty_shard(PG_FUNCTION_ARGS)
 	uint64 shardId = GetNextShardId();
 
 	/* if enough live groups, add an extra candidate node as backup */
-	List *workerNodeList = DistributedTablePlacementNodeList(NoLock);
+	List *workerNodeList = AppendDistributedTablePlacementNodeList(NoLock);
 
 	if (list_length(workerNodeList) > ShardReplicationFactor)
 	{
diff --git a/src/backend/distributed/operations/worker_node_manager.c b/src/backend/distributed/operations/worker_node_manager.c
index a1637bf34..963a4fafa 100644
--- a/src/backend/distributed/operations/worker_node_manager.c
+++ b/src/backend/distributed/operations/worker_node_manager.c
@@ -499,6 +499,19 @@ DistributedTablePlacementNodeList(LOCKMODE lockMode)
 }
 
 
+/*
+ * AppendDistributedTablePlacementNodeList returns a list of all active, primary
+ * worker nodes that can store new data, i.e shouldstoreshards is 'true', and are
+ * not the coordinator
+ */
+List *
+AppendDistributedTablePlacementNodeList(LOCKMODE lockMode)
+{
+	EnsureModificationsCanRun();
+	return FilterActiveNodeListFunc(lockMode, NodeCanHaveAppendDistTablePlacements);
+}
+
+
 /*
  * NodeCanHaveDistTablePlacements returns true if the given node can have
  * shards of a distributed table.
@@ -515,6 +528,22 @@ NodeCanHaveDistTablePlacements(WorkerNode *node)
 }
 
 
+/*
+ * NodeCanHaveAppendDistTablePlacements returns true if the given node can have
+ * shards of a distributed table, and is not the coordinator
+ */
+bool
+NodeCanHaveAppendDistTablePlacements(WorkerNode *node)
+{
+	if (!NodeIsPrimary(node))
+	{
+		return false;
+	}
+
+	return node->shouldHaveShards && (node->groupId != 0);
+}
+
+
 /*
  * ActiveReadableNonCoordinatorNodeList returns a list of all nodes in workerNodeHash
  * that are readable nodes This method excludes coordinator.
diff --git a/src/include/distributed/worker_manager.h b/src/include/distributed/worker_manager.h
index db0045166..13f8ea7d1 100644
--- a/src/include/distributed/worker_manager.h
+++ b/src/include/distributed/worker_manager.h
@@ -79,7 +79,9 @@ extern List * ReferenceTablePlacementNodeList(LOCKMODE lockMode);
 extern WorkerNode * CoordinatorNodeIfAddedAsWorkerOrError(void);
 extern void ErrorIfCoordinatorNotAddedAsWorkerNode(void);
 extern List * DistributedTablePlacementNodeList(LOCKMODE lockMode);
+extern List * AppendDistributedTablePlacementNodeList(LOCKMODE lockMode);
 extern bool NodeCanHaveDistTablePlacements(WorkerNode *node);
+extern bool NodeCanHaveAppendDistTablePlacements(WorkerNode *node);
 extern uint32 ActiveReadableNonCoordinatorNodeCount(void);
 extern List * ActiveReadableNonCoordinatorNodeList(void);
 extern List * ActiveReadableNodeList(void);