v4-0002-Memory-Pools.patch

From 55766fc0354e25621d922bb588a248196d124051 Mon Sep 17 00:00:00 2001
From: Jeff Davis <jeff@j-davis.com>
Date: Mon, 4 May 2026 14:24:40 -0700
Subject: [PATCH v4 2/3] Memory Pools.

A memory context can have a memory pool, which tracks the memory for
that context as well as all subcontexts in a single total, and also
holds an (unenforced) limit. Memory pools are automatically inherited
by child contexts.

Memory Pools solve two problems:

First, the previous implementation in MemoryContextMemAllocated()
recursively walked through subcontexts to total the mem_allocated
fields. For hash aggregates that used a child memory context for each
group, it was too slow. Now, it just reads the single total.

Second, with the block size doubling behavior, it was likely that a
single small chunk allocation could cause a block allocation that
overwhelmed a small value of work_mem. Hash Aggregation protected
against this by constraining maxBlockSize for known contexts, but not
their subcontexts. Now, subcontexts constrain the block size using the
inherited memory pool's limit.

Suggested-by: David Rowley <dgrowleyml@gmail.com>
Discussion: https://postgr.es/m/3579239.1775151498@sss.pgh.pa.us
---
 src/backend/utils/mmgr/aset.c         |   8 +-
 src/backend/utils/mmgr/bump.c         |   7 +-
 src/backend/utils/mmgr/generation.c   |   7 +-
 src/backend/utils/mmgr/mcxt.c         | 110 +++++++++++++++++++++++++-
 src/include/nodes/memnodes.h          |  16 ++++
 src/include/utils/memutils.h          |   1 +
 src/include/utils/memutils_internal.h |  76 +++++++++++++++++-
 src/tools/pgindent/typedefs.list      |   1 +
 8 files changed, 218 insertions(+), 8 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 3ddef26a3d3..8981c128672 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -867,6 +867,7 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	Size		blksize;
 	Size		required_size;
 	Size		chunk_size;
+	Size		blockSizeLimit;
 
 	/* due to the keeper block set->blocks should always be valid */
 	Assert(set->blocks != NULL);
@@ -931,8 +932,11 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	 */
 	blksize = set->nextBlockSize;
 	set->nextBlockSize <<= 1;
-	if (set->nextBlockSize > set->maxBlockSize)
-		set->nextBlockSize = set->maxBlockSize;
+
+	blockSizeLimit = MemoryContextBlockSizeLimit(context, set->initBlockSize,
+												 set->maxBlockSize);
+	if (set->nextBlockSize > blockSizeLimit)
+		set->nextBlockSize = blockSizeLimit;
 
 	/* Choose the actual chunk size to allocate */
 	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
diff --git a/src/backend/utils/mmgr/bump.c b/src/backend/utils/mmgr/bump.c
index 1deb73f451a..154ebdebd12 100644
--- a/src/backend/utils/mmgr/bump.c
+++ b/src/backend/utils/mmgr/bump.c
@@ -457,15 +457,18 @@ BumpAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	BumpBlock  *block;
 	Size		blksize;
 	Size		required_size;
+	Size		blockSizeLimit;
 
 	/*
 	 * The first such block has size initBlockSize, and we double the space in
 	 * each succeeding block, but not more than maxBlockSize.
 	 */
 	blksize = set->nextBlockSize;
+	blockSizeLimit = MemoryContextBlockSizeLimit(context, set->initBlockSize,
+												 set->maxBlockSize);
 	set->nextBlockSize <<= 1;
-	if (set->nextBlockSize > set->maxBlockSize)
-		set->nextBlockSize = set->maxBlockSize;
+	if (set->nextBlockSize > blockSizeLimit)
+		set->nextBlockSize = blockSizeLimit;
 
 	/* we'll need space for the chunk, chunk hdr and block hdr */
 	required_size = chunk_size + Bump_CHUNKHDRSZ + Bump_BLOCKHDRSZ;
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index ddcd25c1e9d..cade3ae3fe9 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -488,15 +488,18 @@ GenerationAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	GenerationBlock *block;
 	Size		blksize;
 	Size		required_size;
