From 72b580e35e91c4fc79a057330c50365eb9e5f1ab Mon Sep 17 00:00:00 2001
From: Daniil Anisimov <d.anisimov@postgrespro.ru>
Date: Thu, 25 May 2023 16:17:11 +0700
Subject: [PATCH] Try reducing the memory overhead. Free some allocated memory
 right after use. Reset AQOPredictMemCtx as soon as posible. Remove learning
 attempts on SubPlan nodes. Bugfix. Free allocated memory on save/load data.
 Add memory context for storage. Change copyright to 2016-2023.

---
 aqo.c                    | 11 +++++++-
 aqo.h                    |  3 +-
 auto_tuning.c            |  2 +-
 cardinality_estimation.c |  2 +-
 cardinality_hooks.c      | 16 +++++++++--
 expected/unsupported.out | 61 ++++++++++++++++++++++++++++++++++++++--
 hash.c                   | 30 ++++++++++++++++----
 machine_learning.c       |  2 +-
 path_utils.c             | 37 +++++++++++-------------
 postprocessing.c         | 10 +++++--
 preprocessing.c          |  2 +-
 sql/unsupported.sql      | 10 +++++++
 storage.c                | 46 ++++++++++++++++++++++--------
 13 files changed, 183 insertions(+), 49 deletions(-)

diff --git a/aqo.c b/aqo.c
index c436f9f7..4cfe0ee4 100644
--- a/aqo.c
+++ b/aqo.c
@@ -2,7 +2,7 @@
  * aqo.c
  *		Adaptive query optimization extension
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/aqo.c
@@ -92,6 +92,9 @@ MemoryContext 		AQOPredictMemCtx = NULL;
 /* Is released at the end of learning */
 MemoryContext 		AQOLearnMemCtx = NULL;
 
+/* Is released at the end of load/store routines */
+MemoryContext 		AQOStorageMemCtx = NULL;
+
 /* Additional plan info */
 int njoins;
 
@@ -348,6 +351,12 @@ _PG_init(void)
 	AQOLearnMemCtx = AllocSetContextCreate(AQOTopMemCtx,
 											 "AQOLearnMemoryContext",
 											 ALLOCSET_DEFAULT_SIZES);
+	/*
+	 * AQOStorageMemoryContext containe data for load/store routines.
+	 */
+	AQOStorageMemCtx = AllocSetContextCreate(AQOTopMemCtx,
+											 "AQOStorageMemoryContext",
+											 ALLOCSET_DEFAULT_SIZES);
 	RegisterResourceReleaseCallback(aqo_free_callback, NULL);
 	RegisterAQOPlanNodeMethods();
 
diff --git a/aqo.h b/aqo.h
index 85c3f3b2..f3275003 100644
--- a/aqo.h
+++ b/aqo.h
@@ -105,7 +105,7 @@
  * Module storage.c is responsible for storage query settings and models
  * (i. e. all information which is used in extension).
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/aqo.h
@@ -232,6 +232,7 @@ extern MemoryContext AQOTopMemCtx;
 extern MemoryContext AQOCacheMemCtx;
 extern MemoryContext AQOPredictMemCtx;
 extern MemoryContext AQOLearnMemCtx;
+extern MemoryContext AQOStorageMemCtx;
 
 extern int aqo_statement_timeout;
 
diff --git a/auto_tuning.c b/auto_tuning.c
index 22e9b4dc..e6f5db83 100644
--- a/auto_tuning.c
+++ b/auto_tuning.c
@@ -8,7 +8,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/auto_tuning.c
diff --git a/cardinality_estimation.c b/cardinality_estimation.c
index f93e0905..8ab98f3c 100644
--- a/cardinality_estimation.c
+++ b/cardinality_estimation.c
@@ -8,7 +8,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/cardinality_estimation.c
diff --git a/cardinality_hooks.c b/cardinality_hooks.c
index f0d745bb..fd2f970c 100644
--- a/cardinality_hooks.c
+++ b/cardinality_hooks.c
@@ -18,7 +18,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/cardinality_hooks.c
@@ -82,6 +82,7 @@ aqo_set_baserel_rows_estimate(PlannerInfo *root, RelOptInfo *rel)
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -100,6 +101,7 @@ aqo_set_baserel_rows_estimate(PlannerInfo *root, RelOptInfo *rel)
 
 	/* Return to the caller's memory context. */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	if (predicted < 0)
 		goto default_estimator;
@@ -191,12 +193,15 @@ aqo_get_parameterized_baserel_size(PlannerInfo *root,
 			cache_selectivity(current_hash, rel->relid, rte->relid,
 							  *((double *) lfirst(l2)));
 		}
