Add palloc_extended for frontend and backend.

This commit also adds pg_malloc_extended for frontend. These interfaces
can be used to control at a lower level memory allocation using an interface
similar to MemoryContextAllocExtended. For example, the callers can specify
MCXT_ALLOC_NO_OOM if they want to suppress the "out of memory" error while
allocating the memory and handle a NULL return value.

Michael Paquier, reviewed by me.

4 files changed, 85 insertions, 12 deletions

diff --git a/src/backend/utils/mmgr/mcxt.c b/src/backend/utils/mmgr/mcxt.c
index e2fbfd420b6..c42a6b60056 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 345221e1317..527f9c8a6bc 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 db7cb3e3318..51f12eb8251 100644
--- a/src/include/common/fe_memutils.h
+++ b/src/include/common/fe_memutils.h
@@ -9,10 +9,23 @@
 #ifndef FE_MEMUTILS_H
 #define FE_MEMUTILS_H
 
-/* "Safe" memory allocation functions --- these exit(1) on failure */
+/*
+ * 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
+ * (except pg_malloc_extended with MCXT_ALLOC_NO_OOM)
+ */
 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 +33,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 2cf51298331..9861f0dac73 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);