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

Author: david-rowley

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;

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);
 		}
 
 		/*

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);
 	}
 
 	/*