From ee6a75179817cbb1223d215a8c0cf612a3d3d573 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Mon, 5 Feb 2018 00:09:53 +0100
Subject: [PATCH] Only copy distributed plan when modifying it

---
 .../executor/multi_router_executor.c          | 22 ++++++++++++++-----
 .../distributed/planner/distributed_planner.c | 14 +++++-------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/backend/distributed/executor/multi_router_executor.c b/src/backend/distributed/executor/multi_router_executor.c
index 04b4beb15..a9c9394bc 100644
--- a/src/backend/distributed/executor/multi_router_executor.c
+++ b/src/backend/distributed/executor/multi_router_executor.c
@@ -395,11 +395,21 @@ void
 CitusModifyBeginScan(CustomScanState *node, EState *estate, int eflags)
 {
 	CitusScanState *scanState = (CitusScanState *) node;
-	DistributedPlan *distributedPlan = scanState->distributedPlan;
-	Job *workerJob = distributedPlan->workerJob;
-	Query *jobQuery = workerJob->jobQuery;
-	List *taskList = workerJob->taskList;
-	bool deferredPruning = workerJob->deferredPruning;
+	DistributedPlan *distributedPlan = NULL;
+	Job *workerJob = NULL;
+	Query *jobQuery = NULL;
+	List *taskList = NIL;
+
+	/*
+	 * We must not change the distributed plan since it may be reused across multiple
+	 * executions of a prepared statement. Instead we create a deep copy that we only
+	 * use for the current execution.
+	 */
+	distributedPlan = scanState->distributedPlan = copyObject(scanState->distributedPlan);
+
+	workerJob = distributedPlan->workerJob;
+	jobQuery = workerJob->jobQuery;
+	taskList = workerJob->taskList;
 
 	if (workerJob->requiresMasterEvaluation)
 	{
@@ -416,7 +426,7 @@ CitusModifyBeginScan(CustomScanState *node, EState *estate, int eflags)
 		 */
 		executorState->es_param_list_info = NULL;
 
-		if (deferredPruning)
+		if (workerJob->deferredPruning)
 		{
 			DeferredErrorMessage *planningError = NULL;
 
diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index 45b2ed79c..17b25bf8a 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/src/backend/distributed/planner/distributed_planner.c
@@ -828,6 +828,9 @@ ResolveExternalParams(Node *inputNode, ParamListInfo boundParams)
 
 /*
  * GetDistributedPlan returns the associated DistributedPlan for a CustomScan.
+ *
+ * Callers should only read from the returned data structure, since it may be
+ * the plan of a prepared statement and may therefore be reused.
  */
 DistributedPlan *
 GetDistributedPlan(CustomScan *customScan)
@@ -840,16 +843,9 @@ GetDistributedPlan(CustomScan *customScan)
 	node = (Node *) linitial(customScan->custom_private);
 	Assert(CitusIsA(node, DistributedPlan));
 
-	node = CheckNodeCopyAndSerialization(node);
+	CheckNodeCopyAndSerialization(node);
 
-	/*
-	 * When using prepared statements the same plan gets reused across
-	 * multiple statements and transactions. We make several modifications
-	 * to the DistributedPlan during execution such as assigning task placements
-	 * and evaluating functions and parameters. These changes should not
-	 * persist, so we always work on a copy.
-	 */
-	distributedPlan = (DistributedPlan *) copyObject(node);
+	distributedPlan = (DistributedPlan *) node;
 
 	return distributedPlan;
 }