Safe memory allocation functions

From b636c809c2f2cb4177bedc2e5a4883a79b61fbc6 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Tue, 13 Jan 2015 15:40:38 +0900
Subject: [PATCH] Add memory allocation APIs able to return NULL instead of
 ERROR

The following functions are added to the existing set for frontend and
backend:
- palloc_noerror
- palloc0_noerror
- repalloc_noerror
---
 src/backend/utils/mmgr/aset.c    | 529 +++++++++++++++++++++++----------------
 src/backend/utils/mmgr/mcxt.c    | 124 +++++----
 src/common/fe_memutils.c         |  72 ++++--
 src/include/common/fe_memutils.h |   3 +
 src/include/nodes/memnodes.h     |   2 +
 src/include/utils/palloc.h       |   3 +
 6 files changed, 451 insertions(+), 282 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 85b3c9a..974e018 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -243,11 +243,22 @@ typedef struct AllocChunkData
 					((AllocPointer)(((char *)(chk)) + ALLOC_CHUNKHDRSZ))
 
 /*
+ * Wrappers for allocation functions.
+ */
+static void *set_alloc_internal(MemoryContext context,
+								Size size, bool noerror);
+static void *set_realloc_internal(MemoryContext context, void *pointer,
+								  Size size, bool noerror);
+
+/*
  * These functions implement the MemoryContext API for AllocSet contexts.
  */
 static void *AllocSetAlloc(MemoryContext context, Size size);
+static void *AllocSetAllocNoError(MemoryContext context, Size size);
 static void AllocSetFree(MemoryContext context, void *pointer);
 static void *AllocSetRealloc(MemoryContext context, void *pointer, Size size);
+static void *AllocSetReallocNoError(MemoryContext context,
+								 void *pointer, Size size);
 static void AllocSetInit(MemoryContext context);
 static void AllocSetReset(MemoryContext context);
 static void AllocSetDelete(MemoryContext context);
@@ -264,8 +275,10 @@ static void AllocSetCheck(MemoryContext context);
  */
 static MemoryContextMethods AllocSetMethods = {
 	AllocSetAlloc,
+	AllocSetAllocNoError,
 	AllocSetFree,
 	AllocSetRealloc,
+	AllocSetReallocNoError,
 	AllocSetInit,
 	AllocSetReset,
 	AllocSetDelete,
@@ -517,140 +530,16 @@ AllocSetContextCreate(MemoryContext parent,
 }
 
 /*
- * AllocSetInit
- *		Context-type-specific initialization routine.
- *
- * This is called by MemoryContextCreate() after setting up the
- * generic MemoryContext fields and before linking the new context
- * into the context tree.  We must do whatever is needed to make the
- * new context minimally valid for deletion.  We must *not* risk
- * failure --- thus, for example, allocating more memory is not cool.
- * (AllocSetContextCreate can allocate memory when it gets control
- * back, however.)
- */
-static void
-AllocSetInit(MemoryContext context)
-{
-	/*
-	 * Since MemoryContextCreate already zeroed the context node, we don't
-	 * have to do anything here: it's already OK.
-	 */
-}
-
-/*
- * AllocSetReset
- *		Frees all memory which is allocated in the given set.
- *
- * Actually, this routine has some discretion about what to do.
- * It should mark all allocated chunks freed, but it need not necessarily
- * give back all the resources the set owns.  Our actual implementation is
- * that we hang onto any "keeper" block specified for the set.  In this way,
- * we don't thrash malloc() when a context is repeatedly reset after small
- * allocations, which is typical behavior for per-tuple contexts.
- */
-static void
-AllocSetReset(MemoryContext context)
-{
-	AllocSet	set = (AllocSet) context;
-	AllocBlock	block;
-
-	AssertArg(AllocSetIsValid(set));
-
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* Check for corruption and leaks before freeing */
-	AllocSetCheck(context);
-#endif
-
-	/* Clear chunk freelists */
-	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
-
-	block = set->blocks;
-
-	/* New blocks list is either empty or just the keeper block */
-	set->blocks = set->keeper;
-
-	while (block != NULL)
-	{
-		AllocBlock	next = block->next;
-
-		if (block == set->keeper)
-		{
-			/* Reset the block, but don't return it to malloc */
-			char	   *datastart = ((char *) block) + ALLOC_BLOCKHDRSZ;
-
-#ifdef CLOBBER_FREED_MEMORY
-			wipe_mem(datastart, block->freeptr - datastart);
-#else
-			/* wipe_mem() would have done this */
-			VALGRIND_MAKE_MEM_NOACCESS(datastart, block->freeptr - datastart);
-#endif
-			block->freeptr = datastart;
-			block->next = NULL;
-		}
-		else
-		{
-			/* Normal case, release the block */
-#ifdef CLOBBER_FREED_MEMORY
-			wipe_mem(block, block->freeptr - ((char *) block));
-#endif
-			free(block);
-		}
-		block = next;
-	}
-
-	/* Reset block size allocation sequence, too */
-	set->nextBlockSize = set->initBlockSize;
-}
-
-/*
- * AllocSetDelete
- *		Frees all memory which is allocated in the given set,
- *		in preparation for deletion of the set.
- *
- * Unlike AllocSetReset, this *must* free all resources of the set.
- * But note we are not responsible for deleting the context node itself.
- */
-static void
-AllocSetDelete(MemoryContext context)
-{
-	AllocSet	set = (AllocSet) context;
-	AllocBlock	block = set->blocks;
-
-	AssertArg(AllocSetIsValid(set));
-
-#ifdef MEMORY_CONTEXT_CHECKING
-	/* Check for corruption and leaks before freeing */
-	AllocSetCheck(context);
-#endif
-
-	/* Make it look empty, just in case... */
-	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
-	set->blocks = NULL;
-	set->keeper = NULL;
-
-	while (block != NULL)
-	{
-		AllocBlock	next = block->next;
-
-#ifdef CLOBBER_FREED_MEMORY
-		wipe_mem(block, block->freeptr - ((char *) block));
-#endif
-		free(block);
-		block = next;
-	}
-}
-
-/*
- * AllocSetAlloc
- *		Returns pointer to allocated memory of given size; memory is added
- *		to the set.
+ * set_alloc_internal
+ *		Wrapper for memory allocation routines.
  *
  * No request may exceed:
  *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
  * All callers use a much-lower limit.
  */
 static void *
-AllocSetAlloc(MemoryContext context, Size size)
+set_alloc_internal(MemoryContext context,
+				   Size size, bool noerror)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocBlock	block;
@@ -673,10 +562,13 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+			if (noerror)
+				return NULL;
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on request of size %zu.", size)));
 		}
 		block->aset = set;
 		block->freeptr = block->endptr = ((char *) block) + blksize;
@@ -867,10 +759,13 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+			if (noerror)
+				return NULL;
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on request of size %zu.", size)));
 		}
 
 		block->aset = set;
