On Fri, 9 Sept 2022 at 11:33, David Rowley <dgrowle...@gmail.com> wrote:
> I really think the Assert only form of MemoryContextContains() is the
> best move, and if it's doing Assert only, then we can do the
> loop-over-the-blocks idea as you described and I drafted in [1].
>
> If the need comes up that we're certain we always have a pointer to
> some allocated chunk, but need to know if it's in some memory context,
> then the proper form of expressing that, I think, should be:
>
> if (GetMemoryChunkContext(pointer) == somecontext)
>
> If we're worried about getting that wrong, we can beef up the
> MemoryChunk struct with a magic_number field in
> MEMORY_CONTEXT_CHECKING builds to ensure we catch any code which
> passes invalid pointers.

I've attached a patch series which is my proposal on what we should do
about MemoryContextContains. In summary, this basically changes
MemoryContextContains() so it's only available in
MEMORY_CONTEXT_CHECKING builds and removes 4 current usages of the
function.

0001: Makes MemoryContextContains work again with the
loop-over-the-blocks method mentioned by Tom.
0002: Adds a new "chunk_magic" field to MemoryChunk.  My thoughts are
that it might be a good idea to do this so that we get Assert failures
if we use functions like GetMemoryChunkContext() on a pointer that's
not a MemoryChunk.
0003: This adjusts aggregate final functions and window functions so
that any byref Datum they return is allocated in CurrentMemoryContext
0004: Makes MemoryContextContains only available in
MEMORY_CONTEXT_CHECKING builds and adjusts our usages of the function
to use GetMemoryChunkContext() instead.

An alternative to 0004, would be more along the lines of what was
mentioned by Andres and just Assert that the returned value is in the
memory context that we expect. I don't think we need to do anything
special with aggregates that take an internal state.  I think the rule
is just as simple as;  all final functions and window functions must
return any byref values in CurrentMemoryContext.  For aggregates
without a finalfn, we can just datumCopy() the returned byref value.
There's no chance for those to be in CurrentMemoryContext anyway. The
return value must be in the aggregate state's context.  The attached
assert.patch shows that this holds true in make check after applying
each of the other patches.

I see that one of the drawbacks from not using MemoryContextContains()
is that window functions such as lead(), lag(), first_value(),
last_value() and nth_value() may now do the datumCopy() when it's not
needed. For example, with a window function call such as
lead(byref_col ), the expression evaluation code in
WinGetFuncArgInPartition() will just return the address in the
tuplestore tuple for "byref_col".  The datumCopy() is needed for that.
However, if the function call was lead(byref_col || 'something') then
we'd have ended up with a new allocation in CurrentMemoryContext to
concatenate the two values.  We'll now do a datumCopy() where we
previously wouldn't have. I don't really see any way around that
without doing some highly invasive surgery to the expression
evaluation code.

None of the attached patches are polished. I can do that once we agree
on the best way to fix the issue.

David
From a9fcf078bc8a87742860abee527199f772f39f5b Mon Sep 17 00:00:00 2001
From: David Rowley <dgrow...@gmail.com>
Date: Thu, 8 Sep 2022 23:32:03 +1200
Subject: [PATCH v1 1/4] Reinstate MemoryContextContains to work with arbitrary
 pointers

---
 src/backend/utils/mmgr/aset.c            | 42 ++++++++++++++++++
 src/backend/utils/mmgr/generation.c      | 45 +++++++++++++++++++
 src/backend/utils/mmgr/mcxt.c            | 56 ++++++++++++++----------
 src/backend/utils/mmgr/slab.c            | 44 +++++++++++++++++++
 src/include/nodes/memnodes.h             |  4 +-
 src/include/utils/memutils_internal.h    |  3 ++
 src/include/utils/memutils_memorychunk.h | 21 +++++++++
 7 files changed, 191 insertions(+), 24 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index ec423375ae..24ab7399a3 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -1333,6 +1333,48 @@ AllocSetGetChunkContext(void *pointer)
        return &set->header;
 }
 
