Merge pull request #1349 from citusdata/fix_check_vanilla

Fix check-vanilla tests

src/backend/distributed/planner/multi_planner.c

@@ -72,8 +72,7 @@ static PlannedStmt * FinalizeRouterPlan(PlannedStmt *localPlan, CustomScan *cust
 static void CheckNodeIsDumpable(Node *node);
 static List * CopyPlanParamList(List *originalPlanParamList);
 static PlannerRestrictionContext * CreateAndPushPlannerRestrictionContext(void);
-static RelationRestrictionContext * CurrentRelationRestrictionContext(void);
+static PlannerRestrictionContext * CurrentPlannerRestrictionContext(void);
-static JoinRestrictionContext * CurrentJoinRestrictionContext(void);
 static void PopPlannerRestrictionContext(void);
 static bool HasUnresolvedExternParamsWalker(Node *expression, ParamListInfo boundParams);
 
@@ -658,20 +657,40 @@ multi_join_restriction_hook(PlannerInfo *root,
 							JoinType jointype,
 							JoinPathExtraData *extra)
 {
-	JoinRestrictionContext *joinContext = NULL;
+	PlannerRestrictionContext *plannerRestrictionContext = NULL;
-	JoinRestriction *joinRestriction = palloc0(sizeof(JoinRestriction));
+	JoinRestrictionContext *joinRestrictionContext = NULL;
+	JoinRestriction *joinRestriction = NULL;
+	MemoryContext restrictionsMemoryContext = NULL;
+	MemoryContext oldMemoryContext = NULL;
 	List *restrictInfoList = NIL;
 
-	restrictInfoList = extra->restrictlist;
+	/*
-	joinContext = CurrentJoinRestrictionContext();
+	 * Use a memory context that's guaranteed to live long enough, could be
-	Assert(joinContext != NULL);
+	 * called in a more shorted lived one (e.g. with GEQO).
+	 */
+	plannerRestrictionContext = CurrentPlannerRestrictionContext();
+	restrictionsMemoryContext = plannerRestrictionContext->memoryContext;
+	oldMemoryContext = MemoryContextSwitchTo(restrictionsMemoryContext);
 
+	/*
+	 * We create a copy of restrictInfoList because it may be created in a memory
+	 * context which will be deleted when we still need it, thus we create a copy
+	 * of it in our memory context.
+	 */
+	restrictInfoList = copyObject(extra->restrictlist);
+
+	joinRestrictionContext = plannerRestrictionContext->joinRestrictionContext;
+	Assert(joinRestrictionContext != NULL);
+
+	joinRestriction = palloc0(sizeof(JoinRestriction));
 	joinRestriction->joinType = jointype;
 	joinRestriction->joinRestrictInfoList = restrictInfoList;
 	joinRestriction->plannerInfo = root;
 
-	joinContext->joinRestrictionList =
+	joinRestrictionContext->joinRestrictionList =
-		lappend(joinContext->joinRestrictionList, joinRestriction);
+		lappend(joinRestrictionContext->joinRestrictionList, joinRestriction);
+
+	MemoryContextSwitchTo(oldMemoryContext);
 }
 
 
@@ -684,7 +703,10 @@ void
 multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index index,
 								RangeTblEntry *rte)
 {
-	RelationRestrictionContext *restrictionContext = NULL;
+	PlannerRestrictionContext *plannerRestrictionContext = NULL;
+	RelationRestrictionContext *relationRestrictionContext = NULL;
+	MemoryContext restrictionsMemoryContext = NULL;
+	MemoryContext oldMemoryContext = NULL;
 	RelationRestriction *relationRestriction = NULL;
 	DistTableCacheEntry *cacheEntry = NULL;
 	bool distributedTable = false;
@@ -695,12 +717,17 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index
 		return;
 	}
 
+	/*
+	 * Use a memory context that's guaranteed to live long enough, could be
+	 * called in a more shorted lived one (e.g. with GEQO).
+	 */
+	plannerRestrictionContext = CurrentPlannerRestrictionContext();
+	restrictionsMemoryContext = plannerRestrictionContext->memoryContext;
+	oldMemoryContext = MemoryContextSwitchTo(restrictionsMemoryContext);
+
 	distributedTable = IsDistributedTable(rte->relid);
 	localTable = !distributedTable;
 
