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

On Wed, 9 Aug 2023 at 20:44, David Rowley <dgrowle...@gmail.com> wrote:
> Here's a review of v2-0001:
> 2. Why do you need to add the NULL check here?
>
>  #ifdef USE_VALGRIND
> - if (method != MCTX_ALIGNED_REDIRECT_ID)
> + if (ret != NULL && method != MCTX_ALIGNED_REDIRECT_ID)
>   VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
>  #endif

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

> 3.
>
> /*
> * XXX: Probably no need to check for huge allocations, we only support
> * one size? Which could theoretically be huge, but that'd not make
> * sense...
> */
>
> They can't be huge per Assert(fullChunkSize <= MEMORYCHUNK_MAX_VALUE)
> in SlabContextCreate().

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

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


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

I've done that too.

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

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

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

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

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

diff --git a/src/backend/utils/mmgr/alignedalloc.c 
b/src/backend/utils/mmgr/alignedalloc.c
index 7204fe64ae..c266fb3dbb 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 2f99fa9a2f..81c3120c2b 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -700,7 +700,7 @@ AllocSetDelete(MemoryContext context)
  * return space that is marked NOACCESS - AllocSetRealloc has to beware!
  */
 void *
-AllocSetAlloc(MemoryContext context, Size size)
+AllocSetAlloc(MemoryContext context, Size size, int flags)
 {
        AllocSet        set = (AllocSet) context;
        AllocBlock      block;
@@ -717,6 +717,9 @@ AllocSetAlloc(MemoryContext context, Size size)
         */
        if (size > set->allocChunkLimit)
        {
+               /* check size, only allocation path where the limits could be 
hit */
+               MemoryContextCheckSize(context, size, flags);
+
 #ifdef MEMORY_CONTEXT_CHECKING
                /* ensure there's always space for the sentinel byte */
                chunk_size = MAXALIGN(size + 1);
@@ -727,7 +730,7 @@ AllocSetAlloc(MemoryContext context, Size size)
                blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
                block = (AllocBlock) malloc(blksize);
                if (block == NULL)
-                       return NULL;
+                       return MemoryContextAllocationFailure(context, size, 
flags);
 
                context->mem_allocated += blksize;
 
@@ -940,7 +943,7 @@ AllocSetAlloc(MemoryContext context, Size size)
                }
 
                if (block == NULL)
-                       return NULL;
+                       return MemoryContextAllocationFailure(context, size, 
flags);
 
                context->mem_allocated += blksize;
 
@@ -1106,7 +1109,7 @@ AllocSetFree(void *pointer)
  * request size.)
  */
 void *