+
+		pfree(args_hash);
+		pfree(eclass_hash);
 	}
 
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(oldctx);
-
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -211,6 +216,7 @@ aqo_get_parameterized_baserel_size(PlannerInfo *root,
 
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(oldctx);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	predicted_ppi_rows = predicted;
 	fss_ppi_hash = fss;
@@ -265,6 +271,7 @@ aqo_set_joinrel_size_estimates(PlannerInfo *root, RelOptInfo *rel,
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -284,6 +291,7 @@ aqo_set_joinrel_size_estimates(PlannerInfo *root, RelOptInfo *rel,
 
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	rel->fss_hash = fss;
 
@@ -343,6 +351,7 @@ aqo_get_parameterized_joinrel_size(PlannerInfo *root,
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -359,6 +368,7 @@ aqo_get_parameterized_joinrel_size(PlannerInfo *root,
 									 &fss);
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	predicted_ppi_rows = predicted;
 	fss_ppi_hash = fss;
@@ -450,6 +460,7 @@ aqo_estimate_num_groups(PlannerInfo *root, List *groupExprs,
 		grouped_rel->rows = predicted;
 		grouped_rel->fss_hash = fss;
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		return predicted;
 	}
 	else
@@ -460,6 +471,7 @@ aqo_estimate_num_groups(PlannerInfo *root, List *groupExprs,
 		grouped_rel->predicted_cardinality = -1;
 
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 default_estimator:
 	if (aqo_estimate_num_groups_next)
diff --git a/expected/unsupported.out b/expected/unsupported.out
index 6e45dcd8..9db07618 100644
--- a/expected/unsupported.out
+++ b/expected/unsupported.out
@@ -311,6 +311,59 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
  JOINS: 0
 (23 rows)
 
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+                        QUERY PLAN                         
+-----------------------------------------------------------
+ Seq Scan on t (actual rows=1000 loops=1)
+   AQO not used
+   Filter: ((x = (SubPlan 1)) AND (SubPlan 2))
+   SubPlan 1
+     ->  Limit (actual rows=1 loops=1000)
+           AQO not used
+           ->  Seq Scan on t t0 (actual rows=1 loops=1000)
+                 AQO not used
+                 Filter: (x = t.x)
+                 Rows Removed by Filter: 475
+   SubPlan 2
+     ->  Seq Scan on t t0_1 (actual rows=1 loops=1000)
+           AQO not used
+           Filter: (x = t.x)
+           Rows Removed by Filter: 475
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(18 rows)
+
+-- No prediction for top SeqScan, because it fss is changed
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+                        QUERY PLAN                         
+-----------------------------------------------------------
+ Seq Scan on t (actual rows=1000 loops=1)
+   AQO not used
+   Filter: ((SubPlan 2) AND (x = (SubPlan 1)))
+   SubPlan 2
+     ->  Seq Scan on t t0_1 (actual rows=1 loops=1000)
+           AQO: rows=1, error=0%
+           Filter: (x = t.x)
+           Rows Removed by Filter: 475
+   SubPlan 1
+     ->  Limit (actual rows=1 loops=1000)
+           AQO not used
+           ->  Seq Scan on t t0 (actual rows=1 loops=1000)
+                 AQO: rows=1, error=0%
+                 Filter: (x = t.x)
+                 Rows Removed by Filter: 475
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(18 rows)
+
 -- It's OK to use the knowledge for a query with different constants.
 EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 	SELECT count(*) FROM t WHERE
@@ -580,6 +633,10 @@ ORDER BY (md5(query_text),error) DESC;
 -------+------------------------------------------------------------------------------------------------
  0.768 | SELECT count(*) FROM (SELECT count(*) FROM t1 GROUP BY (x,y)) AS q1;
  0.070 | SELECT count(*) FROM (SELECT * FROM t GROUP BY (x) HAVING x > 3) AS q1;
+ 1.554 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)                                         +
+       |         SELECT * FROM t WHERE                                                                 +
+       |                         x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND                 +
+       |                         x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
  0.000 | SELECT count(*) FROM t WHERE x < 3 AND mod(x,3) = 1;
  0.000 | SELECT * FROM                                                                                 +
        |         (SELECT * FROM t WHERE x < 0) AS t0                                                   +
@@ -612,13 +669,13 @@ ORDER BY (md5(query_text),error) DESC;
        |                 JOIN                                                                          +
        |         (SELECT * FROM t WHERE x % 3 < (SELECT avg(x) FROM t t0 WHERE t0.x <> t.x)) AS q2     +
        |                 ON q1.x = q2.x+1;
-(13 rows)
+(14 rows)
 
 DROP TABLE t,t1 CASCADE; -- delete all tables used in the test
 SELECT count(*) FROM aqo_data; -- Just to detect some changes in the logic. May some false positives really bother us here?
  count 
 -------
-    44
+    48
 (1 row)
 
 SELECT true AS success FROM aqo_cleanup();
diff --git a/hash.c b/hash.c
index e24d405c..dfb4a55c 100644
--- a/hash.c
+++ b/hash.c
@@ -12,7 +12,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/hash.c
@@ -157,6 +157,8 @@ get_grouped_exprs_hash(int child_fss, List *group_exprs)
 	final_hashes[0] = child_fss;
 	final_hashes[1] = get_int_array_hash(hashes, i);
 
+	pfree(hashes);
+
 	return get_int_array_hash(final_hashes, 2);
 }
 
@@ -224,6 +226,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		clause_has_consts[i] = (args != NULL && has_consts(*args));
 		i++;
 	}
+	pfree(args_hash);
 
 	idx = argsort(clause_hashes, n, sizeof(*clause_hashes), int_cmp);
 	inverse_idx = inverse_permutation(idx, n);
@@ -234,6 +237,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		sorted_clauses[inverse_idx[i]] = clause_hashes[i];
 		i++;
 	}
+	pfree(clause_hashes);
 
 	i = 0;
 	foreach(lc, selectivities)
@@ -249,6 +253,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		}
 		i++;
 	}
