From b46dbb73869067a08746d0aa3bfbdb3a2835a992 Mon Sep 17 00:00:00 2001
From: Andres Freund <andres@anarazel.de>
Date: Mon, 19 Jul 2021 12:55:51 -0700
Subject: [PATCH 2/3] WIP: optimize allocations by separating hot from cold
 paths.

---
 src/backend/utils/mmgr/aset.c       | 429 ++++++++++++++--------------
 src/backend/utils/mmgr/generation.c |  22 +-
 src/backend/utils/mmgr/mcxt.c       | 179 +++---------
 src/backend/utils/mmgr/slab.c       |  14 +-
 src/include/nodes/memnodes.h        |   4 +-
 src/include/utils/memutils.h        |  13 +
 6 files changed, 300 insertions(+), 361 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 77872e77bc..0087835439 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -263,9 +263,9 @@ static AllocSetFreeList context_freelists[2] =
 /*
  * These functions implement the MemoryContext API for AllocSet contexts.
  */
-static void *AllocSetAlloc(MemoryContext context, Size size);
+static void *AllocSetAlloc(MemoryContext context, Size size, int flags);
 static void AllocSetFree(MemoryContext context, void *pointer);
-static void *AllocSetRealloc(MemoryContext context, void *pointer, Size size);
+static void *AllocSetRealloc(MemoryContext context, void *pointer, Size size, int flags);
 static void AllocSetReset(MemoryContext context);
 static void AllocSetDelete(MemoryContext context);
 static Size AllocSetGetChunkSpace(MemoryContext context, void *pointer);
@@ -704,6 +704,208 @@ AllocSetDelete(MemoryContext context)
 	free(set);
 }
 
+static inline void *
+AllocSetAllocReturnChunk(AllocSet set, Size size, AllocChunk chunk, Size chunk_size)
+{
+	chunk->aset = (void *) set;
+#ifdef MEMORY_CONTEXT_CHECKING
+	chunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	if (size < chunk->size)
+		set_sentinel(AllocChunkGetPointer(chunk), size);
+#endif
+#ifdef RANDOMIZE_ALLOCATED_MEMORY
+	/* fill the allocated space with junk */
+	randomize_mem((char *) AllocChunkGetPointer(chunk), size);
+#endif
+
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
+							   chunk_size - size);
+
+	/* Disallow external access to private part of chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+
+	return AllocChunkGetPointer(chunk);
+}
+
+static void * pg_noinline
+AllocSetAllocLarge(AllocSet	set, Size size, int flags)
+{
+	AllocBlock	block;
+	AllocChunk	chunk;
+	Size		chunk_size;
+	Size		blksize;
+
+	/* check size, only allocation path where the limits could be hit */
+	MemoryContextCheckSize(&set->header, size, flags);
+
+	AssertArg(AllocSetIsValid(set));
+
+	chunk_size = MAXALIGN(size);
+	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	block = (AllocBlock) malloc(blksize);
+	if (block == NULL)
+		return NULL;
+
+	set->header.mem_allocated += blksize;
+
+	block->aset = set;
+	block->freeptr = block->endptr = ((char *) block) + blksize;
+
+	/*
+	 * 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;
+	}
+
+	chunk = (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ);
+	chunk->size = chunk_size;
+
+	return AllocSetAllocReturnChunk(set, size, chunk, chunk_size);
+}
+
+static void * pg_noinline
+AllocSetAllocFromNewBlock(AllocSet set, Size size, Size chunk_size)
+{
+	AllocBlock	block;
+	Size		blksize;
+	Size		required_size;
+	AllocChunk	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;
+
+	/*
+	 * 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 NULL;
+
+	set->header.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;
+
+	/*
+	 * OK, do the allocation
+	 */
+	chunk = (AllocChunk) (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);
+
+	chunk->size = chunk_size;
+
+	return AllocSetAllocReturnChunk(set, size, chunk, chunk_size);
+}
+
+static void * pg_noinline
+AllocSetAllocCarveOldAndAlloc(AllocSet set, Size size, Size chunk_size, AllocBlock	block, Size availspace)
+{
+	AllocChunk	chunk;
+
+	/*
+	 * 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))
+	{
+		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 != ((Size) 1 << (a_fidx + ALLOC_MINBITS)))
+		{
+			a_fidx--;
+			Assert(a_fidx >= 0);
+			availchunk = ((Size) 1 << (a_fidx + ALLOC_MINBITS));
+		}
+
+		chunk = (AllocChunk) (block->freeptr);
+
+		/* Prepare to initialize the chunk header. */
+		VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+		block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
+		availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
+
+		chunk->size = availchunk;
+#ifdef MEMORY_CONTEXT_CHECKING
+		chunk->requested_size = 0;	/* mark it free */
+#endif
+		chunk->aset = (void *) set->freelist[a_fidx];
+		set->freelist[a_fidx] = chunk;
+	}
+
+	return AllocSetAllocFromNewBlock(set, size, chunk_size);
+}
+
 /*
  * AllocSetAlloc
  *		Returns pointer to allocated memory of given size or NULL if
@@ -718,14 +920,13 @@ AllocSetDelete(MemoryContext context)
  * return space that is marked NOACCESS - AllocSetRealloc has to beware!
  */
 static void *