-AllocSetRealloc(void *pointer, Size size)
+AllocSetRealloc(void *pointer, Size size, int flags)
 {
        AllocBlock      block;
        AllocSet        set;
@@ -1139,6 +1142,9 @@ AllocSetRealloc(void *pointer, Size size)
 
                set = block->aset;
 
+               /* check size, only allocation path where the limits could be 
hit */
+               MemoryContextCheckSize((MemoryContext) set, size, flags);
+
                oldchksize = block->endptr - (char *) pointer;
 
 #ifdef MEMORY_CONTEXT_CHECKING
@@ -1165,7 +1171,8 @@ AllocSetRealloc(void *pointer, Size size)
                {
                        /* Disallow access to the chunk header. */
                        VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-                       return NULL;
+                       return MemoryContextAllocationFailure(&set->header, 
size, flags);
+
                }
 
                /* updated separately, not to underflow when (oldblksize > 
blksize) */
@@ -1325,15 +1332,15 @@ AllocSetRealloc(void *pointer, Size size)
                AllocPointer newPointer;
                Size            oldsize;
 
-               /* allocate new chunk */
-               newPointer = AllocSetAlloc((MemoryContext) set, size);
+               /* allocate new chunk (also checks size for us) */
+               newPointer = AllocSetAlloc((MemoryContext) set, size, flags);
 
                /* leave immediately if request was not completed */
                if (newPointer == NULL)
                {
                        /* Disallow access to the chunk header. */
                        VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
-                       return NULL;
+                       return MemoryContextAllocationFailure((MemoryContext) 
set, size, flags);
                }
 
                /*
diff --git a/src/backend/utils/mmgr/generation.c 
b/src/backend/utils/mmgr/generation.c
index f9016a7ed7..4b4155b863 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -344,7 +344,7 @@ GenerationDelete(MemoryContext context)
  * return space that is marked NOACCESS - GenerationRealloc has to beware!
  */
 void *
-GenerationAlloc(MemoryContext context, Size size)
+GenerationAlloc(MemoryContext context, Size size, int flags)
 {
        GenerationContext *set = (GenerationContext *) context;
        GenerationBlock *block;
@@ -367,9 +367,11 @@ GenerationAlloc(MemoryContext context, Size size)
        {
                Size            blksize = required_size + Generation_BLOCKHDRSZ;
 
+               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 +474,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 +739,7 @@ GenerationFree(void *pointer)
  *             into the old chunk - in that case we just update chunk header.
  */
 void *
-GenerationRealloc(void *pointer, Size size)
+GenerationRealloc(void *pointer, Size size, int flags)
 {
        MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
        GenerationContext *set;
@@ -840,14 +842,14 @@ GenerationRealloc(void *pointer, Size size)
        }
 
        /* allocate new chunk */
-       newPointer = GenerationAlloc((MemoryContext) set, size);
+       newPointer = GenerationAlloc((MemoryContext) set, size, flags);
 
        /* leave immediately if request was not completed */
        if (newPointer == NULL)
        {
                /* Disallow access to the chunk header. */
                VALGRIND_MAKE_MEM_NOACCESS(chunk, Generation_CHUNKHDRSZ);
-               return NULL;
+               return MemoryContextAllocationFailure((MemoryContext) set, 
size, flags);
        }
 
        /*
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index ad7409a02c..2947b2046b 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));
@@ -1023,6 +1023,27 @@ MemoryContextCreate(MemoryContext node,
        VALGRIND_CREATE_MEMPOOL(node, 0, false);
 }
 
+void *
+MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
+{
+       if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+       {
+               MemoryContextStats(TopMemoryContext);
+               ereport(ERROR,
+                               (errcode(ERRCODE_OUT_OF_MEMORY),
+                                errmsg("out of memory"),
+                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
+                                                  size, context->name)));
+       }
+       return NULL;
+}
+
+void
+MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
+{
+       elog(ERROR, "invalid memory alloc request size %zu", size);
+}
+
 /*
  * MemoryContextAlloc
  *             Allocate space within the specified context.
@@ -1038,28 +1059,9 @@ MemoryContextAlloc(MemoryContext context, Size size)
        Assert(MemoryContextIsValid(context));
        AssertNotInCriticalSection(context);
 
-       if (!AllocSizeIsValid(size))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        context->isReset = false;
 
-       ret = context->methods->alloc(context, size);
-       if (unlikely(ret == NULL))
-       {
-               MemoryContextStats(TopMemoryContext);
-
-               /*
-                * Here, and elsewhere in this module, we show the target 
context's
-                * "name" but not its "ident" (if any) in user-visible error 
messages.
-                * The "ident" string might contain security-sensitive data, 
such as
-                * values in SQL commands.
-                */
-               ereport(ERROR,
-                               (errcode(ERRCODE_OUT_OF_MEMORY),
-                                errmsg("out of memory"),
-                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
-                                                  size, context->name)));
-       }
+       ret = context->methods->alloc(context, size, 0);
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
@@ -1081,24 +1083,16 @@ MemoryContextAllocZero(MemoryContext context, Size size)
        Assert(MemoryContextIsValid(context));
        AssertNotInCriticalSection(context);
 
-       if (!AllocSizeIsValid(size))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        context->isReset = false;
 
-       ret = context->methods->alloc(context, size);
-       if (unlikely(ret == NULL))
-       {
-               MemoryContextStats(TopMemoryContext);
-               ereport(ERROR,
-                               (errcode(ERRCODE_OUT_OF_MEMORY),
-                                errmsg("out of memory"),
-                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
-                                                  size, context->name)));
-       }
+       ret = context->methods->alloc(context, size, 0);
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+       /*
+        * XXX: Should this also be moved into alloc()? We could possibly avoid
+        * zeroing in some cases (e.g. if we used mmap() ourselves.
+        */
        MemSetAligned(ret, 0, size);
 
        return ret;
