Hi,

We have a few routines that are taking up a meaningful share of nearly all
workloads. They are worth micro-optimizing, even though they rarely the most
expensive parts of a profile. The memory allocation infrastructure is an
example of that.

When looking at a profile one can often see that a measurable percentage of
the time is spent doing stack frame setup in functions like palloc(),
AllocSetAlloc(). E.g. here's a perf profile annotating palloc(), the first
column showing the percentage of the time the relevant instruction was
sampled:

       │      void *
       │      palloc(Size size)
       │      {
 11.62 │        push  %rbp
  5.89 │        mov   %rsp,%rbp
 11.79 │        push  %r12
       │        mov   %rdi,%r12
  6.07 │        push  %rbx
       │      /* duplicates MemoryContextAlloc to avoid increased overhead */
       │      void       *ret;
       │      MemoryContext context = CurrentMemoryContext;
       │        mov   CurrentMemoryContext,%rbx
       │
       │      AssertArg(MemoryContextIsValid(context));
       │      AssertNotInCriticalSection(context);
       │
       │      if (!AllocSizeIsValid(size))
  5.86 │        cmp   $0x3fffffff,%rdi
       │      → ja    14fa5b <palloc.cold>
       │      elog(ERROR, "invalid memory alloc request size %zu", size);
       │
       │      context->isReset = false;
 17.71 │        movb  $0x0,0x4(%rbx)
       │
       │      ret = context->methods->alloc(context, size);
  5.98 │        mov   0x10(%rbx),%rax
       │        mov   %rdi,%rsi
       │        mov   %rbx,%rdi
 35.08 │      → callq *(%rax)


The stack frame setup bit is the push ... bit.

At least on x86-64 unixoid systems, that overhead can be avoided in certain
circumstances! The simplest case is if the function doesn't do any function
calls of its own. If simple enough (i.e. no register spilling), the compiler
will just not set up a stack frame - nobody could need it.

The slightly more complicated case is that of a function that only does a
"tail call", i.e. the only function call is just before returning (there can
be multiple such paths though). If the function is simple enough, gcc/clang
will then not use the "call" instruction to call the function (which would
require a proper stack frame being set up), but instead just jump to the other
function. Which ends up reusing the current function's stack frame,
basically. When that called function returns using 'ret', it'll reuse the
location pushed onto the call stack by the caller of the "original" function,
and return to its caller. Having optimized away the need to maintain one stack
frame level, and one indirection when returning from the inner function (which
just would do its own ret).

For that to work, there obviously cannot be any instructions in the function
before calling the inner function. Which brings us back to the palloc example
from above.

As an experiment, if i change the code for palloc() to omit the if (ret == NULL)
check, the assembly (omitting source for brevity) from:

  61c9a0:       55                      push   %rbp
  61c9a1:       48 89 e5                mov    %rsp,%rbp
  61c9a4:       41 54                   push   %r12
  61c9a6:       49 89 fc                mov    %rdi,%r12
  61c9a9:       53                      push   %rbx
  61c9aa:       48 8b 1d 2f f2 2a 00    mov    0x2af22f(%rip),%rbx        # 
8cbbe0 <CurrentMemoryContext>
  61c9b1:       48 81 ff ff ff ff 3f    cmp    $0x3fffffff,%rdi
  61c9b8:       0f 87 9d 30 b3 ff       ja     14fa5b <palloc.cold>
  61c9be:       c6 43 04 00             movb   $0x0,0x4(%rbx)
  61c9c2:       48 8b 43 10             mov    0x10(%rbx),%rax
  61c9c6:       48 89 fe                mov    %rdi,%rsi
  61c9c9:       48 89 df                mov    %rbx,%rdi
  61c9cc:       ff 10                   callq  *(%rax)
  61c9ce:       48 85 c0                test   %rax,%rax
  61c9d1:       0f 84 b9 30 b3 ff       je     14fa90 <palloc.cold+0x35>
  61c9d7:       5b                      pop    %rbx
  61c9d8:       41 5c                   pop    %r12
  61c9da:       5d                      pop    %rbp
  61c9db:       c3                      retq

to

  61c8c0:       48 89 fe                mov    %rdi,%rsi
  61c8c3:       48 8b 3d 16 f3 2a 00    mov    0x2af316(%rip),%rdi        # 
8cbbe0 <CurrentMemoryContext>
  61c8ca:       48 81 fe ff ff ff 3f    cmp    $0x3fffffff,%rsi
  61c8d1:       0f 87 c3 31 b3 ff       ja     14fa9a <palloc.cold>
  61c8d7:       c6 47 04 00             movb   $0x0,0x4(%rdi)
  61c8db:       48 8b 47 10             mov    0x10(%rdi),%rax
  61c8df:       ff 20                   jmpq   *(%rax)

It's not hard to see why that would be faster, I think.


Of course, we cannot just omit that check. But I think this is an argument for
why it is not a great idea to have such a check in palloc() - it prevents the
use of the above optimization, and it adds a branch to a performance critical
function, though there already existing branches in aset.c etc that
specifically know about this case.

The code in palloc() does this check after context->methods->alloc() since
3d6d1b585524: Move out-of-memory error checks from aset.c to mcxt.c

Of course, that commit changed things for a reason: It allows
palloc_extended() to exist.

However, it seems that the above optimization, as well as the desire to avoid
redundant branches (checking for allocation failures in AllocSetAlloc() and
then again in palloc() etc) in critical paths, suggests pushing the handling
of MCXT_ALLOC_NO_OOM (and perhaps others) a layer down, into the memory
context implementations. Which of course means that we would need to pass down
MCXT_ALLOC_NO_OOM into at least MemoryContextMethods->alloc,realloc}. But that
seems like a good idea to me anyway. That way we could pass down further
information as well, e.g. about required alignment.