+	Size		blockSizeLimit;
 
 	/*
 	 * The first such block has size initBlockSize, and we double the space in
 	 * each succeeding block, but not more than maxBlockSize.
 	 */
 	blksize = set->nextBlockSize;
+	blockSizeLimit = MemoryContextBlockSizeLimit(context, set->initBlockSize,
+												 set->maxBlockSize);
 	set->nextBlockSize <<= 1;
-	if (set->nextBlockSize > set->maxBlockSize)
-		set->nextBlockSize = set->maxBlockSize;
+	if (set->nextBlockSize > blockSizeLimit)
+		set->nextBlockSize = blockSizeLimit;
 
 	/* we'll need space for the chunk, chunk hdr and block hdr */
 	required_size = chunk_size + Generation_CHUNKHDRSZ + Generation_BLOCKHDRSZ;
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 073bdb35d2a..472123612f2 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -664,6 +664,86 @@ MemoryContextSetIdentifier(MemoryContext context, const char *id)
 	context->ident = id;
 }
 
+/*
+ * Add a memory pool to an existing context. It may already be a member of an
+ * outer pool owned by an ancestor, but it may not already own its memory
+ * pool.
+ *
+ * Currently, the context must have no child contexts, but that restriction
+ * could be lifted if necessary.
+ */
+MemoryPool *
+MemoryContextCreatePool(MemoryContext context, Size limit)
+{
+	MemoryPool *pool = &context->pool;
+
+	Assert(MemoryContextIsValid(context));
+	Assert(limit > 0);
+	Assert(context->firstchild == NULL);
+	Assert(!MemoryContextOwnsPool(context));
+
+	/* outer may already be set */
+
+	pool->limit = limit;
+	pool->allocated = context->mem_allocated;
+
+	return pool;
+}
+
+/* helper for MemoryContextSetParent() */
+static void
+MemoryContextTransferPool(MemoryContext context, MemoryContext new_parent)
+{
+	MemoryPool *old_outer_pool = context->pool.outer;
+	MemoryPool *new_outer_pool = NULL;
+	Size		subtree_mem;
+
+	if (new_parent)
+		new_outer_pool = MemoryContextGetPool(new_parent);
+
+	if (!old_outer_pool && !new_outer_pool)
+		return;
+
+	if (old_outer_pool == new_outer_pool)
+		return;
+
+	/*
+	 * If the context's pool is inherited from some ancestor context, we need
+	 * to calculate the total for this subtree. If it's owned, we can just use
+	 * the pool's total.
+	 */
+	if (MemoryContextOwnsPool(context))
+		subtree_mem = context->pool.allocated;
+	else
+		subtree_mem = MemoryContextMemAllocated(context, true);
+
+	/* subtract from old subtree */
+	for (MemoryPool *pool = old_outer_pool; pool != NULL; pool = pool->outer)
+		pool->allocated -= subtree_mem;
+
+	/* add to new subtree */
+	for (MemoryPool *pool = new_outer_pool; pool != NULL; pool = pool->outer)
+		pool->allocated += subtree_mem;
+
+	/*
+	 * If it's not the owner of its pool, subcontexts could refer to pools
+	 * owned by the context's old ancestors. Search subtree for references to
+	 * old_outer_pool and replace with new_outer_pool. NB: old_outer_pool or
+	 * new_outer_pool might be NULL here.
+	 */
+	if (!MemoryContextOwnsPool(context))
+	{
+		for (MemoryContext cxt = context->firstchild; cxt != NULL;
+			 cxt = MemoryContextTraverseNext(cxt, context))
+		{
+			if (cxt->pool.outer == old_outer_pool)
+				cxt->pool.outer = new_outer_pool;
+		}
+	}
+
+	context->pool.outer = new_outer_pool;
+}
+
 /*
  * MemoryContextSetParent
  *		Change a context to belong to a new parent (or no parent).
@@ -692,6 +772,12 @@ MemoryContextSetParent(MemoryContext context, MemoryContext new_parent)
 	if (new_parent == context->parent)
 		return;
 
+	/*
+	 * This must happen here while we still have information about the
+	 * existing ancestor.
+	 */
+	MemoryContextTransferPool(context, new_parent);
+
 	/* Delink from existing parent, if any */
 	if (context->parent)
 	{
@@ -726,6 +812,10 @@ MemoryContextSetParent(MemoryContext context, MemoryContext new_parent)
 		context->prevchild = NULL;
 		context->nextchild = NULL;
 	}
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	MemoryContextCheck(context);
+#endif
 }
 
 /*
@@ -1090,18 +1180,30 @@ MemoryContextStatsPrint(MemoryContext context, void *passthru,
 								 level, name, stats_string, truncated_ident)));
 }
 
+#ifdef MEMORY_CONTEXT_CHECKING
+
+static void
+MemoryContextCheckPool(MemoryContext context)
+{
+	if (MemoryContextOwnsPool(context))
+		Assert(context->pool.allocated == MemoryContextMemAllocated(context, true));
+
+	if (context->parent)
+		Assert(context->pool.outer == MemoryContextGetPool(context->parent));
+}
+
 /*
  * MemoryContextCheck
  *		Check all chunks in the named context and its children.
  *
  * This is just a debugging utility, so it's not fancy.
  */
