Adjust memory allocation functions to allow sibling calls

Many modern compilers are able to optimize function calls to functions
where the parameters of the called function match a leading subset of
the calling function's parameters.  If there are no instructions in the
calling function after the function is called, then the compiler is free
to avoid any stack frame setup and implement the function call as a
"jmp" rather than a "call".  This is called sibling call optimization.

Here we adjust the memory allocation functions in mcxt.c to allow this
optimization.  This requires moving some responsibility into the memory
context implementations themselves.  It's now the responsibility of the
MemoryContext to check for malloc failures.  This is good as it both
allows the sibling call optimization, but also because most small and
medium allocations won't call malloc and just allocate memory to an
existing block.  That can't fail, so checking for NULLs in that case
isn't required.

Also, traditionally it's been the responsibility of palloc and the other
allocation functions in mcxt.c to check for invalid allocation size
requests.  Here we also move the responsibility of checking that into the
MemoryContext.  This isn't to allow the sibling call optimization, but
more because most of our allocators handle large allocations separately
and we can just add the size check when doing large allocations.  We no
longer check this for non-large allocations at all.

To make checking the allocation request sizes and ERROR handling easier,
add some helper functions to mcxt.c for the allocators to use.

Author: Andres Freund
Reviewed-by: David Rowley
Discussion: https://postgr.es/m/20210719195950.gavgs6ujzmjfaiig@alap3.anarazel.de

7 files changed, 196 insertions, 166 deletions

diff --git a/src/backend/utils/mmgr/alignedalloc.c b/src/backend/utils/mmgr/alignedalloc.c
index 7204fe64ae6..c266fb3dbb1 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 2f99fa9a2f6..24708f67bb9 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)
 	{
+		/* only check size in paths 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;
 
+		/* only check size in paths where the limits could be hit */
+		MemoryContextCheckSize((MemoryContext) set, size, flags);
+
 		oldchksize = block->endptr - (char *) pointer;
 
 #ifdef MEMORY_CONTEXT_CHECKING
@@ -1165,7 +1171,7 @@ 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 +1331,15 @@ AllocSetRealloc(void *pointer, Size size)
 		AllocPointer newPointer;
 		Size		oldsize;
 
-		/* allocate new chunk */
-		newPointer = AllocSetAlloc((MemoryContext) set, size);
+		/* allocate new chunk (this also checks size is valid) */
+		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 f9016a7ed72..ae4a7c999ed 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,12 @@ GenerationAlloc(MemoryContext context, Size size)
 	{
 		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 NULL;
+			return MemoryContextAllocationFailure(context, size, flags);
 
 		context->mem_allocated += blksize;
 
@@ -472,7 +475,7 @@ GenerationAlloc(MemoryContext context, Size size)
 			block = (GenerationBlock *) malloc(blksize);
 
 			if (block == NULL)
-				return NULL;
+				return MemoryContextAllocationFailure(context, size, flags);
 
 			context->mem_allocated += blksize;
 
@@ -737,7 +740,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;
@@ -839,15 +842,15 @@ GenerationRealloc(void *pointer, Size size)
 		return pointer;
 	}
 
-	/* allocate new chunk */
-	newPointer = GenerationAlloc((MemoryContext) set, size);
+	/* allocate new chunk (this also checks size is valid) */
+	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 ad7409a02cd..41f2390fb8f 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));
@@ -1024,6 +1024,38 @@ MemoryContextCreate(MemoryContext node,
 }
 
 /*
+ * MemoryContextAllocationFailure
+ *		For use by MemoryContextMethods implementations to handle when malloc
+ *		returns NULL.  The behavior is specific to whether MCXT_ALLOC_NO_OOM
+ *		is in 'flags'.
+ */
+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;
+}
+
+/*
+ * MemoryContextSizeFailure
+ *		For use by MemoryContextMethods implementations to handle invalid
+ *		memory allocation request sizes.
+ */
+void
+MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
+{
+	elog(ERROR, "invalid memory alloc request size %zu", size);
+}
+
+/*
  * MemoryContextAlloc
  *		Allocate space within the specified context.
  *
@@ -1038,28 +1070,19 @@ 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)));
-	}
+	/*
+	 * 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 sibling call 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);
 
@@ -1081,21 +1104,9 @@ 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);
 
@@ -1122,20 +1133,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);
 
@@ -1207,22 +1207,21 @@ 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)));
-	}
-
+	/*
+	 * 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 sibling call 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);
+	/* We expect OOM to be handled by the alloc function */
+	Assert(ret != NULL);
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -1238,21 +1237,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);
 
@@ -1271,24 +1258,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;
 	}
 
@@ -1458,26 +1432,22 @@ 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)));
-	}
+	/*
+	 * 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 sibling call 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 (method != MCTX_ALIGNED_REDIRECT_ID)
@@ -1500,31 +1470,24 @@ 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);
+	/*
+	 * 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 sibling call 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))
-	{
-		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);
 
@@ -1565,21 +1528,19 @@ 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)));
-	}
+	/*
+	 * 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 sibling call 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 b8f00cfc7cc..bc91446cb3a 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;
@@ -508,7 +508,10 @@ SlabAlloc(MemoryContext context, Size size)
 	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);
@@ -546,7 +549,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 +773,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;
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index a48f7e5a187..edc0257f36d 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
 {
-	void	   *(*alloc) (MemoryContext context, Size size);
+	/*
+	 * 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);
-	void	   *(*realloc) (void *pointer, Size size);
+
+	/*
+	 * 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 e0c4f3d5af8..ad1048fd829 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,20 @@ 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 */