Subject: Re: Avoid stack frame setup in performance critical routines using tail calls

From: Andres Freund <[email-masked]>
Date: Wed, 19 Jul 2023 01:52:36 -0700

Hi,

David and I were chatting about this patch, in the context of his bump
allocator patch.  Attached is a rebased version that is also split up into two
steps, and a bit more polished.

I wasn't sure what a good test was. I ended up measuring
  COPY pgbench_accounts TO '/dev/null' WITH (FORMAT 'binary');
of a scale 1 database with pgbench:

c=1;pgbench -q -i -s1 && pgbench  -n -c$c -j$c -t100 -f <(echo "COPY pgbench_accounts TO '/dev/null' WITH (FORMAT 'binary');")

	average latency
HEAD:   33.865 ms
01:     32.820 ms
02:     29.934 ms

The server was pinned to the one core, turbo mode disabled.  That's a pretty
nice win, I'd say.  And I don't think this is actually the most allocator
bound workload, I just tried something fairly random...

Greetings,

Andres Freund
>From 4b97070e32ed323ae0f083bf865717c6d271ce53 Mon Sep 17 00:00:00 2001
From: Andres Freund <[email-masked]>
Date: Tue, 18 Jul 2023 18:55:58 -0700
Subject: [PATCH v2 1/4] Optimize palloc() etc to allow sibling calls

The reason palloc() etc previously couldn't use sibling calls is that they did
check the returned value (to e.g. raise an error when the allocation
fails). Push the error handling down into the memory context implementation -
they can avoid performing the check at all in the hot code paths.
---
 src/include/nodes/memnodes.h          |   4 +-
 src/include/utils/memutils_internal.h |  29 +++-
 src/backend/utils/mmgr/alignedalloc.c |  11 +-
 src/backend/utils/mmgr/aset.c         |  23 ++--
 src/backend/utils/mmgr/generation.c   |  14 +-
 src/backend/utils/mmgr/mcxt.c         | 186 ++++++--------------------
 src/backend/utils/mmgr/slab.c         |  12 +-
 7 files changed, 102 insertions(+), 177 deletions(-)

diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index ff6453bb7ac..47d2932e6ed 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -57,10 +57,10 @@ typedef void (*MemoryStatsPrintFunc) (MemoryContext context, void *passthru,
 
 typedef struct MemoryContextMethods
 {
-	void	   *(*alloc) (MemoryContext context, Size size);
+	void	   *(*alloc) (MemoryContext context, Size size, int flags);
 	/* call this free_p in case someone #define's free() */
 	void		(*free_p) (void *pointer);
-	void	   *(*realloc) (void *pointer, Size size);
+	void	   *(*realloc) (void *pointer, Size size, int flags);
 	void		(*reset) (MemoryContext context);
 	void		(*delete_context) (MemoryContext context);
 	MemoryContext (*get_chunk_context) (void *pointer);
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index 2d107bbf9d4..3a050819263 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -19,9 +19,9 @@
 #include "utils/memutils.h"
 
 /* These functions implement the MemoryContext API for AllocSet context. */
-extern void *AllocSetAlloc(MemoryContext context, Size size);
+extern void *AllocSetAlloc(MemoryContext context, Size size, int flags);
 extern void AllocSetFree(void *pointer);
-extern void *AllocSetRealloc(void *pointer, Size size);
+extern void *AllocSetRealloc(void *pointer, Size size, int flags);
 extern void AllocSetReset(MemoryContext context);
 extern void AllocSetDelete(MemoryContext context);
 extern MemoryContext AllocSetGetChunkContext(void *pointer);
@@ -36,9 +36,9 @@ extern void AllocSetCheck(MemoryContext context);
 #endif
 
 /* These functions implement the MemoryContext API for Generation context. */
-extern void *GenerationAlloc(MemoryContext context, Size size);
+extern void *GenerationAlloc(MemoryContext context, Size size, int flags);
 extern void GenerationFree(void *pointer);
-extern void *GenerationRealloc(void *pointer, Size size);
+extern void *GenerationRealloc(void *pointer, Size size, int flags);
 extern void GenerationReset(MemoryContext context);
 extern void GenerationDelete(MemoryContext context);
 extern MemoryContext GenerationGetChunkContext(void *pointer);
@@ -54,9 +54,9 @@ extern void GenerationCheck(MemoryContext context);

 /* These functions implement the MemoryContext API for Slab context. */
-extern void *SlabAlloc(MemoryContext context, Size size);
+extern void *SlabAlloc(MemoryContext context, Size size, int flags);
 extern void SlabFree(void *pointer);
-extern void *SlabRealloc(void *pointer, Size size);
+extern void *SlabRealloc(void *pointer, Size size, int flags);
 extern void SlabReset(MemoryContext context);
 extern void SlabDelete(MemoryContext context);
 extern MemoryContext SlabGetChunkContext(void *pointer);
@@ -75,7 +75,7 @@ extern void SlabCheck(MemoryContext context);
  * part of a fully-fledged MemoryContext type.
  */
 extern void AlignedAllocFree(void *pointer);
-extern void *AlignedAllocRealloc(void *pointer, Size size);
+extern void *AlignedAllocRealloc(void *pointer, Size size, int flags);
 extern MemoryContext AlignedAllocGetChunkContext(void *pointer);
 extern Size AlignedAllocGetChunkSpace(void *pointer);
 
@@ -133,4 +133,19 @@ extern void MemoryContextCreate(MemoryContext node,
 								MemoryContext parent,
 								const char *name);
 
+extern void *MemoryContextAllocationFailure(MemoryContext context, Size size, int flags);
+
+extern void MemoryContextSizeFailure(MemoryContext context, Size size, int flags) pg_attribute_noreturn();
+
+static inline void
+MemoryContextCheckSize(MemoryContext context, Size size, int flags)
+{
+	if (unlikely(!AllocSizeIsValid(size)))
+	{
+		if (!(flags & MCXT_ALLOC_HUGE) || !AllocHugeSizeIsValid(size))
+			MemoryContextSizeFailure(context, size, flags);
+	}
+}
+
+
 #endif							/* MEMUTILS_INTERNAL_H */
diff --git a/src/backend/utils/mmgr/alignedalloc.c b/src/backend/utils/mmgr/alignedalloc.c
index 627e988852b..83c665675b8 100644
--- a/src/backend/utils/mmgr/alignedalloc.c
+++ b/src/backend/utils/mmgr/alignedalloc.c
@@ -57,7 +57,7 @@ AlignedAllocFree(void *pointer)
  *		memory will be uninitialized.
  */
 void *
-AlignedAllocRealloc(void *pointer, Size size)
+AlignedAllocRealloc(void *pointer, Size size, int flags)
 {
 	MemoryChunk *redirchunk = PointerGetMemoryChunk(pointer);
 	Size		alignto;
@@ -97,14 +97,17 @@ AlignedAllocRealloc(void *pointer, Size size)
 #endif
 
 	ctx = GetMemoryChunkContext(unaligned);
-	newptr = MemoryContextAllocAligned(ctx, size, alignto, 0);
+	newptr = MemoryContextAllocAligned(ctx, size, alignto, flags);
 
 	/*
 	 * We may memcpy beyond the end of the original allocation request size,
 	 * so we must mark the entire allocation as defined.
 	 */
-	VALGRIND_MAKE_MEM_DEFINED(pointer, old_size);
-	memcpy(newptr, pointer, Min(size, old_size));
+	if (likely(newptr != NULL))
+	{
+		VALGRIND_MAKE_MEM_DEFINED(pointer, old_size);
+		memcpy(newptr, pointer, Min(size, old_size));
+	}
 	pfree(unaligned);
 
 	return newptr;
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index c3affaf5a8a..b72ec68f7dd 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -700,7 +700,7 @@ AllocSetDelete(MemoryContext context)
  * return space that is marked NOACCESS - AllocSetRealloc has to beware!
  */
 void *
-AllocSetAlloc(MemoryContext context, Size size)
+AllocSetAlloc(MemoryContext context, Size size, int flags)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block;
@@ -717,6 +717,9 @@ AllocSetAlloc(MemoryContext context, Size size)
 	 */
 	if (size > set->allocChunkLimit)
 	{
+		/* check size, only allocation path where the limits could be hit */
+		MemoryContextCheckSize(context, size, flags);
+
 #ifdef MEMORY_CONTEXT_CHECKING
 		/* ensure there's always space for the sentinel byte */
 		chunk_size = MAXALIGN(size + 1);
@@ -727,7 +730,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
 		block = (AllocBlock) malloc(blksize);
 		if (block == NULL)
-			return NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -940,7 +943,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		}
 
 		if (block == NULL)
-			return NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -1106,7 +1109,7 @@ AllocSetFree(void *pointer)
  * request size.)
  */
 void *
-AllocSetRealloc(void *pointer, Size size)
+AllocSetRealloc(void *pointer, Size size, int flags)
 {
 	AllocBlock	block;
 	AllocSet	set;
@@ -1139,6 +1142,9 @@ AllocSetRealloc(void *pointer, Size size)
 
 		set = block->aset;
 
+		/* check size, only allocation path where the limits could be hit */
+		MemoryContextCheckSize((MemoryContext) set, size, flags);
+
 		oldchksize = block->endptr - (char *) pointer;
 
 #ifdef MEMORY_CONTEXT_CHECKING
@@ -1165,7 +1171,8 @@ AllocSetRealloc(void *pointer, Size size)
 		{
 			/* Disallow access to the chunk header. */
 			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-			return NULL;
+			return MemoryContextAllocationFailure(&set->header, size, flags);
+
 		}
 
 		/* updated separately, not to underflow when (oldblksize > blksize) */
@@ -1325,15 +1332,15 @@ AllocSetRealloc(void *pointer, Size size)
 		AllocPointer newPointer;
 		Size		oldsize;
 
-		/* allocate new chunk */
-		newPointer = AllocSetAlloc((MemoryContext) set, size);
+		/* allocate new chunk (also checks size for us) */
+		newPointer = AllocSetAlloc((MemoryContext) set, size, flags);
 
 		/* leave immediately if request was not completed */
 		if (newPointer == NULL)
 		{
 			/* Disallow access to the chunk header. */
 			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-			return NULL;
+			return MemoryContextAllocationFailure((MemoryContext) set, size, flags);
 		}
 
 		/*
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 92401ccf738..e932bd10277 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -342,7 +342,7 @@ GenerationDelete(MemoryContext context)
  * return space that is marked NOACCESS - GenerationRealloc has to beware!
  */
 void *
-GenerationAlloc(MemoryContext context, Size size)
+GenerationAlloc(MemoryContext context, Size size, int flags)
 {
 	GenerationContext *set = (GenerationContext *) context;
 	GenerationBlock *block;
@@ -365,9 +365,11 @@ GenerationAlloc(MemoryContext context, Size size)
 	{
 		Size		blksize = required_size + Generation_BLOCKHDRSZ;
 
+		MemoryContextCheckSize((MemoryContext) set, size, flags);
+
 		block = (GenerationBlock *) malloc(blksize);
 		if (block == NULL)
-			return NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -470,7 +472,7 @@ GenerationAlloc(MemoryContext context, Size size)
 			block = (GenerationBlock *) malloc(blksize);
 
 			if (block == NULL)
-				return NULL;
+				return MemoryContextAllocationFailure(context, size, flags);
 
 			context->mem_allocated += blksize;
 
@@ -735,7 +737,7 @@ GenerationFree(void *pointer)
  *		into the old chunk - in that case we just update chunk header.
  */
 void *
-GenerationRealloc(void *pointer, Size size)
+GenerationRealloc(void *pointer, Size size, int flags)
 {
 	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
 	GenerationContext *set;
@@ -838,14 +840,14 @@ GenerationRealloc(void *pointer, Size size)
 	}
 
 	/* allocate new chunk */
-	newPointer = GenerationAlloc((MemoryContext) set, size);
+	newPointer = GenerationAlloc((MemoryContext) set, size, flags);
 
 	/* leave immediately if request was not completed */
 	if (newPointer == NULL)
 	{
 		/* Disallow access to the chunk header. */
 		VALGRIND_MAKE_MEM_NOACCESS(chunk, Generation_CHUNKHDRSZ);
-		return NULL;
+		return MemoryContextAllocationFailure((MemoryContext) set, size, flags);
 	}
 
 	/*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 9fc83f11f6f..943f041af61 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -34,7 +34,7 @@

 static void BogusFree(void *pointer);
-static void *BogusRealloc(void *pointer, Size size);
+static void *BogusRealloc(void *pointer, Size size, int flags);
 static MemoryContext BogusGetChunkContext(void *pointer);
 static Size BogusGetChunkSpace(void *pointer);
 
@@ -237,7 +237,7 @@ BogusFree(void *pointer)
 }
 
 static void *
-BogusRealloc(void *pointer, Size size)
+BogusRealloc(void *pointer, Size size, int flags)
 {
 	elog(ERROR, "repalloc called with invalid pointer %p (header 0x%016llx)",
 		 pointer, (unsigned long long) GetMemoryChunkHeader(pointer));
@@ -1010,6 +1010,27 @@ MemoryContextCreate(MemoryContext node,
 	VALGRIND_CREATE_MEMPOOL(node, 0, false);
 }
 
+void *
+MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
+{
+	if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+	{
+		MemoryContextStats(TopMemoryContext);
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu in memory context \"%s\".",
+						   size, context->name)));
+	}
+	return NULL;
+}
+
+void
+MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
+{
+	elog(ERROR, "invalid memory alloc request size %zu", size);
+}
+
 /*
  * MemoryContextAlloc
  *		Allocate space within the specified context.
@@ -1025,28 +1046,9 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-
-		/*
-		 * Here, and elsewhere in this module, we show the target context's
-		 * "name" but not its "ident" (if any) in user-visible error messages.
-		 * The "ident" string might contain security-sensitive data, such as
-		 * values in SQL commands.
-		 */
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, 0);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1068,24 +1070,16 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, 0);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	/*
+	 * XXX: Should this also be moved into alloc()? We could possibly avoid
+	 * zeroing in some cases (e.g. if we used mmap() ourselves.
+	 */
 	MemSetAligned(ret, 0, size);
 
 	return ret;
@@ -1106,21 +1100,9 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, 0);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1147,20 +1129,9 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
+	ret = context->methods->alloc(context, size, flags);
 	if (unlikely(ret == NULL))
-	{
-		if ((flags & MCXT_ALLOC_NO_OOM) == 0)
-		{
-			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu in memory context \"%s\".",
-							   size, context->name)));
-		}
 		return NULL;
