Subject: Re: Short-living Memory Context in the optimiser

From: Michael Paquier <[email-masked]>
Date: Wed, 25 Dec 2024 14:30:33 +0900

On Wed, Dec 25, 2024 at 11:56:24AM +0700, Andrei Lepikhov wrote:
> Here is a rebased and slightly commented Tom's patch separated from the
> discussion [1].

Seems to me that there is no patch?
--
Michael

---

From: Andrei Lepikhov <[email-masked]>
Date: Wed, 25 Dec 2024 13:06:52 +0700

> Seems to me that there is no patch?
Pardon me, patch attached.

-- 
regards, Andrei Lepikhov

From 16b438566085e4ceecf358c942075a5cbd4dd7ac Mon Sep 17 00:00:00 2001
From: "Andrei V. Lepikhov" <[email-masked]>
Date: Wed, 25 Dec 2024 13:02:56 +0700
Subject: [PATCH v0] Introduce planner temporary memory context.

This memory context is dedicated to cutting off memory consumption peaks in
accidentally scaling operations, like huge array estimations or massive clause
estimations caused by a parent expansion.

It includes 'depth' and 'usage' parameters. The first one is used to determine
whether it is really safe to reset the context or whether it is still used
somewhere in the call stack.
Another one counts the number of operations it was involved in. 'Operation'
here means not a memory allocation call but a logical operation, like a single
clause estimation. So, after the end of the operation, If 'depth' = 0,
the optimiser can clean up used memory if the counter is big enough.

We can't track that cleaning memory isn't referenced by an external object.
So, it is the developer's responsibility to put any product of this context
usage in a longer-living memory context.

The first obvious usage of this context is selectivity estimation because
the operation's result is just a floating-point value, and only cached numbers
in the RestrictInfo structure created outside the context may be a by-product
of this operation.
---
 src/backend/optimizer/path/clausesel.c    | 31 +++++++++++++++++++
 src/backend/optimizer/plan/planner.c      | 37 +++++++++++++++++++++--
 src/backend/optimizer/prep/prepjointree.c |  2 +-
 src/include/nodes/pathnodes.h             | 21 ++++++++++++-
 4 files changed, 86 insertions(+), 5 deletions(-)

diff --git a/src/backend/optimizer/path/clausesel.c b/src/backend/optimizer/path/clausesel.c
index 0ab021c1e8..144bffd5a7 100644
--- a/src/backend/optimizer/path/clausesel.c
+++ b/src/backend/optimizer/path/clausesel.c
@@ -22,6 +22,7 @@
 #include "statistics/statistics.h"
 #include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
+#include "utils/memutils.h"
 #include "utils/selfuncs.h"
 
 /*
@@ -691,6 +692,7 @@ clause_selectivity_ext(PlannerInfo *root,
 	Selectivity s1 = 0.5;		/* default for any unhandled clause type */
 	RestrictInfo *rinfo = NULL;
 	bool		cacheable = false;
+	MemoryContext saved_cxt;
 
 	if (clause == NULL)			/* can this still happen? */
 		return s1;
@@ -754,6 +756,21 @@ clause_selectivity_ext(PlannerInfo *root,
 			clause = (Node *) rinfo->clause;
 	}
 
+	/*
+	 * Run all the remaining work in the short-lived planner_tmp_cxt, which
+	 * we'll reset at the bottom.  This allows selectivity-related code to not
+	 * be too concerned about leaking memory.
+	 */
+	saved_cxt = MemoryContextSwitchTo(root->glob->planner_tmp_cxt);
+
+	/*
+	 * This function can be called recursively, in which case we'd better not
+	 * reset planner_tmp_cxt until we exit the topmost level.  Use of
+	 * planner_tmp_cxt_depth also makes it safe for other places to use and
+	 * reset planner_tmp_cxt in the same fashion.
+	 */
+	root->glob->planner_tmp_cxt_depth++;
+
 	if (IsA(clause, Var))
 	{
 		Var		   *var = (Var *) clause;
@@ -965,6 +982,20 @@ clause_selectivity_ext(PlannerInfo *root,
 			rinfo->outer_selec = s1;
 	}
 
+	/* Exit and optionally clean up the planner_tmp_cxt */
+	MemoryContextSwitchTo(saved_cxt);
+	root->glob->planner_tmp_cxt_usage++;
+	Assert(root->glob->planner_tmp_cxt_depth > 0);
+	if (--root->glob->planner_tmp_cxt_depth == 0)
+	{
+		/* It's safe to reset the tmp context, but do we want to? */
+		if (root->glob->planner_tmp_cxt_usage >= PLANNER_TMP_CXT_USAGE_RESET)
+		{
+			MemoryContextReset(root->glob->planner_tmp_cxt);
+			root->glob->planner_tmp_cxt_usage = 0;
+		}
+	}
+
 #ifdef SELECTIVITY_DEBUG
 	elog(DEBUG4, "clause_selectivity: s1 %f", s1);
 #endif							/* SELECTIVITY_DEBUG */
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 7468961b01..9f2c5e2134 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -316,6 +316,12 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	glob = makeNode(PlannerGlobal);
 
 	glob->boundParams = boundParams;
+	glob->planner_cxt = CurrentMemoryContext;
+	glob->planner_tmp_cxt = AllocSetContextCreate(glob->planner_cxt,
+												  "Planner temp context",
+												  ALLOCSET_DEFAULT_SIZES);
+	glob->planner_tmp_cxt_depth = 0;
+	glob->planner_tmp_cxt_usage = 0;
 	glob->subplans = NIL;
 	glob->subpaths = NIL;
 	glob->subroots = NIL;
@@ -595,6 +601,9 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 			result->jitFlags |= PGJIT_DEFORM;
 	}
 