+/*
+ * AllocSetContains
+ *             Attempt to determine if 'pointer' is memory which was allocated 
by
+ *             'context'.  Return true if it is, otherwise false.
+ */
+bool
+AllocSetContains(MemoryContext context, void *pointer)
+{
+       MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+       AllocBlock      block;
+       AllocSet        set = (AllocSet) context;
+
+       /*
+        * We must use MemoryChunkIsExternalUnSafePointer() instead of
+        * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
+        * 'pointer' is not pointing to palloced memory.  Below we're careful
+        * never to dereference 'block' as it could point to memory not owned by
+        * this process.
+        */
+       if (MemoryChunkIsExternalUnSafePointer(chunk))
+               block = ExternalChunkGetBlock(chunk);
+       else
+               block = (AllocBlock) MemoryChunkGetBlock(chunk);
+
+       for (AllocBlock blk = set->blocks; blk != NULL; blk = blk->next)
+       {
+               if (block == blk)
+               {
+                       /*
+                        * The block matches. Now check if 'pointer' falls 
within the
+                        * block's memory.  We don't detect if the pointer is 
pointing to
+                        * an allocated chunk as that would require looping 
over the
+                        * freelist for this chunk's size.
+                        */
+                       if ((char *) pointer >= (char *) blk + ALLOC_BLOCKHDRSZ 
+ ALLOC_CHUNKHDRSZ &&
+                               (char *) pointer < blk->endptr)
+                               return true;
+               }
+       }
+       return false;
+}
+
 /*
  * AllocSetGetChunkSpace
  *             Given a currently-allocated chunk, determine the total space
diff --git a/src/backend/utils/mmgr/generation.c 
b/src/backend/utils/mmgr/generation.c
index c743b24fa7..212aba6bf3 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -848,6 +848,51 @@ GenerationGetChunkContext(void *pointer)
        return &block->context->header;
 }
 
+/*
+ * GenerationContains
+ *             Attempt to determine if 'pointer' is memory which was allocated 
by
+ *             'context'.  Return true if it is, otherwise false.
+ */
+bool
+GenerationContains(MemoryContext context, void *pointer)
+{
+       MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+       GenerationBlock *block;
+       GenerationContext *set = (GenerationContext *) context;
+       dlist_iter      iter;
+
+       /*
+        * We must use MemoryChunkIsExternalUnSafePointer() instead of
+        * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
+        * 'pointer' is not pointing to palloced memory.  Below we're careful
+        * never to dereference 'block' as it could point to memory not owned by
+        * this process.
+        */
+       if (MemoryChunkIsExternalUnSafePointer(chunk))
+               block = ExternalChunkGetBlock(chunk);
+       else
+               block = (GenerationBlock *) MemoryChunkGetBlock(chunk);
+
+       dlist_foreach(iter, &set->blocks)
+       {
+               GenerationBlock *blk = dlist_container(GenerationBlock, node, 
iter.cur);
+
+               if (block == blk)
+               {
+                       /*
+                        * The block matches. Now check if 'pointer' falls 
within the
+                        * block's memory.  We don't detect if the pointer is 
pointing to
+                        * an allocated chunk as that would require looping 
over the
+                        * freelist for this chunk's size.
+                        */
+                       if ((char *) pointer >= (char *) blk + 
Generation_BLOCKHDRSZ + Generation_CHUNKHDRSZ &&
+                               (char *) pointer < blk->endptr)
+                               return true;
+               }
+       }
+       return false;
+}
+
 /*
  * GenerationGetChunkSpace
  *             Given a currently-allocated chunk, determine the total space
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 115a64cfe4..3615a83d53 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -44,6 +44,7 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_ASET_ID].reset = AllocSetReset,
        [MCTX_ASET_ID].delete_context = AllocSetDelete,
        [MCTX_ASET_ID].get_chunk_context = AllocSetGetChunkContext,
+       [MCTX_ASET_ID].contains = AllocSetContains,
        [MCTX_ASET_ID].get_chunk_space = AllocSetGetChunkSpace,
        [MCTX_ASET_ID].is_empty = AllocSetIsEmpty,
        [MCTX_ASET_ID].stats = AllocSetStats,
@@ -58,6 +59,7 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_GENERATION_ID].reset = GenerationReset,
        [MCTX_GENERATION_ID].delete_context = GenerationDelete,
        [MCTX_GENERATION_ID].get_chunk_context = GenerationGetChunkContext,
+       [MCTX_GENERATION_ID].contains = GenerationContains,
        [MCTX_GENERATION_ID].get_chunk_space = GenerationGetChunkSpace,
        [MCTX_GENERATION_ID].is_empty = GenerationIsEmpty,
        [MCTX_GENERATION_ID].stats = GenerationStats,
@@ -72,6 +74,7 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_SLAB_ID].reset = SlabReset,
        [MCTX_SLAB_ID].delete_context = SlabDelete,
        [MCTX_SLAB_ID].get_chunk_context = SlabGetChunkContext,
+       [MCTX_SLAB_ID].contains = SlabContains,
        [MCTX_SLAB_ID].get_chunk_space = SlabGetChunkSpace,
        [MCTX_SLAB_ID].is_empty = SlabIsEmpty,
        [MCTX_SLAB_ID].stats = SlabStats
@@ -818,28 +821,16 @@ MemoryContextCheck(MemoryContext context)
  *             Detect whether an allocated chunk of memory belongs to a given
  *             context or not.
  *
- * Caution: 'pointer' must point to a pointer which was allocated by a
- * MemoryContext.  It's not safe or valid to use this function on arbitrary
- * pointers as obtaining the MemoryContext which 'pointer' belongs to requires
- * possibly several pointer dereferences.
+ * Caution: this test is reliable as long as 'pointer' does point to
+ * a chunk of memory allocated from *some* context.  If 'pointer' points
+ * at memory obtained in some other way, there is a chance of a segmentation
+ * violation due to accessing the chunk header, which look for in the 8 bytes
+ * prior to the pointer.
  */
 bool
 MemoryContextContains(MemoryContext context, void *pointer)
 {
        /*
-        * Temporarily make this always return false as we don't yet have a 
fully
-        * baked idea on how to make it work correctly with the new MemoryChunk
-        * code.
-        */
-       return false;
-
-#ifdef NOT_USED
-       MemoryContext ptr_context;
-
-       /*
-        * NB: We must perform run-time checks here which 
GetMemoryChunkContext()
-        * does as assertions before calling GetMemoryChunkContext().
-        *
         * Try to detect bogus pointers handed to us, poorly though we can.
         * Presumably, a pointer that isn't MAXALIGNED isn't pointing at an
         * allocated chunk.
@@ -847,13 +838,10 @@ MemoryContextContains(MemoryContext context, void 
*pointer)
        if (pointer == NULL || pointer != (void *) MAXALIGN(pointer))
                return false;
 
-       /*
-        * OK, it's probably safe to look at the context.
-        */
-       ptr_context = GetMemoryChunkContext(pointer);
+       if (GetMemoryChunkMethodID(pointer) != context->method_id)
+               return false;
 
-       return ptr_context == context;
-#endif
+       return context->methods->contains(context, pointer);
 }
 
 /*
@@ -900,6 +888,8 @@ MemoryContextCreate(MemoryContext node,
 
        /* Initialize all standard fields of memory context header */
        node->type = tag;
+       /* Use uint8 to prevent increasing sizeof(MemoryContextData) */
+       node->method_id = (uint8) method_id;
        node->isReset = true;
        node->methods = &mcxt_methods[method_id];
        node->parent = parent;