-	}
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1232,22 +1203,11 @@ palloc(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, 0);
 
+	Assert(ret != 0);
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -1263,21 +1223,9 @@ palloc0(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, 0);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1296,24 +1244,11 @@ palloc_extended(Size size, int flags)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!((flags & MCXT_ALLOC_HUGE) != 0 ? AllocHugeSizeIsValid(size) :
-		  AllocSizeIsValid(size)))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
+	ret = context->methods->alloc(context, size, flags);
 	if (unlikely(ret == NULL))
 	{
-		if ((flags & MCXT_ALLOC_NO_OOM) == 0)
-		{
-			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu in memory context \"%s\".",
-							   size, context->name)));
-		}
 		return NULL;
 	}
 
@@ -1483,29 +1418,15 @@ repalloc(void *pointer, Size size)
 #endif
 	void	   *ret;
 
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	AssertNotInCriticalSection(context);
 
 	/* isReset must be false already */
 	Assert(!context->isReset);
 
-	ret = MCXT_METHOD(pointer, realloc) (pointer, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContext cxt = GetMemoryChunkContext(pointer);
-
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, cxt->name)));
-	}
+	ret = MCXT_METHOD(pointer, realloc) (pointer, size, 0);
 
 #ifdef USE_VALGRIND
-	if (method != MCTX_ALIGNED_REDIRECT_ID)
+	if (ret != NULL && method != MCTX_ALIGNED_REDIRECT_ID)
 		VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 #endif
 
@@ -1525,31 +1446,14 @@ repalloc_extended(void *pointer, Size size, int flags)
 #endif
 	void	   *ret;
 
-	if (!((flags & MCXT_ALLOC_HUGE) != 0 ? AllocHugeSizeIsValid(size) :
-		  AllocSizeIsValid(size)))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	AssertNotInCriticalSection(context);
 
 	/* isReset must be false already */
 	Assert(!context->isReset);
 
-	ret = MCXT_METHOD(pointer, realloc) (pointer, size);
+	ret = MCXT_METHOD(pointer, realloc) (pointer, size, flags);
 	if (unlikely(ret == NULL))
-	{
-		if ((flags & MCXT_ALLOC_NO_OOM) == 0)
-		{
-			MemoryContext cxt = GetMemoryChunkContext(pointer);
-
-			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu in memory context \"%s\".",
-							   size, cxt->name)));
-		}
 		return NULL;
-	}
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
@@ -1590,21 +1494,9 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocHugeSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	context->isReset = false;
 
-	ret = context->methods->alloc(context, size);
-	if (unlikely(ret == NULL))
-	{
-		MemoryContextStats(TopMemoryContext);
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu in memory context \"%s\".",
-						   size, context->name)));
-	}
+	ret = context->methods->alloc(context, size, MCXT_ALLOC_HUGE);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 40c1d401c4c..9b54b9546e6 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -496,7 +496,7 @@ SlabDelete(MemoryContext context)
  *		request could not be completed; memory is added to the slab.
  */
 void *
-SlabAlloc(MemoryContext context, Size size)
+SlabAlloc(MemoryContext context, Size size, int flags)
 {
 	SlabContext *slab = (SlabContext *) context;
 	SlabBlock  *block;
@@ -504,6 +504,12 @@ SlabAlloc(MemoryContext context, Size size)
 
 	Assert(SlabIsValid(slab));
 
+	/*
+	 * XXX: Probably no need to check for huge allocations, we only support
+	 * one size? Which could theoretically be huge, but that'd not make
+	 * sense...
+	 */
+
 	/* sanity check that this is pointing to a valid blocklist */
 	Assert(slab->curBlocklistIndex >= 0);
 	Assert(slab->curBlocklistIndex <= SlabBlocklistIndex(slab, slab->chunksPerBlock));
@@ -546,7 +552,7 @@ SlabAlloc(MemoryContext context, Size size)
 			block = (SlabBlock *) malloc(slab->blockSize);
 
 			if (unlikely(block == NULL))
-				return NULL;
+				return MemoryContextAllocationFailure(context, size, flags);
 
 			block->slab = slab;
 			context->mem_allocated += slab->blockSize;
@@ -770,7 +776,7 @@ SlabFree(void *pointer)
  * realloc is usually used to enlarge the chunk.
  */
 void *
-SlabRealloc(void *pointer, Size size)
+SlabRealloc(void *pointer, Size size, int flags)
 {
 	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
 	SlabBlock  *block;
-- 
2.38.0

>From 853a494562f2765275bfdfab2dfeed05293a9502 Mon Sep 17 00:00:00 2001
From: Andres Freund <[email-masked]>
Date: Tue, 18 Jul 2023 20:16:18 -0700
Subject: [PATCH v2 2/4] Optimize AllocSetAlloc() by separating hot from cold
 paths.

With gcc the common paths of AllocSetAlloc() now don't need to initialize a
stack frame when compiling with gcc.
---
 src/backend/utils/mmgr/aset.c | 488 +++++++++++++++++++---------------
 1 file changed, 272 insertions(+), 216 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index b72ec68f7dd..79e04b07c84 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -686,6 +686,260 @@ AllocSetDelete(MemoryContext context)
 	free(set);
 }
 
+/*
+ * Helper for AllocSetAlloc() that allocates an entire block for the chunk.
+ *
+ * AllocSetAlloc()'s comment explains why this is separate.
+ */
+pg_noinline
+static void *
+AllocSetAllocLarge(MemoryContext context, Size size, int flags)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock  block;
+	MemoryChunk *chunk;
+	Size        chunk_size;
+	Size        blksize;
+
+	/* check size, only allocation path where the limits could be hit */
+	MemoryContextCheckSize(context, size, flags);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* ensure there's always space for the sentinel byte */
+	chunk_size = MAXALIGN(size + 1);
+#else
+	chunk_size = MAXALIGN(size);
+#endif
+
+	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	block = (AllocBlock) malloc(blksize);
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
+
+	context->mem_allocated += blksize;
+
+	block->aset = set;
+	block->freeptr = block->endptr = ((char *) block) + blksize;
+
+	chunk = (MemoryChunk *) (((char *) block) + ALLOC_BLOCKHDRSZ);
+
+	/* mark the MemoryChunk as externally managed */
+	MemoryChunkSetHdrMaskExternal(chunk, MCTX_ASET_ID);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	chunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	Assert(size < chunk_size);
+	set_sentinel(MemoryChunkGetPointer(chunk), size);
+#endif
+#ifdef RANDOMIZE_ALLOCATED_MEMORY
+	/* fill the allocated space with junk */
+	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
+#endif
+
+	/*
+	 * Stick the new block underneath the active allocation block, if any,
+	 * so that we don't lose the use of the space remaining therein.
+	 */
+	if (set->blocks != NULL)
+	{
+		block->prev = set->blocks;
+		block->next = set->blocks->next;
+		if (block->next)
+			block->next->prev = block;
+		set->blocks->next = block;
+	}
+	else
+	{
+		block->prev = NULL;
+		block->next = NULL;
+		set->blocks = block;
+	}
+
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+							   chunk_size - size);
+
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
+	return MemoryChunkGetPointer(chunk);
+}
+
+/*
+ * Small helper for allocating a new chunk from a chunk, to avoid duplicating
+ * the code between AllocSetAlloc() and AllocSetAllocFromNewBlock().
+ */
+static inline void *
+AllocSetAllocChunkFromBlock(MemoryContext context, AllocBlock block,
+							Size size, Size chunk_size, int fidx)
+{
+	MemoryChunk *chunk;
+
+	chunk = (MemoryChunk *) (block->freeptr);
+
+	/* Prepare to initialize the chunk header. */
+	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
+	Assert(block->freeptr <= block->endptr);
+
+	/* store the free list index in the value field */
+	MemoryChunkSetHdrMask(chunk, block, fidx, MCTX_ASET_ID);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	chunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	if (size < chunk_size)
+		set_sentinel(MemoryChunkGetPointer(chunk), size);
+#endif
+#ifdef RANDOMIZE_ALLOCATED_MEMORY
+	/* fill the allocated space with junk */
+	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
+#endif
+
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+							   chunk_size - size);
+
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
+	return MemoryChunkGetPointer(chunk);
+}
+
+/*
+ * Helper for AllocSetAlloc() that allocates a new block and returns a chunk
+ * allocated from it.
+ *
+ * AllocSetAlloc()'s comment explains why this is separate.
+ */
+pg_noinline
+static void *
+AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
+						  int fidx)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+	Size		availspace;
+	Size		blksize;
+	Size		required_size;
+	Size		chunk_size;
+
+	/* due to the keeper block set->blocks should always be valid */
+	Assert(set->blocks != NULL);
+	block = set->blocks;
+	availspace = block->endptr - block->freeptr;
+
+	/*
+	 * The existing active (top) block does not have enough room for
+	 * the requested allocation, but it might still have a useful
+	 * amount of space in it.  Once we push it down in the block list,
+	 * we'll never try to allocate more space from it. So, before we
+	 * do that, carve up its free space into chunks that we can put on
+	 * the set's freelists.
+	 *
+	 * Because we can only get here when there's less than
+	 * ALLOC_CHUNK_LIMIT left in the block, this loop cannot iterate
+	 * more than ALLOCSET_NUM_FREELISTS-1 times.
+	 */
+	while (availspace >= ((1 << ALLOC_MINBITS) + ALLOC_CHUNKHDRSZ))
+	{
+		AllocFreeListLink *link;
+		Size		availchunk = availspace - ALLOC_CHUNKHDRSZ;
+		int			a_fidx = AllocSetFreeIndex(availchunk);
+		MemoryChunk *chunk;
+
+		/*
+		 * In most cases, we'll get back the index of the next larger
+		 * freelist than the one we need to put this chunk on.  The
+		 * exception is when availchunk is exactly a power of 2.
+		 */
+		if (availchunk != GetChunkSizeFromFreeListIdx(a_fidx))
+		{
+			a_fidx--;
+			Assert(a_fidx >= 0);
+			availchunk = GetChunkSizeFromFreeListIdx(a_fidx);
+		}
+
+		chunk = (MemoryChunk *) (block->freeptr);
+
+		/* Prepare to initialize the chunk header. */
+		VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+		block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
+		availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
+
+		/* store the freelist index in the value field */
+		MemoryChunkSetHdrMask(chunk, block, a_fidx, MCTX_ASET_ID);
+#ifdef MEMORY_CONTEXT_CHECKING
+		chunk->requested_size = InvalidAllocSize;	/* mark it free */
+#endif
+		/* push this chunk onto the free list */
+		link = GetFreeListLink(chunk);
+
+		VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
+		link->next = set->freelist[a_fidx];
+		VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
+
+		set->freelist[a_fidx] = chunk;
+	}
+
+	/*
+	 * The first such block has size initBlockSize, and we double the
+	 * space in each succeeding block, but not more than maxBlockSize.
+	 */
+	blksize = set->nextBlockSize;
+	set->nextBlockSize <<= 1;
+	if (set->nextBlockSize > set->maxBlockSize)
+		set->nextBlockSize = set->maxBlockSize;
+
+	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
+
+	/*
+	 * If initBlockSize is less than ALLOC_CHUNK_LIMIT, we could need more
+	 * space... but try to keep it a power of 2.
+	 */
+	required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	while (blksize < required_size)
+		blksize <<= 1;
+
+	/* Try to allocate it */
+	block = (AllocBlock) malloc(blksize);
+
+	/*
+	 * We could be asking for pretty big blocks here, so cope if malloc
+	 * fails.  But give up if there's less than 1 MB or so available...
+	 */
+	while (block == NULL && blksize > 1024 * 1024)
+	{
+		blksize >>= 1;
+		if (blksize < required_size)
+			break;
+		block = (AllocBlock) malloc(blksize);
+	}
+
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
+
+	context->mem_allocated += blksize;
+
+	block->aset = set;
+	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
+	block->endptr = ((char *) block) + blksize;
+
+	/* Mark unallocated space NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
+							   blksize - ALLOC_BLOCKHDRSZ);
+
+	block->prev = NULL;
+	block->next = set->blocks;
+	if (block->next)
+		block->next->prev = block;
+	set->blocks = block;
+
+	return AllocSetAllocChunkFromBlock(context, block, size, chunk_size, fidx);
+}
+
 /*
  * AllocSetAlloc
  *		Returns pointer to allocated memory of given size or NULL if
@@ -698,6 +952,13 @@ AllocSetDelete(MemoryContext context)
  * Note: when using valgrind, it doesn't matter how the returned allocation
  * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
  * return space that is marked NOACCESS - AllocSetRealloc has to beware!
+ *
+ * This function should only contain the most common code paths, everything
+ * else should be in pg_noinline helper functions. That allows to avoid the
+ * overhead of creating a stack frame for the common cases - this function is
+ * one of the most common bottlenecks, making this worthwhile. The helper
+ * functions should always directly return the newly allocated memory,
+ * otherwise the stack frame is required after all.
  */
 void *
 AllocSetAlloc(MemoryContext context, Size size, int flags)
@@ -707,80 +968,19 @@ AllocSetAlloc(MemoryContext context, Size size, int flags)
 	MemoryChunk *chunk;
 	int			fidx;
 	Size		chunk_size;
-	Size		blksize;
+	Size		availspace;
 
 	Assert(AllocSetIsValid(set));
 