-AllocSetAlloc(MemoryContext context, Size size)
+AllocSetAlloc(MemoryContext context, Size size, int flags)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block;
 	AllocChunk	chunk;
 	int			fidx;
 	Size		chunk_size;
-	Size		blksize;
 
 	AssertArg(AllocSetIsValid(set));
 
@@ -733,61 +934,8 @@ AllocSetAlloc(MemoryContext context, Size size)
 	 * If requested size exceeds maximum for chunks, allocate an entire block
 	 * for this request.
 	 */
-	if (size > set->allocChunkLimit)
-	{
-		chunk_size = MAXALIGN(size);
-		blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-		block = (AllocBlock) malloc(blksize);
-		if (block == NULL)
-			return NULL;
-
-		context->mem_allocated += blksize;
-
-		block->aset = set;
-		block->freeptr = block->endptr = ((char *) block) + blksize;
-
-		chunk = (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ);
-		chunk->aset = set;
-		chunk->size = chunk_size;
-#ifdef MEMORY_CONTEXT_CHECKING
-		chunk->requested_size = size;
-		/* set mark to catch clobber of "unused" space */
-		if (size < chunk_size)
-			set_sentinel(AllocChunkGetPointer(chunk), size);
-#endif
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* fill the allocated space with junk */
-		randomize_mem((char *) AllocChunkGetPointer(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 *) AllocChunkGetPointer(chunk) + size,
-								   chunk_size - size);
-
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
-
-		return AllocChunkGetPointer(chunk);
-	}
+	if (unlikely(size > set->allocChunkLimit))
+		return AllocSetAllocLarge(set, size, flags);
 
 	/*
 	 * Request is small enough to be treated as a chunk.  Look in the
@@ -803,27 +951,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 
 		set->freelist[fidx] = (AllocChunk) chunk->aset;
 
-		chunk->aset = (void *) set;
-
-#ifdef MEMORY_CONTEXT_CHECKING
-		chunk->requested_size = size;
-		/* set mark to catch clobber of "unused" space */
-		if (size < chunk->size)
-			set_sentinel(AllocChunkGetPointer(chunk), size);
-#endif
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* fill the allocated space with junk */
-		randomize_mem((char *) AllocChunkGetPointer(chunk), size);
-#endif
-
-		/* Ensure any padding bytes are marked NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
-								   chunk->size - size);
-
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
-
-		return AllocChunkGetPointer(chunk);
+		return AllocSetAllocReturnChunk(set, size, chunk, chunk->size);
 	}
 
 	/*
@@ -840,115 +968,16 @@ AllocSetAlloc(MemoryContext context, Size size)
 	{
 		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))
-			{
-				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 != ((Size) 1 << (a_fidx + ALLOC_MINBITS)))
-				{
-					a_fidx--;
-					Assert(a_fidx >= 0);
-					availchunk = ((Size) 1 << (a_fidx + ALLOC_MINBITS));
-				}
-
-				chunk = (AllocChunk) (block->freeptr);
-
-				/* Prepare to initialize the chunk header. */
-				VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
-
-				block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
-				availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
-
-				chunk->size = availchunk;
-#ifdef MEMORY_CONTEXT_CHECKING
-				chunk->requested_size = 0;	/* mark it free */
-#endif
-				chunk->aset = (void *) set->freelist[a_fidx];
-				set->freelist[a_fidx] = chunk;
-			}
-
-			/* Mark that we need to create a new block */
-			block = NULL;
-		}
+		if (unlikely(availspace < (chunk_size + ALLOC_CHUNKHDRSZ)))
+			return AllocSetAllocCarveOldAndAlloc(set, size, chunk_size,
+												 block, availspace);
 	}