@@ -969,6 +959,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1007,6 +999,8 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 
        MemSetAligned(ret, 0, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1045,6 +1039,8 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size 
size)
 
        MemSetLoop(ret, 0, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1086,6 +1082,8 @@ MemoryContextAllocExtended(MemoryContext context, Size 
size, int flags)
        if ((flags & MCXT_ALLOC_ZERO) != 0)
                MemSetAligned(ret, 0, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1169,6 +1167,8 @@ palloc(Size size)
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1202,6 +1202,8 @@ palloc0(Size size)
 
        MemSetAligned(ret, 0, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1241,6 +1243,8 @@ palloc_extended(Size size, int flags)
        if ((flags & MCXT_ALLOC_ZERO) != 0)
                MemSetAligned(ret, 0, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1294,6 +1298,8 @@ repalloc(void *pointer, Size size)
 
        VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1329,6 +1335,8 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 
        VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
@@ -1368,6 +1376,8 @@ repalloc_huge(void *pointer, Size size)
 
        VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
+       Assert(MemoryContextContains(context, ret));
+
        return ret;
 }
 
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index 9149aaafcb..b5c801eb01 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -567,6 +567,50 @@ SlabGetChunkContext(void *pointer)
        return &slab->header;
 }
 
+/*
+ * SlabContains
+ *             Attempt to determine if 'pointer' is memory which was allocated 
by
+ *             'context'.  Return true if it is, otherwise false.
+ */
+bool
+SlabContains(MemoryContext context, void *pointer)
+{
+       MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+       SlabBlock  *block;
+       SlabContext *set = (SlabContext *) context;
+
+       /*
+        * Careful not to dereference anything in 'block' as if 'pointer' is 
not a
+        * pointer to an allocated chunk then 'block' could be pointing to about
+        * anything.
+        */
+       block = (SlabBlock *) MemoryChunkGetBlock(chunk);
+
+       for (int i = 0; i <= set->chunksPerBlock; i++)
+       {
+               dlist_iter      iter;
+
+               dlist_foreach(iter, &set->freelist[i])
+               {
+                       SlabBlock  *blk = dlist_container(SlabBlock, node, 
iter.cur);
+
+                       if (block == blk)
+                       {
+                               /*
+                                * The block matches. Now check if 'pointer' 
falls within the
+                                * block's memory.  We don't detect if the 
pointer is pointing
+                                * to an allocated chunk as that would require 
looping over
+                                * the freelist for this chunk's size.
+                                */
+                               if ((char *) pointer >= (char *) blk + 
Slab_BLOCKHDRSZ + Slab_CHUNKHDRSZ &&
+                                       (char *) pointer < (char *) blk + 
set->blockSize)
+                                       return true;
+                       }
+               }
+       }
+       return false;
+}
+
 /*
  * SlabGetChunkSpace
  *             Given a currently-allocated chunk, determine the total space
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index 63d07358cd..6a8d4f4e4e 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -64,6 +64,7 @@ typedef struct MemoryContextMethods
        void            (*reset) (MemoryContext context);
        void            (*delete_context) (MemoryContext context);
        MemoryContext (*get_chunk_context) (void *pointer);
+       bool            (*contains) (MemoryContext context, void *pointer);
        Size            (*get_chunk_space) (void *pointer);
        bool            (*is_empty) (MemoryContext context);
        void            (*stats) (MemoryContext context,
@@ -81,7 +82,8 @@ typedef struct MemoryContextData
        pg_node_attr(abstract)          /* there are no nodes of this type */
 
        NodeTag         type;                   /* identifies exact kind of 
context */
-       /* these two fields are placed here to minimize alignment wastage: */
+       /* these three fields are placed here to minimize alignment wastage: */
+       uint8           method_id;              /* MemoryContextMethodID for 
this context */
        bool            isReset;                /* T = no space alloced since 
last reset */
        bool            allowInCritSection; /* allow palloc in critical section 
*/
        Size            mem_allocated;  /* track memory allocated for this 
context */
diff --git a/src/include/utils/memutils_internal.h 
b/src/include/utils/memutils_internal.h
index 377cee7a84..ebc1d60393 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -25,6 +25,7 @@ extern void *AllocSetRealloc(void *pointer, Size size);
 extern void AllocSetReset(MemoryContext context);
 extern void AllocSetDelete(MemoryContext context);
 extern MemoryContext AllocSetGetChunkContext(void *pointer);
+extern bool AllocSetContains(MemoryContext context, void *pointer);
 extern Size AllocSetGetChunkSpace(void *pointer);
 extern bool AllocSetIsEmpty(MemoryContext context);
 extern void AllocSetStats(MemoryContext context,
@@ -42,6 +43,7 @@ extern void *GenerationRealloc(void *pointer, Size size);
 extern void GenerationReset(MemoryContext context);
 extern void GenerationDelete(MemoryContext context);
 extern MemoryContext GenerationGetChunkContext(void *pointer);
+extern bool GenerationContains(MemoryContext context, void *pointer);
 extern Size GenerationGetChunkSpace(void *pointer);
 extern bool GenerationIsEmpty(MemoryContext context);
 extern void GenerationStats(MemoryContext context,
@@ -60,6 +62,7 @@ extern void *SlabRealloc(void *pointer, Size size);
 extern void SlabReset(MemoryContext context);
 extern void SlabDelete(MemoryContext context);
 extern MemoryContext SlabGetChunkContext(void *pointer);
+extern bool SlabContains(MemoryContext context, void *pointer);
 extern Size SlabGetChunkSpace(void *pointer);
 extern bool SlabIsEmpty(MemoryContext context);
 extern void SlabStats(MemoryContext context,
diff --git a/src/include/utils/memutils_memorychunk.h 
b/src/include/utils/memutils_memorychunk.h
index 2eefc138e3..721c7d004c 100644
--- a/src/include/utils/memutils_memorychunk.h
+++ b/src/include/utils/memutils_memorychunk.h
@@ -54,6 +54,12 @@
  *             Determine if the given MemoryChunk is externally managed, i.e.
  *             MemoryChunkSetHdrMaskExternal() was called on the chunk.
  *
+ * MemoryChunkIsExternalUnSafePointer:
+ *             As MemoryChunkIsExternal but safe to use on pointers that we're
+ *             uncertain are pointers to MemoryChunks.  Callers must be careful
+ *             not to put too much trust into the return value.  The primary 
usecase
+ *             for this is to implement MemoryContextContains.
+ *
  * MemoryChunkGetValue:
  *             For non-external chunks, return the stored 30-bit value as it 
was set
  *             in the call to MemoryChunkSetHdrMask().
@@ -198,6 +204,21 @@ MemoryChunkIsExternal(MemoryChunk *chunk)
        return HdrMaskIsExternal(chunk->hdrmask);
 }
 
+/*
+ * MemoryChunkIsExternalUnSafePointer
+ *             As MemoryChunkIsExternal only without any Assert checking.  This
+ *             version should only be used when we're uncertain of 'chunk' is
+ *             actually a pointer to a MemoryChunk.
+ *
+ * Note: Callers must be careful not to put too much trust into the return
+ * value as 'chunk' may not be a MemoryChunk.
+ */
+static inline bool
+MemoryChunkIsExternalUnSafePointer(MemoryChunk *chunk)
+{
+       return HdrMaskIsExternal(chunk->hdrmask);
+}
+
 /*
  * MemoryChunkGetValue
  *             For non-external chunks, returns the value field as it was set 
in
-- 
2.34.1

From 1de691bf2421c3c87fc513f4c2155f140d21e2c0 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrow...@gmail.com>
Date: Fri, 9 Sep 2022 13:11:32 +1200
Subject: [PATCH v1 2/4] Add a magic number to all MemoryChunks and verify

We only do this for MEMORY_CONTEXT_CHECKING builds.  The aim here is to
give us more confidence that we're dealing with chunks allocated by a
MemoryContext in various cases where there's a chance that we might not
be.
---
 src/backend/utils/mmgr/aset.c            |  6 +--
 src/backend/utils/mmgr/generation.c      |  6 +--
 src/include/utils/memutils_memorychunk.h | 48 ++++++++++++++++++++++--
 3 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 24ab7399a3..b435624d8f 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -1346,16 +1346,16 @@ AllocSetContains(MemoryContext context, void *pointer)
        AllocSet        set = (AllocSet) context;
 
        /*
-        * We must use MemoryChunkIsExternalUnSafePointer() instead of
+        * We must use MemoryChunkIsExternalUnsafePointer() instead of
         * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
         * 'pointer' is not pointing to palloced memory.  Below we're careful
         * never to dereference 'block' as it could point to memory not owned by
         * this process.
         */
-       if (MemoryChunkIsExternalUnSafePointer(chunk))
+       if (MemoryChunkIsExternalUnsafePointer(chunk))
                block = ExternalChunkGetBlock(chunk);
        else
-               block = (AllocBlock) MemoryChunkGetBlock(chunk);
+               block = (AllocBlock) MemoryChunkGetBlockUnsafePointer(chunk);
 
        for (AllocBlock blk = set->blocks; blk != NULL; blk = blk->next)
        {
diff --git a/src/backend/utils/mmgr/generation.c 
b/src/backend/utils/mmgr/generation.c
index 212aba6bf3..f6a3a01912 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -862,16 +862,16 @@ GenerationContains(MemoryContext context, void *pointer)
        dlist_iter      iter;
 
        /*
-        * We must use MemoryChunkIsExternalUnSafePointer() instead of
+        * We must use MemoryChunkIsExternalUnsafePointer() instead of
         * MemoryChunkIsExternal() here as 'chunk' may not be a MemoryChunk if
         * 'pointer' is not pointing to palloced memory.  Below we're careful
         * never to dereference 'block' as it could point to memory not owned by
         * this process.
         */
-       if (MemoryChunkIsExternalUnSafePointer(chunk))
+       if (MemoryChunkIsExternalUnsafePointer(chunk))
                block = ExternalChunkGetBlock(chunk);
        else
-               block = (GenerationBlock *) MemoryChunkGetBlock(chunk);
+               block = (GenerationBlock *) 
MemoryChunkGetBlockUnsafePointer(chunk);
 
        dlist_foreach(iter, &set->blocks)
        {
diff --git a/src/include/utils/memutils_memorychunk.h 
b/src/include/utils/memutils_memorychunk.h
index 721c7d004c..9b267e9ce4 100644
--- a/src/include/utils/memutils_memorychunk.h
+++ b/src/include/utils/memutils_memorychunk.h
@@ -54,7 +54,7 @@
  *             Determine if the given MemoryChunk is externally managed, i.e.
  *             MemoryChunkSetHdrMaskExternal() was called on the chunk.
  *
- * MemoryChunkIsExternalUnSafePointer:
+ * MemoryChunkIsExternalUnsafePointer:
  *             As MemoryChunkIsExternal but safe to use on pointers that we're
  *             uncertain are pointers to MemoryChunks.  Callers must be careful
  *             not to put too much trust into the return value.  The primary 
usecase
@@ -68,6 +68,12 @@
  *             For non-external chunks, return a pointer to the block as it 
was set
  *             in the call to MemoryChunkSetHdrMask().
  *
+ * MemoryChunkGetBlockUnsafePointer:
+ *             As MemoryChunkGetBlock but safe to use on pointers that we're
+ *             uncertain are pointers to MemoryChunks.  Callers must be careful
+ *             not to put too much trust into the return value.  The primary 
usecase
+ *             for this is to implement MemoryContextContains.
+ *
  * Also exports:
  *             MEMORYCHUNK_MAX_VALUE
  *             MEMORYCHUNK_MAX_BLOCKOFFSET
@@ -113,9 +119,18 @@
                                                                 
MEMORYCHUNK_VALUE_BASEBIT << \
                                                                 
MEMORYCHUNK_VALUE_BASEBIT)
 
+#ifdef MEMORY_CONTEXT_CHECKING
+#if SIZEOF_SIZE_T >= 8
+#define MEMORYCHUNK_MAGIC2             UINT64CONST(0xB1A8DB858EB6EFBA)
+#else
+#define MEMORYCHUNK_MAGIC2             0x8EB6EFBA
+#endif
+#endif
+
 typedef struct MemoryChunk
 {
 #ifdef MEMORY_CONTEXT_CHECKING
+       Size            chunk_magic;    /* magic number */
        Size            requested_size;
 #endif
 
@@ -167,6 +182,9 @@ MemoryChunkSetHdrMask(MemoryChunk *chunk, void *block,
        Assert(value <= MEMORYCHUNK_MAX_VALUE);
        Assert((int) methodid <= MEMORY_CONTEXT_METHODID_MASK);
 
+#ifdef MEMORY_CONTEXT_CHECKING
+       chunk->chunk_magic = MEMORYCHUNK_MAGIC2;
+#endif
        chunk->hdrmask = (((uint64) blockoffset) << 
MEMORYCHUNK_BLOCKOFFSET_BASEBIT) |
                (((uint64) value) << MEMORYCHUNK_VALUE_BASEBIT) |
                methodid;
@@ -183,6 +201,9 @@ MemoryChunkSetHdrMaskExternal(MemoryChunk *chunk,
 {
        Assert((int) methodid <= MEMORY_CONTEXT_METHODID_MASK);
 
+#ifdef MEMORY_CONTEXT_CHECKING
+       chunk->chunk_magic = MEMORYCHUNK_MAGIC2;
+#endif
        chunk->hdrmask = MEMORYCHUNK_MAGIC | (((uint64) 1) << 
MEMORYCHUNK_EXTERNAL_BASEBIT) |
                methodid;
 }
@@ -200,12 +221,13 @@ MemoryChunkIsExternal(MemoryChunk *chunk)
         */
        Assert(!HdrMaskIsExternal(chunk->hdrmask) ||
                   HdrMaskCheckMagic(chunk->hdrmask));
+       Assert(chunk->chunk_magic == MEMORYCHUNK_MAGIC2);
 
        return HdrMaskIsExternal(chunk->hdrmask);
 }
 
 /*
- * MemoryChunkIsExternalUnSafePointer
+ * MemoryChunkIsExternalUnsafePointer
  *             As MemoryChunkIsExternal only without any Assert checking.  This
  *             version should only be used when we're uncertain of 'chunk' is
  *             actually a pointer to a MemoryChunk.
@@ -214,7 +236,7 @@ MemoryChunkIsExternal(MemoryChunk *chunk)
  * value as 'chunk' may not be a MemoryChunk.
  */
 static inline bool
-MemoryChunkIsExternalUnSafePointer(MemoryChunk *chunk)
+MemoryChunkIsExternalUnsafePointer(MemoryChunk *chunk)
 {
        return HdrMaskIsExternal(chunk->hdrmask);
 }
@@ -228,6 +250,7 @@ static inline Size
 MemoryChunkGetValue(MemoryChunk *chunk)
 {
        Assert(!HdrMaskIsExternal(chunk->hdrmask));
+       Assert(chunk->chunk_magic == MEMORYCHUNK_MAGIC2);
 
        return HdrMaskGetValue(chunk->hdrmask);
 }
@@ -241,15 +264,34 @@ static inline void *
 MemoryChunkGetBlock(MemoryChunk *chunk)
 {
        Assert(!HdrMaskIsExternal(chunk->hdrmask));
+       Assert(chunk->chunk_magic == MEMORYCHUNK_MAGIC2);
 
        return (void *) ((char *) chunk - HdrMaskBlockOffset(chunk->hdrmask));
 }
 
+/*
+ * MemoryChunkGetBlockUnsafePointer
+ *             For non-external chunks, returns the pointer to the block as 
was set
+ *             in MemoryChunkSetHdrMask.  This differs from 
MemoryChunkGetBlock as
+ *             the coding here does not assume 'chunk' is a valid MemoryChunk.
+ *             We so still assume that the caller checked if the external bit 
was not
+ *             set with MemoryChunkIsExternalUnsafePointer.
+ */
+static inline void *
+MemoryChunkGetBlockUnsafePointer(MemoryChunk *chunk)
+{
+       Assert(!HdrMaskIsExternal(chunk->hdrmask));
+
+       return (void *) ((char *) chunk - HdrMaskBlockOffset(chunk->hdrmask));
+}
+
+
 /* cleanup all internal definitions */
 #undef MEMORYCHUNK_EXTERNAL_BASEBIT
 #undef MEMORYCHUNK_VALUE_BASEBIT
 #undef MEMORYCHUNK_BLOCKOFFSET_BASEBIT
 #undef MEMORYCHUNK_MAGIC
+#undef MEMORYCHUNK_MAGIC2
 #undef HdrMaskIsExternal
 #undef HdrMaskGetValue
 #undef HdrMaskBlockOffset
-- 
2.34.1

From 1a1d69334b9cbecca5038528bc7a5bddcc5f05ca Mon Sep 17 00:00:00 2001
From: David Rowley <dgrow...@gmail.com>
Date: Fri, 9 Sep 2022 13:16:46 +1200
Subject: [PATCH v1 3/4] Adjust memory context API with aggregate and window
 functions

This makes it so an aggregate function's final function always returns
any byref types in memory allocated in CurrentMemoryContext.  The same is
also done here for window functions.  A window function may now call the
PG_WINDOW_RESULTINFO() macro to obtain a WindowResultInfo in order to
determine the result type details of the return value.  Functions which do
not currently return in CurrentMemoryContext must perform a datumCopy().

Extension authors who have created aggregates with a final function or
created their own window functions must ensure their C code conforms.
---
 src/backend/executor/nodeWindowAgg.c | 10 +++++++---
 src/backend/utils/adt/numeric.c      |  2 +-
 src/backend/utils/adt/windowfuncs.c  | 21 +++++++++++++++++++++
 src/include/windowapi.h              | 12 ++++++++++++
 4 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/backend/executor/nodeWindowAgg.c 
b/src/backend/executor/nodeWindowAgg.c
index 8b0858e9f5..ba119046f4 100644
--- a/src/backend/executor/nodeWindowAgg.c
+++ b/src/backend/executor/nodeWindowAgg.c
@@ -1034,9 +1034,13 @@ eval_windowfunction(WindowAggState *winstate, 
WindowStatePerFunc perfuncstate,
 {
        LOCAL_FCINFO(fcinfo, FUNC_MAX_ARGS);
        MemoryContext oldContext;
+       WindowResultInfo resultinfo;
 
        oldContext = 
MemoryContextSwitchTo(winstate->ss.ps.ps_ExprContext->ecxt_per_tuple_memory);
 
+       resultinfo.resulttypeLen = perfuncstate->resulttypeLen;
+       resultinfo.resulttypeByVal = perfuncstate->resulttypeByVal;
+
        /*
         * We don't pass any normal arguments to a window function, but we do 
pass
         * it the number of arguments, in order to permit window function
@@ -1046,7 +1050,8 @@ eval_windowfunction(WindowAggState *winstate, 
WindowStatePerFunc perfuncstate,
        InitFunctionCallInfoData(*fcinfo, &(perfuncstate->flinfo),
                                                         
perfuncstate->numArguments,
                                                         
perfuncstate->winCollation,
-                                                        (void *) 
perfuncstate->winobj, NULL);
+                                                        (void *) 
perfuncstate->winobj,
+                                                        (void *) &resultinfo);
        /* Just in case, make all the regular argument slots be null */
        for (int argno = 0; argno < perfuncstate->numArguments; argno++)
                fcinfo->args[argno].isnull = true;
@@ -1062,8 +1067,7 @@ eval_windowfunction(WindowAggState *winstate, 
WindowStatePerFunc perfuncstate,
         * tuple, as is entirely possible.)
         */
        if (!perfuncstate->resulttypeByVal && !fcinfo->isnull &&
-               !MemoryContextContains(CurrentMemoryContext,
-                                                          
DatumGetPointer(*result)))
+               GetMemoryChunkContext(DatumGetPointer(*result)) != 
CurrentMemoryContext)
                *result = datumCopy(*result,
                                                        
perfuncstate->resulttypeByVal,
                                                        
perfuncstate->resulttypeLen);
diff --git a/src/backend/utils/adt/numeric.c b/src/backend/utils/adt/numeric.c
index 920a63b008..2ac7de70c5 100644
--- a/src/backend/utils/adt/numeric.c
+++ b/src/backend/utils/adt/numeric.c
@@ -6643,7 +6643,7 @@ int2int4_sum(PG_FUNCTION_ARGS)
        if (transdata->count == 0)
                PG_RETURN_NULL();
 
-       PG_RETURN_DATUM(Int64GetDatumFast(transdata->sum));
+       PG_RETURN_DATUM(Int64GetDatum(transdata->sum));
 }
 
 
diff --git a/src/backend/utils/adt/windowfuncs.c 
b/src/backend/utils/adt/windowfuncs.c
index 596564fa15..40f0a08123 100644
--- a/src/backend/utils/adt/windowfuncs.c
+++ b/src/backend/utils/adt/windowfuncs.c
@@ -15,6 +15,7 @@
 
 #include "nodes/supportnodes.h"
 #include "utils/builtins.h"
+#include "utils/datum.h"
 #include "windowapi.h"
 
 /*
@@ -350,6 +351,7 @@ leadlag_common(FunctionCallInfo fcinfo,
                           bool forward, bool withoffset, bool withdefault)
 {
        WindowObject winobj = PG_WINDOW_OBJECT();
+       WindowResultInfo *wresinfo = PG_WINDOW_RESULTINFO();
        int32           offset;
        bool            const_offset;
        Datum           result;
@@ -388,6 +390,10 @@ leadlag_common(FunctionCallInfo fcinfo,
        if (isnull)
                PG_RETURN_NULL();
 
+       if (!wresinfo->resulttypeByVal)
+               result = datumCopy(result, wresinfo->resulttypeByVal,
+                                                  wresinfo->resulttypeLen);
+
        PG_RETURN_DATUM(result);
 }
 
@@ -470,6 +476,7 @@ Datum
 window_first_value(PG_FUNCTION_ARGS)
 {
        WindowObject winobj = PG_WINDOW_OBJECT();
+       WindowResultInfo *wresinfo = PG_WINDOW_RESULTINFO();
        Datum           result;
        bool            isnull;
 
@@ -479,6 +486,10 @@ window_first_value(PG_FUNCTION_ARGS)
        if (isnull)
                PG_RETURN_NULL();
 
+       if (!wresinfo->resulttypeByVal)
+               result = datumCopy(result, wresinfo->resulttypeByVal,
+                                                  wresinfo->resulttypeLen);
+
        PG_RETURN_DATUM(result);
 }
 
@@ -491,6 +502,7 @@ Datum
 window_last_value(PG_FUNCTION_ARGS)
 {
        WindowObject winobj = PG_WINDOW_OBJECT();
+       WindowResultInfo *wresinfo = PG_WINDOW_RESULTINFO();
        Datum           result;
        bool            isnull;
 
@@ -500,6 +512,10 @@ window_last_value(PG_FUNCTION_ARGS)
        if (isnull)
                PG_RETURN_NULL();
 
+       if (!wresinfo->resulttypeByVal)
+               result = datumCopy(result, wresinfo->resulttypeByVal,
+                                                  wresinfo->resulttypeLen);
+
        PG_RETURN_DATUM(result);
 }
 
@@ -512,6 +528,7 @@ Datum
 window_nth_value(PG_FUNCTION_ARGS)
 {
        WindowObject winobj = PG_WINDOW_OBJECT();
+       WindowResultInfo *wresinfo = PG_WINDOW_RESULTINFO();
        bool            const_offset;
        Datum           result;
        bool            isnull;
@@ -533,5 +550,9 @@ window_nth_value(PG_FUNCTION_ARGS)
        if (isnull)
                PG_RETURN_NULL();
 
+       if (!wresinfo->resulttypeByVal)
+               result = datumCopy(result, wresinfo->resulttypeByVal,
+                                                  wresinfo->resulttypeLen);
+
        PG_RETURN_DATUM(result);
 }
diff --git a/src/include/windowapi.h b/src/include/windowapi.h
index 5a620a2e42..c3d65b161c 100644
--- a/src/include/windowapi.h
+++ b/src/include/windowapi.h
@@ -36,7 +36,19 @@
 /* this struct is private in nodeWindowAgg.c */
 typedef struct WindowObjectData *WindowObject;
 
+/*
+ * Evaluation of window functions are passed this object and it's available
+ * via fcinfo->resultinfo and can be accessed in the window function via the
+ * PG_WINDOW_RESULTINFO macro.
+ */
+typedef struct WindowResultInfo
+{
+       int16           resulttypeLen;  /* type length of wfunc's result type */
+       bool            resulttypeByVal;        /* true if wfunc's result is 
byval */
+} WindowResultInfo;
+
 #define PG_WINDOW_OBJECT() ((WindowObject) fcinfo->context)
+#define PG_WINDOW_RESULTINFO() ((WindowResultInfo *) fcinfo->resultinfo)
 
 #define WindowObjectIsValid(winobj) \
        ((winobj) != NULL && IsA(winobj, WindowObjectData))
-- 
2.34.1

From 28f287cdd32f4d3bdde0fd803c9357646e904a17 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrow...@gmail.com>
Date: Fri, 9 Sep 2022 16:19:13 +1200
Subject: [PATCH v1 4/4] Demote MemoryContextContains to
 MEMORY_CONTEXT_CHECKING builds only

Making MemoryContextContains work like it did prior to c6e0fe1f2 required
making it perform a loop over each block that the MemoryContext had
allocated and check if the pointer is in the address-space range of that
block.  This was far too expensive to be used as a general-purpose
function.

Here we instead opt to use GetMemoryChunkContext().  Previously we
couldn't use that function as it only works when given a pointer as it was
returned by palloc and co.  For some use cases of MemoryContextContains
they could receive a pointer into a tuple (as was the case for lead/lag
window functions, and a pointer to a field in an internal aggregate state,
as was the case for int2int4_sum().
---
 src/backend/executor/nodeAgg.c        |  6 ++----
 src/backend/executor/nodeWindowAgg.c  |  3 +--
 src/backend/utils/mmgr/aset.c         |  2 ++
 src/backend/utils/mmgr/generation.c   |  3 +++
 src/backend/utils/mmgr/mcxt.c         | 10 ++++++----
 src/backend/utils/mmgr/slab.c         |  2 ++
 src/include/nodes/memnodes.h          |  2 +-
 src/include/utils/memutils.h          |  2 +-
 src/include/utils/memutils_internal.h |  6 +++---
 9 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 933c304901..323ae94938 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -1122,8 +1122,7 @@ finalize_aggregate(AggState *aggstate,
         * If result is pass-by-ref, make sure it is in the right context.
         */
        if (!peragg->resulttypeByVal && !*resultIsNull &&
-               !MemoryContextContains(CurrentMemoryContext,
-                                                          
DatumGetPointer(*resultVal)))
+               GetMemoryChunkContext(DatumGetPointer(*resultVal)) != 
CurrentMemoryContext)
                *resultVal = datumCopy(*resultVal,
                                                           
peragg->resulttypeByVal,
                                                           
peragg->resulttypeLen);
@@ -1184,8 +1183,7 @@ finalize_partialaggregate(AggState *aggstate,
 
        /* If result is pass-by-ref, make sure it is in the right context. */
        if (!peragg->resulttypeByVal && !*resultIsNull &&
-               !MemoryContextContains(CurrentMemoryContext,
-                                                          
DatumGetPointer(*resultVal)))
+               GetMemoryChunkContext(DatumGetPointer(*resultVal)) != 
CurrentMemoryContext)
                *resultVal = datumCopy(*resultVal,
                                                           
peragg->resulttypeByVal,
                                                           
peragg->resulttypeLen);
diff --git a/src/backend/executor/nodeWindowAgg.c 
b/src/backend/executor/nodeWindowAgg.c
index ba119046f4..2820044cd0 100644
--- a/src/backend/executor/nodeWindowAgg.c
+++ b/src/backend/executor/nodeWindowAgg.c
@@ -639,8 +639,7 @@ finalize_windowaggregate(WindowAggState *winstate,
         * If result is pass-by-ref, make sure it is in the right context.
         */
        if (!peraggstate->resulttypeByVal && !*isnull &&
-               !MemoryContextContains(CurrentMemoryContext,
-                                                          
DatumGetPointer(*result)))
+               GetMemoryChunkContext(DatumGetPointer(*result)) != 
CurrentMemoryContext)
                *result = datumCopy(*result,
                                                        
peraggstate->resulttypeByVal,
                                                        
peraggstate->resulttypeLen);
diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index b435624d8f..3a6f12be32 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -1333,6 +1333,7 @@ AllocSetGetChunkContext(void *pointer)
        return &set->header;
 }
 
+#ifdef MEMORY_CONTEXT_CHECKING
 /*
  * AllocSetContains
  *             Attempt to determine if 'pointer' is memory which was allocated 
by
@@ -1374,6 +1375,7 @@ AllocSetContains(MemoryContext context, void *pointer)
        }
        return false;
 }
+#endif                                                 /* 
MEMORY_CONTEXT_CHECKING */
 
 /*
  * AllocSetGetChunkSpace
diff --git a/src/backend/utils/mmgr/generation.c 
b/src/backend/utils/mmgr/generation.c
index f6a3a01912..9c0d04de08 100644
--- a/src/backend/utils/mmgr/generation.c
+++ b/src/backend/utils/mmgr/generation.c
@@ -848,6 +848,8 @@ GenerationGetChunkContext(void *pointer)
        return &block->context->header;
 }
 
+
+#ifdef MEMORY_CONTEXT_CHECKING
 /*
  * GenerationContains
  *             Attempt to determine if 'pointer' is memory which was allocated 
by
@@ -892,6 +894,7 @@ GenerationContains(MemoryContext context, void *pointer)
        }
        return false;
 }
+#endif                                                 /* 
MEMORY_CONTEXT_CHECKING */
 
 /*
  * GenerationGetChunkSpace
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index 3615a83d53..1097c1129c 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -44,12 +44,12 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_ASET_ID].reset = AllocSetReset,
        [MCTX_ASET_ID].delete_context = AllocSetDelete,
        [MCTX_ASET_ID].get_chunk_context = AllocSetGetChunkContext,
-       [MCTX_ASET_ID].contains = AllocSetContains,
        [MCTX_ASET_ID].get_chunk_space = AllocSetGetChunkSpace,
        [MCTX_ASET_ID].is_empty = AllocSetIsEmpty,
        [MCTX_ASET_ID].stats = AllocSetStats,
 #ifdef MEMORY_CONTEXT_CHECKING
        [MCTX_ASET_ID].check = AllocSetCheck,
+       [MCTX_ASET_ID].contains = AllocSetContains,
 #endif
 
        /* generation.c */
@@ -59,12 +59,12 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_GENERATION_ID].reset = GenerationReset,
        [MCTX_GENERATION_ID].delete_context = GenerationDelete,
        [MCTX_GENERATION_ID].get_chunk_context = GenerationGetChunkContext,
-       [MCTX_GENERATION_ID].contains = GenerationContains,
        [MCTX_GENERATION_ID].get_chunk_space = GenerationGetChunkSpace,
        [MCTX_GENERATION_ID].is_empty = GenerationIsEmpty,
        [MCTX_GENERATION_ID].stats = GenerationStats,
 #ifdef MEMORY_CONTEXT_CHECKING
        [MCTX_GENERATION_ID].check = GenerationCheck,
+       [MCTX_GENERATION_ID].contains = GenerationContains,
 #endif
 
        /* slab.c */
@@ -74,12 +74,12 @@ static const MemoryContextMethods mcxt_methods[] = {
        [MCTX_SLAB_ID].reset = SlabReset,
        [MCTX_SLAB_ID].delete_context = SlabDelete,
        [MCTX_SLAB_ID].get_chunk_context = SlabGetChunkContext,
-       [MCTX_SLAB_ID].contains = SlabContains,
        [MCTX_SLAB_ID].get_chunk_space = SlabGetChunkSpace,
        [MCTX_SLAB_ID].is_empty = SlabIsEmpty,
        [MCTX_SLAB_ID].stats = SlabStats
 #ifdef MEMORY_CONTEXT_CHECKING
-       ,[MCTX_SLAB_ID].check = SlabCheck
+       ,[MCTX_SLAB_ID].check = SlabCheck,
+       [MCTX_SLAB_ID].contains = SlabContains
 #endif
 };
 
@@ -816,6 +816,7 @@ MemoryContextCheck(MemoryContext context)
 }
 #endif
 
+#ifdef MEMORY_CONTEXT_CHECKING
 /*
  * MemoryContextContains
  *             Detect whether an allocated chunk of memory belongs to a given
@@ -843,6 +844,7 @@ MemoryContextContains(MemoryContext context, void *pointer)
 
        return context->methods->contains(context, pointer);
 }
+#endif                                                 /* 
MEMORY_CONTEXT_CHECKING */
 
 /*
  * MemoryContextCreate
diff --git a/src/backend/utils/mmgr/slab.c b/src/backend/utils/mmgr/slab.c
index b5c801eb01..e46f3611ac 100644
--- a/src/backend/utils/mmgr/slab.c
+++ b/src/backend/utils/mmgr/slab.c
@@ -567,6 +567,7 @@ SlabGetChunkContext(void *pointer)
        return &slab->header;
 }
 
+#ifdef MEMORY_CONTEXT_CHECKING
 /*
  * SlabContains
  *             Attempt to determine if 'pointer' is memory which was allocated 
by
@@ -610,6 +611,7 @@ SlabContains(MemoryContext context, void *pointer)
        }
        return false;
 }
+#endif                                                 /* 
MEMORY_CONTEXT_CHECKING */
 
 /*
  * SlabGetChunkSpace
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index 6a8d4f4e4e..6929cc8739 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -64,7 +64,6 @@ typedef struct MemoryContextMethods
        void            (*reset) (MemoryContext context);
        void            (*delete_context) (MemoryContext context);
        MemoryContext (*get_chunk_context) (void *pointer);
-       bool            (*contains) (MemoryContext context, void *pointer);
        Size            (*get_chunk_space) (void *pointer);
        bool            (*is_empty) (MemoryContext context);
        void            (*stats) (MemoryContext context,
@@ -73,6 +72,7 @@ typedef struct MemoryContextMethods
                                                  bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
        void            (*check) (MemoryContext context);
+       bool            (*contains) (MemoryContext context, void *pointer);
 #endif
 } MemoryContextMethods;
 
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index 52bc41ec53..d7144d7307 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -95,8 +95,8 @@ extern void MemoryContextAllowInCriticalSection(MemoryContext 
context,
 
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void MemoryContextCheck(MemoryContext context);
-#endif
 extern bool MemoryContextContains(MemoryContext context, void *pointer);
+#endif
 
 /* Handy macro for copying and assigning context ID ... but note double eval */
 #define MemoryContextCopyAndSetIdentifier(cxt, id) \
diff --git a/src/include/utils/memutils_internal.h 
b/src/include/utils/memutils_internal.h
index ebc1d60393..b767982896 100644
--- a/src/include/utils/memutils_internal.h
+++ b/src/include/utils/memutils_internal.h
@@ -25,7 +25,6 @@ extern void *AllocSetRealloc(void *pointer, Size size);
 extern void AllocSetReset(MemoryContext context);
 extern void AllocSetDelete(MemoryContext context);
 extern MemoryContext AllocSetGetChunkContext(void *pointer);
-extern bool AllocSetContains(MemoryContext context, void *pointer);
 extern Size AllocSetGetChunkSpace(void *pointer);
 extern bool AllocSetIsEmpty(MemoryContext context);
 extern void AllocSetStats(MemoryContext context,
@@ -34,6 +33,7 @@ extern void AllocSetStats(MemoryContext context,
                                                  bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void AllocSetCheck(MemoryContext context);
+extern bool AllocSetContains(MemoryContext context, void *pointer);
 #endif
 
 /* These functions implement the MemoryContext API for Generation context. */
@@ -43,7 +43,6 @@ extern void *GenerationRealloc(void *pointer, Size size);
 extern void GenerationReset(MemoryContext context);
 extern void GenerationDelete(MemoryContext context);
 extern MemoryContext GenerationGetChunkContext(void *pointer);
-extern bool GenerationContains(MemoryContext context, void *pointer);
 extern Size GenerationGetChunkSpace(void *pointer);
 extern bool GenerationIsEmpty(MemoryContext context);
 extern void GenerationStats(MemoryContext context,
@@ -52,6 +51,7 @@ extern void GenerationStats(MemoryContext context,
                                                        bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void GenerationCheck(MemoryContext context);
+extern bool GenerationContains(MemoryContext context, void *pointer);
 #endif
 
 
@@ -62,7 +62,6 @@ extern void *SlabRealloc(void *pointer, Size size);
 extern void SlabReset(MemoryContext context);
 extern void SlabDelete(MemoryContext context);
 extern MemoryContext SlabGetChunkContext(void *pointer);
-extern bool SlabContains(MemoryContext context, void *pointer);
 extern Size SlabGetChunkSpace(void *pointer);
 extern bool SlabIsEmpty(MemoryContext context);
 extern void SlabStats(MemoryContext context,
@@ -71,6 +70,7 @@ extern void SlabStats(MemoryContext context,
                                          bool print_to_stderr);
 #ifdef MEMORY_CONTEXT_CHECKING
 extern void SlabCheck(MemoryContext context);
+extern bool SlabContains(MemoryContext context, void *pointer);
 #endif
 
 /*
-- 
2.34.1

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 323ae94938..a6428a77c4 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -1108,6 +1108,9 @@ finalize_aggregate(AggState *aggstate,
                {
                        *resultVal = FunctionCallInvoke(fcinfo);
                        *resultIsNull = fcinfo->isnull;
+                       Assert(peragg->resulttypeByVal ||
+                                  *resultIsNull ||
+                                  
GetMemoryChunkContext(DatumGetPointer(*resultVal)) == CurrentMemoryContext);
                }
                aggstate->curperagg = NULL;
        }
@@ -1116,6 +1119,9 @@ finalize_aggregate(AggState *aggstate,
                /* Don't need MakeExpandedObjectReadOnly; datumCopy will copy 
it */
                *resultVal = pergroupstate->transValue;
                *resultIsNull = pergroupstate->transValueIsNull;
+               Assert(peragg->resulttypeByVal ||
+                       *resultIsNull ||
+                       GetMemoryChunkContext(DatumGetPointer(*resultVal)) != 
CurrentMemoryContext);
        }
 
        /*
@@ -1172,6 +1178,10 @@ finalize_partialaggregate(AggState *aggstate,
 
                        *resultVal = FunctionCallInvoke(fcinfo);
                        *resultIsNull = fcinfo->isnull;
+
+                       Assert(peragg->resulttypeByVal ||
+                                  *resultIsNull ||
+                                  
GetMemoryChunkContext(DatumGetPointer(*resultVal)) == CurrentMemoryContext);
                }
        }
        else
@@ -1179,6 +1189,10 @@ finalize_partialaggregate(AggState *aggstate,
                /* Don't need MakeExpandedObjectReadOnly; datumCopy will copy 
it */
                *resultVal = pergroupstate->transValue;
                *resultIsNull = pergroupstate->transValueIsNull;
+               Assert(peragg->resulttypeByVal ||
+                          *resultIsNull ||
+                          GetMemoryChunkContext(DatumGetPointer(*resultVal)) 
!= CurrentMemoryContext);
+
        }
 
        /* If result is pass-by-ref, make sure it is in the right context. */
diff --git a/src/backend/executor/nodeWindowAgg.c 
b/src/backend/executor/nodeWindowAgg.c
index 2820044cd0..3e4462b4c4 100644
--- a/src/backend/executor/nodeWindowAgg.c
+++ b/src/backend/executor/nodeWindowAgg.c
@@ -626,6 +626,10 @@ finalize_windowaggregate(WindowAggState *winstate,
                        *result = FunctionCallInvoke(fcinfo);
                        winstate->curaggcontext = NULL;
                        *isnull = fcinfo->isnull;
+
+                       Assert(peraggstate->resulttypeByVal ||
+                                  *isnull ||
+                                  
GetMemoryChunkContext(DatumGetPointer(*result)) == CurrentMemoryContext);
                }
        }
        else
@@ -633,6 +637,9 @@ finalize_windowaggregate(WindowAggState *winstate,
                /* Don't need MakeExpandedObjectReadOnly; datumCopy will copy 
it */
                *result = peraggstate->transValue;
                *isnull = peraggstate->transValueIsNull;
+               Assert(peraggstate->resulttypeByVal ||
+                               *isnull ||
+                               GetMemoryChunkContext(DatumGetPointer(*result)) 
!= CurrentMemoryContext);
        }
 
        /*
@@ -1060,6 +1067,10 @@ eval_windowfunction(WindowAggState *winstate, 
WindowStatePerFunc perfuncstate,
        *result = FunctionCallInvoke(fcinfo);
        *isnull = fcinfo->isnull;
 
+       Assert(perfuncstate->resulttypeByVal ||
+                  *isnull ||
+                  GetMemoryChunkContext(DatumGetPointer(*result)) == 
CurrentMemoryContext);
+
        /*
         * Make sure pass-by-ref data is allocated in the appropriate context. 
(We
         * need this in case the function returns a pointer into some 
short-lived

Reply via email to