@@ -1122,20 +1116,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 +1190,11 @@ palloc(Size size)
        Assert(MemoryContextIsValid(context));
        AssertNotInCriticalSection(context);
 
-       if (!AllocSizeIsValid(size))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        context->isReset = false;
 
-       ret = context->methods->alloc(context, size);
-       if (unlikely(ret == NULL))
-       {
-               MemoryContextStats(TopMemoryContext);
-               ereport(ERROR,
-                               (errcode(ERRCODE_OUT_OF_MEMORY),
-                                errmsg("out of memory"),
-                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
-                                                  size, context->name)));
-       }
+       ret = context->methods->alloc(context, size, 0);
 
+       Assert(ret != 0);
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
        return ret;
@@ -1238,21 +1210,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 +1231,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,29 +1405,15 @@ repalloc(void *pointer, Size size)
 #endif
        void       *ret;
 
-       if (!AllocSizeIsValid(size))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        AssertNotInCriticalSection(context);
 
        /* isReset must be false already */
        Assert(!context->isReset);
 
-       ret = MCXT_METHOD(pointer, realloc) (pointer, size);
-       if (unlikely(ret == NULL))
-       {
-               MemoryContext cxt = GetMemoryChunkContext(pointer);
-
-               MemoryContextStats(TopMemoryContext);
-               ereport(ERROR,
-                               (errcode(ERRCODE_OUT_OF_MEMORY),
-                                errmsg("out of memory"),
-                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
-                                                  size, cxt->name)));
-       }
+       ret = MCXT_METHOD(pointer, realloc) (pointer, size, 0);
 
 #ifdef USE_VALGRIND
-       if (method != MCTX_ALIGNED_REDIRECT_ID)
+       if (ret != NULL && method != MCTX_ALIGNED_REDIRECT_ID)
                VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 #endif
 
@@ -1500,31 +1433,14 @@ repalloc_extended(void *pointer, Size size, int flags)
 #endif
        void       *ret;
 
-       if (!((flags & MCXT_ALLOC_HUGE) != 0 ? AllocHugeSizeIsValid(size) :
-                 AllocSizeIsValid(size)))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        AssertNotInCriticalSection(context);
 
        /* isReset must be false already */
        Assert(!context->isReset);
 
-       ret = MCXT_METHOD(pointer, realloc) (pointer, size);
+       ret = MCXT_METHOD(pointer, realloc) (pointer, size, flags);
        if (unlikely(ret == NULL))
-       {
-               if ((flags & MCXT_ALLOC_NO_OOM) == 0)
-               {
-                       MemoryContext cxt = GetMemoryChunkContext(pointer);
-
-                       MemoryContextStats(TopMemoryContext);
-                       ereport(ERROR,
-                                       (errcode(ERRCODE_OUT_OF_MEMORY),
-                                        errmsg("out of memory"),
-                                        errdetail("Failed on request of size 
%zu in memory context \"%s\".",
-                                                          size, cxt->name)));
-               }
                return NULL;
-       }
 
        VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
@@ -1565,21 +1481,9 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
        Assert(MemoryContextIsValid(context));
        AssertNotInCriticalSection(context);
 
-       if (!AllocHugeSizeIsValid(size))
-               elog(ERROR, "invalid memory alloc request size %zu", size);
-
        context->isReset = false;
 