-
-	/*
-	 * Time to create a new regular (multi-chunk) block?
-	 */
-	if (block == NULL)
+	else if (unlikely(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...
+		 * Time to create a new regular (multi-chunk) block.
 		 */
-		while (block == NULL && blksize > 1024 * 1024)
-		{
-			blksize >>= 1;
-			if (blksize < required_size)
-				break;
-			block = (AllocBlock) malloc(blksize);
-		}
-
-		if (block == NULL)
-			return NULL;
-
-		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 AllocSetAllocFromNewBlock(set, size, chunk_size);
 	}
 
 	/*
@@ -959,30 +988,12 @@ AllocSetAlloc(MemoryContext context, Size size)
 	/* Prepare to initialize the chunk header. */
 	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
 
-	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
-	Assert(block->freeptr <= block->endptr);
-
-	chunk->aset = (void *) set;
 	chunk->size = chunk_size;
-#ifdef MEMORY_CONTEXT_CHECKING
-	chunk->requested_size = size;
-	/* set mark to catch clobber of "unused" space */
-	if (size < chunk->size)
-		set_sentinel(AllocChunkGetPointer(chunk), size);
-#endif
-#ifdef RANDOMIZE_ALLOCATED_MEMORY
-	/* fill the allocated space with junk */
-	randomize_mem((char *) AllocChunkGetPointer(chunk), size);
-#endif
-
-	/* Ensure any padding bytes are marked NOACCESS. */
-	VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
-							   chunk_size - size);
 
-	/* Disallow external access to private part of chunk header. */
-	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
+	Assert(block->freeptr <= block->endptr);
 
-	return AllocChunkGetPointer(chunk);
+	return AllocSetAllocReturnChunk(set, size, chunk, chunk_size);
 }
 
 /*
@@ -1072,7 +1083,7 @@ AllocSetFree(MemoryContext context, void *pointer)
  * request size.)
  */
 static void *
-AllocSetRealloc(MemoryContext context, void *pointer, Size size)
+AllocSetRealloc(MemoryContext context, void *pointer, Size size, int flags)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocChunk	chunk = AllocPointerGetChunk(pointer);
@@ -1081,6 +1092,8 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 	/* Allow access to private part of chunk header. */
 	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOCCHUNK_PRIVATE_LEN);
 
+	MemoryContextCheckSize(context, size, flags);
+
 	oldsize = chunk->size;
 
 #ifdef MEMORY_CONTEXT_CHECKING
@@ -1260,7 +1273,7 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		AllocPointer newPointer;
 
 		/* allocate new chunk */
-		newPointer = AllocSetAlloc((MemoryContext) set, size);
+		newPointer = AllocSetAlloc((MemoryContext) set, size, flags);
 
 		/* leave immediately if request was not completed */
 		if (newPointer == NULL)
diff --git a/src/backend/utils/mmgr/generation.c b/src/backend/utils/mmgr/generation.c
index 584cd614da..5dde15654d 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -146,9 +146,9 @@ struct GenerationChunk
 /*
  * These functions implement the MemoryContext API for Generation contexts.
  */