@@ -928,83 +823,8 @@ AllocSetAlloc(MemoryContext context, Size size)
 }
 
 /*
- * AllocSetFree
- *		Frees allocated memory; memory is removed from the set.
- */
-static void
-AllocSetFree(MemoryContext context, void *pointer)
-{
-	AllocSet	set = (AllocSet) context;
-	AllocChunk	chunk = AllocPointerGetChunk(pointer);
-
-	AllocFreeInfo(set, chunk);
-
-#ifdef MEMORY_CONTEXT_CHECKING
-	VALGRIND_MAKE_MEM_DEFINED(&chunk->requested_size,
-							  sizeof(chunk->requested_size));
-	/* 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
-
-	if (chunk->size > set->allocChunkLimit)
-	{
-		/*
-		 * Big chunks are certain to have been allocated as single-chunk
-		 * blocks.  Find the containing block and return it to malloc().
-		 */
-		AllocBlock	block = set->blocks;
-		AllocBlock	prevblock = NULL;
-
-		while (block != NULL)
-		{
-			if (chunk == (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ))
-				break;
-			prevblock = block;
-			block = block->next;
-		}
-		if (block == NULL)
-			elog(ERROR, "could not find block containing chunk %p", chunk);
-		/* let's just make sure chunk is the only one in the block */
-		Assert(block->freeptr == ((char *) block) +
-			   (chunk->size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ));
-
-		/* OK, remove block from aset's list and free it */
-		if (prevblock == NULL)
-			set->blocks = block->next;
-		else
-			prevblock->next = block->next;
-#ifdef CLOBBER_FREED_MEMORY
-		wipe_mem(block, block->freeptr - ((char *) block));
-#endif
-		free(block);
-	}
-	else
-	{
-		/* Normal case, put the chunk into appropriate freelist */
-		int			fidx = AllocSetFreeIndex(chunk->size);
-
-		chunk->aset = (void *) set->freelist[fidx];
-
-#ifdef CLOBBER_FREED_MEMORY
-		wipe_mem(pointer, chunk->size);
-#endif
-
-#ifdef MEMORY_CONTEXT_CHECKING
-		/* Reset requested_size to 0 in chunks that are on freelist */
-		chunk->requested_size = 0;
-#endif
-		set->freelist[fidx] = chunk;
-	}
-}
-
-/*
- * AllocSetRealloc
- *		Returns new pointer to allocated memory of given size; this memory
- *		is added to the set.  Memory associated with given pointer is copied
- *		into the new memory, and the old memory is freed.
+ * set_realloc_internal
+ *		Wrapper for memory reallocation routines.
  *
  * Without MEMORY_CONTEXT_CHECKING, we don't know the old request size.  This
  * makes our Valgrind client requests less-precise, hazarding false negatives.
@@ -1012,7 +832,8 @@ AllocSetFree(MemoryContext context, void *pointer)
  * request size.)
  */
 static void *
-AllocSetRealloc(MemoryContext context, void *pointer, Size size)
+set_realloc_internal(MemoryContext context, void *pointer,
+					 Size size, bool noerror)
 {
 	AllocSet	set = (AllocSet) context;
 	AllocChunk	chunk = AllocPointerGetChunk(pointer);
@@ -1029,8 +850,8 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 #endif
 
 	/*
-	 * Chunk sizes are aligned to power of 2 in AllocSetAlloc(). Maybe the
-	 * allocated area already is >= the new size.  (In particular, we always
+	 * Chunk sizes are aligned to power of 2 in set_alloc_internal(). Maybe
+	 * the allocated area already is >= the new size.  (In particular, we always
 	 * fall out here if the requested size is a decrease.)
 	 */
 	if (oldsize >= size)
@@ -1109,10 +930,15 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+
+			/* allocation failed */
+			if (noerror)
+				return NULL;
+			else
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on request of size %zu.", size)));
 		}
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
@@ -1177,10 +1003,17 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		AllocPointer newPointer;
 
 		/* allocate new chunk */
