pg_query 4.2.1 → 6.1.0

data/ext/pg_query/src_backend_utils_mmgr_aset.c

@@ -1,18 +1,21 @@
 /*--------------------------------------------------------------------
  * Symbols referenced in this file:
- * - AllocSetContextCreateInternal
- * - context_freelists
- * - AllocSetMethods
  * - AllocSetAlloc
+ * - AllocSetAllocLarge
  * - AllocSetFreeIndex
+ * - AllocSetAllocFromNewBlock
+ * - AllocSetAllocChunkFromBlock
  * - AllocSetFree
  * - AllocSetRealloc
  * - AllocSetReset
  * - AllocSetDelete
+ * - context_freelists
+ * - AllocSetGetChunkContext
  * - AllocSetGetChunkSpace
  * - AllocSetIsEmpty
  * - AllocSetStats
  * - AllocSetCheck
+ * - AllocSetContextCreateInternal
  * - AllocSetDeleteFreeList
  *--------------------------------------------------------------------
  */
@@ -26,7 +29,7 @@
  * type.
  *
  *
- * Portions Copyright (c) 1996-2022, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
  * IDENTIFICATION
@@ -68,6 +71,8 @@
 #include "port/pg_bitutils.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
+#include "utils/memutils_internal.h"
+#include "utils/memutils_memorychunk.h"
 
 /*--------------------
  * Chunk freelist k holds chunks of size 1 << (k + ALLOC_MINBITS),
@@ -85,7 +90,9 @@
  * CAUTION: ALLOC_MINBITS must be large enough so that
  * 1<<ALLOC_MINBITS is at least MAXALIGN,
  * or we may fail to align the smallest chunks adequately.
- * 8-byte alignment is enough on all currently known machines.
+ * 8-byte alignment is enough on all currently known machines.  This 8-byte
+ * minimum also allows us to store a pointer to the next freelist item within
+ * the chunk of memory itself.
  *
  * With the current parameters, request sizes up to 8K are treated as chunks,
  * larger requests go into dedicated blocks.  Change ALLOCSET_NUM_FREELISTS
@@ -117,10 +124,9 @@
  */
 
 #define ALLOC_BLOCKHDRSZ	MAXALIGN(sizeof(AllocBlockData))
-#define ALLOC_CHUNKHDRSZ	sizeof(struct AllocChunkData)
+#define ALLOC_CHUNKHDRSZ	sizeof(MemoryChunk)
 
 typedef struct AllocBlockData *AllocBlock;	/* forward reference */
-typedef struct AllocChunkData *AllocChunk;
 
 /*
  * AllocPointer
@@ -128,6 +134,34 @@ typedef struct AllocChunkData *AllocChunk;
  */
 typedef void *AllocPointer;
 
+/*
+ * AllocFreeListLink
+ *		When pfreeing memory, if we maintain a freelist for the given chunk's
+ *		size then we use a AllocFreeListLink to point to the current item in
+ *		the AllocSetContext's freelist and then set the given freelist element
+ *		to point to the chunk being freed.
+ */
+typedef struct AllocFreeListLink
+{
+	MemoryChunk *next;
+} AllocFreeListLink;
+
+/*
+ * Obtain a AllocFreeListLink for the given chunk.  Allocation sizes are
+ * always at least sizeof(AllocFreeListLink), so we reuse the pointer's memory
+ * itself to store the freelist link.
+ */
+#define GetFreeListLink(chkptr) \
+	(AllocFreeListLink *) ((char *) (chkptr) + ALLOC_CHUNKHDRSZ)
+
+/* Validate a freelist index retrieved from a chunk header */
+#define FreeListIdxIsValid(fidx) \
+	((fidx) >= 0 && (fidx) < ALLOCSET_NUM_FREELISTS)
+
+/* Determine the size of the chunk based on the freelist index */
+#define GetChunkSizeFromFreeListIdx(fidx) \
+	((((Size) 1) << ALLOC_MINBITS) << (fidx))
+
 /*
  * AllocSetContext is our standard implementation of MemoryContext.
  *
@@ -142,13 +176,12 @@ typedef struct AllocSetContext
 	MemoryContextData header;	/* Standard memory-context fields */
 	/* Info about storage allocated in this context: */
 	AllocBlock	blocks;			/* head of list of blocks in this set */
-	AllocChunk	freelist[ALLOCSET_NUM_FREELISTS];	/* free chunk lists */
+	MemoryChunk *freelist[ALLOCSET_NUM_FREELISTS];	/* free chunk lists */
 	/* Allocation parameters for this context: */
-	Size		initBlockSize;	/* initial block size */
-	Size		maxBlockSize;	/* maximum block size */
-	Size		nextBlockSize;	/* next block size to allocate */
-	Size		allocChunkLimit;	/* effective chunk size limit */
-	AllocBlock	keeper;			/* keep this block over resets */
+	uint32		initBlockSize;	/* initial block size */
+	uint32		maxBlockSize;	/* maximum block size */
+	uint32		nextBlockSize;	/* next block size to allocate */
+	uint32		allocChunkLimit;	/* effective chunk size limit */
 	/* freelist this context could be put in, or -1 if not a candidate: */
 	int			freeListIndex;	/* index in context_freelists[], or -1 */
 } AllocSetContext;
@@ -158,8 +191,8 @@ typedef AllocSetContext *AllocSet;
 /*
  * AllocBlock
  *		An AllocBlock is the unit of memory that is obtained by aset.c
- *		from malloc().  It contains one or more AllocChunks, which are
- *		the units requested by palloc() and freed by pfree().  AllocChunks
+ *		from malloc().  It contains one or more MemoryChunks, which are
+ *		the units requested by palloc() and freed by pfree(). MemoryChunks
  *		cannot be returned to malloc() individually, instead they are put
  *		on freelists by pfree() and re-used by the next palloc() that has
  *		a matching request size.
@@ -176,50 +209,6 @@ typedef struct AllocBlockData
 	char	   *endptr;			/* end of space in this block */
 }			AllocBlockData;
 
-/*
- * AllocChunk
- *		The prefix of each piece of memory in an AllocBlock
- *
- * Note: to meet the memory context APIs, the payload area of the chunk must
- * be maxaligned, and the "aset" link must be immediately adjacent to the
- * payload area (cf. GetMemoryChunkContext).  We simplify matters for this
- * module by requiring sizeof(AllocChunkData) to be maxaligned, and then
- * we can ensure things work by adding any required alignment padding before
- * the "aset" field.  There is a static assertion below that the alignment
- * is done correctly.
- */
-typedef struct AllocChunkData
-{
-	/* size is always the size of the usable space in the chunk */
-	Size		size;
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* when debugging memory usage, also store actual requested size */
-	/* this is zero in a free chunk */
-	Size		requested_size;
-
-#define ALLOCCHUNK_RAWSIZE  (SIZEOF_SIZE_T * 2 + SIZEOF_VOID_P)
-#else
-#define ALLOCCHUNK_RAWSIZE  (SIZEOF_SIZE_T + SIZEOF_VOID_P)
-#endif							/* MEMORY_CONTEXT_CHECKING */
-
-	/* ensure proper alignment by adding padding if needed */
-#if (ALLOCCHUNK_RAWSIZE % MAXIMUM_ALIGNOF) != 0
-	char		padding[MAXIMUM_ALIGNOF - ALLOCCHUNK_RAWSIZE % MAXIMUM_ALIGNOF];
-#endif
-
-	/* aset is the owning aset if allocated, or the freelist link if free */
-	void	   *aset;
-	/* there must not be any padding to reach a MAXALIGN boundary here! */
-}			AllocChunkData;
-
-/*
- * Only the "aset" field should be accessed outside this module.
- * We keep the rest of an allocated chunk's header marked NOACCESS when using
- * valgrind.  But note that chunk headers that are in a freelist are kept
- * accessible, for simplicity.
- */
-#define ALLOCCHUNK_PRIVATE_LEN	offsetof(AllocChunkData, aset)
-
 /*
  * AllocPointerIsValid
  *		True iff pointer is valid allocation pointer.
@@ -230,12 +219,23 @@ typedef struct AllocChunkData
  * AllocSetIsValid
  *		True iff set is valid allocation set.
  */
-#define AllocSetIsValid(set) PointerIsValid(set)
+#define AllocSetIsValid(set) \
+	(PointerIsValid(set) && IsA(set, AllocSetContext))
 