+	/* due to the keeper block set->blocks should always be valid */
+	Assert(set->blocks != NULL);
+
 	/*
 	 * If requested size exceeds maximum for chunks, allocate an entire block
 	 * for this request.
 	 */
 	if (size > set->allocChunkLimit)
-	{
-		/* check size, only allocation path where the limits could be hit */
-		MemoryContextCheckSize(context, size, flags);
-
-#ifdef MEMORY_CONTEXT_CHECKING
-		/* ensure there's always space for the sentinel byte */
-		chunk_size = MAXALIGN(size + 1);
-#else
-		chunk_size = MAXALIGN(size);
-#endif
-
-		blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-		block = (AllocBlock) malloc(blksize);
-		if (block == NULL)
-			return MemoryContextAllocationFailure(context, size, flags);
-
-		context->mem_allocated += blksize;
-
-		block->aset = set;
-		block->freeptr = block->endptr = ((char *) block) + blksize;
-
-		chunk = (MemoryChunk *) (((char *) block) + ALLOC_BLOCKHDRSZ);
-
-		/* mark the MemoryChunk as externally managed */
-		MemoryChunkSetHdrMaskExternal(chunk, MCTX_ASET_ID);
-
-#ifdef MEMORY_CONTEXT_CHECKING
-		chunk->requested_size = size;
-		/* set mark to catch clobber of "unused" space */
-		Assert(size < chunk_size);
-		set_sentinel(MemoryChunkGetPointer(chunk), size);
-#endif
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* fill the allocated space with junk */
-		randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
-#endif
-
-		/*
-		 * Stick the new block underneath the active allocation block, if any,
-		 * so that we don't lose the use of the space remaining therein.
-		 */
-		if (set->blocks != NULL)
-		{
-			block->prev = set->blocks;
-			block->next = set->blocks->next;
-			if (block->next)
-				block->next->prev = block;
-			set->blocks->next = block;
-		}
-		else
-		{
-			block->prev = NULL;
-			block->next = NULL;
-			set->blocks = block;
-		}
-
-		/* Ensure any padding bytes are marked NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
-								   chunk_size - size);
-
-		/* Disallow access to the chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-
-		return MemoryChunkGetPointer(chunk);
-	}
+		return AllocSetAllocLarge(context, size, flags);
 
 	/*
 	 * Request is small enough to be treated as a chunk.  Look in the
@@ -837,164 +1037,20 @@ AllocSetAlloc(MemoryContext context, Size size, int flags)
 	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
 	Assert(chunk_size >= size);
 
+	block = set->blocks;
+	availspace = block->endptr - block->freeptr;
+
 	/*
 	 * If there is enough room in the active allocation block, we will put the
 	 * chunk into that block.  Else must start a new one.
 	 */
-	if ((block = set->blocks) != NULL)
-	{
-		Size		availspace = block->endptr - block->freeptr;
-
-		if (availspace < (chunk_size + ALLOC_CHUNKHDRSZ))
-		{
-			/*
-			 * The existing active (top) block does not have enough room for
-			 * the requested allocation, but it might still have a useful
-			 * amount of space in it.  Once we push it down in the block list,
-			 * we'll never try to allocate more space from it. So, before we
-			 * do that, carve up its free space into chunks that we can put on
-			 * the set's freelists.
-			 *
-			 * Because we can only get here when there's less than
-			 * ALLOC_CHUNK_LIMIT left in the block, this loop cannot iterate
-			 * more than ALLOCSET_NUM_FREELISTS-1 times.
-			 */
-			while (availspace >= ((1 << ALLOC_MINBITS) + ALLOC_CHUNKHDRSZ))
-			{
-				AllocFreeListLink *link;
-				Size		availchunk = availspace - ALLOC_CHUNKHDRSZ;
-				int			a_fidx = AllocSetFreeIndex(availchunk);
-
-				/*
-				 * In most cases, we'll get back the index of the next larger
-				 * freelist than the one we need to put this chunk on.  The
-				 * exception is when availchunk is exactly a power of 2.
-				 */
-				if (availchunk != GetChunkSizeFromFreeListIdx(a_fidx))
-				{
-					a_fidx--;
-					Assert(a_fidx >= 0);
-					availchunk = GetChunkSizeFromFreeListIdx(a_fidx);
-				}
-
-				chunk = (MemoryChunk *) (block->freeptr);
-
-				/* Prepare to initialize the chunk header. */
-				VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
-				block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
-				availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
-
-				/* store the freelist index in the value field */
-				MemoryChunkSetHdrMask(chunk, block, a_fidx, MCTX_ASET_ID);
-#ifdef MEMORY_CONTEXT_CHECKING
-				chunk->requested_size = InvalidAllocSize;	/* mark it free */
-#endif
-				/* push this chunk onto the free list */
-				link = GetFreeListLink(chunk);
-
-				VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
-				link->next = set->freelist[a_fidx];
-				VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
-
-				set->freelist[a_fidx] = chunk;
-			}
-			/* Mark that we need to create a new block */
-			block = NULL;
-		}
-	}
-
-	/*
-	 * Time to create a new regular (multi-chunk) block?
-	 */
-	if (block == NULL)
-	{
-		Size		required_size;
-
-		/*
-		 * The first such block has size initBlockSize, and we double the
-		 * space in each succeeding block, but not more than maxBlockSize.
-		 */
-		blksize = set->nextBlockSize;
-		set->nextBlockSize <<= 1;
-		if (set->nextBlockSize > set->maxBlockSize)
-			set->nextBlockSize = set->maxBlockSize;
-
-		/*
-		 * If initBlockSize is less than ALLOC_CHUNK_LIMIT, we could need more
-		 * space... but try to keep it a power of 2.
-		 */
-		required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-		while (blksize < required_size)
-			blksize <<= 1;
-
-		/* Try to allocate it */
-		block = (AllocBlock) malloc(blksize);
-
-		/*
-		 * We could be asking for pretty big blocks here, so cope if malloc
-		 * fails.  But give up if there's less than 1 MB or so available...
-		 */
-		while (block == NULL && blksize > 1024 * 1024)
-		{
-			blksize >>= 1;
-			if (blksize < required_size)
-				break;
-			block = (AllocBlock) malloc(blksize);
-		}
-
-		if (block == NULL)
-			return MemoryContextAllocationFailure(context, size, flags);
-
-		context->mem_allocated += blksize;
-
-		block->aset = set;
-		block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
-		block->endptr = ((char *) block) + blksize;
-
-		/* Mark unallocated space NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
-								   blksize - ALLOC_BLOCKHDRSZ);
-
-		block->prev = NULL;
-		block->next = set->blocks;
-		if (block->next)
-			block->next->prev = block;
-		set->blocks = block;
-	}
+	if (unlikely(availspace < (chunk_size + ALLOC_CHUNKHDRSZ)))
+		return AllocSetAllocFromNewBlock(context, size, flags, fidx);
 
 	/*
 	 * OK, do the allocation
 	 */
-	chunk = (MemoryChunk *) (block->freeptr);
-
-	/* Prepare to initialize the chunk header. */
-	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
-
-	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
-	Assert(block->freeptr <= block->endptr);
-
-	/* store the free list index in the value field */
-	MemoryChunkSetHdrMask(chunk, block, fidx, MCTX_ASET_ID);
-
-#ifdef MEMORY_CONTEXT_CHECKING
-	chunk->requested_size = size;
-	/* set mark to catch clobber of "unused" space */
-	if (size < chunk_size)
-		set_sentinel(MemoryChunkGetPointer(chunk), size);
-#endif
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-	/* fill the allocated space with junk */
-	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
-#endif
-
-	/* Ensure any padding bytes are marked NOACCESS. */
-	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
-							   chunk_size - size);
-
-	/* Disallow access to the chunk header. */
-	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-
-	return MemoryChunkGetPointer(chunk);
+	return AllocSetAllocChunkFromBlock(context, block, size, chunk_size, fidx);
 }
 
 /*
-- 
2.38.0

---

From: David Rowley <[email-masked]>
Date: Fri, 21 Jul 2023 14:03:46 +1200

> steps, and a bit more polished.

I've only just briefly read through the updated patch, but I did take
it for a spin to see what sort of improvements I can get from it.

The attached graph shows the time in seconds that it took for each
allocator to allocate 10GBs of memory resetting the context once 1MB
is allocated.  The data point for aset with 32-byte chunks takes
master 1.697 seconds and with both patches, it goes down to 1.264,
which is a 34% increase in performance.

It's pretty nice that we can hide the AllocSizeIsValid tests inside
the allocChunkLimit path and pretty good that we can skip the NULL
checks in most cases since we're not having to check for malloc
failure unless we malloc a new block.

I'll reply back with a more detailed review next week.

David

---

From: John Naylor <[email-masked]>
Date: Tue, 8 Aug 2023 18:18:51 +0700

> steps, and a bit more polished.

Here is a quick test -- something similar was used to measure the slab
improvements last cycle. With radix tree v37 0001-0011 from [1],

create extension bench_radix_tree;
select avg(load_ms) from generate_series(1,100) x(x), lateral (select *
from bench_load_random_int(100 * 1000 * (1+x-x))) a;

The backend was pinned and turbo off. Perf runs were separate from timed
runs. I included 0002 for completeness.

v37
         avg
---------------------
 27.0400000000000000

  32.42%  postgres  bench_radix_tree.so  [.] rt_recursive_set
  21.60%  postgres  postgres             [.] SlabAlloc
  11.06%  postgres  [unknown]            [k] 0xffffffff930018f7
  10.49%  postgres  bench_radix_tree.so  [.] rt_extend_down
   7.07%  postgres  postgres             [.] MemoryContextAlloc
   4.83%  postgres  bench_radix_tree.so  [.] rt_node_insert_inner
   2.19%  postgres  bench_radix_tree.so  [.] rt_grow_node_48
   2.16%  postgres  bench_radix_tree.so  [.] rt_set.isra.0
   1.50%  postgres  bench_radix_tree.so  [.] MemoryContextAlloc@plt

v37 + palloc sibling calls
         avg
---------------------
 26.0700000000000000

v37 + palloc sibling calls + opt aset
         avg
---------------------
 26.0900000000000000

  33.78%  postgres  bench_radix_tree.so  [.] rt_recursive_set
  23.04%  postgres  postgres             [.] SlabAlloc
  11.43%  postgres  [unknown]            [k] 0xffffffff930018f7
  11.05%  postgres  bench_radix_tree.so  [.] rt_extend_down
   5.52%  postgres  bench_radix_tree.so  [.] rt_node_insert_inner
   2.47%  postgres  bench_radix_tree.so  [.] rt_set.isra.0
   2.30%  postgres  bench_radix_tree.so  [.] rt_grow_node_48
   1.88%  postgres  postgres             [.] MemoryContextAlloc
   1.44%  postgres  bench_radix_tree.so  [.] MemoryContextAlloc@plt

It's nice to see MemoryContextAlloc go down in the profile.

[1]
https://www.postgresql.org/message-id/[email-masked]

--
John Naylor
EDB: http://www.enterprisedb.com

---

From: David Rowley <[email-masked]>
Date: Wed, 9 Aug 2023 20:44:43 +1200

> I'll reply back with a more detailed review next week.

Here's a review of v2-0001:

1.

/*
* XXX: Should this also be moved into alloc()? We could possibly avoid
* zeroing in some cases (e.g. if we used mmap() ourselves.
*/
MemSetAligned(ret, 0, size);

Maybe this should be moved to the alloc function.  It would allow us
to get rid of this:

#define palloc0fast(sz) \
( MemSetTest(0, sz) ? \
MemoryContextAllocZeroAligned(CurrentMemoryContext, sz) : \
MemoryContextAllocZero(CurrentMemoryContext, sz) )

If we do the zeroing inside the alloc function then it can always use
the MemoryContextAllocZeroAligned version providing we zero before
setting the sentinel byte.

It would allow the tail call in the palloc0() case, but the drawback
would be having to check for the MCXT_ALLOC_ZERO flag in the alloc
function. I wonder if that branch would be predictable in most cases,
e.g. the parser will be making lots of nodes and want to zero all
allocations, but the executor won't be doing much of that. There will
be a mix of zeroing and not zeroing in the planner, mostly not, I
think.

2. Why do you need to add the NULL check here?

 #ifdef USE_VALGRIND

I know it's just valgrind code and performance does not matter, but
the realloc flags are being passed as 0, so allocation failures won't
return.

3.

/*
* XXX: Probably no need to check for huge allocations, we only support
* one size? Which could theoretically be huge, but that'd not make
* sense...
*/

They can't be huge per Assert(fullChunkSize <= MEMORYCHUNK_MAX_VALUE)
in SlabContextCreate().

4. It would be good to see some API documentation in the
MemoryContextMethods struct.  This adds a lot of responsibility onto
the context implementation without any extra documentation to explain
what, for example, palloc is responsible for and what the alloc
function needs to do itself.

David

---

From: David Rowley <[email-masked]>
Date: Fri, 23 Feb 2024 00:46:26 +1300

I've rebased the 0001 patch and gone over it again and made a few
additional changes besides what I mentioned in my review.

>   VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);

I removed this NULL check as we're calling the realloc function with
no flags, so it shouldn't return NULL as it'll error out from any OOM
errors.

> They can't be huge per Assert(fullChunkSize <= MEMORYCHUNK_MAX_VALUE)

I removed this comment and adjusted the comment just below that which
checks the 'size' matches the expected slab chunk size. i.e.

/*
* Make sure we only allow correct request size.  This doubles as the
* MemoryContextCheckSize check.
*/
if (unlikely(size != slab->chunkSize))

> function needs to do itself.

I've done that too.

I also added header comments for MemoryContextAllocationFailure and
MemoryContextSizeFailure and added some comments to explain in places
like palloc() to warn people not to add checks after the 'alloc' call.

The rebased patch is 0001 and all of my changes are in 0002.  I will
rebase your original 0002 patch later.  I think 0001 is much more
important, as evident by the reported benchmarks on this thread.

In absence of anyone else looking at this, I think it's ready to go.
If anyone is following along and wants to review or test it, please do
so soon.