-		newPointer = AllocSetAlloc((MemoryContext) set, size);
+		newPointer = set_alloc_internal((MemoryContext) set, size, noerror);
+
+		/* leave if allocation did not complete properly */
+		if (newPointer == NULL)
+		{
+			Assert(noerror);
+			return NULL;
+		}
 
 		/*
-		 * AllocSetAlloc() just made the region NOACCESS.  Change it to
+		 * set_alloc_internal() just made the region NOACCESS.  Change it to
 		 * UNDEFINED for the moment; memcpy() will then transfer definedness
 		 * from the old allocation to the new.  If we know the old allocation,
 		 * copy just that much.  Otherwise, make the entire old chunk defined
@@ -1203,6 +1036,258 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 	}
 }
 
+
+/*
+ * AllocSetInit
+ *		Context-type-specific initialization routine.
+ *
+ * This is called by MemoryContextCreate() after setting up the
+ * generic MemoryContext fields and before linking the new context
+ * into the context tree.  We must do whatever is needed to make the
+ * new context minimally valid for deletion.  We must *not* risk
+ * failure --- thus, for example, allocating more memory is not cool.
+ * (AllocSetContextCreate can allocate memory when it gets control
+ * back, however.)
+ */
+static void
+AllocSetInit(MemoryContext context)
+{
+	/*
+	 * Since MemoryContextCreate already zeroed the context node, we don't
+	 * have to do anything here: it's already OK.
+	 */
+}
+
+/*
+ * AllocSetReset
+ *		Frees all memory which is allocated in the given set.
+ *
+ * Actually, this routine has some discretion about what to do.
+ * It should mark all allocated chunks freed, but it need not necessarily
+ * give back all the resources the set owns.  Our actual implementation is
+ * that we hang onto any "keeper" block specified for the set.  In this way,
+ * we don't thrash malloc() when a context is repeatedly reset after small
+ * allocations, which is typical behavior for per-tuple contexts.
+ */
+static void
+AllocSetReset(MemoryContext context)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block;
+
+	AssertArg(AllocSetIsValid(set));
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* Check for corruption and leaks before freeing */
+	AllocSetCheck(context);
+#endif
+
+	/* Clear chunk freelists */
+	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
+
+	block = set->blocks;
+
+	/* New blocks list is either empty or just the keeper block */
+	set->blocks = set->keeper;
+
+	while (block != NULL)
+	{
+		AllocBlock	next = block->next;
+
+		if (block == set->keeper)
+		{
+			/* Reset the block, but don't return it to malloc */
+			char	   *datastart = ((char *) block) + ALLOC_BLOCKHDRSZ;
+
+#ifdef CLOBBER_FREED_MEMORY
+			wipe_mem(datastart, block->freeptr - datastart);
+#else
+			/* wipe_mem() would have done this */
+			VALGRIND_MAKE_MEM_NOACCESS(datastart, block->freeptr - datastart);
+#endif
+			block->freeptr = datastart;
+			block->next = NULL;
+		}
+		else
+		{
+			/* Normal case, release the block */
+#ifdef CLOBBER_FREED_MEMORY
+			wipe_mem(block, block->freeptr - ((char *) block));
+#endif
+			free(block);
+		}
+		block = next;
+	}
+
+	/* Reset block size allocation sequence, too */
+	set->nextBlockSize = set->initBlockSize;
+}
+
+/*
+ * AllocSetDelete
+ *		Frees all memory which is allocated in the given set,
+ *		in preparation for deletion of the set.
+ *
+ * Unlike AllocSetReset, this *must* free all resources of the set.
+ * But note we are not responsible for deleting the context node itself.
+ */
+static void
+AllocSetDelete(MemoryContext context)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocBlock	block = set->blocks;
+
+	AssertArg(AllocSetIsValid(set));
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	/* Check for corruption and leaks before freeing */
+	AllocSetCheck(context);
+#endif
+
+	/* Make it look empty, just in case... */
+	MemSetAligned(set->freelist, 0, sizeof(set->freelist));
+	set->blocks = NULL;
+	set->keeper = NULL;
+
+	while (block != NULL)
+	{
+		AllocBlock	next = block->next;
+
+#ifdef CLOBBER_FREED_MEMORY
+		wipe_mem(block, block->freeptr - ((char *) block));
+#endif
+		free(block);
+		block = next;
+	}
+}
+
+/*
+ * AllocSetAlloc
+ *		Returns pointer to allocated memory of given size; memory is added
+ *		to the set. This fails with an out-of-memory error if request cannot
+ *		be completed properly.
+ */
+static void *
+AllocSetAlloc(MemoryContext context, Size size)
+{
+	return set_alloc_internal(context, size, false);
+}
+
+/*
+ * AllocSetAllocNoError
+ *		Returns pointer to allocated memory of given size; memory is added
+ *		to the set. This returns NULL if request cannot be completed
+ *		properly.
+ *
+ * No request may exceed:
+ *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
+ * All callers use a much-lower limit.
+ */
+static void *
+AllocSetAllocNoError(MemoryContext context, Size size)
+{
+	return set_alloc_internal(context, size, true);
+}
+
+/*
+ * AllocSetFree
+ *		Frees allocated memory; memory is removed from the set.
+ */
+static void
+AllocSetFree(MemoryContext context, void *pointer)
+{
+	AllocSet	set = (AllocSet) context;
+	AllocChunk	chunk = AllocPointerGetChunk(pointer);
+
+	AllocFreeInfo(set, chunk);
+
+#ifdef MEMORY_CONTEXT_CHECKING
+	VALGRIND_MAKE_MEM_DEFINED(&chunk->requested_size,
+							  sizeof(chunk->requested_size));
+	/* 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
+
+	if (chunk->size > set->allocChunkLimit)
+	{
+		/*
+		 * Big chunks are certain to have been allocated as single-chunk
+		 * blocks.  Find the containing block and return it to malloc().
+		 */
+		AllocBlock	block = set->blocks;
+		AllocBlock	prevblock = NULL;
+
+		while (block != NULL)
+		{
+			if (chunk == (AllocChunk) (((char *) block) + ALLOC_BLOCKHDRSZ))
+				break;
+			prevblock = block;
+			block = block->next;
+		}
+		if (block == NULL)
+			elog(ERROR, "could not find block containing chunk %p", chunk);
+		/* let's just make sure chunk is the only one in the block */
+		Assert(block->freeptr == ((char *) block) +
+			   (chunk->size + ALLOC_BLOCKHDRSZ + ALLOC_CHUNKHDRSZ));
+
+		/* OK, remove block from aset's list and free it */
+		if (prevblock == NULL)
+			set->blocks = block->next;
+		else
+			prevblock->next = block->next;
+#ifdef CLOBBER_FREED_MEMORY
+		wipe_mem(block, block->freeptr - ((char *) block));
+#endif
+		free(block);
+	}
+	else
+	{
+		/* Normal case, put the chunk into appropriate freelist */
+		int			fidx = AllocSetFreeIndex(chunk->size);
+
+		chunk->aset = (void *) set->freelist[fidx];
+
+#ifdef CLOBBER_FREED_MEMORY
+		wipe_mem(pointer, chunk->size);
+#endif
+
+#ifdef MEMORY_CONTEXT_CHECKING
+		/* Reset requested_size to 0 in chunks that are on freelist */
+		chunk->requested_size = 0;
+#endif
+		set->freelist[fidx] = chunk;
+	}
+}
+
+/*
+ * AllocSetRealloc
+ *		Returns new pointer to allocated memory of given size; this memory
+ *		is added to the set.  Memory associated with given pointer is copied
+ *		into the new memory, and the old memory is freed. If request cannot
+ *		be completed, this fails with an out-of-memory error.
+ */
+static void *
+AllocSetRealloc(MemoryContext context, void *pointer, Size size)
+{
+	return set_realloc_internal(context, pointer, size, false);
+}
+
+/*
+ * AllocSetReallocNoError
+ *		Returns new pointer to allocated memory of given size; this memory
+ *		is added to the set.  Memory associated with given pointer is copied
+ *		into the new memory, and the old memory is freed. If request cannot
+ *		be completed, this returns NULL.
+ */
+static void *
+AllocSetReallocNoError(MemoryContext context, void *pointer, Size size)
+{
+	return set_realloc_internal(context, pointer, size, true);
+}
+
 /*
  * AllocSetGetChunkSpace
  *		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 aa0d458..37c0669 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -56,6 +56,10 @@ MemoryContext PortalContext = NULL;
 
 static void MemoryContextStatsInternal(MemoryContext context, int level);
 
+/* wrapper routines for allocation */
+static void* palloc_internal(Size size, bool noerror);
+static void* repalloc_internal(void *pointer, Size size, bool noerror);
+
 /*
  * You should not do memory allocations within a critical section, because
  * an out-of-memory error will be escalated to a PANIC. To enforce that
@@ -684,8 +688,8 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	return ret;
 }
 
-void *
-palloc(Size size)
+static void*
+palloc_internal(Size size, bool noerror)
 {
 	/* duplicates MemoryContextAlloc to avoid increased overhead */
 	void	   *ret;
@@ -698,31 +702,85 @@ palloc(Size size)
 
 	CurrentMemoryContext->isReset = false;
 
-	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
+	if (noerror)
+		ret = (*CurrentMemoryContext->methods->alloc_noerror)
+			(CurrentMemoryContext, size);
+	else
+		ret = (*CurrentMemoryContext->methods->alloc)
+			(CurrentMemoryContext, size);
 	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
 
 	return ret;
 }
 
-void *
-palloc0(Size size)
+static void*
+repalloc_internal(void *pointer, Size size, bool noerror)
 {
-	/* duplicates MemoryContextAllocZero to avoid increased overhead */
+	MemoryContext context;
 	void	   *ret;
 
-	AssertArg(MemoryContextIsValid(CurrentMemoryContext));
-	AssertNotInCriticalSection(CurrentMemoryContext);
-
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
-	CurrentMemoryContext->isReset = false;
+	/*
+	 * 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.
+	 */
+	Assert(pointer != NULL);
+	Assert(pointer == (void *) MAXALIGN(pointer));
 
-	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
-	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
+	/*
+	 * OK, it's probably safe to look at the chunk header.
+	 */
+	context = ((StandardChunkHeader *)
+			   ((char *) pointer - STANDARDCHUNKHEADERSIZE))->context;
 
+	AssertArg(MemoryContextIsValid(context));
+	AssertNotInCriticalSection(context);
+
+	/* isReset must be false already */
+	Assert(!context->isReset);
+
+	if (noerror)
+		ret = (*context->methods->realloc_noerror) (context, pointer, size);
+	else
+		ret = (*context->methods->realloc) (context, pointer, size);
+
+	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+
+	return ret;
+}
+
+void *
+palloc(Size size)
+{
+	return palloc_internal(size, false);
+}
+
+void *
+palloc_noerror(Size size)
+{
+	return palloc_internal(size, true);
+}
+
+void *
+palloc0(Size size)
+{
+	void	   *ret;
+
+	ret = palloc_internal(size, false);
 	MemSetAligned(ret, 0, size);
+	return ret;
+}
 
+void *
+palloc0_noerror(Size size)
+{
+	void	   *ret;
+
+	ret = palloc_internal(size, true);
+	MemSetAligned(ret, 0, size);
 	return ret;
 }
 