-#define AllocPointerGetChunk(ptr)	\
-					((AllocChunk)(((char *)(ptr)) - ALLOC_CHUNKHDRSZ))
-#define AllocChunkGetPointer(chk)	\
-					((AllocPointer)(((char *)(chk)) + ALLOC_CHUNKHDRSZ))
+/*
+ * AllocBlockIsValid
+ *		True iff block is valid block of allocation set.
+ */
+#define AllocBlockIsValid(block) \
+	(PointerIsValid(block) && AllocSetIsValid((block)->aset))
+
+/*
+ * We always store external chunks on a dedicated block.  This makes fetching
+ * the block from an external chunk easy since it's always the first and only
+ * chunk on the block.
+ */
+#define ExternalChunkGetBlock(chunk) \
+	(AllocBlock) ((char *) chunk - ALLOC_BLOCKHDRSZ)
 
 /*
  * Rather than repeatedly creating and deleting memory contexts, we keep some
@@ -262,6 +262,13 @@ typedef struct AllocChunkData
  */
 #define MAX_FREE_CONTEXTS 100	/* arbitrary limit on freelist length */
 
+/* Obtain the keeper block for an allocation set */
+#define KeeperBlock(set) \
+	((AllocBlock) (((char *) set) + MAXALIGN(sizeof(AllocSetContext))))
+
+/* Check if the block is the keeper block of the given allocation set */
+#define IsKeeperBlock(set, block) ((block) == (KeeperBlock(set)))
+
 typedef struct AllocSetFreeList
 {
 	int			num_free;		/* current list length */
@@ -280,42 +287,6 @@ static __thread AllocSetFreeList context_freelists[2] =
 };
 
 
-/*
- * These functions implement the MemoryContext API for AllocSet contexts.
- */
-static void *AllocSetAlloc(MemoryContext context, Size size);
-static void AllocSetFree(MemoryContext context, void *pointer);
-static void *AllocSetRealloc(MemoryContext context, void *pointer, Size size);
-static void AllocSetReset(MemoryContext context);
-static void AllocSetDelete(MemoryContext context);
-static Size AllocSetGetChunkSpace(MemoryContext context, void *pointer);
-static bool AllocSetIsEmpty(MemoryContext context);
-static void AllocSetStats(MemoryContext context,
-						  MemoryStatsPrintFunc printfunc, void *passthru,
-						  MemoryContextCounters *totals,
-						  bool print_to_stderr);
-
-#ifdef MEMORY_CONTEXT_CHECKING
-static void AllocSetCheck(MemoryContext context);
-#endif
-
-/*
- * This is the virtual function table for AllocSet contexts.
- */
-static const MemoryContextMethods AllocSetMethods = {
-	AllocSetAlloc,
-	AllocSetFree,
-	AllocSetRealloc,
-	AllocSetReset,
-	AllocSetDelete,
-	AllocSetGetChunkSpace,
-	AllocSetIsEmpty,
-	AllocSetStats
-#ifdef MEMORY_CONTEXT_CHECKING
-	,AllocSetCheck
-#endif
-};
-
 
 /* ----------
  * AllocSetFreeIndex -
@@ -339,7 +310,7 @@ AllocSetFreeIndex(Size size)
 		 * or equivalently
 		 *		pg_leftmost_one_pos32(size - 1) - ALLOC_MINBITS + 1
 		 *
-		 * However, rather than just calling that function, we duplicate the
+		 * However, for platforms without intrinsic support, we duplicate the
 		 * logic here, allowing an additional optimization.  It's reasonable
 		 * to assume that ALLOC_CHUNK_LIMIT fits in 16 bits, so we can unroll
 		 * the byte-at-a-time loop in pg_leftmost_one_pos32 and just handle
@@ -349,14 +320,14 @@ AllocSetFreeIndex(Size size)
 		 * much trouble.
 		 *----------
 		 */
-#ifdef HAVE__BUILTIN_CLZ
-		idx = 31 - __builtin_clz((uint32) size - 1) - ALLOC_MINBITS + 1;
+#ifdef HAVE_BITSCAN_REVERSE
+		idx = pg_leftmost_one_pos32((uint32) size - 1) - ALLOC_MINBITS + 1;
 #else
 		uint32		t,
 					tsize;
 
 		/* Statically assert that we only have a 16-bit input value. */