David
From 854751aeb1d0b917bf6e96c30e80f8480f96e883 Mon Sep 17 00:00:00 2001
From: Andres Freund <[email-masked]>
Date: Tue, 18 Jul 2023 18:55:58 -0700
Subject: [PATCH v3 1/2] Optimize palloc() etc to allow sibling calls

The reason palloc() etc previously couldn't use sibling calls is that they did
check the returned value (to e.g. raise an error when the allocation
fails). Push the error handling down into the memory context implementation -
they can avoid performing the check at all in the hot code paths.
---
 src/backend/utils/mmgr/alignedalloc.c |  11 +-
 src/backend/utils/mmgr/aset.c         |  23 ++--
 src/backend/utils/mmgr/generation.c   |  14 ++-
 src/backend/utils/mmgr/mcxt.c         | 172 ++++++--------------------
 src/backend/utils/mmgr/slab.c         |  12 +-
 src/include/nodes/memnodes.h          |   4 +-
 src/include/utils/memutils_internal.h |  29 +++--
 7 files changed, 101 insertions(+), 164 deletions(-)

index 7204fe64ae..c266fb3dbb 100644

 	/*
 
 	return newptr;
index 2f99fa9a2f..81c3120c2b 100644

 
 		if (block == NULL)
-			return NULL;

 
 		set = block->aset;

 

 
 		/*
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index f9016a7ed7..4b4155b863 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -344,7 +344,7 @@ GenerationDelete(MemoryContext context)
  * return space that is marked NOACCESS - GenerationRealloc has to beware!
  */
 void *
