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

Alena Rybakina · Alena Rybakina · commit 10745e30f3eb · 2022-09-08T13:43:44.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
@@ -334,7 +334,7 @@ _PG_init(void)
 	shmem_request_hook = aqo_shmem_request;
 
 	init_deactivated_queries_storage();
-	
+
 	/*
 	 * Create own Top memory Context for reporting AQO memory in the future.
 	 */
diff --git a/postprocessing.c b/postprocessing.c
@@ -617,6 +617,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))
@@ -629,6 +630,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
@@ -196,7 +196,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);
@@ -497,6 +497,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)
@@ -538,6 +539,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
@@ -2042,6 +2042,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 */
@@ -2050,6 +2051,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);
 				}
@@ -2100,8 +2102,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
`@@ -617,6 +617,7 @@ static int exec_nested_level = 0;`
`617`	`617`	`static void`
`618`	`618`	`aqo_timeout_handler(void)`
`619`	`619`	`{`
	`620`	`+MemoryContext oldctx = MemoryContextSwitchTo(AQOLearnMemCtx);`
`620`	`621`	`aqo_obj_stat ctx = {NIL, NIL, NIL, false, false};`
`621`	`622`
`622`	`623`	`if (!timeoutCtl.queryDesc \|\| !ExtractFromQueryEnv(timeoutCtl.queryDesc))`
`@@ -629,6 +630,7 @@ aqo_timeout_handler(void)`
`629`	`630`
`630`	`631`	`elog(NOTICE, "[AQO] Time limit for execution of the statement was expired. AQO tried to learn on partial data.");`
`631`	`632`	`learnOnPlanState(timeoutCtl.queryDesc->planstate, (void *) &ctx);`
	`633`	`+MemoryContextSwitchTo(oldctx);`
`632`	`634`	`}`
`633`	`635`
`634`	`636`	`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
`@@ -2042,6 +2042,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2042`	`2042`	`for(i = 0; i < dentry->nrels; i++)`
`2043`	`2043`	`{`
`2044`	`2044`	`Oid reloid = ObjectIdGetDatum((Oid )ptr);`
	`2045`	`+MemoryContext oldctx = MemoryContextSwitchTo(AQOUtilityMemCtx);`
`2045`	`2046`
`2046`	`2047`	`if (!SearchSysCacheExists1(RELOID, reloid))`
`2047`	`2048`	`/* Remember this value */`
`@@ -2050,6 +2051,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2050`	`2051`	`else`
`2051`	`2052`	`actual_fss = list_append_unique_int(actual_fss,`
`2052`	`2053`	`dentry->key.fss);`
	`2054`	`+MemoryContextSwitchTo(oldctx);`
`2053`	`2055`
`2054`	`2056`	`ptr += sizeof(Oid);`
`2055`	`2057`	`}`
`@@ -2100,8 +2102,7 @@ cleanup_aqo_database(bool gentle, int fs_num, int fss_num)`
`2100`	`2102`	`(*fs_num) += (int) _aqo_queries_remove(entry->queryid);`
`2101`	`2103`	`}`
`2102`	`2104`
`2103`		`-list_free(junk_fss);`
`2104`		`-list_free(actual_fss);`
	`2105`	`+MemoryContextReset(AQOUtilityMemCtx);`
`2105`	`2106`	`}`
`2106`	`2107`
`2107`	`2108`	`/*`