From f43fc9a5a5e0ffb246782010b8860829b8304593 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@otacoo.com>
Date: Fri, 3 Apr 2015 14:21:12 +0900
Subject: [PATCH 1/2] Add palloc_extended and pg_malloc_extended for frontend
 and backend

This interface can be used to control at a lower level memory allocation
using an interface similar to MemoryContextAllocExtended, but this time
applied to CurrentMemoryContext on backend side, while on frontend side
a similar interface is available.
---
 src/backend/utils/mmgr/mcxt.c    | 37 ++++++++++++++++++++++++++++++++++
 src/common/fe_memutils.c         | 43 ++++++++++++++++++++++++++++++----------
 src/include/common/fe_memutils.h | 11 ++++++++++
 src/include/utils/palloc.h       |  1 +
 4 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index e2fbfd4..c42a6b6 100644
--- a/src/backend/utils/mmgr/mcxt.c
+++ b/src/backend/utils/mmgr/mcxt.c
@@ -864,6 +864,43 @@ palloc0(Size size)
 	return ret;
 }
 
+void *
+palloc_extended(Size size, int flags)
+{
+	/* duplicates MemoryContextAllocExtended to avoid increased overhead */
+	void	   *ret;
+
+	AssertArg(MemoryContextIsValid(CurrentMemoryContext));
+	AssertNotInCriticalSection(CurrentMemoryContext);
+
+	if (((flags & MCXT_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		((flags & MCXT_ALLOC_HUGE) == 0 && !AllocSizeIsValid(size)))
+		elog(ERROR, "invalid memory alloc request size %zu", size);
+
+	CurrentMemoryContext->isReset = false;
+
+	ret = (*CurrentMemoryContext->methods->alloc) (CurrentMemoryContext, size);
+	if (ret == NULL)
+	{
+		if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+		{
+			MemoryContextStats(TopMemoryContext);
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory"),
+					 errdetail("Failed on request of size %zu.", size)));
+		}
+		return NULL;
+	}
+
+	VALGRIND_MEMPOOL_ALLOC(CurrentMemoryContext, ret, size);
+
+	if ((flags & MCXT_ALLOC_ZERO) != 0)
+		MemSetAligned(ret, 0, size);
+
+	return ret;
+}
+
 /*
  * pfree
  *		Release an allocated chunk.
diff --git a/src/common/fe_memutils.c b/src/common/fe_memutils.c
index 345221e..527f9c8 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, int flags)
 {
 	void	   *tmp;
 
@@ -28,22 +28,37 @@ pg_malloc(size_t size)
 	if (size == 0)
 		size = 1;
 	tmp = malloc(size);
-	if (!tmp)
+	if (tmp == NULL)
 	{
-		fprintf(stderr, _("out of memory\n"));
-		exit(EXIT_FAILURE);
+		if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+		{
+			fprintf(stderr, _("out of memory\n"));
+			exit(EXIT_FAILURE);
+		}
+		return NULL;
 	}
+
+	if ((flags & MCXT_ALLOC_ZERO) != 0)
+		MemSet(tmp, 0, size);
 	return tmp;
 }
 
 void *
+pg_malloc(size_t size)
+{
+	return pg_malloc_internal(size, 0);
+}
+
+void *
 pg_malloc0(size_t size)
 {
-	void	   *tmp;
+	return pg_malloc_internal(size, MCXT_ALLOC_ZERO);
+}
 
-	tmp = pg_malloc(size);
-	MemSet(tmp, 0, size);
-	return tmp;
+void *
+pg_malloc_extended(size_t size, int flags)
+{
+	return pg_malloc_internal(size, flags);
 }
 
 void *
@@ -100,13 +115,19 @@ pg_free(void *ptr)
 void *
 palloc(Size size)
 {
-	return pg_malloc(size);
+	return pg_malloc_internal(size, 0);
 }
 
 void *
 palloc0(Size size)
 {
-	return pg_malloc0(size);
+	return pg_malloc_internal(size, MCXT_ALLOC_ZERO);
+}
+
+void *
+palloc_extended(Size size, int flags)
+{
+	return pg_malloc_internal(size, flags);
 }
 
 void
diff --git a/src/include/common/fe_memutils.h b/src/include/common/fe_memutils.h
index db7cb3e..6466167 100644
--- a/src/include/common/fe_memutils.h
+++ b/src/include/common/fe_memutils.h
@@ -9,10 +9,20 @@
 #ifndef FE_MEMUTILS_H
 #define FE_MEMUTILS_H
 
+/*
+ * Flags for pg_malloc_extended and palloc_extended, deliberately named
+ * the same as the backend flags.
+ */
+#define MCXT_ALLOC_HUGE			0x01	/* allow huge allocation (> 1 GB)
+										 * not actually used for frontends */
+#define MCXT_ALLOC_NO_OOM		0x02	/* no failure if out-of-memory */
+#define MCXT_ALLOC_ZERO			0x04	/* zero allocated memory */
+
 /* "Safe" memory allocation functions --- these exit(1) on failure */
 extern char *pg_strdup(const char *in);
 extern void *pg_malloc(size_t size);
 extern void *pg_malloc0(size_t size);
+extern void *pg_malloc_extended(size_t size, int flags);
 extern void *pg_realloc(void *pointer, size_t size);
 extern void pg_free(void *pointer);
 
@@ -20,6 +30,7 @@ extern void pg_free(void *pointer);
 extern char *pstrdup(const char *in);
 extern void *palloc(Size size);
 extern void *palloc0(Size size);
+extern void *palloc_extended(Size size, int flags);
 extern void *repalloc(void *pointer, Size size);
 extern void pfree(void *pointer);
 
diff --git a/src/include/utils/palloc.h b/src/include/utils/palloc.h
index 2cf5129..9861f0d 100644
--- a/src/include/utils/palloc.h
+++ b/src/include/utils/palloc.h
@@ -76,6 +76,7 @@ extern void *MemoryContextAllocExtended(MemoryContext context,
 
 extern void *palloc(Size size);
 extern void *palloc0(Size size);
+extern void *palloc_extended(Size size, int flags);
 extern void *repalloc(void *pointer, Size size);
 extern void pfree(void *pointer);
 
-- 
2.3.5