-GenerationAlloc(MemoryContext context, Size size)
+GenerationAlloc(MemoryContext context, Size size, int flags)
 {
 	GenerationContext *set = (GenerationContext *) context;
 	GenerationBlock *block;
@@ -367,9 +367,11 @@ GenerationAlloc(MemoryContext context, Size size)
 	{
 		Size		blksize = required_size + Generation_BLOCKHDRSZ;

 
@@ -472,7 +474,7 @@ GenerationAlloc(MemoryContext context, Size size)
 			block = (GenerationBlock *) malloc(blksize);
 
 			if (block == NULL)
-				return NULL;

@@ -737,7 +739,7 @@ GenerationFree(void *pointer)
  *		into the old chunk - in that case we just update chunk header.
  */
 void *
-GenerationRealloc(void *pointer, Size size)
+GenerationRealloc(void *pointer, Size size, int flags)
 {
 	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
 	GenerationContext *set;
@@ -840,14 +842,14 @@ GenerationRealloc(void *pointer, Size size)
 	}
 
 	/* allocate new chunk */

 	/*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index ad7409a02c..2947b2046b 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -34,7 +34,7 @@

 
 static void *
@@ -1023,6 +1023,27 @@ MemoryContextCreate(MemoryContext node,
 
+void *
@@ -1038,28 +1059,9 @@ MemoryContextAlloc(MemoryContext context, Size size)

 
 
@@ -1081,24 +1083,16 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
+	/*
 
 	return ret;
@@ -1122,20 +1116,9 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)

 
 
@@ -1207,22 +1190,11 @@ palloc(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
+	Assert(ret != 0);
 
 	return ret;
@@ -1238,21 +1210,9 @@ palloc0(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
@@ -1271,24 +1231,11 @@ palloc_extended(Size size, int flags)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
@@ -1458,29 +1405,15 @@ repalloc(void *pointer, Size size)
 #endif
 	void	   *ret;

 
 
 #ifdef USE_VALGRIND
 
@@ -1500,31 +1433,14 @@ repalloc_extended(void *pointer, Size size, int flags)
 #endif
 	void	   *ret;

 

@@ -1565,21 +1481,9 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
index b8f00cfc7c..1345c8d500 100644

+	/*

 			block->slab = slab;
index a48f7e5a18..cc8cc58a0c 100644
 
index e0c4f3d5af..e5cecd7122 100644

 
 
2.40.1

From 9b6ffdceb3c6f592d3d1597b19f1240d4a78d848 Mon Sep 17 00:00:00 2001
From: David Rowley <[email-masked]>
Date: Fri, 23 Feb 2024 00:27:07 +1300
Subject: [PATCH v3 2/2] fixup! Optimize palloc() etc to allow sibling calls

---
 src/backend/utils/mmgr/mcxt.c         | 67 +++++++++++++++++++++++++--
 src/backend/utils/mmgr/slab.c         | 11 ++---
 src/include/nodes/memnodes.h          | 38 +++++++++++++++
 src/include/utils/memutils_internal.h |  7 +--
 4 files changed, 110 insertions(+), 13 deletions(-)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 2947b2046b..c214d323c8 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -1023,6 +1023,12 @@ MemoryContextCreate(MemoryContext node,
 	VALGRIND_CREATE_MEMPOOL(node, 0, false);
 }
 
+/*
+ * MemoryContextAllocationFailure
+ *		For use by MemoryContextMethods implementations to handle when malloc
+ *		returns NULL.  The bahavior is specific to whether MCXT_ALLOC_NO_OOM
+ *		is in 'flags'.
+ */
 void *
 MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
 {
@@ -1038,6 +1044,11 @@ MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
 	return NULL;
 }
 
+/*
+ * MemoryContextSizeFailure
+ *		For use by MemoryContextMethods implementations to handle invalid
+ *		memory allocation request sizes.
+ */
 void
 MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
 {
@@ -1061,6 +1072,16 @@ MemoryContextAlloc(MemoryContext context, Size size)

+	/*
+	 * For efficiency reasons, we purposefully offload the handling of
+	 * allocation failures to the MemoryContextMethods implementation as this
+	 * allows these checks to be performed only when an actual malloc needs to
+	 * be done to request more memory from the OS.  Additionally, not having
+	 * to execute any instructions after this call allows the compiler to use
+	 * the tailcall optimization.  If you're considering adding code after
+	 * this call, consider making it the responsibility of the 'alloc'
+	 * function instead.
+	 */
 	ret = context->methods->alloc(context, size, 0);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
@@ -1192,9 +1213,19 @@ palloc(Size size)

+	/*
+	 * For efficiency reasons, we purposefully offload the handling of
+	 * allocation failures to the MemoryContextMethods implementation as this
+	 * allows these checks to be performed only when an actual malloc needs to
+	 * be done to request more memory from the OS.  Additionally, not having
+	 * to execute any instructions after this call allows the compiler to use
+	 * the tailcall optimization.  If you're considering adding code after
+	 * this call, consider making it the responsibility of the 'alloc'
+	 * function instead.
+	 */
 	ret = context->methods->alloc(context, size, 0);
-
-	Assert(ret != 0);
+	/* We expect OOM to be handled by the alloc function */
+	Assert(ret != NULL);
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -1410,10 +1441,20 @@ repalloc(void *pointer, Size size)
 	/* isReset must be false already */
 	Assert(!context->isReset);
 
+	/*
+	 * For efficiency reasons, we purposefully offload the handling of
+	 * allocation failures to the MemoryContextMethods implementation as this
+	 * allows these checks to be performed only when an actual malloc needs to
+	 * be done to request more memory from the OS.  Additionally, not having
+	 * to execute any instructions after this call allows the compiler to use
+	 * the tailcall optimization.  If you're considering adding code after
+	 * this call, consider making it the responsibility of the 'realloc'
+	 * function instead.
+	 */
 	ret = MCXT_METHOD(pointer, realloc) (pointer, size, 0);
 
 #ifdef USE_VALGRIND
-	if (ret != NULL && method != MCTX_ALIGNED_REDIRECT_ID)
+	if (method != MCTX_ALIGNED_REDIRECT_ID)
 		VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 #endif
 
@@ -1438,6 +1479,16 @@ repalloc_extended(void *pointer, Size size, int flags)
 	/* isReset must be false already */
 	Assert(!context->isReset);
 
+	/*
+	 * For efficiency reasons, we purposefully offload the handling of
+	 * allocation failures to the MemoryContextMethods implementation as this
+	 * allows these checks to be performed only when an actual malloc needs to
+	 * be done to request more memory from the OS.  Additionally, not having
+	 * to execute any instructions after this call allows the compiler to use
+	 * the tailcall optimization.  If you're considering adding code after
+	 * this call, consider making it the responsibility of the 'realloc'
+	 * function instead.
+	 */
 	ret = MCXT_METHOD(pointer, realloc) (pointer, size, flags);
 	if (unlikely(ret == NULL))
 		return NULL;
@@ -1483,6 +1534,16 @@ MemoryContextAllocHuge(MemoryContext context, Size size)

+	/*
+	 * For efficiency reasons, we purposefully offload the handling of
+	 * allocation failures to the MemoryContextMethods implementation as this
+	 * allows these checks to be performed only when an actual malloc needs to
+	 * be done to request more memory from the OS.  Additionally, not having
+	 * to execute any instructions after this call allows the compiler to use
+	 * the tailcall optimization.  If you're considering adding code after
+	 * this call, consider making it the responsibility of the 'alloc'
+	 * function instead.
+	 */
 	ret = context->methods->alloc(context, size, MCXT_ALLOC_HUGE);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 1345c8d500..bc91446cb3 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -504,17 +504,14 @@ SlabAlloc(MemoryContext context, Size size, int flags)

-	/*
-	 * XXX: Probably no need to check for huge allocations, we only support
-	 * one size? Which could theoretically be huge, but that'd not make
-	 * sense...
-	 */
-
 	/* sanity check that this is pointing to a valid blocklist */
 	Assert(slab->curBlocklistIndex >= 0);
 	Assert(slab->curBlocklistIndex <= SlabBlocklistIndex(slab, slab->chunksPerBlock));
 
-	/* make sure we only allow correct request size */
+	/*
+	 * Make sure we only allow correct request size.  This doubles as the
+	 * MemoryContextCheckSize check.
+	 */
 	if (unlikely(size != slab->chunkSize))
 		elog(ERROR, "unexpected alloc chunk size %zu (expected %u)",
 			 size, slab->chunkSize);
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index cc8cc58a0c..edc0257f36 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -57,20 +57,58 @@ typedef void (*MemoryStatsPrintFunc) (MemoryContext context, void *passthru,
 
 typedef struct MemoryContextMethods
 {
+	/*
+	 * Function to handle memory allocation requests of 'size' to allocate
+	 * memory into the given 'context'.  The function must handle flags
+	 * MCXT_ALLOC_HUGE and MCXT_ALLOC_NO_OOM.  MCXT_ALLOC_ZERO is handled by
+	 * the calling function.
+	 */
 	void	   *(*alloc) (MemoryContext context, Size size, int flags);
+
 	/* call this free_p in case someone #define's free() */
 	void		(*free_p) (void *pointer);
+
+	/*
+	 * Function to handle a size change request for an existing allocation.
+	 * The implementation must handle flags MCXT_ALLOC_HUGE and
+	 * MCXT_ALLOC_NO_OOM.  MCXT_ALLOC_ZERO is handled by the calling function.
+	 */
 	void	   *(*realloc) (void *pointer, Size size, int flags);
+
+	/*
+	 * Invalidate all previous allocations in the given memory context and
+	 * prepare the context for a new set of allocations.  Implementations may
+	 * optionally free() excess memory back to the OS during this time.
+	 */
 	void		(*reset) (MemoryContext context);
+
+	/* Free all memory consumed by the given MemoryContext. */
 	void		(*delete_context) (MemoryContext context);
+
+	/* Return the MemoryContext that the given pointer belongs to. */
 	MemoryContext (*get_chunk_context) (void *pointer);
+
+	/*
+	 * Return the number of bytes consumed by the given pointer within its
+	 * memory context, including the overhead of alignment and chunk headers.
+	 */
 	Size		(*get_chunk_space) (void *pointer);
+
+	/*
+	 * Return true if the given MemoryContext has not had any allocations
+	 * since it was created or last reset.
+	 */
 	bool		(*is_empty) (MemoryContext context);
 	void		(*stats) (MemoryContext context,
 						  MemoryStatsPrintFunc printfunc, void *passthru,
 						  MemoryContextCounters *totals,
 						  bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
+
+	/*
+	 * Perform validation checks on the given context and raise any discovered
+	 * anomalies as WARNINGs.
+	 */
 	void		(*check) (MemoryContext context);
 #endif
 } MemoryContextMethods;
diff --git a/src/include/utils/memutils_internal.h b/src/include/utils/memutils_internal.h
index e5cecd7122..ad1048fd82 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -133,9 +133,11 @@ extern void MemoryContextCreate(MemoryContext node,
 								MemoryContext parent,
 								const char *name);
 
-extern void *MemoryContextAllocationFailure(MemoryContext context, Size size, int flags);
+extern void *MemoryContextAllocationFailure(MemoryContext context, Size size,
+											int flags);
 
-extern void MemoryContextSizeFailure(MemoryContext context, Size size, int flags) pg_attribute_noreturn();
+extern void MemoryContextSizeFailure(MemoryContext context, Size size,
+									 int flags) pg_attribute_noreturn();
 
 static inline void
 MemoryContextCheckSize(MemoryContext context, Size size, int flags)
@@ -147,5 +149,4 @@ MemoryContextCheckSize(MemoryContext context, Size size, int flags)
 	}
 }
 
-
2.40.1

---

From: Andres Freund <[email-masked]>
Date: Thu, 22 Feb 2024 14:53:43 -0800

Hi,

> no flags, so it shouldn't return NULL as it'll error out from any OOM

That was probably a copy-paste issue...

> rebase your original 0002 patch later.

Thanks!

> I think 0001 is much more important, as evident by the reported benchmarks
> on this thread.

I agree that it's good to tackle 0001 first.

I don't understand the benchmark point though. Your benchmark seems to suggest
that 0002 improves aset performance by *more* than 0001: for 8 byte aset
allocs:

        time
master: 8.86
0001:   8.12
0002:   7.02

So 0001 reduces time by 0.92x and 0002 by 0.86x.

John's test shows basically no change for 0002 - which is unsurprising, as
0002 changes aset.c, but the test seems to solely exercise slab, as only
SlabAlloc() shows up in the profile. As 0002 only touches aset.c it couldn't
really have affected that test.

> If anyone is following along and wants to review or test it, please do

Makes sense!

>  	context->isReset = false;

For a moment this made me wonder if we could move the isReset handling into
the allocator slow paths as well - it's annoying to write that bit (and thus
dirty the cacheline) over and ove. But it'd be somewhat awkward due to
pre-allocated blocks. So that'd be a larger change better done separately.

Greetings,

Andres Freund

---

From: David Rowley <[email-masked]>
Date: Mon, 26 Feb 2024 20:42:34 +1300

> pre-allocated blocks. So that'd be a larger change better done separately.

It makes sense to do this, but on looking closer for aset.c, it seems
like the only time we can avoid un-setting the isReset flag is when
allocating from the freelist.  We must unset it for large allocations
and for allocations that don't fit onto the existing block (the
exiting block could be the keeper block) and for allocations that
require a new block.

With the current arrangement of code in generation.c, I didn't see any
path we could skip doing context->isReset = false.

For slab.c, it's very easy and we can skip setting the isReset in most cases.

I've attached the patches I benchmarked against 449e798c7 and also the
patch I used to add a function to exercise palloc.

The query I ran was:

select chksz,mtype,pg_allocate_memory_test_reset(chksz, 1024*1024,
1024*1024*1024, mtype)
from (values(8),(16),(32),(64)) sizes(chksz),
(values('aset'),('generation'),('slab')) cxt(mtype)
order by mtype,chksz;

David
From 14c1a7dd73c8b011cb2c986aaefdc7588a44c84a Mon Sep 17 00:00:00 2001
From: Andres Freund <[email-masked]>
Date: Tue, 18 Jul 2023 18:55:58 -0700
Subject: [PATCH v4 1/3] Optimize palloc() etc to allow sibling calls

The reason palloc() etc previously couldn't use sibling calls is that they did
check the returned value (to e.g. raise an error when the allocation
fails). Push the error handling down into the memory context implementation -
they can avoid performing the check at all in the hot code paths.
---
 src/backend/utils/mmgr/alignedalloc.c |  11 +-
 src/backend/utils/mmgr/aset.c         |  23 ++-
 src/backend/utils/mmgr/generation.c   |  14 +-
 src/backend/utils/mmgr/mcxt.c         | 233 +++++++++++---------------
 src/backend/utils/mmgr/slab.c         |  11 +-
 src/include/nodes/memnodes.h          |  42 ++++-
 src/include/utils/memutils_internal.h |  30 +++-
 7 files changed, 199 insertions(+), 165 deletions(-)

 
 	/*
 
 	return newptr;

 
 		if (block == NULL)
-			return NULL;

 
 		set = block->aset;

 

 
 		/*

 
 
 			if (block == NULL)
-				return NULL;

 
 	/* allocate new chunk */

 	/*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index ad7409a02c..c214d323c8 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -34,7 +34,7 @@

 
 static void *
@@ -1023,6 +1023,38 @@ MemoryContextCreate(MemoryContext node,
 
+/*
@@ -1038,28 +1070,19 @@ MemoryContextAlloc(MemoryContext context, Size size)

 
 
@@ -1081,24 +1104,16 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
+	/*
 
 	return ret;
@@ -1122,20 +1137,9 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)

 
 
@@ -1207,22 +1211,21 @@ palloc(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 	return ret;
@@ -1238,21 +1241,9 @@ palloc0(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
@@ -1271,24 +1262,11 @@ palloc_extended(Size size, int flags)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
@@ -1458,26 +1436,22 @@ repalloc(void *pointer, Size size)
 #endif
 	void	   *ret;

 
 
 #ifdef USE_VALGRIND
 	if (method != MCTX_ALIGNED_REDIRECT_ID)
@@ -1500,31 +1474,24 @@ repalloc_extended(void *pointer, Size size, int flags)
 #endif
 	void	   *ret;

 

@@ -1565,21 +1532,19 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);

 
 
index b8f00cfc7c..bc91446cb3 100644
@@ -508,7 +508,10 @@ SlabAlloc(MemoryContext context, Size size)
 
@@ -546,7 +549,7 @@ SlabAlloc(MemoryContext context, Size size)

 			block->slab = slab;
@@ -770,7 +773,7 @@ SlabFree(void *pointer)
index a48f7e5a18..edc0257f36 100644
 
index e0c4f3d5af..ad1048fd82 100644

 
@@ -133,4 +133,20 @@ extern void MemoryContextCreate(MemoryContext node,
 								MemoryContext parent,
 								const char *name);
 
2.40.1.windows.1

From 239295dc3dca9ad1f05a71a2169c89405f4ce041 Mon Sep 17 00:00:00 2001
From: Andres Freund <[email-masked]>
Date: Tue, 18 Jul 2023 20:16:18 -0700
Subject: [PATCH v4 2/3] Optimize AllocSetAlloc() by separating hot from cold
 paths.

With gcc the common paths of AllocSetAlloc() now don't need to initialize a
stack frame when compiling with gcc.
---
 src/backend/utils/mmgr/aset.c | 482 +++++++++++++++++++---------------
 1 file changed, 269 insertions(+), 213 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 81c3120c2b..4e52eb063c 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -687,101 +687,301 @@ AllocSetDelete(MemoryContext context)
 }
 
 /*
- * AllocSetAlloc
- *		Returns pointer to allocated memory of given size or NULL if
- *		request could not be completed; memory is added to the set.
+ * Helper for AllocSetAlloc() that allocates an entire block for the chunk.
  *
- * No request may exceed:
- *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
- * All callers use a much-lower limit.
- *
- * Note: when using valgrind, it doesn't matter how the returned allocation
- * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
- * return space that is marked NOACCESS - AllocSetRealloc has to beware!
+ * AllocSetAlloc()'s comment explains why this is separate.
  */
-void *
-AllocSetAlloc(MemoryContext context, Size size, int flags)
+pg_noinline
+static void *
+AllocSetAllocLarge(MemoryContext context, Size size, int flags)
 {
 	AllocSet	set = (AllocSet) context;
-	AllocBlock	block;
+	AllocBlock  block;
 	MemoryChunk *chunk;
-	int			fidx;
-	Size		chunk_size;
-	Size		blksize;
+	Size        chunk_size;
+	Size        blksize;
 
-	Assert(AllocSetIsValid(set));
 
 	/*
-	 * If requested size exceeds maximum for chunks, allocate an entire block
-	 * for this request.
-	if (size > set->allocChunkLimit)

 
 
-		block->aset = set;

 
 		set_sentinel(MemoryChunkGetPointer(chunk), size);
 #ifdef RANDOMIZE_ALLOCATED_MEMORY

 

+	/*
+	 * The first such block has size initBlockSize, and we double the
+	 * space in each succeeding block, but not more than maxBlockSize.
+	 */
+	blksize = set->nextBlockSize;
+	set->nextBlockSize <<= 1;
+	if (set->nextBlockSize > set->maxBlockSize)
+		set->nextBlockSize = set->maxBlockSize;
+
+	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
+
+	/*
+	 * If initBlockSize is less than ALLOC_CHUNK_LIMIT, we could need more
+	 * space... but try to keep it a power of 2.
+	 */
+	required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	while (blksize < required_size)
+		blksize <<= 1;
+
+	/* Try to allocate it */
+	block = (AllocBlock) malloc(blksize);
+
+	/*
+	 * We could be asking for pretty big blocks here, so cope if malloc
+	 * fails.  But give up if there's less than 1 MB or so available...
+	 */
+	while (block == NULL && blksize > 1024 * 1024)
+	{
+		blksize >>= 1;
+		if (blksize < required_size)
+			break;
+		block = (AllocBlock) malloc(blksize);
+	}
+
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
+
+	context->mem_allocated += blksize;
+
+	block->aset = set;
+	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
+	block->endptr = ((char *) block) + blksize;
+
+	/* Mark unallocated space NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
+							   blksize - ALLOC_BLOCKHDRSZ);
+
+	block->prev = NULL;
+	block->next = set->blocks;
+	if (block->next)
+		block->next->prev = block;
+	set->blocks = block;
+
+	return AllocSetAllocChunkFromBlock(context, block, size, chunk_size, fidx);
+}
+
+/*
+ * AllocSetAlloc
+ *		Returns pointer to allocated memory of given size or NULL if
+ *		request could not be completed; memory is added to the set.
+ *
+ * No request may exceed:
+ *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
+ * All callers use a much-lower limit.
+ *
+ * Note: when using valgrind, it doesn't matter how the returned allocation
+ * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
+ * return space that is marked NOACCESS - AllocSetRealloc has to beware!
+ *
+ * This function should only contain the most common code paths, everything
+ * else should be in pg_noinline helper functions. That allows to avoid the
+ * overhead of creating a stack frame for the common cases - this function is
+ * one of the most common bottlenecks, making this worthwhile. The helper
+ * functions should always directly return the newly allocated memory,
+ * otherwise the stack frame is required after all.
+ */
+void *
+AllocSetAlloc(MemoryContext context, Size size, int flags)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+	MemoryChunk *chunk;
+	int			fidx;
+	Size		chunk_size;
+	Size		availspace;
+
+	Assert(AllocSetIsValid(set));
+
+	/* due to the keeper block set->blocks should always be valid */
+	Assert(set->blocks != NULL);
+
+	/*
+	 * If requested size exceeds maximum for chunks, allocate an entire block
+	 * for this request.
+	 */
+	if (size > set->allocChunkLimit)
+		return AllocSetAllocLarge(context, size, flags);
+
 	/*
 	 * Request is small enough to be treated as a chunk.  Look in the
 	 * corresponding free list to see if there is a free chunk we could reuse.
@@ -837,164 +1037,20 @@ AllocSetAlloc(MemoryContext context, Size size, int flags)
 	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
 	Assert(chunk_size >= size);
 
+	block = set->blocks;
 
 	/*
 	 * OK, do the allocation
 	 */
 
 /*
-- 
2.40.1.windows.1

From c69840cd85767fa7ee4d1688a605a0978c3fc39f Mon Sep 17 00:00:00 2001
From: David Rowley <[email-masked]>
Date: Mon, 26 Feb 2024 20:05:34 +1300
Subject: [PATCH v4 3/3] Make unsetting the isReset flag the MemoryContext's
 job

Up until now, setting this boolean flag to false has been the job of the
allocation function in mcxt.c.  Making this the job of the MemoryContext
implementation can make things more efficient as adjusting the flag can
just be put in code paths that are only possible to reach in cases where
we may not have already done an allocation since the context was created
or reset.

Discussion: https://postgr.es/m/[email-masked]
---
 src/backend/utils/mmgr/aset.c       | 11 ++++++++++
 src/backend/utils/mmgr/generation.c |  2 ++
 src/backend/utils/mmgr/mcxt.c       | 31 ++++++++++++++++-------------
 src/backend/utils/mmgr/slab.c       | 14 +++++++++++++
 src/include/nodes/memnodes.h        |  5 ++++-
 5 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 4e52eb063c..a1c59381d5 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -711,6 +711,7 @@ AllocSetAllocLarge(MemoryContext context, Size size, int flags)
 	chunk_size = MAXALIGN(size);
 #endif
 
+	context->isReset = false;
 	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
 	block = (AllocBlock) malloc(blksize);
 	if (block == NULL)
@@ -1000,6 +1001,9 @@ AllocSetAlloc(MemoryContext context, Size size, int flags)
 	{
 		AllocFreeListLink *link = GetFreeListLink(chunk);
 
+		/* should already be unset if we've something in the freelist */
+		Assert(context->isReset == false);
+
 		/* Allow access to the chunk header. */
 		VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
 
@@ -1040,6 +1044,13 @@ AllocSetAlloc(MemoryContext context, Size size, int flags)
 	block = set->blocks;
 	availspace = block->endptr - block->freeptr;
 
+	/*
+	 * We must un-reset the context here as lack of space for this allocation
+	 * on a block is no guarantee that we've already seen an allocation as
+	 * the given block could be the keeper block.
+	 */
+	context->isReset = false;
+
 	/*
 	 * If there is enough room in the active allocation block, we will put the
 	 * chunk into that block.  Else must start a new one.
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 4b4155b863..5a351c26df 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -352,6 +352,8 @@ GenerationAlloc(MemoryContext context, Size size, int flags)
 	Size		chunk_size;
 	Size		required_size;
 
+	context->isReset = false;
+
 	Assert(GenerationIsValid(set));
 
 #ifdef MEMORY_CONTEXT_CHECKING
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index c214d323c8..01f3bec028 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -1070,8 +1070,6 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	/*
 	 * For efficiency reasons, we purposefully offload the handling of
 	 * allocation failures to the MemoryContextMethods implementation as this
@@ -1084,6 +1082,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	 */
 	ret = context->methods->alloc(context, size, 0);
 
+	Assert(!context->isReset);
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -1104,10 +1104,10 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	ret = context->methods->alloc(context, size, 0);
 
+	Assert(!context->isReset);
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	/*
@@ -1135,9 +1135,10 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 		  AllocSizeIsValid(size)))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
-	context->isReset = false;
-
 	ret = context->methods->alloc(context, size, flags);
+
+	Assert(!context->isReset);
+
 	if (unlikely(ret == NULL))
 		return NULL;
 
@@ -1211,8 +1212,6 @@ palloc(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	/*
 	 * For efficiency reasons, we purposefully offload the handling of
 	 * allocation failures to the MemoryContextMethods implementation as this
@@ -1224,6 +1223,9 @@ palloc(Size size)
 	 * function instead.
 	 */
 	ret = context->methods->alloc(context, size, 0);
+
+	Assert(!context->isReset);
+
 	/* We expect OOM to be handled by the alloc function */
 	Assert(ret != NULL);
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
@@ -1241,10 +1243,10 @@ palloc0(Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	ret = context->methods->alloc(context, size, 0);
 
+	Assert(!context->isReset);
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	MemSetAligned(ret, 0, size);
@@ -1262,9 +1264,10 @@ palloc_extended(Size size, int flags)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	ret = context->methods->alloc(context, size, flags);
+
+	Assert(!context->isReset);
+
 	if (unlikely(ret == NULL))
 	{
 		return NULL;
@@ -1532,8 +1535,6 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	Assert(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	context->isReset = false;
-
 	/*
 	 * For efficiency reasons, we purposefully offload the handling of
 	 * allocation failures to the MemoryContextMethods implementation as this
@@ -1546,6 +1547,8 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	 */
 	ret = context->methods->alloc(context, size, MCXT_ALLOC_HUGE);
 
+	Assert(!context->isReset);
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index bc91446cb3..1160aa7e20 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -533,6 +533,12 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 		{
 			dlist_node *node = dclist_pop_head_node(&slab->emptyblocks);
 
+			/*
+			 * We should already have already un-reset the context if we've
+			 * got a block on the emptyblocks list.
+			 */
+			Assert(!context->isReset);
+
 			block = dlist_container(SlabBlock, node, node);
 
 			/*
@@ -546,6 +552,8 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 		}
 		else
 		{
+			context->isReset = false;
+
 			block = (SlabBlock *) malloc(slab->blockSize);
 
 			if (unlikely(block == NULL))
@@ -581,6 +589,12 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 
 		Assert(!dlist_is_empty(blocklist));
 
+		/*
+		 * We should already have already un-reset the context if we've got a
+		 * non-zero curBlocklistIndex
+		 */
+		Assert(!context->isReset);
+
 		/* grab the block from the blocklist */
 		block = dlist_head_element(SlabBlock, node, blocklist);
 
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index edc0257f36..72539cfc5d 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -61,7 +61,10 @@ typedef struct MemoryContextMethods
 	 * Function to handle memory allocation requests of 'size' to allocate
 	 * memory into the given 'context'.  The function must handle flags
 	 * MCXT_ALLOC_HUGE and MCXT_ALLOC_NO_OOM.  MCXT_ALLOC_ZERO is handled by
-	 * the calling function.
+	 * the calling function.  The implementation must handle setting
+	 * MemoryContext.isReset to false before returning.  Leaving this up to
+	 * the implementing function can allow this flag to be unset more
+	 * efficiently.
 	 */
 	void	   *(*alloc) (MemoryContext context, Size size, int flags);
 
-- 
2.40.1.windows.1

From 519ba53c785e1858952c2fa2692dc34aed54b76a Mon Sep 17 00:00:00 2001
From: David Rowley <[email-masked]>
Date: Wed, 25 May 2022 14:24:37 +1200
Subject: [PATCH v99] Function to test palloc/pfree performance

---
 src/backend/utils/adt/mcxtfuncs.c | 213 ++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.dat   |  12 ++
 2 files changed, 225 insertions(+)

diff --git a/src/backend/utils/adt/mcxtfuncs.c b/src/backend/utils/adt/mcxtfuncs.c
index 4708d73f5f..416235b473 100644
--- a/src/backend/utils/adt/mcxtfuncs.c
+++ b/src/backend/utils/adt/mcxtfuncs.c
@@ -15,6 +15,8 @@
 
 #include "postgres.h"
 
+#include <time.h>
+
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "mb/pg_wchar.h"
@@ -193,3 +195,214 @@ pg_log_backend_memory_contexts(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(true);
 }
+
+typedef struct AllocateTestNext
+{
+	struct AllocateTestNext *next;		/* ptr to the next allocation */
+} AllocateTestNext;
+
+/* #define ALLOCATE_TEST_DEBUG */
+/*
+ * pg_allocate_memory_test
+ *		Used to test the performance of a memory context types
+ */
+Datum
+pg_allocate_memory_test(PG_FUNCTION_ARGS)
+{
+	int32	chunk_size = PG_GETARG_INT32(0);
+	int64	keep_memory = PG_GETARG_INT64(1);
+	int64	total_alloc = PG_GETARG_INT64(2);
+	text   *context_type_text = PG_GETARG_TEXT_PP(3);
+	char   *context_type;
+	int64	curr_memory_use = 0;
+	int64	remaining_alloc_bytes = total_alloc;
+	MemoryContext context;
+	MemoryContext oldContext;
+	AllocateTestNext	   *next_free_ptr = NULL;
+	AllocateTestNext	   *last_alloc = NULL;
+	clock_t	start, end;
+
+	if (chunk_size < sizeof(AllocateTestNext))
+		elog(ERROR, "chunk_size (%d) must be at least %ld bytes", chunk_size,
+			 sizeof(AllocateTestNext));
+	if (keep_memory > total_alloc)
+		elog(ERROR, "keep_memory (" INT64_FORMAT ") must be less than total_alloc (" INT64_FORMAT ")",
+			 keep_memory, total_alloc);
+
+	context_type = text_to_cstring(context_type_text);
+
+	start = clock();
+
+	if (strcmp(context_type, "generation") == 0)
+		context = GenerationContextCreate(CurrentMemoryContext,
+										  "pg_allocate_memory_test",
+										  ALLOCSET_DEFAULT_SIZES);
+	else if (strcmp(context_type, "aset") == 0)
+		context = AllocSetContextCreate(CurrentMemoryContext,
+										"pg_allocate_memory_test",
+										ALLOCSET_DEFAULT_SIZES);
+	else if (strcmp(context_type, "slab") == 0)
+		context = SlabContextCreate(CurrentMemoryContext,
+									"pg_allocate_memory_test",
+									ALLOCSET_DEFAULT_MAXSIZE,
+									chunk_size);
+	else
+		elog(ERROR, "context_type must be \"generation\", \"aset\" or \"slab\"");
+
+	oldContext = MemoryContextSwitchTo(context);
+
+	while (remaining_alloc_bytes > 0)
+	{
+		AllocateTestNext *curr_alloc;
+
+		CHECK_FOR_INTERRUPTS();
+
+		/* Allocate the memory and update the counters */
+		curr_alloc = (AllocateTestNext *) palloc(chunk_size);
+		remaining_alloc_bytes -= chunk_size;
+		curr_memory_use += chunk_size;
+
+#ifdef ALLOCATE_TEST_DEBUG
+		elog(NOTICE, "alloc %p (curr_memory_use " INT64_FORMAT " bytes, remaining_alloc_bytes " INT64_FORMAT ")", curr_alloc, curr_memory_use, remaining_alloc_bytes);
+#endif
+
+		/*
+		 * Point the last allocate to this one so that we can free allocations
+		 * starting with the oldest first.
+		 */
+		curr_alloc->next = NULL;
+		if (last_alloc != NULL)
+			last_alloc->next = curr_alloc;
+
+		if (next_free_ptr == NULL)
+		{
+			/*
+			 * Remember the first chunk to free. We will follow the ->next
+			 * pointers to find the next chunk to free when freeing memory
+			 */
+			next_free_ptr = curr_alloc;
+		}
+
+		/*
+		 * If the currently allocated memory has reached or exceeded the amount
+		 * of memory we want to keep allocated at once then we'd better free
+		 * some.  Since all allocations are the same size we only need to free
+		 * one allocation per loop.
+		 */
+		if (curr_memory_use >= keep_memory)
+		{
+			AllocateTestNext	 *next = next_free_ptr->next;
+
+			/* free the memory and update the current memory usage */
+			pfree(next_free_ptr);
+			curr_memory_use -= chunk_size;
+
+#ifdef ALLOCATE_TEST_DEBUG
+			elog(NOTICE, "free %p (curr_memory_use " INT64_FORMAT " bytes, remaining_alloc_bytes " INT64_FORMAT ")", next_free_ptr, curr_memory_use, remaining_alloc_bytes);
+#endif
+			/* get the next chunk to free */
+			next_free_ptr = next;
+		}
+
+		if (curr_memory_use > 0)
+			last_alloc = curr_alloc;
+		else
+			last_alloc = NULL;
+	}
+
+	/* cleanup loop -- pfree remaining memory */
+	while (next_free_ptr != NULL)
+	{
+		AllocateTestNext	 *next = next_free_ptr->next;
+
+		/* free the memory and update the current memory usage */
+		pfree(next_free_ptr);
+		curr_memory_use -= chunk_size;
+
+#ifdef ALLOCATE_TEST_DEBUG
+		elog(NOTICE, "free %p (curr_memory_use " INT64_FORMAT " bytes, remaining_alloc_bytes " INT64_FORMAT ")", next_free_ptr, curr_memory_use, remaining_alloc_bytes);
+#endif
+
+		next_free_ptr = next;
+	}
+
+	MemoryContextSwitchTo(oldContext);
+
+	end = clock();
+
+	PG_RETURN_FLOAT8((double) (end - start) / CLOCKS_PER_SEC);
+}
+
+Datum
+pg_allocate_memory_test_reset(PG_FUNCTION_ARGS)
+{
+	int32	chunk_size = PG_GETARG_INT32(0);
+	int64	keep_memory = PG_GETARG_INT64(1);
+	int64	total_alloc = PG_GETARG_INT64(2);
+	text   *context_type_text = PG_GETARG_TEXT_PP(3);
+	char   *context_type;
+	int64	curr_memory_use = 0;
+	int64	remaining_alloc_bytes = total_alloc;
+	MemoryContext context;
+	MemoryContext oldContext;
+	clock_t	start, end;
+
+	if (chunk_size < 1)
+		elog(ERROR, "size of chunk must be above 0");
+	if (keep_memory > total_alloc)
+		elog(ERROR, "keep_memory (" INT64_FORMAT ") must be less than total_alloc (" INT64_FORMAT ")",
+			 keep_memory, total_alloc);
+
+	context_type = text_to_cstring(context_type_text);
+
+	start = clock();
+
+	if (strcmp(context_type, "generation") == 0)
+		context = GenerationContextCreate(CurrentMemoryContext,
+										  "pg_allocate_memory_test",
+										  ALLOCSET_DEFAULT_SIZES);
+	else if (strcmp(context_type, "aset") == 0)
+		context = AllocSetContextCreate(CurrentMemoryContext,
+										"pg_allocate_memory_test",
+										ALLOCSET_DEFAULT_SIZES);
+	else if (strcmp(context_type, "slab") == 0)
+		context = SlabContextCreate(CurrentMemoryContext,
+									"pg_allocate_memory_test",
+									ALLOCSET_DEFAULT_MAXSIZE,
+									chunk_size);
+	else
+		elog(ERROR, "context_type must be \"generation\", \"aset\" or \"slab\"");
+
+	oldContext = MemoryContextSwitchTo(context);
+
+	while (remaining_alloc_bytes > 0)
+	{
+		CHECK_FOR_INTERRUPTS();
+
+		/* Allocate the memory and update the counters */
+		(void) palloc(chunk_size);
+		remaining_alloc_bytes -= chunk_size;
+		curr_memory_use += chunk_size;
+
+#ifdef ALLOCATE_TEST_DEBUG
+		elog(NOTICE, "alloc %p (curr_memory_use " INT64_FORMAT " bytes, remaining_alloc_bytes " INT64_FORMAT ")", curr_alloc, curr_memory_use, remaining_alloc_bytes);
+#endif
+
+		/*
+		 * If the currently allocated memory has reached or exceeded the amount
+		 * of memory we want to keep allocated at once then reset the context.
+		 */
+		if (curr_memory_use >= keep_memory)
+		{
+			curr_memory_use = 0;
+			MemoryContextReset(context);
+		}
+	}
+
+	MemoryContextSwitchTo(oldContext);
+	MemoryContextDelete(context);
+
+	end = clock();
+
+	PG_RETURN_FLOAT8((double) (end - start) / CLOCKS_PER_SEC);
+}
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 9c120fc2b7..f307229c60 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -8274,6 +8274,18 @@
   prorettype => 'bool', proargtypes => 'int4',
   prosrc => 'pg_log_backend_memory_contexts' },
 
+# just for testing memory context allocation speed
+{ oid => '9319', descr => 'for testing performance of allocation and freeing',
+  proname => 'pg_allocate_memory_test', provolatile => 'v',
+  prorettype => 'float8', proargtypes => 'int4 int8 int8 text',
+  prosrc => 'pg_allocate_memory_test' },
+
+# just for testing memory context allocation speed
+{ oid => '9320', descr => 'for testing performance of allocation and resetting',
+  proname => 'pg_allocate_memory_test_reset', provolatile => 'v',
+  prorettype => 'float8', proargtypes => 'int4 int8 int8 text',
+  prosrc => 'pg_allocate_memory_test_reset' },
+
 # non-persistent series generator
 { oid => '1066', descr => 'non-persistent series generator',
   proname => 'generate_series', prorows => '1000',
-- 
2.40.1.windows.1

---

From: David Rowley <[email-masked]>
Date: Thu, 29 Feb 2024 00:29:17 +1300

> > If anyone is following along and wants to review or test it, please do

I pushed the 0001 and 0002 patches today.

I switched over to working on doing what you did in 0002 for
generation.c and slab.c.

See the attached patch which runs the same test as in [1] (aset.c is
just there for comparisons between slab and generation)

The attached includes some additional tuning to generation.c:

1) Changed GenerationFree() to not free() the current block when it
becomes empty. The code now just marks it as empty and reuses it.
Saves free()/malloc() cycle. Also means we can get rid of a NULL check
in GenerationAlloc().

2) Removed code in GenerationAlloc() which I felt was trying too hard
to fill the keeper, free and current block.  The changes I made here
do mean that once the keeper block becomes empty, it won't be used
again until the context is reset and gets a new allocation.  I don't
see this as a big issue as the keeper block is small anyway.

generation.c is now ~30% faster for the 8-byte test.

David

[1] https://postgr.es/m/[email-masked]
From 2f15c7ef04ebe5af539cfa2e3e05f2d491d4fc97 Mon Sep 17 00:00:00 2001
From: David Rowley <[email-masked]>
Date: Wed, 28 Feb 2024 17:08:58 +1300
Subject: [PATCH v1] Optimize GenerationAlloc() and SlabAlloc()

In a similar effort to 413c18401, separate out the hot and cold paths in
GenerationAlloc() and SlabAlloc() to avoid having to setup the stack frame
for the hot path.

This additionally adjusts how we adjust the GenerationContext's
freeblock.  freeblock, when set is now always empty and we only switch
to using it when the current allocation request finds the current block
does not have enough space and the freeblock is large enough to
accomodate the allocation.

This commit also adjusts GenerationFree() so that if we pfree a pointer
that's the current generation block, we now mark that block as empty and
keep it as the current block.  Previously we free'd that block and set
the current block to NULL.  Doing that meant we needed a special case in
GenerationAlloc to check if GenerationContext.block was NULL.

Author: David Rowley
---
 src/backend/utils/mmgr/aset.c       |   5 +-
 src/backend/utils/mmgr/generation.c | 389 ++++++++++++++++------------
 src/backend/utils/mmgr/slab.c       | 230 +++++++++-------
 3 files changed, 368 insertions(+), 256 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 0cfee52274..e5eb7a1f39 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -943,8 +943,9 @@ AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
 
 /*
  * AllocSetAlloc
- *		Returns pointer to allocated memory of given size or NULL if
- *		request could not be completed; memory is added to the set.
+ *		Returns a pointer to allocated memory of given size or raises an ERROR
+ *		on allocation failure, or returns NULL when flags contains
+ *		MCXT_ALLOC_NO_OOM.
  *
  * No request may exceed:
  *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index ae4a7c999e..8e9289c659 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -69,8 +69,8 @@ typedef struct GenerationContext
 	GenerationBlock *block;		/* current (most recently allocated) block, or
 								 * NULL if we've just freed the most recent
 								 * block */
-	GenerationBlock *freeblock; /* pointer to a block that's being recycled,
-								 * or NULL if there's no such block. */
+	GenerationBlock *freeblock; /* pointer to empty block that's being
+								 * recycled, or NULL if there's no such block. */
 	dlist_head	blocks;			/* list of blocks */
 } GenerationContext;
 
@@ -331,28 +331,23 @@ GenerationDelete(MemoryContext context)
 }
 
 /*
- * GenerationAlloc
- *		Returns pointer to allocated memory of given size or NULL if
- *		request could not be completed; memory is added to the set.
- *
- * No request may exceed:
- *		MAXALIGN_DOWN(SIZE_MAX) - Generation_BLOCKHDRSZ - Generation_CHUNKHDRSZ
- * All callers use a much-lower limit.
+ * Helper for GenerationAlloc() that allocates an entire block for the chunk.
  *
- * Note: when using valgrind, it doesn't matter how the returned allocation
- * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
- * return space that is marked NOACCESS - GenerationRealloc has to beware!
+ * GenerationAlloc()'s comment explains why this is separate.
  */
-void *
-GenerationAlloc(MemoryContext context, Size size, int flags)
+pg_noinline
+static void *
+GenerationAllocLarge(MemoryContext context, Size size, int flags)
 {
 	GenerationContext *set = (GenerationContext *) context;
 	GenerationBlock *block;
 	MemoryChunk *chunk;
 	Size		chunk_size;
 	Size		required_size;
+	Size		blksize;
 
-	Assert(GenerationIsValid(set));
+	/* validate 'size' is within the limits for the given 'flags' */
+	MemoryContextCheckSize(context, size, flags);
 
 #ifdef MEMORY_CONTEXT_CHECKING
 	/* ensure there's always space for the sentinel byte */
@@ -361,141 +356,66 @@ GenerationAlloc(MemoryContext context, Size size, int flags)
 	chunk_size = MAXALIGN(size);
 #endif
 	required_size = chunk_size + Generation_CHUNKHDRSZ;
+	blksize = required_size + Generation_BLOCKHDRSZ;
 
-	/* is it an over-sized chunk? if yes, allocate special block */
-	if (chunk_size > set->allocChunkLimit)
-	{
-		Size		blksize = required_size + Generation_BLOCKHDRSZ;
-
-		/* only check size in paths where the limits could be hit */
-		MemoryContextCheckSize((MemoryContext) set, size, flags);
-
-		block = (GenerationBlock *) malloc(blksize);
-		if (block == NULL)
-			return MemoryContextAllocationFailure(context, size, flags);
+	block = (GenerationBlock *) malloc(blksize);
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);

-		/* block with a single (used) chunk */
-		block->context = set;
-		block->blksize = blksize;
-		block->nchunks = 1;
-		block->nfree = 0;
+	/* block with a single (used) chunk */
+	block->context = set;
+	block->blksize = blksize;
+	block->nchunks = 1;
+	block->nfree = 0;
 
-		/* the block is completely full */
-		block->freeptr = block->endptr = ((char *) block) + blksize;
+	/* the block is completely full */
+	block->freeptr = block->endptr = ((char *) block) + blksize;
 
-		chunk = (MemoryChunk *) (((char *) block) + Generation_BLOCKHDRSZ);
+	chunk = (MemoryChunk *) (((char *) block) + Generation_BLOCKHDRSZ);
 
-		/* mark the MemoryChunk as externally managed */
-		MemoryChunkSetHdrMaskExternal(chunk, MCTX_GENERATION_ID);
+	/* mark the MemoryChunk as externally managed */
+	MemoryChunkSetHdrMaskExternal(chunk, MCTX_GENERATION_ID);

-		/* add the block to the list of allocated blocks */
-		dlist_push_head(&set->blocks, &block->node);
-
-		/* Ensure any padding bytes are marked NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
-								   chunk_size - size);
-
-		/* Disallow access to the chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, Generation_CHUNKHDRSZ);
-
-		return MemoryChunkGetPointer(chunk);
-	}
-
-	/*
-	 * Not an oversized chunk.  We try to first make use of the current block,
-	 * but if there's not enough space in it, instead of allocating a new
-	 * block, we look to see if the freeblock is empty and has enough space.
-	 * If not, we'll also try the same using the keeper block.  The keeper
-	 * block may have become empty and we have no other way to reuse it again
-	 * if we don't try to use it explicitly here.
-	 *
-	 * We don't want to start filling the freeblock before the current block
-	 * is full, otherwise we may cause fragmentation in FIFO type workloads.
-	 * We only switch to using the freeblock or keeper block if those blocks
-	 * are completely empty.  If we didn't do that we could end up fragmenting
-	 * consecutive allocations over multiple blocks which would be a problem
-	 * that would compound over time.
-	 */
-	block = set->block;
-
-	if (block == NULL ||
-		GenerationBlockFreeBytes(block) < required_size)
-	{
-		Size		blksize;
-		GenerationBlock *freeblock = set->freeblock;
-
-		if (freeblock != NULL &&
-			GenerationBlockIsEmpty(freeblock) &&
-			GenerationBlockFreeBytes(freeblock) >= required_size)
-		{
-			block = freeblock;
-
-			/*
-			 * Zero out the freeblock as we'll set this to the current block
-			 * below
-			 */
-			set->freeblock = NULL;
-		}
-		else if (GenerationBlockIsEmpty(KeeperBlock(set)) &&
-				 GenerationBlockFreeBytes(KeeperBlock(set)) >= required_size)
-		{
-			block = KeeperBlock(set);
-		}
-		else
-		{
-			/*
-			 * The first such block has size initBlockSize, and we double the
-			 * space in each succeeding block, but not more than maxBlockSize.
-			 */
-			blksize = set->nextBlockSize;
-			set->nextBlockSize <<= 1;
-			if (set->nextBlockSize > set->maxBlockSize)
-				set->nextBlockSize = set->maxBlockSize;
-
-			/* we'll need a block hdr too, so add that to the required size */
-			required_size += Generation_BLOCKHDRSZ;
-
-			/* round the size up to the next power of 2 */
-			if (blksize < required_size)
-				blksize = pg_nextpower2_size_t(required_size);
-
-			block = (GenerationBlock *) malloc(blksize);
-
-			if (block == NULL)
-				return MemoryContextAllocationFailure(context, size, flags);
-
-			context->mem_allocated += blksize;
+	/* add the block to the list of allocated blocks */
+	dlist_push_head(&set->blocks, &block->node);
 
-			/* initialize the new block */
-			GenerationBlockInit(set, block, blksize);
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+							   chunk_size - size);
 
-			/* add it to the doubly-linked list of blocks */
-			dlist_push_head(&set->blocks, &block->node);
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, Generation_CHUNKHDRSZ);
 
-			/* Zero out the freeblock in case it's become full */
-			set->freeblock = NULL;
-		}
+	return MemoryChunkGetPointer(chunk);
+}
 
-		/* and also use it as the current allocation block */
-		set->block = block;
-	}
+/*
+ * Small helper for allocating a new chunk from a chunk, to avoid duplicating
+ * the code between GenerationAlloc() and GenerationAllocFromNewBlock().
+ */
+static inline void *
+GenerationAllocChunkFromBlock(MemoryContext context, GenerationBlock *block,
+							  Size size, Size chunk_size)
+{
+	MemoryChunk *chunk = (MemoryChunk *) (block->freeptr);
 
-	/* we're supposed to have a block with enough free space now */
+	/* validate we've been given a block with enough free space */
 	Assert(block != NULL);
-	Assert((block->endptr - block->freeptr) >= Generation_CHUNKHDRSZ + chunk_size);
+	Assert((block->endptr - block->freeptr) >=
+		   Generation_CHUNKHDRSZ + chunk_size);
 
 	chunk = (MemoryChunk *) block->freeptr;
 
@@ -529,6 +449,155 @@ GenerationAlloc(MemoryContext context, Size size, int flags)
 	return MemoryChunkGetPointer(chunk);
 }
 
