Hi all,

For the last couple of weeks it has been mentioned a couple of times
that it would be useful to have a set of palloc APIs able to return
NULL on OOM to allow certain code paths to not ERROR and to take
another route when memory is under pressure. This has been for example
mentioned on the FPW compression thread or here:
http://www.postgresql.org/message-id/cab7npqrbewhsbj_tkaogtpcmrxyjsvkkb9p030d0tpijb4t...@mail.gmail.com

Attached is a patch adding the following set of functions for frontend
and backends returning NULL instead of reporting ERROR when allocation
fails:
- palloc_safe
- palloc0_safe
- repalloc_safe
This has simply needed some refactoring in aset.c to set up the new
functions by passing an additional control flag, and I didn't think
that adding a new safe version for AllocSetContextCreate was worth it.
Those APIs are not called anywhere yet, but I could for example write
a small extension for that that could be put in src/test/modules or
publish on github in my plugin repo. Also, I am not sure if this is
material for 9.5, even if the patch is not complicated, but let me
know if you are interested in it and I'll add it to the next CF.
Regards,
-- 
Michael
From 008f6bf5a1691fbdf59004157ed521d3b8c41eaf Mon Sep 17 00:00:00 2001
From: Michael Paquier <mich...@otacoo.com>
Date: Tue, 13 Jan 2015 15:40:38 +0900
Subject: [PATCH] Add safe memory allocation APIs able to return NULL on OOM

The following functions are added to the existing set for frontend and
backend:
- palloc_safe
- palloc0_safe
- repalloc_safe
---
 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..5911c53 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 is_safe);
+static void *set_realloc_internal(MemoryContext context, void *pointer,
+								  Size size, bool is_safe);
+
+/*
  * These functions implement the MemoryContext API for AllocSet contexts.
  */
 static void *AllocSetAlloc(MemoryContext context, Size size);
+static void *AllocSetAllocSafe(MemoryContext context, Size size);
 static void AllocSetFree(MemoryContext context, void *pointer);
 static void *AllocSetRealloc(MemoryContext context, void *pointer, Size size);
+static void *AllocSetReallocSafe(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,
+	AllocSetAllocSafe,
 	AllocSetFree,
 	AllocSetRealloc,
+	AllocSetReallocSafe,
 	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 is_safe)
 {
 	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 (is_safe)
+				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 (is_safe)
+				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 is_safe)
 {
 	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 (is_safe)
+				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, is_safe);
+
+		/* leave if allocation did not complete properly */
+		if (newPointer == NULL)
+		{
+			Assert(is_safe);
+			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);
+}
+
+/*
+ * AllocSetAllocSafe
+ *		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 *
+AllocSetAllocSafe(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);
+}
+
+/*
+ * AllocSetReallocSafe
+ *		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 *
+AllocSetReallocSafe(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..7484770 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 is_safe);
+static void* repalloc_internal(void *pointer, Size size, bool is_safe);
+
 /*
  * 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 is_safe)
 {
 	/* 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 (is_safe)
+		ret = (*CurrentMemoryContext->methods->alloc_safe)
+			(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 is_safe)
 {
-	/* 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 (is_safe)
+		ret = (*context->methods->realloc_safe) (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_safe(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_safe(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_safe(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..3bba799 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 is_safe)
 {
 	void	   *tmp;
 
@@ -28,7 +28,24 @@ pg_malloc(size_t size)
 	if (size == 0)
 		size = 1;
 	tmp = malloc(size);
-	if (!tmp)
+	if (!tmp && !is_safe)
+	{
+		fprintf(stderr, _("out of memory\n"));
+		exit(EXIT_FAILURE);
+	}
+	return tmp;
+}
+
+static void *
+pg_realloc_internal(void *ptr, size_t size, bool is_safe)
+{
+	void	   *tmp;
+
+	/* Avoid unportable behavior of realloc(NULL, 0) */
+	if (ptr == NULL && size == 0)
+		size = 1;
+	tmp = realloc(ptr, size);
+	if (!tmp && !is_safe)
 	{
 		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_safe(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_safe(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_safe(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..977e35d 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_safe(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_safe(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_safe(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..df7de1e 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_safe) (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_safe) (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..96cfa77 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_safe(Size size);
 extern void *palloc0(Size size);
+extern void *palloc0_safe(Size size);
 extern void *repalloc(void *pointer, Size size);
+extern void *repalloc_safe(void *pointer, Size size);
 extern void pfree(void *pointer);
 
 /*
-- 
2.2.1

-- 
Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org)
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers

Reply via email to