Re: Optimize planner memory consumption for huge arrays

I wrote:
> Tomas Vondra <tomas.vondra@enterprisedb.com> writes:
>> On 2/19/24 16:45, Tom Lane wrote:
>>> Tomas Vondra <tomas.vondra@enterprisedb.com> writes:
>>>> For example, I don't think we expect selectivity functions to allocate
>>>> long-lived objects, right? So maybe we could run them in a dedicated
>>>> memory context, and reset it aggressively (after each call).

>>> That could eliminate a whole lot of potential leaks.  I'm not sure 
>>> though how much it moves the needle in terms of overall planner
>>> memory consumption.

>> It was an ad hoc thought, inspired by the issue at hand. Maybe it would
>> be possible to find similar "boundaries" in other parts of the planner.

> Here's a quick and probably-incomplete implementation of that idea.
> I've not tried to study its effects on memory consumption, just made
> sure it passes check-world.

I spent a bit more time on this patch.  One thing I was concerned
about was whether it causes any noticeable slowdown, and it seems that
it does: testing with "pgbench -S" I observe perhaps 1% slowdown.
However, we don't necessarily need to reset the temp context after
every single usage.  I experimented with resetting it every tenth
time, and that got me from 1% slower than HEAD to 1% faster.  Of
course "every tenth time" is very ad hoc.  I wondered if we could
make it somehow conditional on how much memory had been consumed
in the temp context, but there doesn't seem to be any cheap way
to get that.  Applying something like MemoryContextMemConsumed
would surely be a loser.  I'm not sure if it'd be worth extending
the mcxt.c API to provide something like "MemoryContextResetIfBig",
with some internal rule that would be cheap to apply like "reset
if we have any non-keeper blocks".

I also looked into whether it really does reduce overall memory
consumption noticeably, by collecting stats about planner memory
consumption during the core regression tests.  The answer is that
it barely helps.  I see the average used space across all planner
invocations drop from 23344 bytes to 23220, and the worst-case
numbers hardly move at all.  So that's a little discouraging.
But of course the regression tests prefer not to deal in very
large/expensive test cases, so maybe it's not surprising that
I don't see much win in this test.

Anyway, 0001 attached is a cleaned-up patch with the every-tenth-
time rule, and 0002 (not meant for commit) is the quick and
dirty instrumentation patch I used for collecting usage stats.

Even though this seems of only edge-case value, I'd much prefer
to do this than the sort of ad-hoc patching initially proposed
in this thread.

            regards, tom lane

diff --git a/src/backend/optimizer/path/clausesel.c b/src/backend/optimizer/path/clausesel.c
index c949dc1866..4678c778b0 100644
--- a/src/backend/optimizer/path/clausesel.c
+++ b/src/backend/optimizer/path/clausesel.c
@@ -24,6 +24,7 @@
 #include "statistics/statistics.h"
 #include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
+#include "utils/memutils.h"
 #include "utils/selfuncs.h"

 /*
@@ -693,6 +694,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;
@@ -756,6 +758,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;
@@ -967,6 +984,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 be4e182869..df6a1fd330 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -307,6 +307,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->subroots = NIL;
     glob->rewindPlanIDs = NULL;
@@ -586,6 +592,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);

@@ -642,7 +651,7 @@ subquery_planner(PlannerGlobal *glob, Query *parse,
     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;
@@ -6509,11 +6518,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));

@@ -6552,7 +6568,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
@@ -6584,6 +6603,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);
 }

@@ -6631,11 +6653,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));

@@ -6725,6 +6754,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 aa83dd3636..6a36110df2 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -987,7 +987,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 534692bee1..92f465c55e 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -101,6 +101,18 @@ 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 */
+#define PLANNER_TMP_CXT_USAGE_RESET 10
+
     /* Plans for SubPlan nodes */
     List       *subplans;

@@ -468,7 +480,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 */
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index df6a1fd330..db1d0fd7de 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -297,6 +297,10 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
     Plan       *top_plan;
     ListCell   *lp,
                *lr;
+    MemoryContextCounters before_counters,
+                after_counters;
+
+    MemoryContextMemConsumed(CurrentMemoryContext, &before_counters);

     /*
      * Set up global state for this planner invocation.  This data is needed
@@ -598,6 +602,13 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
     if (glob->partition_directory != NULL)
         DestroyPartitionDirectory(glob->partition_directory);

+    MemoryContextMemConsumed(CurrentMemoryContext, &after_counters);
+    after_counters.freespace -= before_counters.freespace;
+    after_counters.totalspace -= before_counters.totalspace;
+    elog(LOG, "planning memory consumption: used=%lld allocated=%lld",
+         (long long) (after_counters.totalspace - after_counters.freespace),
+         (long long) (after_counters.totalspace));
+
     return result;
 }