-		StaticAssertStmt(ALLOC_CHUNK_LIMIT < (1 << 16),
+		StaticAssertDecl(ALLOC_CHUNK_LIMIT < (1 << 16),
 						 "ALLOC_CHUNK_LIMIT must be less than 64kB");
 
 		tsize = size - 1;
@@ -407,18 +378,24 @@ AllocSetContextCreateInternal(MemoryContext parent,
 	AllocSet	set;
 	AllocBlock	block;
 
-	/* Assert we padded AllocChunkData properly */
-	StaticAssertStmt(ALLOC_CHUNKHDRSZ == MAXALIGN(ALLOC_CHUNKHDRSZ),
-					 "sizeof(AllocChunkData) is not maxaligned");
-	StaticAssertStmt(offsetof(AllocChunkData, aset) + sizeof(MemoryContext) ==
-					 ALLOC_CHUNKHDRSZ,
-					 "padding calculation in AllocChunkData is wrong");
+	/* ensure MemoryChunk's size is properly maxaligned */
+	StaticAssertDecl(ALLOC_CHUNKHDRSZ == MAXALIGN(ALLOC_CHUNKHDRSZ),
+					 "sizeof(MemoryChunk) is not maxaligned");
+	/* check we have enough space to store the freelist link */
+	StaticAssertDecl(sizeof(AllocFreeListLink) <= (1 << ALLOC_MINBITS),
+					 "sizeof(AllocFreeListLink) larger than minimum allocation size");
 
 	/*
 	 * First, validate allocation parameters.  Once these were regular runtime
-	 * test and elog's, but in practice Asserts seem sufficient because nobody
-	 * varies their parameters at runtime.  We somewhat arbitrarily enforce a
-	 * minimum 1K block size.
+	 * tests and elog's, but in practice Asserts seem sufficient because
+	 * nobody varies their parameters at runtime.  We somewhat arbitrarily
+	 * enforce a minimum 1K block size.  We restrict the maximum block size to
+	 * MEMORYCHUNK_MAX_BLOCKOFFSET as MemoryChunks are limited to this in
+	 * regards to addressing the offset between the chunk and the block that
+	 * the chunk is stored on.  We would be unable to store the offset between
+	 * the chunk and block for any chunks that were beyond
+	 * MEMORYCHUNK_MAX_BLOCKOFFSET bytes into the block if the block was to be
+	 * larger than this.
 	 */
 	Assert(initBlockSize == MAXALIGN(initBlockSize) &&
 		   initBlockSize >= 1024);
@@ -429,6 +406,7 @@ AllocSetContextCreateInternal(MemoryContext parent,
 		   (minContextSize == MAXALIGN(minContextSize) &&
 			minContextSize >= 1024 &&
 			minContextSize <= maxBlockSize));
+	Assert(maxBlockSize <= MEMORYCHUNK_MAX_BLOCKOFFSET);
 
 	/*
 	 * Check whether the parameters match either available freelist.  We do
@@ -463,12 +441,12 @@ AllocSetContextCreateInternal(MemoryContext parent,
 			/* Reinitialize its header, installing correct name and parent */
 			MemoryContextCreate((MemoryContext) set,
 								T_AllocSetContext,
-								&AllocSetMethods,
+								MCTX_ASET_ID,
 								parent,
 								name);
 
 			((MemoryContext) set)->mem_allocated =
-				set->keeper->endptr - ((char *) set);
+				KeeperBlock(set)->endptr - ((char *) set);
 
 			return (MemoryContext) set;
 		}
@@ -504,7 +482,7 @@ AllocSetContextCreateInternal(MemoryContext parent,
 	 */
 
 	/* Fill in the initial block's block header */
-	block = (AllocBlock) (((char *) set) + MAXALIGN(sizeof(AllocSetContext)));
+	block = KeeperBlock(set);
 	block->aset = set;
 	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
 	block->endptr = ((char *) set) + firstBlockSize;
@@ -516,15 +494,13 @@ AllocSetContextCreateInternal(MemoryContext parent,
 
 	/* Remember block as part of block list */
 	set->blocks = block;
-	/* Mark block as not to be released at reset time */
-	set->keeper = block;
 
 	/* Finish filling in aset-specific parts of the context header */
 	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
 
-	set->initBlockSize = initBlockSize;
-	set->maxBlockSize = maxBlockSize;
-	set->nextBlockSize = initBlockSize;
+	set->initBlockSize = (uint32) initBlockSize;
+	set->maxBlockSize = (uint32) maxBlockSize;
+	set->nextBlockSize = (uint32) initBlockSize;
 	set->freeListIndex = freeListIndex;
 
 	/*
@@ -537,15 +513,20 @@ AllocSetContextCreateInternal(MemoryContext parent,
 	 * requests that are all the maximum chunk size we will waste at most
 	 * 1/8th of the allocated space.
 	 *
-	 * We have to have allocChunkLimit a power of two, because the requested
-	 * and actually-allocated sizes of any chunk must be on the same side of
-	 * the limit, else we get confused about whether the chunk is "big".
-	 *
 	 * Also, allocChunkLimit must not exceed ALLOCSET_SEPARATE_THRESHOLD.
 	 */
 	StaticAssertStmt(ALLOC_CHUNK_LIMIT == ALLOCSET_SEPARATE_THRESHOLD,
 					 "ALLOC_CHUNK_LIMIT != ALLOCSET_SEPARATE_THRESHOLD");
 
+	/*
+	 * Determine the maximum size that a chunk can be before we allocate an
+	 * entire AllocBlock dedicated for that chunk.  We set the absolute limit
+	 * of that size as ALLOC_CHUNK_LIMIT but we reduce it further so that we
+	 * can fit about ALLOC_CHUNK_FRACTION chunks this size on a maximally
+	 * sized block.  (We opt to keep allocChunkLimit a power-of-2 value
+	 * primarily for legacy reasons rather than calculating it so that exactly
+	 * ALLOC_CHUNK_FRACTION chunks fit on a maximally sized block.)
+	 */
 	set->allocChunkLimit = ALLOC_CHUNK_LIMIT;
 	while ((Size) (set->allocChunkLimit + ALLOC_CHUNKHDRSZ) >
 		   (Size) ((maxBlockSize - ALLOC_BLOCKHDRSZ) / ALLOC_CHUNK_FRACTION))
@@ -554,7 +535,7 @@ AllocSetContextCreateInternal(MemoryContext parent,
 	/* Finally, do the type-independent part of context creation */
 	MemoryContextCreate((MemoryContext) set,
 						T_AllocSetContext,
-						&AllocSetMethods,
+						MCTX_ASET_ID,
 						parent,
 						name);
 
@@ -575,34 +556,36 @@ AllocSetContextCreateInternal(MemoryContext parent,
  * thrash malloc() when a context is repeatedly reset after small allocations,
  * which is typical behavior for per-tuple contexts.
  */
-static void
+void
 AllocSetReset(MemoryContext context)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block;
-	Size		keepersize PG_USED_FOR_ASSERTS_ONLY
-	= set->keeper->endptr - ((char *) set);
+	Size		keepersize PG_USED_FOR_ASSERTS_ONLY;
 
-	AssertArg(AllocSetIsValid(set));
+	Assert(AllocSetIsValid(set));
 
 #ifdef MEMORY_CONTEXT_CHECKING
 	/* Check for corruption and leaks before freeing */
 	AllocSetCheck(context);
 #endif
 
+	/* Remember keeper block size for Assert below */
+	keepersize = KeeperBlock(set)->endptr - ((char *) set);
+
 	/* Clear chunk freelists */
 	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
 
 	block = set->blocks;
 
 	/* New blocks list will be just the keeper block */
-	set->blocks = set->keeper;
+	set->blocks = KeeperBlock(set);
 
 	while (block != NULL)
 	{
 		AllocBlock	next = block->next;
 
-		if (block == set->keeper)
+		if (IsKeeperBlock(set, block))
 		{
 			/* Reset the block, but don't return it to malloc */
 			char	   *datastart = ((char *) block) + ALLOC_BLOCKHDRSZ;
@@ -643,21 +626,23 @@ AllocSetReset(MemoryContext context)
  *
  * Unlike AllocSetReset, this *must* free all resources of the set.
  */
-static void
+void
 AllocSetDelete(MemoryContext context)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block = set->blocks;
-	Size		keepersize PG_USED_FOR_ASSERTS_ONLY
-	= set->keeper->endptr - ((char *) set);
+	Size		keepersize PG_USED_FOR_ASSERTS_ONLY;
 
-	AssertArg(AllocSetIsValid(set));
+	Assert(AllocSetIsValid(set));
 
 #ifdef MEMORY_CONTEXT_CHECKING
 	/* Check for corruption and leaks before freeing */
 	AllocSetCheck(context);
 #endif
 
+	/* Remember keeper block size for Assert below */
+	keepersize = KeeperBlock(set)->endptr - ((char *) set);
+
 	/*
 	 * If the context is a candidate for a freelist, put it into that freelist
 	 * instead of destroying it.
@@ -705,14 +690,14 @@ AllocSetDelete(MemoryContext context)
 	{
 		AllocBlock	next = block->next;
 
-		if (block != set->keeper)
+		if (!IsKeeperBlock(set, block))
 			context->mem_allocated -= block->endptr - ((char *) block);
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
 #endif
 
-		if (block != set->keeper)
+		if (!IsKeeperBlock(set, block))
 			free(block);
 
 		block = next;
@@ -725,326 +710,410 @@ AllocSetDelete(MemoryContext context)
 }
 
 /*
- * AllocSetAlloc
- *		Returns pointer to allocated memory of given size or NULL if
- *		request could not be completed; memory is added to the set.
+ * Helper for AllocSetAlloc() that allocates an entire block for the chunk.
  *
- * No request may exceed:
- *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
- * All callers use a much-lower limit.
- *
- * Note: when using valgrind, it doesn't matter how the returned allocation
- * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
- * return space that is marked NOACCESS - AllocSetRealloc has to beware!
+ * AllocSetAlloc()'s comment explains why this is separate.
  */
+pg_noinline
 static void *
-AllocSetAlloc(MemoryContext context, Size size)
+AllocSetAllocLarge(MemoryContext context, Size size, int flags)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block;
-	AllocChunk	chunk;
-	int			fidx;
+	MemoryChunk *chunk;
 	Size		chunk_size;
 	Size		blksize;
 
-	AssertArg(AllocSetIsValid(set));
+	/* validate 'size' is within the limits for the given 'flags' */
+	MemoryContextCheckSize(context, size, flags);
 
-	/*
-	 * If requested size exceeds maximum for chunks, allocate an entire block
-	 * for this request.
-	 */
-	if (size > set->allocChunkLimit)
-	{
-		chunk_size = MAXALIGN(size);
-		blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-		block = (AllocBlock) malloc(blksize);
-		if (block == NULL)
-			return NULL;
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* ensure there's always space for the sentinel byte */
+	chunk_size = MAXALIGN(size + 1);
+#else
+	chunk_size = MAXALIGN(size);
+#endif
 
-		context->mem_allocated += blksize;
+	blksize = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	block = (AllocBlock) malloc(blksize);
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
 
-		block->aset = set;
-		block->freeptr = block->endptr = ((char *) block) + blksize;
+	context->mem_allocated += blksize;
+
+	block->aset = set;
+	block->freeptr = block->endptr = ((char *) block) + blksize;
+
+	chunk = (MemoryChunk *) (((char *) block) + ALLOC_BLOCKHDRSZ);
+
+	/* mark the MemoryChunk as externally managed */
+	MemoryChunkSetHdrMaskExternal(chunk, MCTX_ASET_ID);
 
-		chunk = (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ);
-		chunk->aset = set;
-		chunk->size = chunk_size;
 #ifdef MEMORY_CONTEXT_CHECKING
-		chunk->requested_size = size;
-		/* set mark to catch clobber of "unused" space */
-		if (size < chunk_size)
-			set_sentinel(AllocChunkGetPointer(chunk), size);
+	chunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	Assert(size < chunk_size);
+	set_sentinel(MemoryChunkGetPointer(chunk), size);
 #endif
 #ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* fill the allocated space with junk */
-		randomize_mem((char *) AllocChunkGetPointer(chunk), size);
+	/* fill the allocated space with junk */
+	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
 #endif
 
-		/*
-		 * Stick the new block underneath the active allocation block, if any,
-		 * so that we don't lose the use of the space remaining therein.
-		 */
-		if (set->blocks != NULL)
-		{
-			block->prev = set->blocks;
-			block->next = set->blocks->next;
-			if (block->next)
-				block->next->prev = block;
-			set->blocks->next = block;
-		}
-		else
-		{
-			block->prev = NULL;
-			block->next = NULL;
-			set->blocks = block;
-		}
+	/*
+	 * Stick the new block underneath the active allocation block, if any, so
+	 * that we don't lose the use of the space remaining therein.
+	 */
+	if (set->blocks != NULL)
+	{
+		block->prev = set->blocks;
+		block->next = set->blocks->next;
+		if (block->next)
+			block->next->prev = block;
+		set->blocks->next = block;
+	}
+	else
+	{
+		block->prev = NULL;
+		block->next = NULL;
+		set->blocks = block;
+	}
 
-		/* Ensure any padding bytes are marked NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
-								   chunk_size - size);
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+							   chunk_size - size);
 
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
-		return AllocChunkGetPointer(chunk);
-	}
+	return MemoryChunkGetPointer(chunk);
+}
 
-	/*
-	 * Request is small enough to be treated as a chunk.  Look in the
-	 * corresponding free list to see if there is a free chunk we could reuse.
-	 * If one is found, remove it from the free list, make it again a member
-	 * of the alloc set and return its data address.
-	 */
-	fidx = AllocSetFreeIndex(size);
-	chunk = set->freelist[fidx];
-	if (chunk != NULL)
-	{
-		Assert(chunk->size >= size);
+/*
+ * Small helper for allocating a new chunk from a chunk, to avoid duplicating
+ * the code between AllocSetAlloc() and AllocSetAllocFromNewBlock().
+ */
+static inline void *
+AllocSetAllocChunkFromBlock(MemoryContext context, AllocBlock block,
+							Size size, Size chunk_size, int fidx)
+{
+	MemoryChunk *chunk;
+
+	chunk = (MemoryChunk *) (block->freeptr);
+
+	/* Prepare to initialize the chunk header. */
+	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
 
-		set->freelist[fidx] = (AllocChunk) chunk->aset;
+	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
+	Assert(block->freeptr <= block->endptr);
 
-		chunk->aset = (void *) set;
+	/* store the free list index in the value field */
+	MemoryChunkSetHdrMask(chunk, block, fidx, MCTX_ASET_ID);
 
 #ifdef MEMORY_CONTEXT_CHECKING
-		chunk->requested_size = size;
-		/* set mark to catch clobber of "unused" space */
-		if (size < chunk->size)
-			set_sentinel(AllocChunkGetPointer(chunk), size);
+	chunk->requested_size = size;
+	/* set mark to catch clobber of "unused" space */
+	if (size < chunk_size)
+		set_sentinel(MemoryChunkGetPointer(chunk), size);
 #endif
 #ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* fill the allocated space with junk */
-		randomize_mem((char *) AllocChunkGetPointer(chunk), size);
+	/* fill the allocated space with junk */
+	randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
 #endif
 
-		/* Ensure any padding bytes are marked NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
-								   chunk->size - size);
+	/* Ensure any padding bytes are marked NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+							   chunk_size - size);
 
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
-		return AllocChunkGetPointer(chunk);
-	}
+	return MemoryChunkGetPointer(chunk);
+}
 
-	/*
-	 * Choose the actual chunk size to allocate.
-	 */
-	chunk_size = (1 << ALLOC_MINBITS) << fidx;
-	Assert(chunk_size >= size);
+/*
+ * Helper for AllocSetAlloc() that allocates a new block and returns a chunk
+ * allocated from it.
+ *
+ * AllocSetAlloc()'s comment explains why this is separate.
+ */
+pg_noinline
+static void *
+AllocSetAllocFromNewBlock(MemoryContext context, Size size, int flags,
+						  int fidx)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+	Size		availspace;
+	Size		blksize;
+	Size		required_size;
+	Size		chunk_size;
+
+	/* due to the keeper block set->blocks should always be valid */
+	Assert(set->blocks != NULL);
+	block = set->blocks;
+	availspace = block->endptr - block->freeptr;
 
 	/*
-	 * If there is enough room in the active allocation block, we will put the
-	 * chunk into that block.  Else must start a new one.
+	 * The existing active (top) block does not have enough room for the
+	 * requested allocation, but it might still have a useful amount of space
+	 * in it.  Once we push it down in the block list, we'll never try to
+	 * allocate more space from it. So, before we do that, carve up its free
+	 * space into chunks that we can put on the set's freelists.
+	 *
+	 * Because we can only get here when there's less than ALLOC_CHUNK_LIMIT
+	 * left in the block, this loop cannot iterate more than
+	 * ALLOCSET_NUM_FREELISTS-1 times.
 	 */
-	if ((block = set->blocks) != NULL)
+	while (availspace >= ((1 << ALLOC_MINBITS) + ALLOC_CHUNKHDRSZ))
 	{
-		Size		availspace = block->endptr - block->freeptr;
+		AllocFreeListLink *link;
+		MemoryChunk *chunk;
+		Size		availchunk = availspace - ALLOC_CHUNKHDRSZ;
+		int			a_fidx = AllocSetFreeIndex(availchunk);
 
-		if (availspace < (chunk_size + ALLOC_CHUNKHDRSZ))
+		/*
+		 * In most cases, we'll get back the index of the next larger freelist
+		 * than the one we need to put this chunk on.  The exception is when
+		 * availchunk is exactly a power of 2.
+		 */
+		if (availchunk != GetChunkSizeFromFreeListIdx(a_fidx))
 		{
-			/*
-			 * The existing active (top) block does not have enough room for
-			 * the requested allocation, but it might still have a useful
-			 * amount of space in it.  Once we push it down in the block list,
-			 * we'll never try to allocate more space from it. So, before we
-			 * do that, carve up its free space into chunks that we can put on
-			 * the set's freelists.
-			 *
-			 * Because we can only get here when there's less than
-			 * ALLOC_CHUNK_LIMIT left in the block, this loop cannot iterate
-			 * more than ALLOCSET_NUM_FREELISTS-1 times.
-			 */
-			while (availspace >= ((1 << ALLOC_MINBITS) + ALLOC_CHUNKHDRSZ))
-			{
-				Size		availchunk = availspace - ALLOC_CHUNKHDRSZ;
-				int			a_fidx = AllocSetFreeIndex(availchunk);
-
-				/*
-				 * In most cases, we'll get back the index of the next larger
-				 * freelist than the one we need to put this chunk on.  The
-				 * exception is when availchunk is exactly a power of 2.
-				 */
-				if (availchunk != ((Size) 1 << (a_fidx + ALLOC_MINBITS)))
-				{
-					a_fidx--;
-					Assert(a_fidx >= 0);
-					availchunk = ((Size) 1 << (a_fidx + ALLOC_MINBITS));
-				}
-
-				chunk = (AllocChunk) (block->freeptr);
+			a_fidx--;
+			Assert(a_fidx >= 0);
+			availchunk = GetChunkSizeFromFreeListIdx(a_fidx);
+		}
 
-				/* Prepare to initialize the chunk header. */
-				VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+		chunk = (MemoryChunk *) (block->freeptr);
 
-				block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
-				availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
+		/* Prepare to initialize the chunk header. */
+		VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+		block->freeptr += (availchunk + ALLOC_CHUNKHDRSZ);
+		availspace -= (availchunk + ALLOC_CHUNKHDRSZ);
 
-				chunk->size = availchunk;
+		/* store the freelist index in the value field */
+		MemoryChunkSetHdrMask(chunk, block, a_fidx, MCTX_ASET_ID);
 #ifdef MEMORY_CONTEXT_CHECKING
-				chunk->requested_size = 0;	/* mark it free */
+		chunk->requested_size = InvalidAllocSize;	/* mark it free */
 #endif
-				chunk->aset = (void *) set->freelist[a_fidx];
-				set->freelist[a_fidx] = chunk;
-			}
+		/* push this chunk onto the free list */
+		link = GetFreeListLink(chunk);
 
-			/* Mark that we need to create a new block */
-			block = NULL;
-		}
+		VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
+		link->next = set->freelist[a_fidx];
+		VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
+
+		set->freelist[a_fidx] = chunk;
 	}
 
 	/*
-	 * Time to create a new regular (multi-chunk) block?
+	 * The first such block has size initBlockSize, and we double the space in
+	 * each succeeding block, but not more than maxBlockSize.
 	 */
-	if (block == NULL)
-	{
-		Size		required_size;
+	blksize = set->nextBlockSize;
+	set->nextBlockSize <<= 1;
+	if (set->nextBlockSize > set->maxBlockSize)
+		set->nextBlockSize = set->maxBlockSize;
 
-		/*
-		 * The first such block has size initBlockSize, and we double the
-		 * space in each succeeding block, but not more than maxBlockSize.
-		 */
-		blksize = set->nextBlockSize;
-		set->nextBlockSize <<= 1;
-		if (set->nextBlockSize > set->maxBlockSize)
-			set->nextBlockSize = set->maxBlockSize;
+	/* Choose the actual chunk size to allocate */
+	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
+	Assert(chunk_size >= size);
 
-		/*
-		 * If initBlockSize is less than ALLOC_CHUNK_LIMIT, we could need more
-		 * space... but try to keep it a power of 2.
-		 */
-		required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
-		while (blksize < required_size)
-			blksize <<= 1;
+	/*
+	 * If initBlockSize is less than ALLOC_CHUNK_LIMIT, we could need more
+	 * space... but try to keep it a power of 2.
+	 */
+	required_size = chunk_size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
+	while (blksize < required_size)
+		blksize <<= 1;
+
+	/* Try to allocate it */
+	block = (AllocBlock) malloc(blksize);
 
-		/* Try to allocate it */
+	/*
+	 * We could be asking for pretty big blocks here, so cope if malloc fails.
+	 * But give up if there's less than 1 MB or so available...
+	 */
+	while (block == NULL && blksize > 1024 * 1024)
+	{
+		blksize >>= 1;
+		if (blksize < required_size)
+			break;
 		block = (AllocBlock) malloc(blksize);
+	}
 
-		/*
-		 * We could be asking for pretty big blocks here, so cope if malloc
-		 * fails.  But give up if there's less than 1 MB or so available...
-		 */
-		while (block == NULL && blksize > 1024 * 1024)
-		{
-			blksize >>= 1;
-			if (blksize < required_size)
-				break;
-			block = (AllocBlock) malloc(blksize);
-		}
+	if (block == NULL)
+		return MemoryContextAllocationFailure(context, size, flags);
 
-		if (block == NULL)
-			return NULL;
+	context->mem_allocated += blksize;
 
-		context->mem_allocated += blksize;
+	block->aset = set;
+	block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
+	block->endptr = ((char *) block) + blksize;
 
-		block->aset = set;
-		block->freeptr = ((char *) block) + ALLOC_BLOCKHDRSZ;
-		block->endptr = ((char *) block) + blksize;
+	/* Mark unallocated space NOACCESS. */
+	VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
+							   blksize - ALLOC_BLOCKHDRSZ);
 
-		/* Mark unallocated space NOACCESS. */
-		VALGRIND_MAKE_MEM_NOACCESS(block->freeptr,
-								   blksize - ALLOC_BLOCKHDRSZ);
+	block->prev = NULL;
+	block->next = set->blocks;
+	if (block->next)
+		block->next->prev = block;
+	set->blocks = block;
 
-		block->prev = NULL;
-		block->next = set->blocks;
-		if (block->next)
-			block->next->prev = block;
-		set->blocks = block;
-	}
+	return AllocSetAllocChunkFromBlock(context, block, size, chunk_size, fidx);
+}
+
+/*
+ * AllocSetAlloc
+ *		Returns a pointer to allocated memory of given size or raises an ERROR
+ *		on allocation failure, or returns NULL when flags contains
+ *		MCXT_ALLOC_NO_OOM.
+ *
+ * No request may exceed:
+ *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
+ * All callers use a much-lower limit.
+ *
+ * Note: when using valgrind, it doesn't matter how the returned allocation
+ * is marked, as mcxt.c will set it to UNDEFINED.  In some paths we will
+ * return space that is marked NOACCESS - AllocSetRealloc has to beware!
+ *
+ * This function should only contain the most common code paths.  Everything
+ * else should be in pg_noinline helper functions, thus avoiding the overhead
+ * of creating a stack frame for the common cases.  Allocating memory is often
+ * a bottleneck in many workloads, so avoiding stack frame setup is
+ * worthwhile.  Helper functions should always directly return the newly
+ * allocated memory so that we can just return that address directly as a tail
+ * call.
+ */
+void *
+AllocSetAlloc(MemoryContext context, Size size, int flags)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+	MemoryChunk *chunk;
+	int			fidx;
+	Size		chunk_size;
+	Size		availspace;
+
+	Assert(AllocSetIsValid(set));
+
+	/* due to the keeper block set->blocks should never be NULL */
+	Assert(set->blocks != NULL);
 
 	/*
-	 * OK, do the allocation
+	 * If requested size exceeds maximum for chunks we hand the request off to
+	 * AllocSetAllocLarge().
 	 */
-	chunk = (AllocChunk) (block->freeptr);
+	if (size > set->allocChunkLimit)
+		return AllocSetAllocLarge(context, size, flags);
 
-	/* Prepare to initialize the chunk header. */
-	VALGRIND_MAKE_MEM_UNDEFINED(chunk, ALLOC_CHUNKHDRSZ);
+	/*
+	 * Request is small enough to be treated as a chunk.  Look in the
+	 * corresponding free list to see if there is a free chunk we could reuse.
+	 * If one is found, remove it from the free list, make it again a member
+	 * of the alloc set and return its data address.
+	 *
+	 * Note that we don't attempt to ensure there's space for the sentinel
+	 * byte here.  We expect a large proportion of allocations to be for sizes
+	 * which are already a power of 2.  If we were to always make space for a
+	 * sentinel byte in MEMORY_CONTEXT_CHECKING builds, then we'd end up
+	 * doubling the memory requirements for such allocations.
+	 */
+	fidx = AllocSetFreeIndex(size);
+	chunk = set->freelist[fidx];
+	if (chunk != NULL)
+	{
+		AllocFreeListLink *link = GetFreeListLink(chunk);
 
-	block->freeptr += (chunk_size + ALLOC_CHUNKHDRSZ);
-	Assert(block->freeptr <= block->endptr);
+		/* Allow access to the chunk header. */
+		VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+		Assert(fidx == MemoryChunkGetValue(chunk));
+
+		/* pop this chunk off the freelist */
+		VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
+		set->freelist[fidx] = link->next;
+		VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
 
-	chunk->aset = (void *) set;
-	chunk->size = chunk_size;
 #ifdef MEMORY_CONTEXT_CHECKING
-	chunk->requested_size = size;
-	/* set mark to catch clobber of "unused" space */
-	if (size < chunk->size)
-		set_sentinel(AllocChunkGetPointer(chunk), size);
+		chunk->requested_size = size;
+		/* set mark to catch clobber of "unused" space */
+		if (size < GetChunkSizeFromFreeListIdx(fidx))
+			set_sentinel(MemoryChunkGetPointer(chunk), size);
 #endif
 #ifdef RANDOMIZE_ALLOCATED_MEMORY
-	/* fill the allocated space with junk */
-	randomize_mem((char *) AllocChunkGetPointer(chunk), size);
+		/* fill the allocated space with junk */
+		randomize_mem((char *) MemoryChunkGetPointer(chunk), size);
 #endif
 
-	/* Ensure any padding bytes are marked NOACCESS. */
-	VALGRIND_MAKE_MEM_NOACCESS((char *) AllocChunkGetPointer(chunk) + size,
-							   chunk_size - size);
+		/* Ensure any padding bytes are marked NOACCESS. */
+		VALGRIND_MAKE_MEM_NOACCESS((char *) MemoryChunkGetPointer(chunk) + size,
+								   GetChunkSizeFromFreeListIdx(fidx) - size);
 
-	/* Disallow external access to private part of chunk header. */
-	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+		/* Disallow access to the chunk header. */
+		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
-	return AllocChunkGetPointer(chunk);
+		return MemoryChunkGetPointer(chunk);
+	}
+
+	/*
+	 * Choose the actual chunk size to allocate.
+	 */
+	chunk_size = GetChunkSizeFromFreeListIdx(fidx);
+	Assert(chunk_size >= size);
+
+	block = set->blocks;
+	availspace = block->endptr - block->freeptr;
+
+	/*
+	 * If there is enough room in the active allocation block, we will put the
+	 * chunk into that block.  Else must start a new one.
+	 */
+	if (unlikely(availspace < (chunk_size + ALLOC_CHUNKHDRSZ)))
+		return AllocSetAllocFromNewBlock(context, size, flags, fidx);
+
+	/* There's enough space on the current block, so allocate from that */
+	return AllocSetAllocChunkFromBlock(context, block, size, chunk_size, fidx);
 }
 
 /*
  * AllocSetFree
  *		Frees allocated memory; memory is removed from the set.
  */
-static void
-AllocSetFree(MemoryContext context, void *pointer)
+void
+AllocSetFree(void *pointer)
 {
-	AllocSet	set = (AllocSet) context;
-	AllocChunk	chunk = AllocPointerGetChunk(pointer);
-
-	/* Allow access to private part of chunk header. */
-	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOCCHUNK_PRIVATE_LEN);
+	AllocSet	set;
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
 
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* Test for someone scribbling on unused space in chunk */
-	if (chunk->requested_size < chunk->size)
-		if (!sentinel_ok(pointer, chunk->requested_size))
-			elog(WARNING, "detected write past chunk end in %s %p",
-				 set->header.name, chunk);
-#endif
+	/* Allow access to the chunk header. */
+	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
 
-	if (chunk->size > set->allocChunkLimit)
+	if (MemoryChunkIsExternal(chunk))
 	{
-		/*
-		 * Big chunks are certain to have been allocated as single-chunk
-		 * blocks.  Just unlink that block and return it to malloc().
-		 */
-		AllocBlock	block = (AllocBlock) (((char *) chunk) - ALLOC_BLOCKHDRSZ);
+		/* Release single-chunk block. */
+		AllocBlock	block = ExternalChunkGetBlock(chunk);
 
 		/*
-		 * Try to verify that we have a sane block pointer: it should
-		 * reference the correct aset, and freeptr and endptr should point
-		 * just past the chunk.
+		 * Try to verify that we have a sane block pointer: the block header
+		 * should reference an aset and the freeptr should match the endptr.
 		 */
-		if (block->aset != set ||
-			block->freeptr != block->endptr ||
-			block->freeptr != ((char *) block) +
-			(chunk->size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ))
+		if (!AllocBlockIsValid(block) || block->freeptr != block->endptr)
 			elog(ERROR, "could not find block containing chunk %p", chunk);
 
+		set = block->aset;
+
+#ifdef MEMORY_CONTEXT_CHECKING
+		{
+			/* Test for someone scribbling on unused space in chunk */
+			Assert(chunk->requested_size < (block->endptr - (char *) pointer));
+			if (!sentinel_ok(pointer, chunk->requested_size))
+				elog(WARNING, "detected write past chunk end in %s %p",
+					 set->header.name, chunk);
+		}
+#endif
+
 		/* OK, remove block from aset's list and free it */
 		if (block->prev)
 			block->prev->next = block->next;
@@ -1053,7 +1122,7 @@ AllocSetFree(MemoryContext context, void *pointer)
 		if (block->next)
 			block->next->prev = block->prev;
 
-		context->mem_allocated -= block->endptr - ((char *) block);
+		set->header.mem_allocated -= block->endptr - ((char *) block);
 
 #ifdef CLOBBER_FREED_MEMORY
 		wipe_mem(block, block->freeptr - ((char *) block));
@@ -1062,20 +1131,48 @@ AllocSetFree(MemoryContext context, void *pointer)
 	}
 	else
 	{
-		/* Normal case, put the chunk into appropriate freelist */
-		int			fidx = AllocSetFreeIndex(chunk->size);
+		AllocBlock	block = MemoryChunkGetBlock(chunk);
+		int			fidx;
+		AllocFreeListLink *link;
+
+		/*
+		 * In this path, for speed reasons we just Assert that the referenced
+		 * block is good.  We can also Assert that the value field is sane.
+		 * Future field experience may show that these Asserts had better
+		 * become regular runtime test-and-elog checks.
+		 */
+		Assert(AllocBlockIsValid(block));
+		set = block->aset;
+
+		fidx = MemoryChunkGetValue(chunk);
+		Assert(FreeListIdxIsValid(fidx));
+		link = GetFreeListLink(chunk);
 
-		chunk->aset = (void *) set->freelist[fidx];
+#ifdef MEMORY_CONTEXT_CHECKING
+		/* Test for someone scribbling on unused space in chunk */
+		if (chunk->requested_size < GetChunkSizeFromFreeListIdx(fidx))
+			if (!sentinel_ok(pointer, chunk->requested_size))
+				elog(WARNING, "detected write past chunk end in %s %p",
+					 set->header.name, chunk);
+#endif
 
 #ifdef CLOBBER_FREED_MEMORY
-		wipe_mem(pointer, chunk->size);
+		wipe_mem(pointer, GetChunkSizeFromFreeListIdx(fidx));
 #endif
+		/* push this chunk onto the top of the free list */
+		VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
+		link->next = set->freelist[fidx];
+		VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
+		set->freelist[fidx] = chunk;
 
 #ifdef MEMORY_CONTEXT_CHECKING
-		/* Reset requested_size to 0 in chunks that are on freelist */
-		chunk->requested_size = 0;
+
+		/*
+		 * Reset requested_size to InvalidAllocSize in chunks that are on free
+		 * list.
+		 */
+		chunk->requested_size = InvalidAllocSize;
 #endif
-		set->freelist[fidx] = chunk;
 	}
 }
 
@@ -1091,57 +1188,59 @@ AllocSetFree(MemoryContext context, void *pointer)
  * (In principle, we could use VALGRIND_GET_VBITS() to rediscover the old
  * request size.)
  */
-static void *
-AllocSetRealloc(MemoryContext context, void *pointer, Size size)
+void *
+AllocSetRealloc(void *pointer, Size size, int flags)
 {
-	AllocSet	set = (AllocSet) context;
-	AllocChunk	chunk = AllocPointerGetChunk(pointer);
-	Size		oldsize;
-
-	/* Allow access to private part of chunk header. */
-	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOCCHUNK_PRIVATE_LEN);
-
-	oldsize = chunk->size;
+	AllocBlock	block;
+	AllocSet	set;
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	Size		oldchksize;
+	int			fidx;
 
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* Test for someone scribbling on unused space in chunk */
-	if (chunk->requested_size < oldsize)
-		if (!sentinel_ok(pointer, chunk->requested_size))
-			elog(WARNING, "detected write past chunk end in %s %p",
-				 set->header.name, chunk);
-#endif
+	/* Allow access to the chunk header. */
+	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
 
-	if (oldsize > set->allocChunkLimit)
+	if (MemoryChunkIsExternal(chunk))
 	{
 		/*
 		 * The chunk must have been allocated as a single-chunk block.  Use
 		 * realloc() to make the containing block bigger, or smaller, with
 		 * minimum space wastage.
 		 */
-		AllocBlock	block = (AllocBlock) (((char *) chunk) - ALLOC_BLOCKHDRSZ);
 		Size		chksize;
 		Size		blksize;
 		Size		oldblksize;
 
+		block = ExternalChunkGetBlock(chunk);
+
 		/*
-		 * Try to verify that we have a sane block pointer: it should
-		 * reference the correct aset, and freeptr and endptr should point
-		 * just past the chunk.
+		 * Try to verify that we have a sane block pointer: the block header
+		 * should reference an aset and the freeptr should match the endptr.
 		 */
-		if (block->aset != set ||
-			block->freeptr != block->endptr ||
-			block->freeptr != ((char *) block) +
-			(oldsize + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ))
+		if (!AllocBlockIsValid(block) || block->freeptr != block->endptr)
 			elog(ERROR, "could not find block containing chunk %p", chunk);
 
-		/*
-		 * Even if the new request is less than set->allocChunkLimit, we stick
-		 * with the single-chunk block approach.  Therefore we need
-		 * chunk->size to be bigger than set->allocChunkLimit, so we don't get
-		 * confused about the chunk's status in future calls.
-		 */
-		chksize = Max(size, set->allocChunkLimit + 1);
-		chksize = MAXALIGN(chksize);
+		set = block->aset;
+
+		/* only check size in paths where the limits could be hit */
+		MemoryContextCheckSize((MemoryContext) set, size, flags);
+
+		oldchksize = block->endptr - (char *) pointer;
+
+#ifdef MEMORY_CONTEXT_CHECKING
+		/* Test for someone scribbling on unused space in chunk */
+		Assert(chunk->requested_size < oldchksize);
+		if (!sentinel_ok(pointer, chunk->requested_size))
+			elog(WARNING, "detected write past chunk end in %s %p",
+				 set->header.name, chunk);
+#endif
+
+#ifdef MEMORY_CONTEXT_CHECKING
+		/* ensure there's always space for the sentinel byte */
+		chksize = MAXALIGN(size + 1);
+#else
+		chksize = MAXALIGN(size);
+#endif
 
 		/* Do the realloc */
 		blksize = chksize + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ;
@@ -1150,77 +1249,109 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		block = (AllocBlock) realloc(block, blksize);
 		if (block == NULL)
 		{
-			/* Disallow external access to private part of chunk header. */
-			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
-			return NULL;
+			/* Disallow access to the chunk header. */
+			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+			return MemoryContextAllocationFailure(&set->header, size, flags);
 		}
 
 		/* updated separately, not to underflow when (oldblksize > blksize) */
-		context->mem_allocated -= oldblksize;
-		context->mem_allocated += blksize;
+		set->header.mem_allocated -= oldblksize;
+		set->header.mem_allocated += blksize;
 
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
 		/* Update pointers since block has likely been moved */
-		chunk = (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ);
-		pointer = AllocChunkGetPointer(chunk);
+		chunk = (MemoryChunk *) (((char *) block) + ALLOC_BLOCKHDRSZ);
+		pointer = MemoryChunkGetPointer(chunk);
 		if (block->prev)
 			block->prev->next = block;
 		else
 			set->blocks = block;
 		if (block->next)
 			block->next->prev = block;
-		chunk->size = chksize;
 
 #ifdef MEMORY_CONTEXT_CHECKING
 #ifdef RANDOMIZE_ALLOCATED_MEMORY
-		/* We can only fill the extra space if we know the prior request */
+
+		/*
+		 * We can only randomize the extra space if we know the prior request.
+		 * When using Valgrind, randomize_mem() also marks memory UNDEFINED.
+		 */
 		if (size > chunk->requested_size)
 			randomize_mem((char *) pointer + chunk->requested_size,
 						  size - chunk->requested_size);
-#endif
+#else
 
 		/*
-		 * realloc() (or randomize_mem()) will have left any newly-allocated
-		 * part UNDEFINED, but we may need to adjust trailing bytes from the
-		 * old allocation.
+		 * If this is an increase, realloc() will have marked any
+		 * newly-allocated part (from oldchksize to chksize) UNDEFINED, but we
+		 * also need to adjust trailing bytes from the old allocation (from
+		 * chunk->requested_size to oldchksize) as they are marked NOACCESS.
+		 * Make sure not to mark too many bytes in case chunk->requested_size
+		 * < size < oldchksize.
 		 */
 #ifdef USE_VALGRIND
-		if (oldsize > chunk->requested_size)
+		if (Min(size, oldchksize) > chunk->requested_size)
 			VALGRIND_MAKE_MEM_UNDEFINED((char *) pointer + chunk->requested_size,
-										oldsize - chunk->requested_size);
+										Min(size, oldchksize) - chunk->requested_size);
+#endif
 #endif
 
 		chunk->requested_size = size;
-
 		/* set mark to catch clobber of "unused" space */
-		if (size < chunk->size)
-			set_sentinel(pointer, size);
+		Assert(size < chksize);
+		set_sentinel(pointer, size);
 #else							/* !MEMORY_CONTEXT_CHECKING */
 
 		/*
-		 * We don't know how much of the old chunk size was the actual
-		 * allocation; it could have been as small as one byte.  We have to be
-		 * conservative and just mark the entire old portion DEFINED.
+		 * We may need to adjust marking of bytes from the old allocation as
+		 * some of them may be marked NOACCESS.  We don't know how much of the
+		 * old chunk size was the requested size; it could have been as small
+		 * as one byte.  We have to be conservative and just mark the entire
+		 * old portion DEFINED.  Make sure not to mark memory beyond the new
+		 * allocation in case it's smaller than the old one.
 		 */
-		VALGRIND_MAKE_MEM_DEFINED(pointer, oldsize);
+		VALGRIND_MAKE_MEM_DEFINED(pointer, Min(size, oldchksize));
 #endif
 
 		/* Ensure any padding bytes are marked NOACCESS. */
 		VALGRIND_MAKE_MEM_NOACCESS((char *) pointer + size, chksize - size);
 
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+		/* Disallow access to the chunk header . */
+		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
 		return pointer;
 	}
 
+	block = MemoryChunkGetBlock(chunk);
+
+	/*
+	 * In this path, for speed reasons we just Assert that the referenced
+	 * block is good. We can also Assert that the value field is sane. Future
+	 * field experience may show that these Asserts had better become regular
+	 * runtime test-and-elog checks.
+	 */
+	Assert(AllocBlockIsValid(block));
+	set = block->aset;
+
+	fidx = MemoryChunkGetValue(chunk);
+	Assert(FreeListIdxIsValid(fidx));
+	oldchksize = GetChunkSizeFromFreeListIdx(fidx);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* Test for someone scribbling on unused space in chunk */
+	if (chunk->requested_size < oldchksize)
+		if (!sentinel_ok(pointer, chunk->requested_size))
+			elog(WARNING, "detected write past chunk end in %s %p",
+				 set->header.name, chunk);
+#endif
+
 	/*
 	 * Chunk sizes are aligned to power of 2 in AllocSetAlloc().  Maybe the
 	 * allocated area already is >= the new size.  (In particular, we will
 	 * fall out here if the requested size is a decrease.)
 	 */
-	else if (oldsize >= size)
+	if (oldchksize >= size)
 	{
 #ifdef MEMORY_CONTEXT_CHECKING
 		Size		oldrequest = chunk->requested_size;
@@ -1243,10 +1374,10 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 										size - oldrequest);
 		else
 			VALGRIND_MAKE_MEM_NOACCESS((char *) pointer + size,
-									   oldsize - size);
+									   oldchksize - size);
 
 		/* set mark to catch clobber of "unused" space */
-		if (size < oldsize)
+		if (size < oldchksize)
 			set_sentinel(pointer, size);
 #else							/* !MEMORY_CONTEXT_CHECKING */
 
@@ -1255,12 +1386,12 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		 * the old request or shrinking it, so we conservatively mark the
 		 * entire new allocation DEFINED.
 		 */
-		VALGRIND_MAKE_MEM_NOACCESS(pointer, oldsize);
+		VALGRIND_MAKE_MEM_NOACCESS(pointer, oldchksize);
 		VALGRIND_MAKE_MEM_DEFINED(pointer, size);
 #endif
 
-		/* Disallow external access to private part of chunk header. */
-		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+		/* Disallow access to the chunk header. */
+		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
 		return pointer;
 	}
@@ -1278,16 +1409,17 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		 * memory indefinitely.  See pgsql-hackers archives for 2007-08-11.)
 		 */
 		AllocPointer newPointer;
+		Size		oldsize;
 
-		/* allocate new chunk */
-		newPointer = AllocSetAlloc((MemoryContext) set, size);
+		/* allocate new chunk (this also checks size is valid) */
+		newPointer = AllocSetAlloc((MemoryContext) set, size, flags);
 
 		/* leave immediately if request was not completed */
 		if (newPointer == NULL)
 		{
-			/* Disallow external access to private part of chunk header. */
-			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
-			return NULL;
+			/* Disallow access to the chunk header. */
+			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+			return MemoryContextAllocationFailure((MemoryContext) set, size, flags);
 		}
 
 		/*
@@ -1302,6 +1434,7 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 #ifdef MEMORY_CONTEXT_CHECKING
 		oldsize = chunk->requested_size;
 #else
+		oldsize = oldchksize;
 		VALGRIND_MAKE_MEM_DEFINED(pointer, oldsize);
 #endif
 
@@ -1309,36 +1442,84 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		memcpy(newPointer, pointer, oldsize);
 
 		/* free old chunk */
-		AllocSetFree((MemoryContext) set, pointer);
+		AllocSetFree(pointer);
 
 		return newPointer;
 	}
 }
 