+	pfree(inverse_idx);
 
 	for (i = 0; i < n;)
 	{
@@ -272,6 +277,8 @@ get_fss_for_object(List *relsigns, List *clauselist,
 				  sizeof(**features), double_cmp);
 		i = j;
 	}
+	pfree(idx);
+	pfree(clause_has_consts);
 
 	/*
 	 * Generate feature subspace hash.
@@ -281,6 +288,8 @@ get_fss_for_object(List *relsigns, List *clauselist,
 	eclasses_hash = get_int_array_hash(eclass_hash, nargs);
 	relations_hash = get_relations_hash(relsigns);
 	fss_hash = get_fss_hash(clauses_hash, eclasses_hash, relations_hash);
+	pfree(sorted_clauses);
+	pfree(eclass_hash);
 
 	if (nfeatures != NULL)
 	{
@@ -340,11 +349,17 @@ static int
 get_node_hash(Node *node)
 {
 	char	   *str;
+	char	   *no_consts;
+	char	   *no_locations;
 	int			hash;
 
-	str = remove_locations(remove_consts(nodeToString(node)));
-	hash = get_str_hash(str);
+	str = nodeToString(node);
+	no_consts = remove_consts(str);
 	pfree(str);
+	no_locations = remove_locations(no_consts);
+	pfree(no_consts);
+	hash = get_str_hash(no_locations);
+	pfree(no_locations);
 	return hash;
 }
 
@@ -467,6 +482,7 @@ get_relations_hash(List *relsigns)
 
 	result = DatumGetInt32(hash_any((const unsigned char *) hashes,
 									nhashes * sizeof(uint32)));
+	pfree(hashes);
 
 	return result;
 }
@@ -479,9 +495,11 @@ static char *
 remove_consts(const char *str)
 {
 	char *res;
+	char *tmp;
 
-	res = replace_patterns(str, "{CONST", is_brace);
-	res = replace_patterns(res, ":stmt_len", is_brace);
+	tmp = replace_patterns(str, "{CONST", is_brace);
+	res = replace_patterns(tmp, ":stmt_len", is_brace);
+	pfree(tmp);
 	return res;
 }
 
@@ -683,6 +701,8 @@ get_eclasses(List *clauselist, int *nargs, int **args_hash, int **eclass_hash)
 
 	for (i = 0; i < *nargs; ++i)
 		(*eclass_hash)[i] = e_hashes[disjoint_set_get_parent(p, i)];
+
+	pfree(e_hashes);
 }
 
 /*
diff --git a/machine_learning.c b/machine_learning.c
index d4f5cbee..bfdf0aaa 100644
--- a/machine_learning.c
+++ b/machine_learning.c
@@ -12,7 +12,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/machine_learning.c
diff --git a/path_utils.c b/path_utils.c
index 5a34b645..7617bfd8 100644
--- a/path_utils.c
+++ b/path_utils.c
@@ -5,7 +5,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/path_utils.c
@@ -51,7 +51,7 @@ static AQOPlanNode DefaultAQOPlanNode =
  */
 static create_plan_hook_type			aqo_create_plan_next					= NULL;
 
