Bugfix. Choose memory context for a query environment piece in more safe way.

danolivo · danolivo · commit 80bcf2d23b78 · 2022-05-19T17:58:31.000+05:00
Sometimes someone can invent a queryEnv and use it with short-lived plans. So,
anyone under the hood should create its queryEnv in the same memory context.
diff --git a/postprocessing.c b/postprocessing.c
@@ -913,11 +913,20 @@ StoreToQueryEnv(QueryDesc *queryDesc)
 	MemoryContext	oldCxt;
 	bool newentry = false;
 
-	oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc->plannedstmt));
-
-	if (queryDesc->queryEnv == NULL)
-			queryDesc->queryEnv = create_queryEnv();
+	/*
+	 * Choose memory context for AQO parameters. Use pre-existed context if
+	 * someone earlier created queryEnv (usually, SPI), or base on the queryDesc
+	 * memory context.
+	 */
+	if (queryDesc->queryEnv != NULL)
+		oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc->queryEnv));
+	else
+	{
+		oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc));
+		queryDesc->queryEnv = create_queryEnv();
+	}
 
+	Assert(queryDesc->queryEnv);
 	enr = get_ENR(queryDesc->queryEnv, AQOPrivateData);
 	if (enr == NULL)
 	{
@@ -965,11 +974,20 @@ StorePlanInternals(QueryDesc *queryDesc)
 	njoins = 0;
 	planstate_tree_walker(queryDesc->planstate, calculateJoinNum, &njoins);
 
-	oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc->plannedstmt));
-
-	if (queryDesc->queryEnv == NULL)
-			queryDesc->queryEnv = create_queryEnv();
+	/*
+	 * Choose memory context for AQO parameters. Use pre-existed context if
+	 * someone earlier created queryEnv (usually, SPI), or base on the queryDesc
+	 * memory context.
+	 */
+	if (queryDesc->queryEnv != NULL)
+		oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc->queryEnv));
+	else
+	{
+		oldCxt = MemoryContextSwitchTo(GetMemoryChunkContext(queryDesc));
+		queryDesc->queryEnv = create_queryEnv();
+	}
 
+	Assert(queryDesc->queryEnv);
 	enr = get_ENR(queryDesc->queryEnv, PlanStateInfo);
 	if (enr == NULL)
 	{
