From 93ab0ff7e857acecdefc0a9c01b2d6d805c90559 Mon Sep 17 00:00:00 2001
From: Jeff Davis <jeff@j-davis.com>
Date: Wed, 29 Apr 2026 11:19:53 -0700
Subject: [PATCH v4 1/3] Refactor memory accounting into inline function.

---
 src/backend/utils/mmgr/aset.c         | 21 ++++++++++-----------
 src/backend/utils/mmgr/bump.c         |  8 ++++----
 src/backend/utils/mmgr/generation.c   |  6 +++---
 src/backend/utils/mmgr/slab.c         |  8 ++++----
 src/include/utils/memutils_internal.h | 12 ++++++++++++
 5 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 6a9ea367107..3ddef26a3d3 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -422,8 +422,8 @@ AllocSetContextCreateInternal(MemoryContext parent,
 								parent,
 								name);
 
-			((MemoryContext) set)->mem_allocated =
-				KeeperBlock(set)->endptr - ((char *) set);
+			MemoryContextUpdateAlloc((MemoryContext) set,
+									 KeeperBlock(set)->endptr - ((char *) set));
 
 			return (MemoryContext) set;
 		}
@@ -525,7 +525,7 @@ AllocSetContextCreateInternal(MemoryContext parent,
 						parent,
 						name);
 
-	((MemoryContext) set)->mem_allocated = firstBlockSize;
+	MemoryContextUpdateAlloc((MemoryContext) set, firstBlockSize);
 
 	return (MemoryContext) set;
 }