-#ifdef MEMORY_CONTEXT_CHECKING
 void
 MemoryContextCheck(MemoryContext context)
 {
 	Assert(MemoryContextIsValid(context));
 	context->methods->check(context);
+	MemoryContextCheckPool(context);
 
 	for (MemoryContext curr = context->firstchild;
 		 curr != NULL;
@@ -1109,8 +1211,10 @@ MemoryContextCheck(MemoryContext context)
 	{
 		Assert(MemoryContextIsValid(curr));
 		curr->methods->check(curr);
+		MemoryContextCheckPool(curr);
 	}
 }
+
 #endif
 
 /*
@@ -1166,6 +1270,9 @@ MemoryContextCreate(MemoryContext node,
 	node->parent = parent;
 	node->firstchild = NULL;
 	node->mem_allocated = 0;
+	node->pool.limit = 0;
+	node->pool.allocated = 0;
+	node->pool.outer = NULL;
 	node->prevchild = NULL;
 	node->name = name;
 	node->ident = NULL;
@@ -1180,6 +1287,7 @@ MemoryContextCreate(MemoryContext node,
 		parent->firstchild = node;
 		/* inherit allowInCritSection flag from parent */
 		node->allowInCritSection = parent->allowInCritSection;
+		node->pool.outer = MemoryContextGetPool(parent);
 	}
 	else
 	{
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index 4b24778fe1e..8c3bc7b7d30 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -114,6 +114,21 @@ typedef struct MemoryContextMethods
 } MemoryContextMethods;
 
 