-       ret = context->methods->alloc(context, size);
-       if (unlikely(ret == NULL))
-       {
-               MemoryContextStats(TopMemoryContext);
-               ereport(ERROR,
-                               (errcode(ERRCODE_OUT_OF_MEMORY),
-                                errmsg("out of memory"),
-                                errdetail("Failed on request of size %zu in 
memory context \"%s\".",
-                                                  size, context->name)));
-       }
+       ret = context->methods->alloc(context, size, MCXT_ALLOC_HUGE);
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index b8f00cfc7c..1345c8d500 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -496,7 +496,7 @@ SlabDelete(MemoryContext context)
  *             request could not be completed; memory is added to the slab.
  */
 void *
-SlabAlloc(MemoryContext context, Size size)
+SlabAlloc(MemoryContext context, Size size, int flags)
 {
        SlabContext *slab = (SlabContext *) context;
        SlabBlock  *block;
@@ -504,6 +504,12 @@ SlabAlloc(MemoryContext context, Size size)
 
        Assert(SlabIsValid(slab));
 
+       /*
+        * XXX: Probably no need to check for huge allocations, we only support
+        * one size? Which could theoretically be huge, but that'd not make
+        * sense...
+        */
+
        /* sanity check that this is pointing to a valid blocklist */
        Assert(slab->curBlocklistIndex >= 0);
        Assert(slab->curBlocklistIndex <= SlabBlocklistIndex(slab, 
slab->chunksPerBlock));
@@ -546,7 +552,7 @@ SlabAlloc(MemoryContext context, Size size)
                        block = (SlabBlock *) malloc(slab->blockSize);
 
                        if (unlikely(block == NULL))
-                               return NULL;
+                               return MemoryContextAllocationFailure(context, 
size, flags);
 
                        block->slab = slab;
                        context->mem_allocated += slab->blockSize;
@@ -770,7 +776,7 @@ SlabFree(void *pointer)
  * realloc is usually used to enlarge the chunk.
  */
 void *
-SlabRealloc(void *pointer, Size size)
+SlabRealloc(void *pointer, Size size, int flags)
 {
        MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
        SlabBlock  *block;
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index a48f7e5a18..cc8cc58a0c 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -57,10 +57,10 @@ typedef void (*MemoryStatsPrintFunc) (MemoryContext 
context, void *passthru,
 
 typedef struct MemoryContextMethods
 {
-       void       *(*alloc) (MemoryContext context, Size size);
+       void       *(*alloc) (MemoryContext context, Size size, int flags);
        /* call this free_p in case someone #define's free() */
        void            (*free_p) (void *pointer);
-       void       *(*realloc) (void *pointer, Size size);
+       void       *(*realloc) (void *pointer, Size size, int flags);
        void            (*reset) (MemoryContext context);
        void            (*delete_context) (MemoryContext context);
        MemoryContext (*get_chunk_context) (void *pointer);
diff --git a/src/include/utils/memutils_internal.h 
b/src/include/utils/memutils_internal.h
index e0c4f3d5af..e5cecd7122 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -19,9 +19,9 @@
 #include "utils/memutils.h"
 
 /* These functions implement the MemoryContext API for AllocSet context. */
-extern void *AllocSetAlloc(MemoryContext context, Size size);
+extern void *AllocSetAlloc(MemoryContext context, Size size, int flags);
 extern void AllocSetFree(void *pointer);
-extern void *AllocSetRealloc(void *pointer, Size size);
+extern void *AllocSetRealloc(void *pointer, Size size, int flags);
 extern void AllocSetReset(MemoryContext context);
 extern void AllocSetDelete(MemoryContext context);
 extern MemoryContext AllocSetGetChunkContext(void *pointer);
@@ -36,9 +36,9 @@ extern void AllocSetCheck(MemoryContext context);
 #endif
 
 /* These functions implement the MemoryContext API for Generation context. */
-extern void *GenerationAlloc(MemoryContext context, Size size);
+extern void *GenerationAlloc(MemoryContext context, Size size, int flags);
 extern void GenerationFree(void *pointer);
-extern void *GenerationRealloc(void *pointer, Size size);
+extern void *GenerationRealloc(void *pointer, Size size, int flags);
 extern void GenerationReset(MemoryContext context);
 extern void GenerationDelete(MemoryContext context);
 extern MemoryContext GenerationGetChunkContext(void *pointer);
@@ -54,9 +54,9 @@ extern void GenerationCheck(MemoryContext context);
 
 
 /* These functions implement the MemoryContext API for Slab context. */
-extern void *SlabAlloc(MemoryContext context, Size size);
+extern void *SlabAlloc(MemoryContext context, Size size, int flags);
 extern void SlabFree(void *pointer);
-extern void *SlabRealloc(void *pointer, Size size);
+extern void *SlabRealloc(void *pointer, Size size, int flags);
 extern void SlabReset(MemoryContext context);
 extern void SlabDelete(MemoryContext context);
 extern MemoryContext SlabGetChunkContext(void *pointer);
@@ -75,7 +75,7 @@ extern void SlabCheck(MemoryContext context);
  * part of a fully-fledged MemoryContext type.
  */
 extern void AlignedAllocFree(void *pointer);
-extern void *AlignedAllocRealloc(void *pointer, Size size);
+extern void *AlignedAllocRealloc(void *pointer, Size size, int flags);
 extern MemoryContext AlignedAllocGetChunkContext(void *pointer);
 extern Size AlignedAllocGetChunkSpace(void *pointer);
 
@@ -133,4 +133,19 @@ extern void MemoryContextCreate(MemoryContext node,
                                                                MemoryContext 
parent,
                                                                const char 
*name);
 
+extern void *MemoryContextAllocationFailure(MemoryContext context, Size size, 
int flags);
+
+extern void MemoryContextSizeFailure(MemoryContext context, Size size, int 
flags) pg_attribute_noreturn();
+
+static inline void
+MemoryContextCheckSize(MemoryContext context, Size size, int flags)
+{
+       if (unlikely(!AllocSizeIsValid(size)))
+       {
+               if (!(flags & MCXT_ALLOC_HUGE) || !AllocHugeSizeIsValid(size))
+                       MemoryContextSizeFailure(context, size, flags);
+       }
+}
+
+
 #endif                                                 /* MEMUTILS_INTERNAL_H 
*/
-- 
2.40.1

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

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

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 2947b2046b..c214d323c8 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -1023,6 +1023,12 @@ MemoryContextCreate(MemoryContext node,
        VALGRIND_CREATE_MEMPOOL(node, 0, false);
 }
 
+/*
+ * MemoryContextAllocationFailure
+ *             For use by MemoryContextMethods implementations to handle when 
malloc
+ *             returns NULL.  The bahavior is specific to whether 
MCXT_ALLOC_NO_OOM
+ *             is in 'flags'.
+ */
 void *
 MemoryContextAllocationFailure(MemoryContext context, Size size, int flags)
 {
@@ -1038,6 +1044,11 @@ MemoryContextAllocationFailure(MemoryContext context, 
Size size, int flags)
        return NULL;
 }
 
+/*
+ * MemoryContextSizeFailure
+ *             For use by MemoryContextMethods implementations to handle 
invalid
+ *             memory allocation request sizes.
+ */
 void
 MemoryContextSizeFailure(MemoryContext context, Size size, int flags)
 {
@@ -1061,6 +1072,16 @@ MemoryContextAlloc(MemoryContext context, Size size)
 
        context->isReset = false;
 
+       /*
+        * For efficiency reasons, we purposefully offload the handling of
+        * allocation failures to the MemoryContextMethods implementation as 
this
+        * allows these checks to be performed only when an actual malloc needs 
to
+        * be done to request more memory from the OS.  Additionally, not having
+        * to execute any instructions after this call allows the compiler to 
use
+        * the tailcall optimization.  If you're considering adding code after
+        * this call, consider making it the responsibility of the 'alloc'
+        * function instead.
+        */
        ret = context->methods->alloc(context, size, 0);
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
@@ -1192,9 +1213,19 @@ palloc(Size size)
 
        context->isReset = false;
 
+       /*
+        * For efficiency reasons, we purposefully offload the handling of
+        * allocation failures to the MemoryContextMethods implementation as 
this
+        * allows these checks to be performed only when an actual malloc needs 
to
+        * be done to request more memory from the OS.  Additionally, not having
+        * to execute any instructions after this call allows the compiler to 
use
+        * the tailcall optimization.  If you're considering adding code after
+        * this call, consider making it the responsibility of the 'alloc'
+        * function instead.
+        */
        ret = context->methods->alloc(context, size, 0);
-
-       Assert(ret != 0);
+       /* We expect OOM to be handled by the alloc function */
+       Assert(ret != NULL);
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
        return ret;
@@ -1410,10 +1441,20 @@ repalloc(void *pointer, Size size)
        /* isReset must be false already */
        Assert(!context->isReset);
 
+       /*
+        * For efficiency reasons, we purposefully offload the handling of
+        * allocation failures to the MemoryContextMethods implementation as 
this
+        * allows these checks to be performed only when an actual malloc needs 
to
+        * be done to request more memory from the OS.  Additionally, not having
+        * to execute any instructions after this call allows the compiler to 
use
+        * the tailcall optimization.  If you're considering adding code after
+        * this call, consider making it the responsibility of the 'realloc'
+        * function instead.
+        */
        ret = MCXT_METHOD(pointer, realloc) (pointer, size, 0);
 
 #ifdef USE_VALGRIND
-       if (ret != NULL && method != MCTX_ALIGNED_REDIRECT_ID)
+       if (method != MCTX_ALIGNED_REDIRECT_ID)
                VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 #endif
 
@@ -1438,6 +1479,16 @@ repalloc_extended(void *pointer, Size size, int flags)
        /* isReset must be false already */
        Assert(!context->isReset);
 
+       /*
+        * For efficiency reasons, we purposefully offload the handling of
+        * allocation failures to the MemoryContextMethods implementation as 
this
+        * allows these checks to be performed only when an actual malloc needs 
to
+        * be done to request more memory from the OS.  Additionally, not having
+        * to execute any instructions after this call allows the compiler to 
use
+        * the tailcall optimization.  If you're considering adding code after
+        * this call, consider making it the responsibility of the 'realloc'
+        * function instead.
+        */
        ret = MCXT_METHOD(pointer, realloc) (pointer, size, flags);
        if (unlikely(ret == NULL))
                return NULL;
@@ -1483,6 +1534,16 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 
        context->isReset = false;
 
+       /*
+        * For efficiency reasons, we purposefully offload the handling of
+        * allocation failures to the MemoryContextMethods implementation as 
this
+        * allows these checks to be performed only when an actual malloc needs 
to
+        * be done to request more memory from the OS.  Additionally, not having
+        * to execute any instructions after this call allows the compiler to 
use
+        * the tailcall optimization.  If you're considering adding code after
+        * this call, consider making it the responsibility of the 'alloc'
+        * function instead.
+        */
        ret = context->methods->alloc(context, size, MCXT_ALLOC_HUGE);
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 1345c8d500..bc91446cb3 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -504,17 +504,14 @@ SlabAlloc(MemoryContext context, Size size, int flags)
 
        Assert(SlabIsValid(slab));
 
-       /*
-        * XXX: Probably no need to check for huge allocations, we only support
-        * one size? Which could theoretically be huge, but that'd not make
-        * sense...
-        */
-
        /* sanity check that this is pointing to a valid blocklist */
        Assert(slab->curBlocklistIndex >= 0);
        Assert(slab->curBlocklistIndex <= SlabBlocklistIndex(slab, 
slab->chunksPerBlock));
 
-       /* make sure we only allow correct request size */
+       /*
+        * Make sure we only allow correct request size.  This doubles as the
+        * MemoryContextCheckSize check.
+        */
        if (unlikely(size != slab->chunkSize))
                elog(ERROR, "unexpected alloc chunk size %zu (expected %u)",
                         size, slab->chunkSize);
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index cc8cc58a0c..edc0257f36 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -57,20 +57,58 @@ typedef void (*MemoryStatsPrintFunc) (MemoryContext 
context, void *passthru,
 
 typedef struct MemoryContextMethods
 {
+       /*
+        * Function to handle memory allocation requests of 'size' to allocate
+        * memory into the given 'context'.  The function must handle flags
+        * MCXT_ALLOC_HUGE and MCXT_ALLOC_NO_OOM.  MCXT_ALLOC_ZERO is handled by
+        * the calling function.
+        */
        void       *(*alloc) (MemoryContext context, Size size, int flags);
+
        /* call this free_p in case someone #define's free() */
        void            (*free_p) (void *pointer);
+
+       /*
+        * Function to handle a size change request for an existing allocation.
+        * The implementation must handle flags MCXT_ALLOC_HUGE and
+        * MCXT_ALLOC_NO_OOM.  MCXT_ALLOC_ZERO is handled by the calling 
function.
+        */
        void       *(*realloc) (void *pointer, Size size, int flags);
+
+       /*
+        * Invalidate all previous allocations in the given memory context and
+        * prepare the context for a new set of allocations.  Implementations 
may
+        * optionally free() excess memory back to the OS during this time.
+        */
        void            (*reset) (MemoryContext context);
+
+       /* Free all memory consumed by the given MemoryContext. */
        void            (*delete_context) (MemoryContext context);
+
+       /* Return the MemoryContext that the given pointer belongs to. */
        MemoryContext (*get_chunk_context) (void *pointer);
+
+       /*
+        * Return the number of bytes consumed by the given pointer within its
+        * memory context, including the overhead of alignment and chunk 
headers.
+        */
        Size            (*get_chunk_space) (void *pointer);
+
+       /*
+        * Return true if the given MemoryContext has not had any allocations
+        * since it was created or last reset.
+        */
        bool            (*is_empty) (MemoryContext context);
        void            (*stats) (MemoryContext context,
                                                  MemoryStatsPrintFunc 
printfunc, void *passthru,
                                                  MemoryContextCounters *totals,
                                                  bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
+
+       /*
+        * Perform validation checks on the given context and raise any 
discovered
+        * anomalies as WARNINGs.
+        */
        void            (*check) (MemoryContext context);
 #endif
 } MemoryContextMethods;
diff --git a/src/include/utils/memutils_internal.h 
b/src/include/utils/memutils_internal.h
index e5cecd7122..ad1048fd82 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -133,9 +133,11 @@ extern void MemoryContextCreate(MemoryContext node,
                                                                MemoryContext 
parent,
                                                                const char 
*name);
 
-extern void *MemoryContextAllocationFailure(MemoryContext context, Size size, 
int flags);
+extern void *MemoryContextAllocationFailure(MemoryContext context, Size size,
+                                                                               
        int flags);
 
-extern void MemoryContextSizeFailure(MemoryContext context, Size size, int 
flags) pg_attribute_noreturn();
+extern void MemoryContextSizeFailure(MemoryContext context, Size size,
+                                                                        int 
flags) pg_attribute_noreturn();
 
 static inline void
 MemoryContextCheckSize(MemoryContext context, Size size, int flags)
@@ -147,5 +149,4 @@ MemoryContextCheckSize(MemoryContext context, Size size, 
int flags)
        }
 }
 
-
 #endif                                                 /* MEMUTILS_INTERNAL_H 
*/
-- 
2.40.1

Reply via email to