-static create_upper_paths_hook_type	aqo_create_upper_paths_next				= NULL;
+/*static create_upper_paths_hook_type	aqo_create_upper_paths_next				= NULL;*/
 
 
 static AQOPlanNode *
@@ -260,7 +260,7 @@ get_list_of_relids(PlannerInfo *root, Relids relids, RelSortOut *rels)
 
 /*
  * Search for any subplans or initplans.
- * if subplan is found, replace it by the feature space value of this subplan.
+ * if subplan is found, replace it by zero Const.
  */
 static Node *
 subplan_hunter(Node *node, void *context)
@@ -271,21 +271,13 @@ subplan_hunter(Node *node, void *context)
 
 	if (IsA(node, SubPlan))
 	{
-		SubPlan		*splan = (SubPlan *) node;
-		PlannerInfo	*root = (PlannerInfo *) context;
-		PlannerInfo	*subroot;
-		RelOptInfo	*upper_rel;
-		A_Const		*fss;
+		A_Const	*fss = makeNode(A_Const);
 
-		subroot = (PlannerInfo *) list_nth(root->glob->subroots,
-										   splan->plan_id - 1);
-		upper_rel = fetch_upper_rel(subroot, UPPERREL_FINAL, NULL);
+		fss->val.ival.type = T_Integer;
+		fss->location = -1;
+		fss->val.ival.ival = 0;
+		return (Node *) fss;
 
-		Assert(list_length(upper_rel->ext_nodes) == 1);
-		Assert(IsA((Node *) linitial(upper_rel->ext_nodes), A_Const));
-
-		fss = (A_Const *) linitial(upper_rel->ext_nodes);
-		return (Node *) copyObject(fss);
 	}
 	return expression_tree_mutator(node, subplan_hunter, context);
 }
@@ -766,11 +758,14 @@ RegisterAQOPlanNodeMethods(void)
 }
 
 /*
+ * Warning! This function does not word properly.
+ * Because value of Const nodes removed by hash routine.
+ *
  * Hook for create_upper_paths_hook
  *
  * Assume, that we are last in the chain of path creators.
  */
-static void
+/*static void
 aqo_store_upper_signature(PlannerInfo *root,
 						  UpperRelationKind stage,
 						  RelOptInfo *input_rel,
@@ -786,7 +781,7 @@ aqo_store_upper_signature(PlannerInfo *root,
 		(*aqo_create_upper_paths_next)(root, stage, input_rel, output_rel, extra);
 
 	if (!query_context.use_aqo && !query_context.learn_aqo && !force_collect_stat)
-		/* Includes 'disabled query' state. */
+		/ * Includes 'disabled query' state. * /
 		return;
 
 	if (stage != UPPERREL_FINAL)
@@ -801,7 +796,7 @@ aqo_store_upper_signature(PlannerInfo *root,
 	fss_node->val.ival.ival = get_fss_for_object(rels.signatures, clauses, NIL,
 												NULL, NULL);
 	output_rel->ext_nodes = lappend(output_rel->ext_nodes, (void *) fss_node);
-}
+}*/
 
 void
 aqo_path_utils_init(void)
@@ -809,6 +804,6 @@ aqo_path_utils_init(void)
 	aqo_create_plan_next				= create_plan_hook;
 	create_plan_hook					= aqo_create_plan;
 
-	aqo_create_upper_paths_next			= create_upper_paths_hook;
-	create_upper_paths_hook				= aqo_store_upper_signature;
+	/*aqo_create_upper_paths_next			= create_upper_paths_hook;
+	create_upper_paths_hook				= aqo_store_upper_signature;*/
 }
diff --git a/postprocessing.c b/postprocessing.c
index 66aca901..a6b6d030 100644
--- a/postprocessing.c
+++ b/postprocessing.c
@@ -9,7 +9,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/postprocessing.c
@@ -224,6 +224,12 @@ restore_selectivities(List *clauselist, List *relidslist, JoinType join_type,
 		lst = lappend(lst, cur_sel);
 	}
 
+	if (parametrized_sel)
+	{
+		pfree(args_hash);
+		pfree(eclass_hash);
+	}
+
 	return lst;
 }
 
@@ -833,11 +839,11 @@ aqo_ExecutorEnd(QueryDesc *queryDesc)
 		}
 	}
 
-	selectivity_cache_clear();
 	cur_classes = ldelete_uint64(cur_classes, query_context.query_hash);
 
 end:
 	/* Release all AQO-specific memory, allocated during learning procedure */
+	selectivity_cache_clear();
 	MemoryContextSwitchTo(oldctx);
 	MemoryContextReset(AQOLearnMemCtx);
 