-	restrictionContext = CurrentRelationRestrictionContext();
-	Assert(restrictionContext != NULL);
-
 	relationRestriction = palloc0(sizeof(RelationRestriction));
 	relationRestriction->index = index;
 	relationRestriction->relationId = rte->relid;
@@ -718,8 +745,9 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index
 			CopyPlanParamList(root->parent_root->plan_params);
 	}
 
-	restrictionContext->hasDistributedRelation |= distributedTable;
+	relationRestrictionContext = plannerRestrictionContext->relationRestrictionContext;
-	restrictionContext->hasLocalRelation |= localTable;
+	relationRestrictionContext->hasDistributedRelation |= distributedTable;
+	relationRestrictionContext->hasLocalRelation |= localTable;
 
 	/*
 	 * We're also keeping track of whether all participant
@@ -729,12 +757,14 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index
 	{
 		cacheEntry = DistributedTableCacheEntry(rte->relid);
 
-		restrictionContext->allReferenceTables &=
+		relationRestrictionContext->allReferenceTables &=
 			(cacheEntry->partitionMethod == DISTRIBUTE_BY_NONE);
 	}
 
-	restrictionContext->relationRestrictionList =
+	relationRestrictionContext->relationRestrictionList =
-		lappend(restrictionContext->relationRestrictionList, relationRestriction);
+		lappend(relationRestrictionContext->relationRestrictionList, relationRestriction);
+
+	MemoryContextSwitchTo(oldMemoryContext);
 }
 
 
@@ -784,6 +814,8 @@ CreateAndPushPlannerRestrictionContext(void)
 	plannerRestrictionContext->joinRestrictionContext =
 		palloc0(sizeof(JoinRestrictionContext));
 
+	plannerRestrictionContext->memoryContext = CurrentMemoryContext;
+
 	/* we'll apply logical AND as we add tables */
 	plannerRestrictionContext->relationRestrictionContext->allReferenceTables = true;
 
@@ -795,44 +827,20 @@ CreateAndPushPlannerRestrictionContext(void)
 
 
 /*
- * CurrentRelationRestrictionContext returns the the last restriction context from the
+ * CurrentRestrictionContext returns the the most recently added
- * relationRestrictionContext list.
+ * PlannerRestrictionContext from the plannerRestrictionContextList list.
  */
-static RelationRestrictionContext *
+static PlannerRestrictionContext *
-CurrentRelationRestrictionContext(void)
+CurrentPlannerRestrictionContext(void)
 {
 	PlannerRestrictionContext *plannerRestrictionContext = NULL;
-	RelationRestrictionContext *relationRestrictionContext = NULL;
 
 	Assert(plannerRestrictionContextList != NIL);
 
 	plannerRestrictionContext =
 		(PlannerRestrictionContext *) linitial(plannerRestrictionContextList);
 
-	relationRestrictionContext = plannerRestrictionContext->relationRestrictionContext;
+	return plannerRestrictionContext;
-
-	return relationRestrictionContext;
-}
-
-
-/*
- * CurrentJoinRestrictionContext returns the the last restriction context from the
- * list.
- */
-static JoinRestrictionContext *
-CurrentJoinRestrictionContext(void)
-{
-	PlannerRestrictionContext *plannerRestrictionContext = NULL;
-	JoinRestrictionContext *joinRestrictionContext = NULL;
-
-	Assert(plannerRestrictionContextList != NIL);
-
-	plannerRestrictionContext =
-		(PlannerRestrictionContext *) linitial(plannerRestrictionContextList);
-
-	joinRestrictionContext = plannerRestrictionContext->joinRestrictionContext;
-
-	return joinRestrictionContext;
 }
 
 

src/include/distributed/multi_planner.h

@@ -58,6 +58,7 @@ typedef struct PlannerRestrictionContext
 {
 	RelationRestrictionContext *relationRestrictionContext;
 	JoinRestrictionContext *joinRestrictionContext;
+	MemoryContext memoryContext;
 } PlannerRestrictionContext;
 
 typedef struct RelationShard