-static void *GenerationAlloc(MemoryContext context, Size size);
+static void *GenerationAlloc(MemoryContext context, Size size, int flags);
 static void GenerationFree(MemoryContext context, void *pointer);
-static void *GenerationRealloc(MemoryContext context, void *pointer, Size size);
+static void *GenerationRealloc(MemoryContext context, void *pointer, Size size, int flags);
 static void GenerationReset(MemoryContext context);
 static void GenerationDelete(MemoryContext context);
 static Size GenerationGetChunkSpace(MemoryContext context, void *pointer);
@@ -323,7 +323,7 @@ GenerationDelete(MemoryContext context)
  * return space that is marked NOACCESS - GenerationRealloc has to beware!
  */
 static void *
-GenerationAlloc(MemoryContext context, Size size)
+GenerationAlloc(MemoryContext context, Size size, int flags)
 {
 	GenerationContext *set = (GenerationContext *) context;
 	GenerationBlock *block;
@@ -335,9 +335,12 @@ GenerationAlloc(MemoryContext context, Size size)
 	{
 		Size		blksize = chunk_size + Generation_BLOCKHDRSZ + Generation_CHUNKHDRSZ;
 
+		/* check size, only allocation path where the limits could be hit */
+		MemoryContextCheckSize(context, size, flags);
+
 		block = (GenerationBlock *) malloc(blksize);
 		if (block == NULL)
-			return NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -392,7 +395,7 @@ GenerationAlloc(MemoryContext context, Size size)
 		block = (GenerationBlock *) malloc(blksize);
 
 		if (block == NULL)
-			return NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -520,13 +523,15 @@ GenerationFree(MemoryContext context, void *pointer)
  *		into the old chunk - in that case we just update chunk header.
  */
 static void *
-GenerationRealloc(MemoryContext context, void *pointer, Size size)
+GenerationRealloc(MemoryContext context, void *pointer, Size size, int flags)
 {
 	GenerationContext *set = (GenerationContext *) context;
 	GenerationChunk *chunk = GenerationPointerGetChunk(pointer);
 	GenerationPointer newPointer;
 	Size		oldsize;
 
+	MemoryContextCheckSize(context, size, flags);
+
 	/* Allow access to private part of chunk header. */
 	VALGRIND_MAKE_MEM_DEFINED(chunk, GENERATIONCHUNK_PRIVATE_LEN);
 
@@ -596,14 +601,15 @@ GenerationRealloc(MemoryContext context, 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 external access to private part of chunk header. */
 		VALGRIND_MAKE_MEM_NOACCESS(chunk, GENERATIONCHUNK_PRIVATE_LEN);
-		return NULL;
+		/* again? */
+		return MemoryContextAllocationFailure(context, size, flags);
 	}
 
 	/*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 6919a73280..13125c4a56 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -867,28 +867,9 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	AssertArg(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);
 
@@ -910,21 +891,10 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	AssertArg(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 != NULL);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -948,21 +918,10 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	AssertArg(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 != NULL);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -983,26 +942,11 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 	AssertArg(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (((flags & MCXT_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
-		((flags & MCXT_ALLOC_HUGE) == 0 && !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;
-	}
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1067,22 +1011,10 @@ palloc(Size size)
 
 	AssertArg(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 != NULL);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1099,21 +1031,10 @@ palloc0(Size size)
 	AssertArg(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 != NULL);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1132,26 +1053,11 @@ palloc_extended(Size size, int flags)
 	AssertArg(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (((flags & MCXT_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
-		((flags & MCXT_ALLOC_HUGE) == 0 && !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, 0);
 	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);
 
@@ -1184,24 +1090,13 @@ repalloc(void *pointer, Size size)
 	MemoryContext context = GetMemoryChunkContext(pointer);
 	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 = context->methods->realloc(context, pointer, 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->realloc(context, pointer, size, 0);
+	Assert(ret != NULL);
 
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
@@ -1222,21 +1117,9 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	AssertArg(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);
+	Assert(ret != NULL);
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1254,16 +1137,23 @@ repalloc_huge(void *pointer, Size size)
 	MemoryContext context = GetMemoryChunkContext(pointer);
 	void	   *ret;
 
-	if (!AllocHugeSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
 	AssertNotInCriticalSection(context);
 
 	/* isReset must be false already */
 	Assert(!context->isReset);
 
-	ret = context->methods->realloc(context, pointer, size);
-	if (unlikely(ret == NULL))
+	ret = context->methods->realloc(context, pointer, size, MCXT_ALLOC_HUGE);
+	Assert(ret != NULL);
+
+	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+
+	return ret;
+}
+
+void *
+MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
+{
+	if ((flags & MCXT_ALLOC_NO_OOM) == 0)
 	{
 		MemoryContextStats(TopMemoryContext);
 		ereport(ERROR,
@@ -1273,9 +1163,13 @@ repalloc_huge(void *pointer, Size size)
 						   size, context->name)));
 	}
 
-	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+	return NULL;
+}
 