diff --git a/preprocessing.c b/preprocessing.c
index ef41ab0e..feb28d39 100644
--- a/preprocessing.c
+++ b/preprocessing.c
@@ -49,7 +49,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/preprocessing.c
diff --git a/sql/unsupported.sql b/sql/unsupported.sql
index 8b36d721..e5853306 100644
--- a/sql/unsupported.sql
+++ b/sql/unsupported.sql
@@ -98,6 +98,16 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 		x = (SELECT avg(x) FROM t t0 WHERE t0.x = t.x + 21) OR
 		x IN (SELECT avg(x) FROM t t0 WHERE t0.x = t.x + 21);
 
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+-- No prediction for top SeqScan, because it fss is changed
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+
 -- It's OK to use the knowledge for a query with different constants.
 EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 	SELECT count(*) FROM t WHERE
diff --git a/storage.c b/storage.c
index bf004199..f71f5207 100644
--- a/storage.c
+++ b/storage.c
@@ -8,7 +8,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/storage.c
@@ -666,11 +666,12 @@ static int
 data_store(const char *filename, form_record_t callback,
 		   long nrecs, void *ctx)
 {
-	FILE   *file;
-	size_t	size;
-	uint32	counter = 0;
-	void   *data;
-	char   *tmpfile;
+	FILE		   *file;
+	size_t			size;
+	uint32			counter = 0;
+	void		   *data;
+	char		   *tmpfile;
+	MemoryContext	old_context = MemoryContextSwitchTo(AQOStorageMemCtx);
 
 	tmpfile = psprintf("%s.tmp", filename);
 	file = AllocateFile(tmpfile, PG_BINARY_W);
@@ -687,7 +688,11 @@ data_store(const char *filename, form_record_t callback,
 		/* TODO: Add CRC code ? */
 		if (fwrite(&size, sizeof(size), 1, file) != 1 ||
 			fwrite(data, size, 1, file) != 1)
+		{
+			pfree(data);
 			goto error;
+		}
+		pfree(data);
 		counter++;
 	}
 
@@ -701,6 +706,9 @@ data_store(const char *filename, form_record_t callback,
 	/* Parallel (re)writing into a file haven't happen. */
 	(void) durable_rename(tmpfile, filename, PANIC);
 	elog(LOG, "[AQO] %d records stored in file %s.", counter, filename);
+
+	MemoryContextSwitchTo(old_context);
+	MemoryContextReset(AQOStorageMemCtx);
 	return 0;
 
 error:
@@ -712,6 +720,9 @@ data_store(const char *filename, form_record_t callback,
 		FreeFile(file);
 	unlink(tmpfile);
 	pfree(tmpfile);
+
+	MemoryContextSwitchTo(old_context);
+	MemoryContextReset(AQOStorageMemCtx);
 	return -1;
 }
 
@@ -936,17 +947,20 @@ aqo_queries_load(void)
 static void
 data_load(const char *filename, deform_record_t callback, void *ctx)
 {
-	FILE   *file;
-	long	i;
-	uint32	header;
-	int32	pgver;
-	long	num;
+	FILE		   *file;
+	long			i;
+	uint32			header;
+	int32			pgver;
+	long			num;
+	MemoryContext	old_context = MemoryContextSwitchTo(AQOStorageMemCtx);
 
 	file = AllocateFile(filename, PG_BINARY_R);
 	if (file == NULL)
 	{
 		if (errno != ENOENT)
 			goto read_error;
+
+		MemoryContextSwitchTo(old_context);
 		return;
 	}
 
@@ -968,8 +982,12 @@ data_load(const char *filename, deform_record_t callback, void *ctx)
 			goto read_error;
 		data = palloc(size);
 		if (fread(data, size, 1, file) != 1)
+		{
+			pfree(data);
 			goto read_error;
+		}
 		res = callback(data, size);
+		pfree(data);
 
 		if (!res)
 		{
@@ -983,6 +1001,9 @@ data_load(const char *filename, deform_record_t callback, void *ctx)
 	FreeFile(file);
 
 	elog(LOG, "[AQO] %ld records loaded from file %s.", num, filename);
+
+	MemoryContextSwitchTo(old_context);
+	MemoryContextReset(AQOStorageMemCtx);
 	return;
 
 read_error:
@@ -998,6 +1019,9 @@ data_load(const char *filename, deform_record_t callback, void *ctx)
 	if (file)
 		FreeFile(file);
 	unlink(filename);
+
+	MemoryContextSwitchTo(old_context);
+	MemoryContextReset(AQOStorageMemCtx);
 }
 
 static void