+/*
+ * Memory Pools account for memory usage.
+ *
+ * If limit > 0, then this pool is owned by the containing context, and
+ * 'allocated' is valid. If limit == 0, the pool was inherited from the parent
+ * context, and 'allocated' is invalid.
+ */
+typedef struct MemoryPool
+{
+	Size		limit;			/* limit (not enforced) */
+	Size		allocated;		/* allocation total for subtree */
+	struct MemoryPool *outer;	/* outer containing pool */
+} MemoryPool;
+
+
 typedef struct MemoryContextData
 {
 	pg_node_attr(abstract)		/* there are no nodes of this type */
@@ -123,6 +138,7 @@ typedef struct MemoryContextData
 	bool		isReset;		/* T = no space allocated since last reset */
 	bool		allowInCritSection; /* allow palloc in critical section */
 	Size		mem_allocated;	/* track memory allocated for this context */
+	MemoryPool	pool;			/* accounting, if enabled */
 	const MemoryContextMethods *methods;	/* virtual function table */
 	MemoryContext parent;		/* NULL if no parent (toplevel context) */
 	MemoryContext firstchild;	/* head of linked list of children */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 11ab1717a16..e531876ba66 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -86,6 +86,7 @@ extern bool MemoryContextIsEmpty(MemoryContext context);
 extern Size MemoryContextMemAllocated(MemoryContext context, bool recurse);
 extern void MemoryContextMemConsumed(MemoryContext context,
 									 MemoryContextCounters *consumed);
+extern MemoryPool *MemoryContextCreatePool(MemoryContext owner, Size limit);
 extern void MemoryContextStats(MemoryContext context);
 extern void MemoryContextStatsDetail(MemoryContext context,
 									 int max_level, int max_children,
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index 8f793f3c781..8cacd8abb8c 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -17,6 +17,7 @@
 #define MEMUTILS_INTERNAL_H
 
 #include "utils/memutils.h"
+#include "port/pg_bitutils.h"
 
 /* These functions implement the MemoryContext API for AllocSet context. */
 extern void *AllocSetAlloc(MemoryContext context, Size size, int flags);
@@ -157,16 +158,89 @@ extern void MemoryContextCreate(MemoryContext node,
 								MemoryContext parent,
 								const char *name);
 
+static inline bool
+MemoryPoolIsValid(MemoryPool *pool)
+{
+	if (pool->limit > 0)
+		return true;
+
+	Assert(pool->allocated == 0);
+	return false;
+}
+
+static inline bool
+MemoryContextOwnsPool(MemoryContext context)
+{
+	return MemoryPoolIsValid(&context->pool);
+}
+
+static inline MemoryPool *
+MemoryPoolOuter(MemoryPool *pool)
+{
+	Assert(pool->outer == NULL || MemoryPoolIsValid(pool->outer));
+	return pool->outer;
+}
+
+/* get owned or inherited pool for memory context, or NULL */
+static inline MemoryPool *
+MemoryContextGetPool(MemoryContext context)
+{
+	MemoryPool *pool;
+
+	if (MemoryContextOwnsPool(context))
+		pool = &context->pool;
+	else
+		pool = context->pool.outer;
+
+	Assert(pool == NULL || MemoryPoolIsValid(pool));
+	return pool;
+}
+
+/*
+ * If the context is part of a memory pool, calculate a reasonable max block
+ * size given the memory pool limit. This avoids a case where a small chunk
+ * request causes a large block allocation that overwhelms the limit.
+ * Calculated dynamically because the limit could change, e.g. during
+ * MemoryContextSetParent().
+ */
+static inline Size
+MemoryContextBlockSizeLimit(MemoryContext context, Size min, Size max)
+{
+	MemoryPool *pool = MemoryContextGetPool(context);
+	Size		maxBlockSize;
+
+	if (pool != NULL)
+	{
+		/* reasonable value chosen as 1/16th of the limit */
+		maxBlockSize = pg_prevpower2_size_t(pool->limit / 16);
+
+		maxBlockSize = Min(maxBlockSize, max);
+		maxBlockSize = Max(maxBlockSize, min);
+	}
+	else
+		maxBlockSize = max;
+
+	return maxBlockSize;
+}
+
 /*
  * MemoryContextUpdateAlloc()
  *
- * Update allocation total.
+ * Update allocation total for context and pool(s).
  */
 static inline void
 MemoryContextUpdateAlloc(MemoryContext context, ssize_t size)
 {
 	Assert(size >= 0 || -size <= context->mem_allocated);
 	context->mem_allocated += size;
+
+	for (MemoryPool *pool = MemoryContextGetPool(context);
+		 pool != NULL;
+		 pool = MemoryPoolOuter(pool))
+	{
+		Assert(size >= 0 || -size <= pool->allocated);
+		pool->allocated += size;
+	}
 }
 
 extern void *MemoryContextAllocationFailure(MemoryContext context, Size size,
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 0abdb2d37e2..9fe2b799476 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1754,6 +1754,7 @@ MemoryContextData
 MemoryContextId
 MemoryContextMethodID
 MemoryContextMethods
+MemoryPool
 MemoryStatsPrintFunc
 MergeAction
 MergeActionState
-- 
2.43.0