-	return ret;
+void
+MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
+{
+	elog(ERROR, "invalid memory alloc request size %zu", size);
 }
 
 /*
@@ -1298,6 +1192,13 @@ MemoryContextStrdup(MemoryContext context, const char *string)
 char *
 pstrdup(const char *in)
 {
+	/*
+	 * 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.
+	 */
+
 	return MemoryContextStrdup(CurrentMemoryContext, in);
 }
 
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 553dd7f667..e4b8275045 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -126,9 +126,9 @@ typedef struct SlabChunk
 /*
  * These functions implement the MemoryContext API for Slab contexts.
  */
-static void *SlabAlloc(MemoryContext context, Size size);
+static void *SlabAlloc(MemoryContext context, Size size, int flags);
 static void SlabFree(MemoryContext context, void *pointer);
-static void *SlabRealloc(MemoryContext context, void *pointer, Size size);
+static void *SlabRealloc(MemoryContext context, void *pointer, Size size, int flags);
 static void SlabReset(MemoryContext context);
 static void SlabDelete(MemoryContext context);
 static Size SlabGetChunkSpace(MemoryContext context, void *pointer);
@@ -337,7 +337,7 @@ SlabDelete(MemoryContext context)
  *		request could not be completed; memory is added to the slab.
  */
 static void *
-SlabAlloc(MemoryContext context, Size size)
+SlabAlloc(MemoryContext context, Size size, int flags)
 {
 	SlabContext *slab = castNode(SlabContext, context);
 	SlabBlock  *block;
@@ -346,6 +346,12 @@ SlabAlloc(MemoryContext context, Size size)
 
 	Assert(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...
+	 */
+
 	Assert((slab->minFreeChunks >= 0) &&
 		   (slab->minFreeChunks < slab->chunksPerBlock));
 
@@ -583,7 +589,7 @@ SlabFree(MemoryContext context, void *pointer)
  * realloc is usually used to enlarge the chunk.
  */
 static void *
-SlabRealloc(MemoryContext context, void *pointer, Size size)
+SlabRealloc(MemoryContext context, void *pointer, Size size, int flags)
 {
 	SlabContext *slab = castNode(SlabContext, context);
 
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index e6a757d6a0..8a42d2ff99 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) (MemoryContext context, void *pointer);
-	void	   *(*realloc) (MemoryContext context, void *pointer, Size size);
+	void	   *(*realloc) (MemoryContext context, void *pointer, Size size, int flags);
 	void		(*reset) (MemoryContext context);
 	void		(*delete_context) (MemoryContext context);
 	Size		(*get_chunk_space) (MemoryContext context, void *pointer);
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index ff872274d4..2f75b4cca4 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -147,6 +147,19 @@ extern void MemoryContextCreate(MemoryContext node,
 
 extern void HandleLogMemoryContextInterrupt(void);
 extern void ProcessLogMemoryContextInterrupt(void);
+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);
+	}
+}
 
 /*
  * Memory-context-type-specific functions
-- 
2.31.1