+/*
+ * Helper for GenerationAlloc() that allocates a new block and returns a chunk
+ * allocated from it.
+ *
+ * GenerationAlloc()'s comment explains why this is separate.
+ */
+pg_noinline
+static void *
+GenerationAllocFromNewBlock(MemoryContext context, Size size, int flags,
+							Size chunk_size)
+{
+	GenerationContext *set = (GenerationContext *) context;
+	GenerationBlock *block;
+	Size		blksize;
+	Size		required_size;
+
+	/*
+	 * The first such block has size initBlockSize, and we double the space in
+	 * each succeeding block, but not more than maxBlockSize.
+	 */
+	blksize = set->nextBlockSize;
+	set->nextBlockSize <<= 1;
+	if (set->nextBlockSize > set->maxBlockSize)
+		set->nextBlockSize = set->maxBlockSize;
+
+	/* we'll need space for the chunk, chunk hdr and block hdr */
+	required_size = chunk_size + Generation_CHUNKHDRSZ + Generation_BLOCKHDRSZ;
+
+	/* round the size up to the next power of 2 */
+	if (blksize < required_size)
+		blksize = pg_nextpower2_size_t(required_size);
+
+	block = (GenerationBlock *) malloc(blksize);
+
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
+
+	context->mem_allocated += blksize;
+
+	/* initialize the new block */
+	GenerationBlockInit(set, block, blksize);
+
+	/* add it to the doubly-linked list of blocks */
+	dlist_push_head(&set->blocks, &block->node);
+
+	/* make this the current block */
+	set->block = block;
+
+	return GenerationAllocChunkFromBlock(context, block, size, chunk_size);
+}
+
+/*
+ * GenerationAlloc
+ *		Returns a pointer to allocated memory of given size or raises an ERROR
+ *		on allocation failure, or returns NULL when flags contains
+ *		MCXT_ALLOC_NO_OOM.
+ *
+ * No request may exceed:
+ *		MAXALIGN_DOWN(SIZE_MAX) - Generation_BLOCKHDRSZ - Generation_CHUNKHDRSZ
+ * All callers use a much-lower limit.
+ *
+ * Note: when using valgrind, it doesn't matter how the returned allocation
+ * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
+ * return space that is marked NOACCESS - GenerationRealloc has to beware!
+ *
+ * This function should only contain the most common code paths.  Everything
+ * else should be in pg_noinline helper functions, thus avoiding the overheads
+ * creating a stack frame for the common cases.  Allocating memory is often a
+ * bottleneck in many workloads, so avoiding stack frame setup is worthwhile.
+ * Helper functions should always directly return the newly allocated memory
+ * so that we can just return that address directly as a tail call.
+ */
+void *
+GenerationAlloc(MemoryContext context, Size size, int flags)
+{
+	GenerationContext *set = (GenerationContext *) context;
+	GenerationBlock *block;
+	Size		chunk_size;
+	Size		required_size;
+
+	Assert(GenerationIsValid(set));
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* ensure there's always space for the sentinel byte */
+	chunk_size = MAXALIGN(size + 1);
+#else
+	chunk_size = MAXALIGN(size);
+#endif
+
+	/*
+	 * If requested size exceeds maximum for chunks we hand the the request
+	 * off to GenerationAllocLarge().
+	 */
+	if (chunk_size > set->allocChunkLimit)
+		return GenerationAllocLarge(context, size, flags);
+
+	required_size = chunk_size + Generation_CHUNKHDRSZ;
+
+	/*
+	 * Not an oversized chunk.  We try to first make use of the current block,
+	 * but if there's not enough space in it, instead of allocating a new
+	 * block, we look to see if the empty freeblock has enough space.  We
+	 * don't try reusing the keeper block.  If it's become empty we'll reuse
+	 * that again only if the context is reset.
+	 *
+	 * We only try reusing the freeblock if we've no space for this allocation
+	 * on the current block.  When a freeblock exists, we'll switch to it once
+	 * the first time we can't fit an allocation in the current block.  We
+	 * avoid ping-ponging between the two as we need to be careful not to
+	 * fragment differently sized consecutive allocations between several
+	 * blocks.  Going between the two could cause fragmentation for FIFO
+	 * workloads, which generation is meant to be good at.
+	 */
+	block = set->block;
+
+	if (unlikely(GenerationBlockFreeBytes(block) < required_size))
+	{
+		GenerationBlock *freeblock = set->freeblock;
+
+		/* freeblock, if set, must be empty */
+		Assert(freeblock == NULL || GenerationBlockIsEmpty(freeblock));
+
+		/* check if we have a freeblock and if it's big enough */
+		if (freeblock != NULL &&
+			GenerationBlockFreeBytes(freeblock) >= required_size)
+		{
+			/* make the freeblock the current block */
+			set->freeblock = NULL;
+			set->block = freeblock;
+
+			return GenerationAllocChunkFromBlock(context,
+												 freeblock,
+												 size,
+												 chunk_size);
+		}
+		else
+		{
+			/*
+			 * No freeblock, or it's not big enough for this allocation.  Make
+			 * a new block.
+			 */
+			return GenerationAllocFromNewBlock(context, size, flags, chunk_size);
+		}
+	}
+
+	/* The current block has space, so just allocate chunk there. */
+	return GenerationAllocChunkFromBlock(context, block, size, chunk_size);
+}
+
 /*
  * GenerationBlockInit
  *		Initializes 'block' assuming 'blksize'.  Does not update the context's
@@ -621,8 +690,8 @@ GenerationBlockFree(GenerationContext *set, GenerationBlock *block)
 
 /*
  * GenerationFree
- *		Update number of chunks in the block, and if all chunks in the block
- *		are now free then discard the block.
+ *		Update number of chunks in the block, and consider freeing the block
+ *		if it's become empty.
  */
 void
 GenerationFree(void *pointer)
