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