@@ -589,7 +589,7 @@ AllocSetReset(MemoryContext context)
 		else
 		{
 			/* Normal case, release the block */
-			context->mem_allocated -= block->endptr - ((char *) block);
+			MemoryContextUpdateAlloc(context, -(block->endptr - ((char *) block)));
 
 #ifdef CLOBBER_FREED_MEMORY
 			wipe_mem(block, block->freeptr - ((char *) block));
@@ -696,7 +696,7 @@ AllocSetDelete(MemoryContext context)
 		AllocBlock	next = block->next;
 
 		if (!IsKeeperBlock(set, block))
-			context->mem_allocated -= block->endptr - ((char *) block);
+			MemoryContextUpdateAlloc(context, -(block->endptr - ((char *) block)));
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
@@ -758,7 +758,7 @@ AllocSetAllocLarge(MemoryContext context, Size size, int flags)
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, ALLOC_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	block->aset = set;
 	block->freeptr = block->endptr = ((char *) block) + blksize;
@@ -967,7 +967,7 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, ALLOC_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	block->aset = set;
 	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
@@ -1144,7 +1144,7 @@ AllocSetFree(void *pointer)
 		if (block->next)
 			block->next->prev = block->prev;
 
-		set->header.mem_allocated -= block->endptr - ((char *) block);
+		MemoryContextUpdateAlloc(&set->header, -(block->endptr - ((char *) block)));
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
@@ -1307,9 +1307,8 @@ AllocSetRealloc(void *pointer, Size size, int flags)
 		VALGRIND_MEMPOOL_CHANGE(set, block, newblock, ALLOC_BLOCKHDRSZ);
 		block = newblock;
 
-		/* updated separately, not to underflow when (oldblksize > blksize) */
-		set->header.mem_allocated -= oldblksize;
-		set->header.mem_allocated += blksize;
+		MemoryContextUpdateAlloc(&set->header,
+								 (ssize_t) blksize - (ssize_t) oldblksize);
 
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
diff --git a/src/backend/utils/mmgr/bump.c b/src/backend/utils/mmgr/bump.c
index bfb5a114147..1deb73f451a 100644
--- a/src/backend/utils/mmgr/bump.c
+++ b/src/backend/utils/mmgr/bump.c
@@ -235,7 +235,7 @@ BumpContextCreate(MemoryContext parent, const char *name, Size minContextSize,
 	MemoryContextCreate((MemoryContext) set, T_BumpContext, MCTX_BUMP_ID,
 						parent, name);
 
-	((MemoryContext) set)->mem_allocated = allocSize;
+	MemoryContextUpdateAlloc((MemoryContext) set, allocSize);
 
 	return (MemoryContext) set;
 }
@@ -341,7 +341,7 @@ BumpAllocLarge(MemoryContext context, Size size, int flags)
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, Bump_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	/* the block is completely full */
 	block->freeptr = block->endptr = ((char *) block) + blksize;
@@ -481,7 +481,7 @@ BumpAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, Bump_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	/* initialize the new block */
 	BumpBlockInit(set, block, blksize);
@@ -626,7 +626,7 @@ BumpBlockFree(BumpContext *set, BumpBlock *block)
 	/* release the block from the list of blocks */
 	dlist_delete(&block->node);
 
-	((MemoryContext) set)->mem_allocated -= ((char *) block->endptr - (char *) block);
+	MemoryContextUpdateAlloc((MemoryContext) set, -((char *) block->endptr - (char *) block));
 
 #ifdef CLOBBER_FREED_MEMORY
 	wipe_mem(block, ((char *) block->endptr - (char *) block));
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 609c9bdc9a6..ddcd25c1e9d 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -388,7 +388,7 @@ GenerationAllocLarge(MemoryContext context, Size size, int flags)
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, Generation_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	/* block with a single (used) chunk */
 	block->context = set;
@@ -513,7 +513,7 @@ GenerationAllocFromNewBlock(MemoryContext context, Size size, int flags,
 	/* Make a vchunk covering the new block's header */
 	VALGRIND_MEMPOOL_ALLOC(set, block, Generation_BLOCKHDRSZ);
 
-	context->mem_allocated += blksize;
+	MemoryContextUpdateAlloc(context, blksize);
 
 	/* initialize the new block */
 	GenerationBlockInit(set, block, blksize);
@@ -697,7 +697,7 @@ GenerationBlockFree(GenerationContext *set, GenerationBlock *block)
 	/* release the block from the list of blocks */
 	dlist_delete(&block->node);
 
-	((MemoryContext) set)->mem_allocated -= block->blksize;
+	MemoryContextUpdateAlloc((MemoryContext) set, -((ssize_t) block->blksize));
 
 #ifdef CLOBBER_FREED_MEMORY
 	wipe_mem(block, block->blksize);
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index dd1db9566d1..f27d9683db6 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -461,7 +461,7 @@ SlabReset(MemoryContext context)
 		VALGRIND_MEMPOOL_FREE(slab, block);
 
 		free(block);
-		context->mem_allocated -= slab->blockSize;
+		MemoryContextUpdateAlloc(context, -((ssize_t) slab->blockSize));
 	}
 
 	/* walk over blocklist and free the blocks */
@@ -481,7 +481,7 @@ SlabReset(MemoryContext context)
 			VALGRIND_MEMPOOL_FREE(slab, block);
 
 			free(block);
-			context->mem_allocated -= slab->blockSize;
+			MemoryContextUpdateAlloc(context, -((ssize_t) slab->blockSize));
 		}
 	}
 
@@ -597,7 +597,7 @@ SlabAllocFromNewBlock(MemoryContext context, Size size, int flags)
 		VALGRIND_MEMPOOL_ALLOC(slab, block, Slab_BLOCKHDRSZ);
 
 		block->slab = slab;
-		context->mem_allocated += slab->blockSize;
+		MemoryContextUpdateAlloc(context, slab->blockSize);
 
 		/* use the first chunk in the new block */
 		chunk = SlabBlockGetChunk(slab, block, 0);
@@ -836,7 +836,7 @@ SlabFree(void *pointer)
 			VALGRIND_MEMPOOL_FREE(slab, block);
 
 			free(block);
-			slab->header.mem_allocated -= slab->blockSize;
+			MemoryContextUpdateAlloc(&slab->header, -((ssize_t) slab->blockSize));
 		}
 
 		/*
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index 475e91b336b..8f793f3c781 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -157,6 +157,18 @@ extern void MemoryContextCreate(MemoryContext node,
 								MemoryContext parent,
 								const char *name);
 
+/*
+ * MemoryContextUpdateAlloc()
+ *
+ * Update allocation total.
+ */
+static inline void
+MemoryContextUpdateAlloc(MemoryContext context, ssize_t size)
+{
+	Assert(size >= 0 || -size <= context->mem_allocated);
+	context->mem_allocated += size;
+}
+
 extern void *MemoryContextAllocationFailure(MemoryContext context, Size size,
 											int flags);
 
-- 
2.43.0