@@ -694,43 +763,37 @@ GenerationFree(void *pointer)
 	Assert(block->nfree <= block->nchunks);
 
 	/* If there are still allocated chunks in the block, we're done. */
-	if (block->nfree < block->nchunks)
+	if (likely(block->nfree < block->nchunks))
 		return;
 
 	set = block->context;
 
-	/* Don't try to free the keeper block, just mark it empty */
-	if (IsKeeperBlock(set, block))
-	{
-		GenerationBlockMarkEmpty(block);
-		return;
-	}
-
-	/*
-	 * If there is no freeblock set or if this is the freeblock then instead
-	 * of freeing this memory, we keep it around so that new allocations have
-	 * the option of recycling it.
+	/*-----------------------
+	 * The block this allocation was on has now become completely empty of
+	 * chunks.  In the general case, we can now return the memory for this
+	 * block back to malloc.  However, there are cases where we don't want to
+	 * do that:
+	 *
+	 * 1)	If it's the keeper block.  This block is malloc'd with the context
+	 *		and can't be free'd without freeing the context itself.
+	 * 2)	If it's the current block.  We could free this, but doing so would
+	 *		leave us nothing to set the current block to, so we just mark the
+	 *		block as empty so new allocations can reuse it again.
+	 * 3)	If we have no "freeblock" set, then to avoid new allocations from
+	 *		having to malloc a new block again, we save a single block to
+	 *		avoid that.  This is especially useful for FIFO workloads as it
+	 *		avoids continual free/malloc cycles.
 	 */