+	/* Clean up. We aren't very thorough here. */
+	MemoryContextDelete(glob->planner_tmp_cxt);
+
 	if (glob->partition_directory != NULL)
 		DestroyPartitionDirectory(glob->partition_directory);
 
@@ -655,7 +664,7 @@ subquery_planner(PlannerGlobal *glob, Query *parse, PlannerInfo *parent_root,
 	root->parent_root = parent_root;
 	root->plan_params = NIL;
 	root->outer_params = NULL;
-	root->planner_cxt = CurrentMemoryContext;
+	root->planner_cxt = glob->planner_cxt;
 	root->init_plans = NIL;
 	root->cte_plan_ids = NIL;
 	root->multiexpr_params = NIL;
@@ -6672,11 +6681,18 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
 
 	glob = makeNode(PlannerGlobal);
 
+	glob->planner_cxt = CurrentMemoryContext;
+	glob->planner_tmp_cxt = AllocSetContextCreate(glob->planner_cxt,
+												  "Planner temp context",
+												  ALLOCSET_DEFAULT_SIZES);
+	glob->planner_tmp_cxt_depth = 0;
+	glob->planner_tmp_cxt_usage = 0;
+
 	root = makeNode(PlannerInfo);
 	root->parse = query;
 	root->glob = glob;
 	root->query_level = 1;
-	root->planner_cxt = CurrentMemoryContext;
+	root->planner_cxt = glob->planner_cxt;
 	root->wt_param_id = -1;
 	root->join_domains = list_make1(makeNode(JoinDomain));
 
@@ -6715,7 +6731,10 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
 	 * trust the index contents but use seqscan-and-sort.
 	 */
 	if (lc == NULL)				/* not in the list? */
+	{
+		MemoryContextDelete(glob->planner_tmp_cxt);
 		return true;			/* use sort */
+	}
 
 	/*
 	 * Rather than doing all the pushups that would be needed to use
@@ -6748,6 +6767,9 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
 									  ForwardScanDirection, false,
 									  NULL, 1.0, false);
 
+	/* We assume this won't free *indexScanPath */
+	MemoryContextDelete(glob->planner_tmp_cxt);
+
 	return (seqScanAndSortPath.total_cost < indexScanPath->path.total_cost);
 }
 
@@ -6796,11 +6818,18 @@ plan_create_index_workers(Oid tableOid, Oid indexOid)
 
 	glob = makeNode(PlannerGlobal);
 
+	glob->planner_cxt = CurrentMemoryContext;
+	glob->planner_tmp_cxt = AllocSetContextCreate(glob->planner_cxt,
+												  "Planner temp context",
+												  ALLOCSET_DEFAULT_SIZES);
+	glob->planner_tmp_cxt_depth = 0;
+	glob->planner_tmp_cxt_usage = 0;
+
 	root = makeNode(PlannerInfo);
 	root->parse = query;
 	root->glob = glob;
 	root->query_level = 1;
-	root->planner_cxt = CurrentMemoryContext;
+	root->planner_cxt = glob->planner_cxt;
 	root->wt_param_id = -1;
 	root->join_domains = list_make1(makeNode(JoinDomain));
 
@@ -6890,6 +6919,8 @@ plan_create_index_workers(Oid tableOid, Oid indexOid)
 		parallel_workers--;
 
 done:
+	MemoryContextDelete(glob->planner_tmp_cxt);
+
 	index_close(index, NoLock);
 	table_close(heap, NoLock);
 
diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index adad7ea9a9..8e12900e56 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -1147,7 +1147,7 @@ pull_up_simple_subquery(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte,
 	subroot->parent_root = root->parent_root;
 	subroot->plan_params = NIL;
 	subroot->outer_params = NULL;
-	subroot->planner_cxt = CurrentMemoryContext;
+	subroot->planner_cxt = root->planner_cxt;
 	subroot->init_plans = NIL;
 	subroot->cte_plan_ids = NIL;
 	subroot->multiexpr_params = NIL;
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 58748d2ca6..226f103c67 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -101,6 +101,25 @@ typedef struct PlannerGlobal
 	/* Param values provided to planner() */
 	ParamListInfo boundParams pg_node_attr(read_write_ignore);
 
+	/* Context holding PlannerGlobal */
+	MemoryContext planner_cxt pg_node_attr(read_write_ignore);
+
+	/* short-lived context for purposes such as calling selectivity functions */
+	MemoryContext planner_tmp_cxt pg_node_attr(read_write_ignore);
+	/* nesting depth of uses of planner_tmp_cxt; reset it only at level 0 */
+	int			planner_tmp_cxt_depth;
+	/* how many times we've used planner_tmp_cxt since last reset */
+	int			planner_tmp_cxt_usage;
+
+	/*
+	 * How often we want to reset it.
+	 * It may be costly to clean up the context used for quick operations
+	 * likewise selectivity estimation causing instensive allocations of small
+	 * pieces of memory. Use it as high as possible to cut off peak memory
+	 * consumptions during the planning.
+	 */
+#define PLANNER_TMP_CXT_USAGE_RESET 1000
+
 	/* Plans for SubPlan nodes */
 	List	   *subplans;
 
@@ -477,7 +496,7 @@ struct PlannerInfo
 	/* List of MinMaxAggInfos */
 	List	   *minmax_aggs;
 
-	/* context holding PlannerInfo */
+	/* context holding PlannerInfo (copied from PlannerGlobal) */
 	MemoryContext planner_cxt pg_node_attr(read_write_ignore);
 
 	/* # of pages in all non-dummy tables of query */
-- 
2.39.5