@@ -757,41 +815,23 @@ pfree(void *pointer)
 
 /*
  * repalloc
- *		Adjust the size of a previously allocated chunk.
+ *		Adjust the size of a previously allocated chunk, failing on OOM.
  */
 void *
 repalloc(void *pointer, Size size)
 {
-	MemoryContext context;
-	void	   *ret;
-
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
-	/*
-	 * 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.
-	 */
-	Assert(pointer != NULL);
-	Assert(pointer == (void *) MAXALIGN(pointer));
-
-	/*
-	 * OK, it's probably safe to look at the chunk header.
-	 */
-	context = ((StandardChunkHeader *)
-			   ((char *) pointer - STANDARDCHUNKHEADERSIZE))->context;
-
-	AssertArg(MemoryContextIsValid(context));
-	AssertNotInCriticalSection(context);
-
-	/* isReset must be false already */
-	Assert(!context->isReset);
-
-	ret = (*context->methods->realloc) (context, pointer, size);
-	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+	return repalloc_internal(pointer, size, false);
+}
 
-	return ret;
+/*
+ * repalloc
+ *		Adjust the size of a previously allocated chunk, returning NULL
+ *		on OOM.
+ */
+void *
+repalloc_noerror(void *pointer, Size size)
+{
+	return repalloc_internal(pointer, size, true);
 }
 
 /*
diff --git a/src/common/fe_memutils.c b/src/common/fe_memutils.c
index 345221e..fd343ba 100644
--- a/src/common/fe_memutils.c
+++ b/src/common/fe_memutils.c
@@ -19,8 +19,8 @@
 
 #include "postgres_fe.h"
 
-void *
-pg_malloc(size_t size)
+static void *
+pg_malloc_internal(size_t size, bool noerror)
 {
 	void	   *tmp;
 
@@ -28,7 +28,24 @@ pg_malloc(size_t size)
 	if (size == 0)
 		size = 1;
 	tmp = malloc(size);
-	if (!tmp)
+	if (!tmp && !noerror)
+	{
+		fprintf(stderr, _("out of memory\n"));
+		exit(EXIT_FAILURE);
+	}
+	return tmp;
+}
+
+static void *
+pg_realloc_internal(void *ptr, size_t size, bool noerror)
+{
+	void	   *tmp;
+
+	/* Avoid unportable behavior of realloc(NULL, 0) */
+	if (ptr == NULL && size == 0)
+		size = 1;
+	tmp = realloc(ptr, size);
+	if (!tmp && !noerror)
 	{
 		fprintf(stderr, _("out of memory\n"));
 		exit(EXIT_FAILURE);
@@ -37,6 +54,12 @@ pg_malloc(size_t size)
 }
 
 void *