+/*
+ * AllocSetGetChunkContext
+ *		Return the MemoryContext that 'pointer' belongs to.
+ */
+MemoryContext
+AllocSetGetChunkContext(void *pointer)
+{
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	AllocBlock	block;
+	AllocSet	set;
+
+	/* Allow access to the chunk header. */
+	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+	if (MemoryChunkIsExternal(chunk))
+		block = ExternalChunkGetBlock(chunk);
+	else
+		block = (AllocBlock) MemoryChunkGetBlock(chunk);
+
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
+	Assert(AllocBlockIsValid(block));
+	set = block->aset;
+
+	return &set->header;
+}
+
 /*
  * AllocSetGetChunkSpace
  *		Given a currently-allocated chunk, determine the total space
  *		it occupies (including all memory-allocation overhead).
  */
-static Size
-AllocSetGetChunkSpace(MemoryContext context, void *pointer)
+Size
+AllocSetGetChunkSpace(void *pointer)
 {
-	AllocChunk	chunk = AllocPointerGetChunk(pointer);
-	Size		result;
+	MemoryChunk *chunk = PointerGetMemoryChunk(pointer);
+	int			fidx;
+
+	/* Allow access to the chunk header. */
+	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+	if (MemoryChunkIsExternal(chunk))
+	{
+		AllocBlock	block = ExternalChunkGetBlock(chunk);
 
-	VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOCCHUNK_PRIVATE_LEN);
-	result = chunk->size + ALLOC_CHUNKHDRSZ;
-	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
-	return result;
+		/* Disallow access to the chunk header. */
+		VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
+		Assert(AllocBlockIsValid(block));
+
+		return block->endptr - (char *) chunk;
+	}
+
+	fidx = MemoryChunkGetValue(chunk);
+	Assert(FreeListIdxIsValid(fidx));
+
+	/* Disallow access to the chunk header. */
+	VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
+	return GetChunkSizeFromFreeListIdx(fidx) + ALLOC_CHUNKHDRSZ;
 }
 
 /*
  * AllocSetIsEmpty
  *		Is an allocset empty of any allocated space?
  */