-	if (set->freeblock == NULL || set->freeblock == block)
+	if (IsKeeperBlock(set, block) || set->block == block)
+		GenerationBlockMarkEmpty(block);	/* case 1 and 2 */
+	else if (set->freeblock == NULL)
 	{
-		/* XXX should we only recycle maxBlockSize sized blocks? */
-		set->freeblock = block;
+		/* case 3 */
 		GenerationBlockMarkEmpty(block);
-		return;
+		set->freeblock = block;
 	}
-
-	/* Also make sure the block is not marked as the current block. */
-	if (set->block == block)
-		set->block = NULL;
-
-	/*
-	 * The block is empty, so let's get rid of it. First remove it from the
-	 * list of blocks, then return it to malloc().
-	 */
-	dlist_delete(&block->node);
-
-	set->header.mem_allocated -= block->blksize;
-	free(block);
+	else
+		GenerationBlockFree(set, block);	/* Otherwise, free it */
 }
 
 /*
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index bc91446cb3..c5f985d5ea 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -490,10 +490,139 @@ SlabDelete(MemoryContext context)
 	free(context);
 }
 
+/*
+ * Small helper for allocating a new chunk from a chunk, to avoid duplicating
+ * the code between SlabAlloc() and SlabAllocFromNewBlock().
+ */
+static inline void *
+SlabAllocSetupNewChunk(MemoryContext context, SlabBlock *block,
+					   MemoryChunk *chunk, Size size)
+{
+	SlabContext *slab = (SlabContext *) context;
+
+	/*
+	 * Check that the chunk pointer is actually somewhere on the block and is
+	 * aligned as expected.
+	 */
+	Assert(chunk >= SlabBlockGetChunk(slab, block, 0));
+	Assert(chunk <= SlabBlockGetChunk(slab, block, slab->chunksPerBlock - 1));
+	Assert(SlabChunkMod(slab, block, chunk) == 0);
+
+	/* Prepare to initialize the chunk header. */
+	VALGRIND_MAKE_MEM_UNDEFINED(chunk, Slab_CHUNKHDRSZ);
+
+	MemoryChunkSetHdrMask(chunk, block, MAXALIGN(slab->chunkSize), MCTX_SLAB_ID);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* slab mark to catch clobber of "unused" space */
+	Assert(slab->chunkSize < (slab->fullChunkSize - Slab_CHUNKHDRSZ));
+	set_sentinel(MemoryChunkGetPointer(chunk), size);
+	VALGRIND_MAKE_MEM_NOACCESS(((char *) chunk) + Slab_CHUNKHDRSZ +
+							   slab->chunkSize,
+							   slab->fullChunkSize -
+							   (slab->chunkSize + Slab_CHUNKHDRSZ));
+#endif
+
+#ifdef RANDOMIZE_ALLOCATED_MEMORY
+	/* fill the allocated space with junk */
+	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
+#endif
+
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, Slab_CHUNKHDRSZ);
+
+	return MemoryChunkGetPointer(chunk);
+}
+
+pg_noinline
+static void *
+SlabAllocFromNewBlock(MemoryContext context, Size size, int flags)
+{
+	SlabContext *slab = (SlabContext *) context;
+	SlabBlock  *block;
+	MemoryChunk *chunk;
+	dlist_head *blocklist;
+	int			blocklist_idx;
+
+	/* to save allocating a new one, first check the empty blocks list */
+	if (dclist_count(&slab->emptyblocks) > 0)
+	{
+		dlist_node *node = dclist_pop_head_node(&slab->emptyblocks);
+
+		block = dlist_container(SlabBlock, node, node);
+
+		/*
+		 * SlabFree() should have left this block in a valid state with all
+		 * chunks free.  Ensure that's the case.
+		 */
+		Assert(block->nfree == slab->chunksPerBlock);
+
+		/* fetch the next chunk from this block */
+		chunk = SlabGetNextFreeChunk(slab, block);
+	}
+	else
+	{
+		block = (SlabBlock *) malloc(slab->blockSize);
+
+		if (unlikely(block == NULL))
+			return MemoryContextAllocationFailure(context, size, flags);
+
+		block->slab = slab;
+		context->mem_allocated += slab->blockSize;
+
+		/* use the first chunk in the new block */
+		chunk = SlabBlockGetChunk(slab, block, 0);
+
+		block->nfree = slab->chunksPerBlock - 1;
+		block->unused = SlabBlockGetChunk(slab, block, 1);
+		block->freehead = NULL;
+		block->nunused = slab->chunksPerBlock - 1;
+	}
+
+	/* find the blocklist element for storing blocks with 1 used chunk */
+	blocklist_idx = SlabBlocklistIndex(slab, block->nfree);
+	blocklist = &slab->blocklist[blocklist_idx];
+
+	/* this better be empty.  We just added a block thinking it was */
+	Assert(dlist_is_empty(blocklist));
+
+	dlist_push_head(blocklist, &block->node);
+
+	slab->curBlocklistIndex = blocklist_idx;
+
+	return SlabAllocSetupNewChunk(context, block, chunk, size);
+}
+
+/*
+ * SlabAllocInvalidSize
+ *		Handle raising an ERROR for an invalid size request.  We don't do this
+ *		in slab alloc as calling the elog functions would force the compiler
+ *		to setup the stack frame in SlabAlloc.  For performance reasons, we
+ *		want to avoid that.
+ */
+pg_noinline
+static void
+pg_attribute_noreturn()
+SlabAllocInvalidSize(MemoryContext context, Size size)
+{
+	SlabContext *slab = (SlabContext *) context;
+
+	elog(ERROR, "unexpected alloc chunk size %zu (expected %u)", size,
+		 slab->chunkSize);
+}
+
 /*
  * SlabAlloc
- *		Returns a pointer to allocated memory of given size or NULL if
- *		request could not be completed; memory is added to the slab.
+ *		Returns a pointer to allocated memory of given size or raises an ERROR
+ *		on allocation failure, or returns NULL when flags contains
+ *		MCXT_ALLOC_NO_OOM.
+ *
+ * This function should only contain the most common code paths.  Everything
+ * else should be in pg_noinline helper functions, thus avoiding the overheads
+ * creating a stack frame for the common cases.  Allocating memory is often a
+ * bottleneck in many workloads, so avoiding stack frame setup is worthwhile.
+ * Helper functions should always directly return the newly allocated memory
+ * so that we can just return that address directly as a tail call.
  */
 void *
 SlabAlloc(MemoryContext context, Size size, int flags)
@@ -513,66 +642,16 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 	 * MemoryContextCheckSize check.
 	 */
 	if (unlikely(size != slab->chunkSize))
-		elog(ERROR, "unexpected alloc chunk size %zu (expected %u)",
-			 size, slab->chunkSize);
+		SlabAllocInvalidSize(context, size);
 
-	/*
-	 * Handle the case when there are no partially filled blocks available.
-	 * SlabFree() will have updated the curBlocklistIndex setting it to zero
-	 * to indicate that it has freed the final block.  Also later in
-	 * SlabAlloc() we will set the curBlocklistIndex to zero if we end up
-	 * filling the final block.
-	 */
 	if (unlikely(slab->curBlocklistIndex == 0))
 	{
-		dlist_head *blocklist;
-		int			blocklist_idx;
-
-		/* to save allocating a new one, first check the empty blocks list */
-		if (dclist_count(&slab->emptyblocks) > 0)
-		{
-			dlist_node *node = dclist_pop_head_node(&slab->emptyblocks);
-
-			block = dlist_container(SlabBlock, node, node);
-
-			/*
-			 * SlabFree() should have left this block in a valid state with
-			 * all chunks free.  Ensure that's the case.
-			 */
-			Assert(block->nfree == slab->chunksPerBlock);
-
-			/* fetch the next chunk from this block */
-			chunk = SlabGetNextFreeChunk(slab, block);
-		}
-		else
-		{
-			block = (SlabBlock *) malloc(slab->blockSize);
-
-			if (unlikely(block == NULL))
-				return MemoryContextAllocationFailure(context, size, flags);
-
-			block->slab = slab;
-			context->mem_allocated += slab->blockSize;
-
-			/* use the first chunk in the new block */
-			chunk = SlabBlockGetChunk(slab, block, 0);
-
-			block->nfree = slab->chunksPerBlock - 1;
-			block->unused = SlabBlockGetChunk(slab, block, 1);
-			block->freehead = NULL;
-			block->nunused = slab->chunksPerBlock - 1;
-		}
-
-		/* find the blocklist element for storing blocks with 1 used chunk */
-		blocklist_idx = SlabBlocklistIndex(slab, block->nfree);
-		blocklist = &slab->blocklist[blocklist_idx];
-
-		/* this better be empty.  We just added a block thinking it was */
-		Assert(dlist_is_empty(blocklist));
-
-		dlist_push_head(blocklist, &block->node);
-
-		slab->curBlocklistIndex = blocklist_idx;
+		/*
+		 * Handle the case when there are no partially filled blocks
+		 * available.  This happens either when the last allocation took the
+		 * last chunk in the block, or when SlabFree() free'd the final block.
+		 */
+		return SlabAllocFromNewBlock(context, size, flags);
 	}
 	else
 	{
@@ -609,38 +688,7 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 		}
 	}
 
-	/*
-	 * Check that the chunk pointer is actually somewhere on the block and is
-	 * aligned as expected.
-	 */
-	Assert(chunk >= SlabBlockGetChunk(slab, block, 0));
-	Assert(chunk <= SlabBlockGetChunk(slab, block, slab->chunksPerBlock - 1));
-	Assert(SlabChunkMod(slab, block, chunk) == 0);
-
-	/* Prepare to initialize the chunk header. */
-	VALGRIND_MAKE_MEM_UNDEFINED(chunk, Slab_CHUNKHDRSZ);
-
-	MemoryChunkSetHdrMask(chunk, block, MAXALIGN(slab->chunkSize),
-						  MCTX_SLAB_ID);
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* slab mark to catch clobber of "unused" space */
-	Assert(slab->chunkSize < (slab->fullChunkSize - Slab_CHUNKHDRSZ));
-	set_sentinel(MemoryChunkGetPointer(chunk), size);
-	VALGRIND_MAKE_MEM_NOACCESS(((char *) chunk) +
-							   Slab_CHUNKHDRSZ + slab->chunkSize,
-							   slab->fullChunkSize -
-							   (slab->chunkSize + Slab_CHUNKHDRSZ));
-#endif
-
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-	/* fill the allocated space with junk */
-	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
-#endif
-
-	/* Disallow access to the chunk header. */
-	VALGRIND_MAKE_MEM_NOACCESS(chunk, Slab_CHUNKHDRSZ);
-
-	return MemoryChunkGetPointer(chunk);
+	return SlabAllocSetupNewChunk(context, block, chunk, size);
 }
 
 /*
-- 
2.40.1

---

From: David Rowley <[email-masked]>
Date: Mon, 4 Mar 2024 17:43:50 +1300

> The attached includes some additional tuning to generation.c:

I've now pushed this.

David

---

From: Andres Freund <[email-masked]>
Date: Mon, 4 Mar 2024 00:20:21 -0800

Hi,

> > The attached includes some additional tuning to generation.c:

Thanks for working on all these, much appreciated!

Greetings,

Andres Freund