+pg_malloc(size_t size)
+{
+	return pg_malloc_internal(size, false);
+}
+
+void *
 pg_malloc0(size_t size)
 {
 	void	   *tmp;
@@ -49,18 +72,7 @@ pg_malloc0(size_t size)
 void *
 pg_realloc(void *ptr, size_t size)
 {
-	void	   *tmp;
-
-	/* Avoid unportable behavior of realloc(NULL, 0) */
-	if (ptr == NULL && size == 0)
-		size = 1;
-	tmp = realloc(ptr, size);
-	if (!tmp)
-	{
-		fprintf(stderr, _("out of memory\n"));
-		exit(EXIT_FAILURE);
-	}
-	return tmp;
+	return pg_realloc_internal(ptr, size, false);
 }
 
 /*
@@ -100,13 +112,31 @@ pg_free(void *ptr)
 void *
 palloc(Size size)
 {
-	return pg_malloc(size);
+	return pg_malloc_internal(size, false);
+}
+
+void *
+palloc_noerror(Size size)
+{
+	return pg_malloc_internal(size, true);
 }
 
 void *
 palloc0(Size size)
 {
-	return pg_malloc0(size);
+	void	   *tmp;
+	tmp = pg_malloc_internal(size, false);
+	MemSet(tmp, 0, size);
+	return tmp;
+}
+
+void *
+palloc0_noeror(Size size)
+{
+	void	   *tmp;
+	tmp = pg_malloc_internal(size, true);
+	MemSet(tmp, 0, size);
+	return tmp;
 }
 
 void
@@ -124,5 +154,11 @@ pstrdup(const char *in)
 void *
 repalloc(void *pointer, Size size)
 {
-	return pg_realloc(pointer, size);
+	return pg_realloc_internal(pointer, size, false);
+}
+
+void *
+repalloc_noerror(void *pointer, Size size)
+{
+	return pg_realloc_internal(pointer, size, true);
 }
diff --git a/src/include/common/fe_memutils.h b/src/include/common/fe_memutils.h
index f7114c2..0151e6e 100644
--- a/src/include/common/fe_memutils.h
+++ b/src/include/common/fe_memutils.h
@@ -19,8 +19,11 @@ extern void pg_free(void *pointer);
 /* Equivalent functions, deliberately named the same as backend functions */
 extern char *pstrdup(const char *in);
 extern void *palloc(Size size);
+extern void *palloc_noerror(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_noerror(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_noerror(void *pointer, Size size);
 extern void pfree(void *pointer);
 
 /* sprintf into a palloc'd buffer --- these are in psprintf.c */
diff --git a/src/include/nodes/memnodes.h b/src/include/nodes/memnodes.h
index ca9c3de..6d61acb 100644
--- a/src/include/nodes/memnodes.h
+++ b/src/include/nodes/memnodes.h
@@ -36,9 +36,11 @@
 typedef struct MemoryContextMethods
 {
 	void	   *(*alloc) (MemoryContext context, Size size);
+	void	   *(*alloc_noerror) (MemoryContext context, Size size);
 	/* call this free_p in case someone #define's free() */
 	void		(*free_p) (MemoryContext context, void *pointer);
 	void	   *(*realloc) (MemoryContext context, void *pointer, Size size);
+	void	   *(*realloc_noerror) (MemoryContext context, void *pointer, Size size);
 	void		(*init) (MemoryContext context);
 	void		(*reset) (MemoryContext context);
 	void		(*delete_context) (MemoryContext context);
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index ca03f2b..3634a7f 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -50,8 +50,11 @@ extern void *MemoryContextAllocZero(MemoryContext context, Size size);
 extern void *MemoryContextAllocZeroAligned(MemoryContext context, Size size);
 
 extern void *palloc(Size size);
+extern void *palloc_noerror(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_noerror(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_noerror(void *pointer, Size size);
 extern void pfree(void *pointer);
 
 /*
-- 
2.2.2

From 337c439554ce66486cf9d29dced6c72d034b8f8d Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Wed, 28 Jan 2015 22:10:13 +0900
Subject: [PATCH 1/3] Make allocation return functions return NULL on OOM

On counterpart, higher-level APIs in mcxt.c return an explicit error
message when memory requests cannot be completed.
---
 src/backend/utils/mmgr/aset.c | 30 ++++++++++++---------------
 src/backend/utils/mmgr/mcxt.c | 48 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/src/backend/utils/mmgr/aset.c b/src/backend/utils/mmgr/aset.c
index 85b3c9a..bf6f09a 100644
--- a/src/backend/utils/mmgr/aset.c
+++ b/src/backend/utils/mmgr/aset.c
@@ -642,8 +642,8 @@ AllocSetDelete(MemoryContext context)
 
 /*
  * AllocSetAlloc
- *		Returns pointer to allocated memory of given size; memory is added
- *		to the set.
+ *		Returns pointer to allocated memory of given size or NULL if
+ *		request could not be completed; memory is added to the set.
  *
  * No request may exceed:
  *		MAXALIGN_DOWN(SIZE_MAX) - ALLOC_BLOCKHDRSZ - ALLOC_CHUNKHDRSZ
@@ -673,10 +673,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+			return NULL;
 		}
 		block->aset = set;
 		block->freeptr = block->endptr = ((char *) block) + blksize;
@@ -867,10 +864,7 @@ AllocSetAlloc(MemoryContext context, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+			return NULL;
 		}
 
 		block->aset = set;
@@ -1002,9 +996,10 @@ AllocSetFree(MemoryContext context, void *pointer)
 
 /*
  * AllocSetRealloc
- *		Returns new pointer to allocated memory of given size; this memory
- *		is added to the set.  Memory associated with given pointer is copied
- *		into the new memory, and the old memory is freed.
+ *		Returns new pointer to allocated memory of given size or NULL if
+ *		request could not be completed; this memory is added to the set.
+ *		Memory associated with given pointer is copied into the new memory,
+ *		and the old memory is freed.
  *
  * Without MEMORY_CONTEXT_CHECKING, we don't know the old request size.  This
  * makes our Valgrind client requests less-precise, hazarding false negatives.
@@ -1109,10 +1104,7 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		if (block == NULL)
 		{
 			MemoryContextStats(TopMemoryContext);
-			ereport(ERROR,
-					(errcode(ERRCODE_OUT_OF_MEMORY),
-					 errmsg("out of memory"),
-					 errdetail("Failed on request of size %zu.", size)));
+			return NULL;
 		}
 		block->freeptr = block->endptr = ((char *) block) + blksize;
 
@@ -1179,6 +1171,10 @@ AllocSetRealloc(MemoryContext context, void *pointer, Size size)
 		/* allocate new chunk */
 		newPointer = AllocSetAlloc((MemoryContext) set, size);
 
+		/* leave immediately if request was not completed */
+		if (newPointer == NULL)
+			return NULL;
+
 		/*
 		 * AllocSetAlloc() just made the region NOACCESS.  Change it to
 		 * UNDEFINED for the moment; memcpy() will then transfer definedness
diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index aa0d458..f6cf0cc 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -623,6 +623,12 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	context->isReset = false;
 
 	ret = (*context->methods->alloc) (context, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -649,6 +655,12 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	context->isReset = false;
 
 	ret = (*context->methods->alloc) (context, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	MemSetAligned(ret, 0, size);
@@ -677,6 +689,12 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	context->isReset = false;
 
 	ret = (*context->methods->alloc) (context, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	MemSetLoop(ret, 0, size);
@@ -699,6 +717,12 @@ palloc(Size size)
 	CurrentMemoryContext->isReset = false;
 
 	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
 
 	return ret;
@@ -719,6 +743,12 @@ palloc0(Size size)
 	CurrentMemoryContext->isReset = false;
 
 	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
 
 	MemSetAligned(ret, 0, size);
@@ -789,6 +819,12 @@ repalloc(void *pointer, Size size)
 	Assert(!context->isReset);
 
 	ret = (*context->methods->realloc) (context, pointer, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
 	return ret;
@@ -814,6 +850,12 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	context->isReset = false;
 
 	ret = (*context->methods->alloc) (context, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
@@ -854,6 +896,12 @@ repalloc_huge(void *pointer, Size size)
 	Assert(!context->isReset);
 
 	ret = (*context->methods->realloc) (context, pointer, size);
+	if (ret == NULL)
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+
 	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
 
 	return ret;
-- 
2.2.2

From 13a05f5b28ab00864a9d96fae5a72402cf43184b Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Wed, 28 Jan 2015 22:46:04 +0900
Subject: [PATCH 2/3] Add *_noerror routines for palloc memory allocation

The following routines are added for frontend and backend:
- palloc_noerror, equivalent of palloc but returning NULL if memory request
could not be completed
- palloc0_noerror, equivalent of palloc0 but returning NULL if memory
request could not be completed
- repalloc_noerror, equivalent of repalloc bur returning NULL if memory
request could not be completed
---
 src/backend/utils/mmgr/mcxt.c    | 129 ++++++++++++++++++++++++---------------
 src/common/fe_memutils.c         |  73 ++++++++++++++++------
 src/include/common/fe_memutils.h |   3 +
 src/include/utils/palloc.h       |   3 +
 4 files changed, 140 insertions(+), 68 deletions(-)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index f6cf0cc..f8907a8 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -56,6 +56,10 @@ MemoryContext PortalContext = NULL;
 
 static void MemoryContextStatsInternal(MemoryContext context, int level);
 
+/* wrapper routines for allocation */
+static inline void* palloc_internal(Size size, bool noerror);
+static inline void* repalloc_internal(void *pointer, Size size, bool noerror);
+
 /*
  * You should not do memory allocations within a critical section, because
  * an out-of-memory error will be escalated to a PANIC. To enforce that
@@ -702,8 +706,8 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	return ret;
 }
 
-void *
-palloc(Size size)
+static inline void *
+palloc_internal(Size size, bool noerror)
 {
 	/* duplicates MemoryContextAlloc to avoid increased overhead */
 	void	   *ret;
@@ -717,42 +721,90 @@ palloc(Size size)
 	CurrentMemoryContext->isReset = false;
 
 	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
-	if (ret == NULL)
+	if (!noerror && ret == NULL)
 		ereport(ERROR,
 				(errcode(ERRCODE_OUT_OF_MEMORY),
 				 errmsg("out of memory"),
 				 errdetail("Failed on request of size %zu.", size)));
 
-	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
+	if (ret != NULL)
+		VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
 
 	return ret;
 }
 
-void *
-palloc0(Size size)
+static inline void *
+repalloc_internal(void *pointer, Size size, bool noerror)
 {
-	/* duplicates MemoryContextAllocZero to avoid increased overhead */
+	MemoryContext context;
 	void	   *ret;
 
-	AssertArg(MemoryContextIsValid(CurrentMemoryContext));
-	AssertNotInCriticalSection(CurrentMemoryContext);
-
 	if (!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
-	CurrentMemoryContext->isReset = false;
+	/*
+	 * 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.
+	 */
+	Assert(pointer != NULL);
+	Assert(pointer == (void *) MAXALIGN(pointer));
 
-	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
-	if (ret == NULL)
+	/*
+	 * OK, it's probably safe to look at the chunk header.
+	 */
+	context = ((StandardChunkHeader *)
+			   ((char *) pointer - STANDARDCHUNKHEADERSIZE))->context;
+
+	AssertArg(MemoryContextIsValid(context));
+	AssertNotInCriticalSection(context);
+
+	/* isReset must be false already */
+	Assert(!context->isReset);
+
+	ret = (*context->methods->realloc) (context, pointer, size);
+	if (!noerror && ret == NULL)
 		ereport(ERROR,
 				(errcode(ERRCODE_OUT_OF_MEMORY),
 				 errmsg("out of memory"),
 				 errdetail("Failed on request of size %zu.", size)));
 
-	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
+	if (ret != NULL)
+		VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+
+	return ret;
+
+}
+
+void *
+palloc(Size size)
+{
+	return palloc_internal(size, false);
+}
+
+void *
+palloc_noerror(Size size)
+{
+	return palloc_internal(size, true);
+}
+
+void *
+palloc0(Size size)
+{
+	void	   *ret;
 
+	ret = palloc_internal(size, false);
 	MemSetAligned(ret, 0, size);
+	return ret;
+}
 
+void *
+palloc0_noerror(Size size)
+{
+	void	   *ret;
+
+	ret = palloc_internal(size, true);
+	MemSetAligned(ret, 0, size);
 	return ret;
 }
 
@@ -787,47 +839,24 @@ pfree(void *pointer)
 
 /*
  * repalloc
- *		Adjust the size of a previously allocated chunk.
+ *		Adjust the size of a previously allocated chunk, failing if
+ *		request could not be completed.
  */
 void *
 repalloc(void *pointer, Size size)
 {
-	MemoryContext context;
-	void	   *ret;
-
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
-	/*
-	 * 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.
-	 */
-	Assert(pointer != NULL);
-	Assert(pointer == (void *) MAXALIGN(pointer));
-
-	/*
-	 * OK, it's probably safe to look at the chunk header.
-	 */
-	context = ((StandardChunkHeader *)
-			   ((char *) pointer - STANDARDCHUNKHEADERSIZE))->context;
-
-	AssertArg(MemoryContextIsValid(context));
-	AssertNotInCriticalSection(context);
-
-	/* isReset must be false already */
-	Assert(!context->isReset);
-
-	ret = (*context->methods->realloc) (context, pointer, size);
-	if (ret == NULL)
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu.", size)));
-
-	VALGRIND_MEMPOOL_CHANGE(context, pointer, ret, size);
+	return repalloc_internal(pointer, size, false);
+}
 
-	return ret;
+/*
+ * repalloc_noerror
+ *		Adjust the size of a previously allocated chunk, returning NULL
+ *		if request could not be completed.
+ */
+void *
+repalloc_noerror(void *pointer, Size size)
+{
+	return repalloc_internal(pointer, size, true);
 }
 
 /*
diff --git a/src/common/fe_memutils.c b/src/common/fe_memutils.c
index 345221e..9a312c4 100644
--- a/src/common/fe_memutils.c
+++ b/src/common/fe_memutils.c
@@ -19,8 +19,8 @@
 
 #include "postgres_fe.h"
 
-void *
-pg_malloc(size_t size)
+static inline void *
+pg_malloc_internal(size_t size, bool noerror)
 {
 	void	   *tmp;
 
@@ -28,7 +28,24 @@ pg_malloc(size_t size)
 	if (size == 0)
 		size = 1;
 	tmp = malloc(size);
-	if (!tmp)
+	if (!tmp && !noerror)
+	{
+		fprintf(stderr, _("out of memory\n"));
+		exit(EXIT_FAILURE);
+	}
+	return tmp;
+}
+
+static inline void *
+pg_realloc_internal(void *ptr, size_t size, bool noerror)
+{
+	void	   *tmp;
+
+	/* Avoid unportable behavior of realloc(NULL, 0) */
+	if (ptr == NULL && size == 0)
+		size = 1;
+	tmp = realloc(ptr, size);
+	if (!tmp && !noerror)
 	{
 		fprintf(stderr, _("out of memory\n"));
 		exit(EXIT_FAILURE);
@@ -37,6 +54,12 @@ pg_malloc(size_t size)
 }
 
 void *
+pg_malloc(size_t size)
+{
+	return pg_malloc_internal(size, false);
+}
+
+void *
 pg_malloc0(size_t size)
 {
 	void	   *tmp;
@@ -49,20 +72,10 @@ pg_malloc0(size_t size)
 void *
 pg_realloc(void *ptr, size_t size)
 {
-	void	   *tmp;
-
-	/* Avoid unportable behavior of realloc(NULL, 0) */
-	if (ptr == NULL && size == 0)
-		size = 1;
-	tmp = realloc(ptr, size);
-	if (!tmp)
-	{
-		fprintf(stderr, _("out of memory\n"));
-		exit(EXIT_FAILURE);
-	}
-	return tmp;
+	return pg_realloc_internal(ptr, size, false);
 }
 
+
 /*
  * "Safe" wrapper around strdup().
  */
@@ -100,13 +113,31 @@ pg_free(void *ptr)
 void *
 palloc(Size size)
 {
-	return pg_malloc(size);
+	return pg_malloc_internal(size, false);
+}
+
+void *
+palloc_noerror(Size size)
+{
+	return pg_malloc_internal(size, true);
 }
 
 void *
 palloc0(Size size)
 {
-	return pg_malloc0(size);
+	void	   *tmp;
+	tmp = pg_malloc_internal(size, false);
+	MemSet(tmp, 0, size);
+	return tmp;
+}
+
+void *
+palloc0_noerror(Size size)
+{
+	void	   *tmp;
+	tmp = pg_malloc_internal(size, true);
+	MemSet(tmp, 0, size);
+	return tmp;
 }
 
 void
@@ -124,5 +155,11 @@ pstrdup(const char *in)
 void *
 repalloc(void *pointer, Size size)
 {
-	return pg_realloc(pointer, size);
+	return pg_realloc_internal(pointer, size, false);
+}
+
+void *
+repalloc_noerror(void *pointer, Size size)
+{
+	return pg_realloc_internal(pointer, size, true);
 }
diff --git a/src/include/common/fe_memutils.h b/src/include/common/fe_memutils.h
index f7114c2..0151e6e 100644
--- a/src/include/common/fe_memutils.h
+++ b/src/include/common/fe_memutils.h
@@ -19,8 +19,11 @@ extern void pg_free(void *pointer);
 /* Equivalent functions, deliberately named the same as backend functions */
 extern char *pstrdup(const char *in);
 extern void *palloc(Size size);
+extern void *palloc_noerror(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_noerror(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_noerror(void *pointer, Size size);
 extern void pfree(void *pointer);
 
 /* sprintf into a palloc'd buffer --- these are in psprintf.c */
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index ca03f2b..3634a7f 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -50,8 +50,11 @@ extern void *MemoryContextAllocZero(MemoryContext context, Size size);
 extern void *MemoryContextAllocZeroAligned(MemoryContext context, Size size);
 
 extern void *palloc(Size size);
+extern void *palloc_noerror(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_noerror(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_noerror(void *pointer, Size size);
 extern void pfree(void *pointer);
 
 /*
-- 
2.2.2

From 1ff0406b11a6e54cdb0a4312c199dbe7800fea82 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Wed, 28 Jan 2015 23:27:54 +0900
Subject: [PATCH 3/3] Create MemoryContextAllocExtended central routine for
 memory allocation

This new routine is the central point of all the MemoryContextAlloc*
functions controlled by a set of flags allowing a far wider control
of how allocation can be done in a memory context.
---
 src/backend/utils/mmgr/mcxt.c | 107 ++++++++++++------------------------------
 src/include/utils/palloc.h    |  11 +++++
 2 files changed, 41 insertions(+), 77 deletions(-)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index f8907a8..01201fb 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -607,21 +607,20 @@ MemoryContextCreate(NodeTag tag, Size size,
 }
 
 /*
- * MemoryContextAlloc
- *		Allocate space within the specified context.
- *
- * This could be turned into a macro, but we'd have to import
- * nodes/memnodes.h into postgres.h which seems a bad idea.
+ * MemoryContextAllocExtended
+ *		Allocation space within the specified context using flag options
+ *		specified by caller.
  */
 void *
-MemoryContextAlloc(MemoryContext context, Size size)
+MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 {
 	void	   *ret;
 
 	AssertArg(MemoryContextIsValid(context));
 	AssertNotInCriticalSection(context);
 
-	if (!AllocSizeIsValid(size))
+	if (((flags & ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		!AllocSizeIsValid(size))
 		elog(ERROR, "invalid memory alloc request size %zu", size);
 
 	context->isReset = false;
@@ -635,75 +634,48 @@ MemoryContextAlloc(MemoryContext context, Size size)
 
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
+	/*
+	 * MemSetAligned and MemSetLoop should not be called in parallel
+	 * (see c.h for more details).
+	 */
+	Assert((flags & ALLOC_ALIGNED) == 0 || (flags & ALLOC_HUGE) == 0);
+
+	if ((flags & ALLOC_ZERO) != 0)
+		MemSetAligned(ret, 0, size);
+	if ((flags & ALLOC_ALIGNED) != 0)
+		MemSetLoop(ret, 0, size);
+
 	return ret;
 }
 
 /*
+ * MemoryContextAlloc
+ *		Allocate space within the specified context.
+ */
+void *
+MemoryContextAlloc(MemoryContext context, Size size)
+{
+	return MemoryContextAllocExtended(context, size, 0);
+}
+
+/*
  * MemoryContextAllocZero
  *		Like MemoryContextAlloc, but clears allocated memory
- *
- *	We could just call MemoryContextAlloc then clear the memory, but this
- *	is a very common combination, so we provide the combined operation.
  */
 void *
 MemoryContextAllocZero(MemoryContext context, Size size)
 {
-	void	   *ret;
-
-	AssertArg(MemoryContextIsValid(context));
-	AssertNotInCriticalSection(context);
-
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
-	context->isReset = false;
-
-	ret = (*context->methods->alloc) (context, size);
-	if (ret == NULL)
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu.", size)));
-
-	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
-
-	MemSetAligned(ret, 0, size);
-
-	return ret;
+	return MemoryContextAllocExtended(context, size, ALLOC_ZERO);
 }
 
 /*
  * MemoryContextAllocZeroAligned
  *		MemoryContextAllocZero where length is suitable for MemSetLoop
- *
- *	This might seem overly specialized, but it's not because newNode()
- *	is so often called with compile-time-constant sizes.
  */
 void *
 MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 {
-	void	   *ret;
-
-	AssertArg(MemoryContextIsValid(context));
-	AssertNotInCriticalSection(context);
-
-	if (!AllocSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
-	context->isReset = false;
-
-	ret = (*context->methods->alloc) (context, size);
-	if (ret == NULL)
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu.", size)));
-
-	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
-
-	MemSetLoop(ret, 0, size);
-
-	return ret;
+	return MemoryContextAllocExtended(context, size, ALLOC_ALIGNED);
 }
 
 static inline void *
@@ -868,26 +840,7 @@ repalloc_noerror(void *pointer, Size size)
 void *
 MemoryContextAllocHuge(MemoryContext context, Size size)
 {
-	void	   *ret;
-
-	AssertArg(MemoryContextIsValid(context));
-	AssertNotInCriticalSection(context);
-
-	if (!AllocHugeSizeIsValid(size))
-		elog(ERROR, "invalid memory alloc request size %zu", size);
-
-	context->isReset = false;
-
-	ret = (*context->methods->alloc) (context, size);
-	if (ret == NULL)
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory"),
-				 errdetail("Failed on request of size %zu.", size)));
-
-	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
-
-	return ret;
+	return MemoryContextAllocExtended(context, size, ALLOC_HUGE);
 }
 
 /*
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index 3634a7f..a8bf9a1 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -43,8 +43,19 @@ typedef struct MemoryContextData *MemoryContext;
 extern PGDLLIMPORT MemoryContext CurrentMemoryContext;
 
 /*
+ * Additional options for MemoryContextAllocExtended()
+ */
+#define ALLOC_HUGE		0x01	/* huge allocation */
+#define ALLOC_ZERO		0x02	/* clear allocated memory */
+#define ALLOC_NO_OOM	0x04	/* no failure if out-of-memory */
+#define ALLOC_ALIGNED	0x08	/* request length suitable for MemSetLoop */
+
+/*
  * Fundamental memory-allocation operations (more are in utils/memutils.h)
  */
+extern void *MemoryContextAllocExtended(MemoryContext context,
+										Size size,
+										int flags);
 extern void *MemoryContextAlloc(MemoryContext context, Size size);
 extern void *MemoryContextAllocZero(MemoryContext context, Size size);
 extern void *MemoryContextAllocZeroAligned(MemoryContext context, Size size);
-- 
2.2.2

From b5d43fd598e177c402c354f0b76aca52305463c6 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Fri, 30 Jan 2015 12:56:21 +0900
Subject: [PATCH] Create MemoryContextAllocExtended central routine for memory
 allocation

This new routine is the central point can be used by extensions and
third-part utilities in a more extensive way than the already present
routines MemoryContextAlloc, one of the control flags introduced being
particularly useful to avoid out-of-memory errors when allocation request
cannot be completed correctly.
---
 src/backend/utils/mmgr/mcxt.c | 57 +++++++++++++++++++++++++++++++++++++++++++
 src/include/utils/palloc.h    | 12 +++++++++
 2 files changed, 69 insertions(+)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index c62922a..26579e3 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -603,6 +603,63 @@ MemoryContextCreate(NodeTag tag, Size size,
 }
 
 /*
+ * MemoryContextAllocExtended
+ *		Allocate space within the specified context using flag options
+ *		defined by caller.
+ *
+ * The following control flags can be used:
+ * - MCXT_ALLOC_HUGE, allocate possibly-expansive space. this is
+ *   equivalent to MemoryContextAllocHuge.
+ * - MCXT_ALLOC_NO_OOM, not fail in case of allocation request
+ *   failure and return NULL.
+ * - MCXT_ALLOC_ZERO, clear allocated memory using MemSetAligned.
+ * - MCXT_ALLOC_ZERO_ALIGNED, clear memory using MemSetLoop.
+ */
+void *
+MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
+{
+	void	   *ret;
+
+	AssertArg(MemoryContextIsValid(context));
+	AssertNotInCriticalSection(context);
+
+	if (((flags & MCXT_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		!AllocSizeIsValid(size))
+		elog(ERROR, "invalid memory alloc request size %zu", size);
+
+	context->isReset = false;
+
+	ret = (*context->methods->alloc) (context, size);
+	if ((flags & MCXT_ALLOC_NO_OOM) == 0 && ret == NULL)
+	{
+		MemoryContextStats(TopMemoryContext);
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+	}
+
+	if (ret == NULL)
+		return NULL;
+
+	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
+
+	/*
+	 * MemSetAligned and MemSetLoop should not be called in the same
+	 * context (see c.h for more details).
+	 */
+	Assert((flags & MCXT_ALLOC_ZERO) == 0 ||
+		   (flags & MCXT_ALLOC_ZERO_ALIGNED) == 0);
+
+	if ((flags & MCXT_ALLOC_ZERO) != 0)
+		MemSetAligned(ret, 0, size);
+	if ((flags & MCXT_ALLOC_ZERO_ALIGNED) != 0)
+		MemSetLoop(ret, 0, size);
+
+	return ret;
+}
+
+/*
  * MemoryContextAlloc
  *		Allocate space within the specified context.
  *
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index ca03f2b..34acabe 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -43,8 +43,20 @@ typedef struct MemoryContextData *MemoryContext;
 extern PGDLLIMPORT MemoryContext CurrentMemoryContext;
 
 /*
+ * Control flags for options of MemoryContextAllocExtended()
+ */
+#define MCXT_ALLOC_HUGE			0x01	/* huge allocation */
+#define MCXT_ALLOC_NO_OOM		0x02	/* no failure if out-of-memory */
+#define MCXT_ALLOC_ZERO			0x04	/* clear allocated memory using
+										 * MemSetAligned */
+#define MCXT_ALLOC_ZERO_ALIGNED	0x08	/* clear allocated memory using
+										 * MemSetLoop */
+
+/*
  * Fundamental memory-allocation operations (more are in utils/memutils.h)
  */
+extern void *MemoryContextAllocExtended(MemoryContext context,
+										Size size, int flags);
 extern void *MemoryContextAlloc(MemoryContext context, Size size);
 extern void *MemoryContextAllocZero(MemoryContext context, Size size);
 extern void *MemoryContextAllocZeroAligned(MemoryContext context, Size size);
-- 
2.2.2

From 09bf9364a80dfe6426c91bdefe6ae2135e6f20c3 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Fri, 30 Jan 2015 12:56:21 +0900
Subject: [PATCH 1/2] Create MemoryContextAllocExtended central routine for
 memory allocation

This new routine is the central point can be used by extensions and
third-part utilities in a more extensive way than the already present
routines MemoryContextAlloc, one of the control flags introduced being
particularly useful to avoid out-of-memory errors when allocation request
cannot be completed correctly.
---
 src/backend/utils/mmgr/mcxt.c | 57 +++++++++++++++++++++++++++++++++++++++++++
 src/include/utils/palloc.h    | 12 +++++++++
 2 files changed, 69 insertions(+)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index c62922a..fce8a0e 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -603,6 +603,63 @@ MemoryContextCreate(NodeTag tag, Size size,
 }
 
 /*
+ * MemoryContextAllocExtended
+ *		Allocate space within the specified context using flag options
+ *		defined by caller.
+ *
+ * The following control flags can be used:
+ * - MCXT_ALLOC_HUGE, allocate possibly-expansive space. this is
+ *   equivalent to MemoryContextAllocHuge.
+ * - MCXT_ALLOC_NO_OOM, not fail in case of allocation request
+ *   failure and return NULL.
+ * - MCXT_ALLOC_ZERO, clear allocated memory using MemSetAligned.
+ * - MCXT_ALLOC_ZERO_ALIGNED, clear memory using MemSetLoop.
+ */
+void *
+MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
+{
+	void	   *ret;
+
+	AssertArg(MemoryContextIsValid(context));
+	AssertNotInCriticalSection(context);
+
+	if (((flags & MCXT_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		((flags & MCXT_ALLOC_HUGE) == 0 && !AllocSizeIsValid(size)))
+		elog(ERROR, "invalid memory alloc request size %zu", size);
+
+	context->isReset = false;
+
+	ret = (*context->methods->alloc) (context, size);
+	if ((flags & MCXT_ALLOC_NO_OOM) == 0 && ret == NULL)
+	{
+		MemoryContextStats(TopMemoryContext);
+		ereport(ERROR,
+				(errcode(ERRCODE_OUT_OF_MEMORY),
+				 errmsg("out of memory"),
+				 errdetail("Failed on request of size %zu.", size)));
+	}
+
+	if (ret == NULL)
+		return NULL;
+
+	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
+
+	/*
+	 * MemSetAligned and MemSetLoop should not be called in the same
+	 * context (see c.h for more details).
+	 */
+	Assert((flags & MCXT_ALLOC_ZERO) == 0 ||
+		   (flags & MCXT_ALLOC_ZERO_ALIGNED) == 0);
+
+	if ((flags & MCXT_ALLOC_ZERO) != 0)
+		MemSetAligned(ret, 0, size);
+	if ((flags & MCXT_ALLOC_ZERO_ALIGNED) != 0)
+		MemSetLoop(ret, 0, size);
+
+	return ret;
+}
+
+/*
  * MemoryContextAlloc
  *		Allocate space within the specified context.
  *
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index ca03f2b..34acabe 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -43,8 +43,20 @@ typedef struct MemoryContextData *MemoryContext;
 extern PGDLLIMPORT MemoryContext CurrentMemoryContext;
 
 /*
+ * Control flags for options of MemoryContextAllocExtended()
+ */
+#define MCXT_ALLOC_HUGE			0x01	/* huge allocation */
+#define MCXT_ALLOC_NO_OOM		0x02	/* no failure if out-of-memory */
+#define MCXT_ALLOC_ZERO			0x04	/* clear allocated memory using
+										 * MemSetAligned */
+#define MCXT_ALLOC_ZERO_ALIGNED	0x08	/* clear allocated memory using
+										 * MemSetLoop */
+
+/*
  * Fundamental memory-allocation operations (more are in utils/memutils.h)
  */
+extern void *MemoryContextAllocExtended(MemoryContext context,
+										Size size, int flags);
 extern void *MemoryContextAlloc(MemoryContext context, Size size);
 extern void *MemoryContextAllocZero(MemoryContext context, Size size);
 extern void *MemoryContextAllocZeroAligned(MemoryContext context, Size size);
-- 
2.2.2

From 88f3fb0129a6fad35114902320b006183ea6da5f Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Fri, 30 Jan 2015 13:14:23 +0900
Subject: [PATCH 2/2] Small hack to test extended function for mcxt

---
 src/backend/utils/mmgr/mcxt.c | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index fce8a0e..efa81fc 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -669,6 +669,8 @@ MemoryContextAllocExtended(MemoryContext context, Size size, int flags)
 void *
 MemoryContextAlloc(MemoryContext context, Size size)
 {
+	return MemoryContextAllocExtended(context, size, 0);
+#ifdef NOT_USED
 	void	   *ret;
 
 	AssertArg(MemoryContextIsValid(context));
@@ -692,6 +694,7 @@ MemoryContextAlloc(MemoryContext context, Size size)
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
+#endif
 }
 
 /*
@@ -704,6 +707,8 @@ MemoryContextAlloc(MemoryContext context, Size size)
 void *
 MemoryContextAllocZero(MemoryContext context, Size size)
 {
+	return MemoryContextAllocExtended(context, size, MCXT_ALLOC_ZERO);
+#ifdef NOT_USED
 	void	   *ret;
 
 	AssertArg(MemoryContextIsValid(context));
@@ -729,6 +734,7 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 	MemSetAligned(ret, 0, size);
 
 	return ret;
+#endif
 }
 
 /*
@@ -741,6 +747,8 @@ MemoryContextAllocZero(MemoryContext context, Size size)
 void *
 MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 {
+	return MemoryContextAllocExtended(context, size, MCXT_ALLOC_ZERO_ALIGNED);
+#ifdef NOT_USED
 	void	   *ret;
 
 	AssertArg(MemoryContextIsValid(context));
@@ -766,6 +774,7 @@ MemoryContextAllocZeroAligned(MemoryContext context, Size size)
 	MemSetLoop(ret, 0, size);
 
 	return ret;
+#endif
 }
 
 void *
@@ -911,6 +920,8 @@ repalloc(void *pointer, Size size)
 void *
 MemoryContextAllocHuge(MemoryContext context, Size size)
 {
+	return MemoryContextAllocExtended(context, size, MCXT_ALLOC_HUGE);
+#ifdef NOT_USED
 	void	   *ret;
 
 	AssertArg(MemoryContextIsValid(context));
@@ -934,6 +945,7 @@ MemoryContextAllocHuge(MemoryContext context, Size size)
 	VALGRIND_MEMPOOL_ALLOC(context, ret, size);
 
 	return ret;
+#endif
 }
 
 /*
-- 
2.2.2