-static bool
+bool
 AllocSetIsEmpty(MemoryContext context)
 {
+	Assert(AllocSetIsValid(context));
+
 	/*
 	 * For now, we say "empty" only if the context is new or just reset. We
 	 * could examine the freelists to determine if all space has been freed,
@@ -1359,7 +1540,7 @@ AllocSetIsEmpty(MemoryContext context)
  * totals: if not NULL, add stats about this context into *totals.
  * print_to_stderr: print stats to stderr if true, elog otherwise.
  */
-static void
+void
 AllocSetStats(MemoryContext context,
 			  MemoryStatsPrintFunc printfunc, void *passthru,
 			  MemoryContextCounters *totals, bool print_to_stderr)
@@ -1372,6 +1553,8 @@ AllocSetStats(MemoryContext context,
 	AllocBlock	block;
 	int			fidx;
 
+	Assert(AllocSetIsValid(set));
+
 	/* Include context header in totalspace */
 	totalspace = MAXALIGN(sizeof(AllocSetContext));
 
@@ -1383,13 +1566,24 @@ AllocSetStats(MemoryContext context,
 	}
 	for (fidx = 0; fidx < ALLOCSET_NUM_FREELISTS; fidx++)
 	{
-		AllocChunk	chunk;
+		Size		chksz = GetChunkSizeFromFreeListIdx(fidx);
+		MemoryChunk *chunk = set->freelist[fidx];
 
-		for (chunk = set->freelist[fidx]; chunk != NULL;
-			 chunk = (AllocChunk) chunk->aset)
+		while (chunk != NULL)
 		{
+			AllocFreeListLink *link = GetFreeListLink(chunk);
+
+			/* Allow access to the chunk header. */
+			VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
+			Assert(MemoryChunkGetValue(chunk) == fidx);
+			VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
+
 			freechunks++;
-			freespace += chunk->size + ALLOC_CHUNKHDRSZ;
+			freespace += chksz + ALLOC_CHUNKHDRSZ;
+
+			VALGRIND_MAKE_MEM_DEFINED(link, sizeof(AllocFreeListLink));
+			chunk = link->next;
+			VALGRIND_MAKE_MEM_NOACCESS(link, sizeof(AllocFreeListLink));
 		}
 	}
 
@@ -1424,7 +1618,7 @@ AllocSetStats(MemoryContext context,
  * find yourself in an infinite loop when trouble occurs, because this
  * routine will be entered again when elog cleanup tries to release memory!
  */
-static void
+void
 AllocSetCheck(MemoryContext context)
 {
 	AllocSet	set = (AllocSet) context;
@@ -1441,8 +1635,9 @@ AllocSetCheck(MemoryContext context)
 		long		blk_used = block->freeptr - bpoz;
 		long		blk_data = 0;
 		long		nchunks = 0;
+		bool		has_external_chunk = false;
 
-		if (set->keeper == block)
+		if (IsKeeperBlock(set, block))
 			total_allocated += block->endptr - ((char *) set);
 		else
 			total_allocated += block->endptr - ((char *) block);
@@ -1452,7 +1647,7 @@ AllocSetCheck(MemoryContext context)
 		 */
 		if (!blk_used)
 		{
-			if (set->keeper != block)
+			if (!IsKeeperBlock(set, block))
 				elog(WARNING, "problem in alloc set %s: empty block %p",
 					 name, block);
 		}
@@ -1472,56 +1667,64 @@ AllocSetCheck(MemoryContext context)
 		 */
 		while (bpoz < block->freeptr)
 		{
-			AllocChunk	chunk = (AllocChunk) bpoz;
+			MemoryChunk *chunk = (MemoryChunk *) bpoz;
 			Size		chsize,
 						dsize;
 
-			/* Allow access to private part of chunk header. */
-			VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOCCHUNK_PRIVATE_LEN);
+			/* Allow access to the chunk header. */
+			VALGRIND_MAKE_MEM_DEFINED(chunk, ALLOC_CHUNKHDRSZ);
+
+			if (MemoryChunkIsExternal(chunk))
+			{
+				chsize = block->endptr - (char *) MemoryChunkGetPointer(chunk); /* aligned chunk size */
+				has_external_chunk = true;
+
+				/* make sure this chunk consumes the entire block */
+				if (chsize + ALLOC_CHUNKHDRSZ != blk_used)
+					elog(WARNING, "problem in alloc set %s: bad single-chunk %p in block %p",
+						 name, chunk, block);
+			}
+			else
+			{
+				int			fidx = MemoryChunkGetValue(chunk);
+
+				if (!FreeListIdxIsValid(fidx))
+					elog(WARNING, "problem in alloc set %s: bad chunk size for chunk %p in block %p",
+						 name, chunk, block);
+
+				chsize = GetChunkSizeFromFreeListIdx(fidx); /* aligned chunk size */
 
-			chsize = chunk->size;	/* aligned chunk size */
+				/*
+				 * Check the stored block offset correctly references this
+				 * block.
+				 */
+				if (block != MemoryChunkGetBlock(chunk))
+					elog(WARNING, "problem in alloc set %s: bad block offset for chunk %p in block %p",
+						 name, chunk, block);
+			}
 			dsize = chunk->requested_size;	/* real data */
 
-			/*
-			 * Check chunk size
-			 */
-			if (dsize > chsize)
+			/* an allocated chunk's requested size must be <= the chsize */
+			if (dsize != InvalidAllocSize && dsize > chsize)
 				elog(WARNING, "problem in alloc set %s: req size > alloc size for chunk %p in block %p",
 					 name, chunk, block);
+
+			/* chsize must not be smaller than the first freelist's size */
 			if (chsize < (1 << ALLOC_MINBITS))
 				elog(WARNING, "problem in alloc set %s: bad size %zu for chunk %p in block %p",
 					 name, chsize, chunk, block);
 
-			/* single-chunk block? */
-			if (chsize > set->allocChunkLimit &&
-				chsize + ALLOC_CHUNKHDRSZ != blk_used)
-				elog(WARNING, "problem in alloc set %s: bad single-chunk %p in block %p",
-					 name, chunk, block);
-
-			/*
-			 * If chunk is allocated, check for correct aset pointer. (If it's
-			 * free, the aset is the freelist pointer, which we can't check as
-			 * easily...)  Note this is an incomplete test, since palloc(0)
-			 * produces an allocated chunk with requested_size == 0.
-			 */
-			if (dsize > 0 && chunk->aset != (void *) set)
-				elog(WARNING, "problem in alloc set %s: bogus aset link in block %p, chunk %p",
-					 name, block, chunk);
-
 			/*
 			 * Check for overwrite of padding space in an allocated chunk.
 			 */
-			if (chunk->aset == (void *) set && dsize < chsize &&
+			if (dsize != InvalidAllocSize && dsize < chsize &&
 				!sentinel_ok(chunk, ALLOC_CHUNKHDRSZ + dsize))
 				elog(WARNING, "problem in alloc set %s: detected write past chunk end in block %p, chunk %p",
 					 name, block, chunk);
 
-			/*
-			 * If chunk is allocated, disallow external access to private part
-			 * of chunk header.
-			 */
-			if (chunk->aset == (void *) set)
-				VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOCCHUNK_PRIVATE_LEN);
+			/* if chunk is allocated, disallow access to the chunk header */
+			if (dsize != InvalidAllocSize)
+				VALGRIND_MAKE_MEM_NOACCESS(chunk, ALLOC_CHUNKHDRSZ);
 
 			blk_data += chsize;
 			nchunks++;
@@ -1532,6 +1735,10 @@ AllocSetCheck(MemoryContext context)
 		if ((blk_data + (nchunks * ALLOC_CHUNKHDRSZ)) != blk_used)
 			elog(WARNING, "problem in alloc set %s: found inconsistent memory block %p",
 				 name, block);
+
+		if (has_external_chunk && nchunks > 1)
+			elog(WARNING, "problem in alloc set %s: external chunk on non-dedicated block %p",
+				 name, block);
 	}
 
 	Assert(total_allocated == context->mem_allocated);