Of course it'd make sense to avoid duplicating the same error message across
all contexts, but that could be addressed using a mcxt.c helper function to
deal with the allocation failure case.

E.g the existing cases like

    block = (AllocBlock) malloc(blksize);
    if (block == NULL)
        return NULL;

could become something like
    block = (AllocBlock) malloc(blksize);
    if (unlikely(block == NULL))
        return MemoryContextAllocationFailure(context, size, flags);


The trick of avoiding stack frame setup does not just apply to wrapper
functions like palloc(). It even can apply to AllocSetAlloc() itself! If one
separates out the "slow paths" from the "fast paths" of AllocSetAlloc(), the
fast path can avoid needing the stack frame, for the price of the slow paths
being a tiny bit slower. Often the generated code turns out to be better,
because the register allocation pressure is lower in the fast path.

For that to work, the paths of AllocSetAlloc() that call malloc() need to be
separated out. As we obviously need to process malloc()'s result, the call to
malloc cannot be a tail call. So we need to split out two paths:
1) handling of large allocations
2) running out of space in the current block / having no block

To actually benefit from the optimization, those paths need to actually return
the allocated memory. And they need to be marked pg_noinline, otherwise the
compiler won't get the message...

I think this actually makes the aset.c code a good bit more readable, and
highlights where in AllocSetAlloc() adding instructions hurts, and where its
fine.

I have *not* carefully benchmarked this, but a quick implementation of this
does seem to increase readonly pgbench tps at a small scale by 2-3% (both
-Mprepared/simple). Despite not being an all that pgbench bound workload.


Rough prototype patch for the above attached.

Comments?


A slightly tangential improvement would be to move the memset() in palloc0()
et al do into a static inline. There's two benefits of that:

1) compilers can generate much better code for memset() if the length is known
   - instead of a function call with length dispatch replace that with a
   handful of instructions doing the zeroing for the precise length.

2) compilers can often optimize away [part of ]the overhead of needing to do
   the memset, as many callers will go on to overwrite a good portion of the
   zeroed data.


Greetings,

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

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

diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index e6a757d6a07..8a42d2ff999 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 ff872274d44..2f75b4cca46 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
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 77872e77bcd..00878354392 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,266 +704,10 @@ AllocSetDelete(MemoryContext context)
 	free(set);
 }
 
-/*
- * 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!
- */
-static void *
-AllocSetAlloc(MemoryContext context, Size size)
+static inline void *
+AllocSetAllocReturnChunk(AllocSet set, Size size, AllocChunk chunk, Size chunk_size)
 {
-	AllocSet	set = (AllocSet) context;
-	AllocBlock	block;
-	AllocChunk	chunk;
-	int			fidx;
-	Size		chunk_size;
-	Size		blksize;
-
-	AssertArg(AllocSetIsValid(set));
-
-	/*
-	 * 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);
-	}
-
-	/*
-	 * 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.
-	 * If one is found, remove it from the free list, make it again a member
-	 * of the alloc set and return its data address.
-	 */
-	fidx = AllocSetFreeIndex(size);
-	chunk = set->freelist[fidx];
-	if (chunk != NULL)
-	{
-		Assert(chunk->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);
-	}
-
-	/*
-	 * Choose the actual chunk size to allocate.
-	 */
-	chunk_size = (1 << ALLOC_MINBITS) << fidx;
-	Assert(chunk_size >= size);
-
-	/*
-	 * 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))
-			{
-				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;
-		}
-	}
-
-	/*
-	 * 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 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;
-	}
-
-	/*
-	 * 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->aset = (void *) set;
-	chunk->size = chunk_size;
 #ifdef MEMORY_CONTEXT_CHECKING
 	chunk->requested_size = size;
 	/* set mark to catch clobber of "unused" space */
@@ -985,6 +729,273 @@ AllocSetAlloc(MemoryContext context, Size size)
 	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
+ *		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!
+ */
+static void *
+AllocSetAlloc(MemoryContext context, Size size, int flags)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+	AllocChunk	chunk;
+	int			fidx;
+	Size		chunk_size;
+
+	AssertArg(AllocSetIsValid(set));
+
+	/*
+	 * If requested size exceeds maximum for chunks, allocate an entire block
+	 * for this request.
+	 */
+	if (unlikely(size > set->allocChunkLimit))
+		return AllocSetAllocLarge(set, 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.
+	 * If one is found, remove it from the free list, make it again a member
+	 * of the alloc set and return its data address.
+	 */
+	fidx = AllocSetFreeIndex(size);
+	chunk = set->freelist[fidx];
+	if (chunk != NULL)
+	{
+		Assert(chunk->size >= size);
+
+		set->freelist[fidx] = (AllocChunk) chunk->aset;
+
+		return AllocSetAllocReturnChunk(set, size, chunk, chunk->size);
+	}
+
+	/*
+	 * Choose the actual chunk size to allocate.
+	 */
+	chunk_size = (1 << ALLOC_MINBITS) << fidx;
+	Assert(chunk_size >= size);
+
+	/*
+	 * 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 (unlikely(availspace < (chunk_size + ALLOC_CHUNKHDRSZ)))
+			return AllocSetAllocCarveOldAndAlloc(set, size, chunk_size,
+												 block, availspace);
+	}
+	else if (unlikely(block == NULL))
+	{
+		/*
+		 * Time to create a new regular (multi-chunk) block.
+		 */
+		return AllocSetAllocFromNewBlock(set, size, chunk_size);
+	}
+
+	/*
+	 * OK, do the allocation
+	 */
+	chunk = (AllocChunk) (block->freeptr);
+
+	/* Prepare to initialize the chunk header. */
+	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+	chunk->size = chunk_size;
+
+	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
+	Assert(block->freeptr <= block->endptr);
+
+	return AllocSetAllocReturnChunk(set, size, chunk, chunk_size);
+}
+
 /*
  * AllocSetFree
  *		Frees allocated memory; memory is removed from the set.
@@ -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 584cd614da0..5dde15654d7 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 6919a732804..13125c4a562 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 553dd7f6674..e4b8275045f 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);
 
-- 
2.32.0.rc2

Reply via email to