Add memory context to cover memory space when applying aqo_timeout handler and applying

Alena Rybakina · Alena Rybakina · commit c297fec3439c · 2022-09-08T13:57:47.000+03:00
isQueryUsingSystemRelation_walker and add UtilityMemCtx for allocation space for junk_fss
and actual_fss list and reset it after cleaning aqo database process.
diff --git a/aqo.c b/aqo.c
@@ -316,7 +316,7 @@ _PG_init(void)
 	create_upper_paths_hook						= aqo_store_upper_signature_hook;
 
 	init_deactivated_queries_storage();
-	
+
 	/*
 	 * Create own Top memory Context for reporting AQO memory in the future.
 	 */
diff --git a/postprocessing.c b/postprocessing.c
@@ -615,6 +615,7 @@ static int exec_nested_level = 0;
 static void
 aqo_timeout_handler(void)
 {
+	MemoryContext oldctx = MemoryContextSwitchTo(AQOLearnMemCtx);
 	aqo_obj_stat ctx = {NIL, NIL, NIL, false, false};
 
 	if (!timeoutCtl.queryDesc || !ExtractFromQueryEnv(timeoutCtl.queryDesc))
@@ -627,6 +628,7 @@ aqo_timeout_handler(void)
 
 	elog(NOTICE, "[AQO] Time limit for execution of the statement was expired. AQO tried to learn on partial data.");
 	learnOnPlanState(timeoutCtl.queryDesc->planstate, (void *) &ctx);
+	MemoryContextSwitchTo(oldctx);
 }
 
 static bool
diff --git a/preprocessing.c b/preprocessing.c
@@ -157,7 +157,7 @@ aqo_planner(Query *parse,
 
 	selectivity_cache_clear();
 	MemoryContextSwitchTo(oldctx);
- 
+
 	oldctx = MemoryContextSwitchTo(AQOUtilityMemCtx);
 	query_context.query_hash = get_query_hash(parse, query_string);
 	MemoryContextSwitchTo(oldctx);
@@ -458,6 +458,7 @@ jointree_walker(Node *jtnode, void *context)
 static bool
 isQueryUsingSystemRelation_walker(Node *node, void *context)
 {
+	MemoryContext oldctx = MemoryContextSwitchTo(AQOLearnMemCtx);
 	AQOPreWalkerCtx   *ctx = (AQOPreWalkerCtx *) context;
 
 	if (node == NULL)
@@ -499,6 +500,7 @@ isQueryUsingSystemRelation_walker(Node *node, void *context)
 		}
 
 		jointree_walker((Node *) query->jointree, context);
+		MemoryContextSwitchTo(oldctx);
 
 		/* Recursively plunge into subqueries and CTEs */
 		return query_tree_walker(query,
diff --git a/selectivity_cache.c b/selectivity_cache.c
@@ -51,7 +51,7 @@ cache_selectivity(int clause_hash,
 		AQOCacheSelectivity = AllocSetContextCreate(AQOTopMemCtx,
 													"AQOCacheSelectivity",
 													ALLOCSET_DEFAULT_SIZES);
- 
+
 
 	foreach(l, objects)
 	{
diff --git a/storage.c b/storage.c
@@ -2026,6 +2026,7 @@ cleanup_aqo_database(bool gentle, int *fs_num, int *fss_num)
 				for(i = 0; i < dentry->nrels; i++)
 				{
 					Oid reloid = ObjectIdGetDatum(*(Oid *)ptr);
+					MemoryContext oldctx = MemoryContextSwitchTo(AQOUtilityMemCtx);
 
 					if (!SearchSysCacheExists1(RELOID, reloid))
 						/* Remember this value */
@@ -2034,6 +2035,7 @@ cleanup_aqo_database(bool gentle, int *fs_num, int *fss_num)
 					else
 						actual_fss = list_append_unique_int(actual_fss,
 															dentry->key.fss);
+					MemoryContextSwitchTo(oldctx);
 
 					ptr += sizeof(Oid);
 				}
@@ -2083,8 +2085,7 @@ cleanup_aqo_database(bool gentle, int *fs_num, int *fss_num)
 			(*fs_num) += (int) _aqo_queries_remove(entry->queryid);
 		}
 
-		list_free(junk_fss);
-		list_free(actual_fss);
+		MemoryContextReset(AQOUtilityMemCtx);
 	}
 
 	/*

Original file line number	Diff line number	Diff line change
`@@ -615,6 +615,7 @@ static int exec_nested_level = 0;`
`615`	`615`	`static void`
`616`	`616`	`aqo_timeout_handler(void)`
`617`	`617`	`{`
	`618`	`+ MemoryContext oldctx = MemoryContextSwitchTo(AQOLearnMemCtx);`
`618`	`619`	`aqo_obj_stat ctx = {NIL, NIL, NIL, false, false};`
`619`	`620`
`620`	`621`	`if (!timeoutCtl.queryDesc \|\| !ExtractFromQueryEnv(timeoutCtl.queryDesc))`
`@@ -627,6 +628,7 @@ aqo_timeout_handler(void)`
`627`	`628`
`628`	`629`	`elog(NOTICE, "[AQO] Time limit for execution of the statement was expired. AQO tried to learn on partial data.");`
`629`	`630`	`learnOnPlanState(timeoutCtl.queryDesc->planstate, (void *) &ctx);`
	`631`	`+ MemoryContextSwitchTo(oldctx);`
`630`	`632`	`}`
`631`	`633`
`632`	`634`	`static bool`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ cache_selectivity(int clause_hash,`
`51`	`51`	`AQOCacheSelectivity = AllocSetContextCreate(AQOTopMemCtx,`
`52`	`52`	`"AQOCacheSelectivity",`
`53`	`53`	`ALLOCSET_DEFAULT_SIZES);`
`54`		`-`
	`54`	`+`
`55`	`55`
`56`	`56`	`foreach(l, objects)`
`57`	`57`	`{`
Original file line number	Diff line number	Diff line change
`@@ -2026,6 +2026,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2026`	`2026`	`for(i = 0; i < dentry->nrels; i++)`
`2027`	`2027`	`{`
`2028`	`2028`	`Oid reloid = ObjectIdGetDatum((Oid )ptr);`
	`2029`	`+ MemoryContext oldctx = MemoryContextSwitchTo(AQOUtilityMemCtx);`
`2029`	`2030`
`2030`	`2031`	`if (!SearchSysCacheExists1(RELOID, reloid))`
`2031`	`2032`	`/* Remember this value */`
`@@ -2034,6 +2035,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2034`	`2035`	`else`
`2035`	`2036`	`actual_fss = list_append_unique_int(actual_fss,`
`2036`	`2037`	`dentry->key.fss);`
	`2038`	`+ MemoryContextSwitchTo(oldctx);`
`2037`	`2039`
`2038`	`2040`	`ptr += sizeof(Oid);`
`2039`	`2041`	`}`
`@@ -2083,8 +2085,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2083`	`2085`	`(*fs_num) += (int) _aqo_queries_remove(entry->queryid);`
`2084`	`2086`	`}`
`2085`	`2087`
`2086`		`- list_free(junk_fss);`
`2087`		`- list_free(actual_fss);`
	`2088`	`+ MemoryContextReset(AQOUtilityMemCtx);`
`2088`	`2089`	`}`
`2089`	`2090`
`2090`	`2091`	`/*`