pgsql: Add new function dsa_allocate0.

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d..5026bc9 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -762,11 +762,10 @@ dsa_pointer
 dsa_allocate0(dsa_area *area, Size size)
 {
 	dsa_pointer dp;
-	char	   *object;
 
 	dp = dsa_allocate(area, size);
-	object = dsa_get_address(area, dp);
-	memset(object, 0, size);
+	if (DsaPointerIsValid(dp))
+		memset(dsa_get_address(area, dp), 0, size);
 
 	return dp;
 }

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d..8ffef37 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -644,13 +644,17 @@ dsa_pin_mapping(dsa_area *area)
  * that can be passed to other processes, and converted to a local pointer
  * with dsa_get_address.  If no memory is available, returns
  * InvalidDsaPointer.
+ *
+ * See also the macros dsa_allocate and dsa_allocate0 which expand to a call
+ * to this function with commonly used flags.
  */
 dsa_pointer
-dsa_allocate(dsa_area *area, Size size)
+dsa_allocate_extended(dsa_area *area, Size size, int flags)
 {
 	uint16		size_class;
 	dsa_pointer start_pointer;
 	dsa_segment_map *segment_map;
+	dsa_pointer result;
 
 	Assert(size > 0);
 
@@ -684,6 +688,14 @@ dsa_allocate(dsa_area *area, Size size)
 			/* Can't make any more segments: game over. */
 			LWLockRelease(DSA_AREA_LOCK(area));
 			dsa_free(area, span_pointer);
+
+			/* Raise error unless asked not to. */
+			if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on DSA request of size %zu.",
+								   size)));
 			return InvalidDsaPointer;
 		}
 
@@ -710,6 +722,10 @@ dsa_allocate(dsa_area *area, Size size)
 		segment_map->pagemap[first_page] = span_pointer;
 		LWLockRelease(DSA_SCLASS_LOCK(area, DSA_SCLASS_SPAN_LARGE));
 
+		/* Zero-initialize the memory if requested. */
+		if ((flags & DSA_ALLOC_ZERO) != 0)
+			memset(dsa_get_address(area, start_pointer), 0, size);
+
 		return start_pointer;
 	}
 
@@ -748,27 +764,28 @@ dsa_allocate(dsa_area *area, Size size)
 	Assert(size <= dsa_size_classes[size_class]);
 	Assert(size_class == 0 || size > dsa_size_classes[size_class - 1]);
 
-	/*
-	 * Attempt to allocate an object from the appropriate pool.  This might
-	 * return InvalidDsaPointer if there's no space available.
-	 */
-	return alloc_object(area, size_class);
-}
+	/* Attempt to allocate an object from the appropriate pool. */
+	result = alloc_object(area, size_class);
 
-/*
- * As dsa_allocate, but zeroes the allocated memory.
- */
-dsa_pointer
-dsa_allocate0(dsa_area *area, Size size)
-{
-	dsa_pointer dp;
-	char	   *object;
+	/* Check for failure to allocate. */
+	if (!DsaPointerIsValid(result))
+	{
+		/* Raise error unless asked not to. */
+		if ((flags & DSA_ALLOC_NO_OOM) == 0)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory"),
+					 errdetail("Failed on DSA request of size %zu.", size)));
+		}
+		return InvalidDsaPointer;
+	}
 
-	dp = dsa_allocate(area, size);
-	object = dsa_get_address(area, dp);
-	memset(object, 0, size);
+	/* Zero-initialize the memory if requested. */
+	if ((flags & DSA_ALLOC_ZERO) != 0)
+		memset(dsa_get_address(area, result), 0, size);
 
-	return dp;
+	return result;
 }
 
 /*
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 3fd9bbd..cccd4d0 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -71,12 +71,24 @@ typedef pg_atomic_uint64 dsa_pointer_atomic;
 #define DSA_POINTER_FORMAT "%016" INT64_MODIFIER "x"
 #endif
 
+/* Flags for dsa_allocate_extended. */
+#define DSA_ALLOC_NO_OOM	0x02	/* no failure if out-of-memory */
+#define DSA_ALLOC_ZERO		0x04	/* zero allocated memory */
+
 /* A sentinel value for dsa_pointer used to indicate failure to allocate. */
 #define InvalidDsaPointer ((dsa_pointer) 0)
 
 /* Check if a dsa_pointer value is valid. */
 #define DsaPointerIsValid(x) ((x) != InvalidDsaPointer)
 
+/* Allocate uninitialized memory with error on out-of-memory. */
+#define dsa_allocate(area, size) \
+	dsa_allocate_extended(area, size, 0)
+
+/* Allocate zero-initialized memory with error on out-of-memory. */
+#define dsa_allocate0(area, size) \
+	dsa_allocate_extended(area, size, DSA_ALLOC_ZERO)
+
 /*
  * The type used for dsa_area handles.  dsa_handle values can be shared with
  * other processes, so that they can attach to them.  This provides a way to
@@ -105,8 +117,7 @@ extern void dsa_unpin(dsa_area *area);
 extern void dsa_set_size_limit(dsa_area *area, Size limit);
 extern Size dsa_minimum_size(void);
 extern dsa_handle dsa_get_handle(dsa_area *area);
-extern dsa_pointer dsa_allocate(dsa_area *area, Size size);
-extern dsa_pointer dsa_allocate0(dsa_area *area, Size size);
+extern dsa_pointer dsa_allocate_extended(dsa_area *area, Size size, int flags);
 extern void dsa_free(dsa_area *area, dsa_pointer dp);
 extern void *dsa_get_address(dsa_area *area, dsa_pointer dp);
 extern void dsa_trim(dsa_area *area);

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d..f91f286 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -644,16 +644,32 @@ dsa_pin_mapping(dsa_area *area)
  * that can be passed to other processes, and converted to a local pointer
  * with dsa_get_address.  If no memory is available, returns
  * InvalidDsaPointer.
+ *
+ * 'flags' is a bitmap which should be constructed from the following values:
+ * DSA_ALLOC_HUGE to allow very large allocations, DSA_ALLOC_NO_OOM to return
+ * InvalidDsaPointer instead of raising an error on failrue to allocate due to
+ * lack of memory (including failure due to exhausting a size limit imposed by
+ * calling dsa_set_size_limit) and DSA_ALLOC_ZERO to request that the memory
+ * be set to all zeroes.  These flags correspond to similarly named flags used
+ * by MemoryContextAllocExtended().  See also the macros dsa_allocate and
+ * dsa_allocate0 which expand to a call to this function with commonly used
+ * flags.
  */
 dsa_pointer
-dsa_allocate(dsa_area *area, Size size)
+dsa_allocate_extended(dsa_area *area, Size size, int flags)
 {
 	uint16		size_class;
 	dsa_pointer start_pointer;
 	dsa_segment_map *segment_map;
+	dsa_pointer result;
 
 	Assert(size > 0);
 
+	/* Sanity check on huge individual allocation size. */
+	if (((flags & DSA_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		((flags & DSA_ALLOC_HUGE) == 0 && !AllocSizeIsValid(size)))
+		elog(ERROR, "invalid DSA memory alloc request size %zu", size);
+
 	/*
 	 * If bigger than the largest size class, just grab a run of pages from
 	 * the free page manager, instead of allocating an object from a pool.
@@ -684,6 +700,14 @@ dsa_allocate(dsa_area *area, Size size)
 			/* Can't make any more segments: game over. */
 			LWLockRelease(DSA_AREA_LOCK(area));
 			dsa_free(area, span_pointer);
+
+			/* Raise error unless asked not to. */
+			if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on DSA request of size %zu.",
+								   size)));
 			return InvalidDsaPointer;
 		}
 
@@ -710,6 +734,10 @@ dsa_allocate(dsa_area *area, Size size)
 		segment_map->pagemap[first_page] = span_pointer;
 		LWLockRelease(DSA_SCLASS_LOCK(area, DSA_SCLASS_SPAN_LARGE));
 
+		/* Zero-initialize the memory if requested. */
+		if ((flags & DSA_ALLOC_ZERO) != 0)
+			memset(dsa_get_address(area, start_pointer), 0, size);
+
 		return start_pointer;
 	}
 
@@ -748,27 +776,28 @@ dsa_allocate(dsa_area *area, Size size)
 	Assert(size <= dsa_size_classes[size_class]);
 	Assert(size_class == 0 || size > dsa_size_classes[size_class - 1]);
 
-	/*
-	 * Attempt to allocate an object from the appropriate pool.  This might
-	 * return InvalidDsaPointer if there's no space available.
-	 */
-	return alloc_object(area, size_class);
-}
+	/* Attempt to allocate an object from the appropriate pool. */
+	result = alloc_object(area, size_class);
 
-/*
- * As dsa_allocate, but zeroes the allocated memory.
- */
-dsa_pointer
-dsa_allocate0(dsa_area *area, Size size)
-{
-	dsa_pointer dp;
-	char	   *object;
+	/* Check for failure to allocate. */
+	if (!DsaPointerIsValid(result))
+	{
+		/* Raise error unless asked not to. */
+		if ((flags & DSA_ALLOC_NO_OOM) == 0)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory"),
+					 errdetail("Failed on DSA request of size %zu.", size)));
+		}
+		return InvalidDsaPointer;
+	}
 
-	dp = dsa_allocate(area, size);
-	object = dsa_get_address(area, dp);
-	memset(object, 0, size);
+	/* Zero-initialize the memory if requested. */
+	if ((flags & DSA_ALLOC_ZERO) != 0)
+		memset(dsa_get_address(area, result), 0, size);
 
-	return dp;
+	return result;
 }
 
 /*
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 3fd9bbd..c4b122e 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -71,12 +71,25 @@ typedef pg_atomic_uint64 dsa_pointer_atomic;
 #define DSA_POINTER_FORMAT "%016" INT64_MODIFIER "x"
 #endif
 
+/* Flags for dsa_allocate_extended. */
+#define DSA_ALLOC_HUGE		0x01	/* allow huge allocation (> 1 GB) */
+#define DSA_ALLOC_NO_OOM	0x02	/* no failure if out-of-memory */
+#define DSA_ALLOC_ZERO		0x04	/* zero allocated memory */
+
 /* A sentinel value for dsa_pointer used to indicate failure to allocate. */
 #define InvalidDsaPointer ((dsa_pointer) 0)
 
 /* Check if a dsa_pointer value is valid. */
 #define DsaPointerIsValid(x) ((x) != InvalidDsaPointer)
 
+/* Allocate uninitialized memory with error on out-of-memory. */
+#define dsa_allocate(area, size) \
+	dsa_allocate_extended(area, size, 0)
+
+/* Allocate zero-initialized memory with error on out-of-memory. */
+#define dsa_allocate0(area, size) \
+	dsa_allocate_extended(area, size, DSA_ALLOC_ZERO)
+
 /*
  * The type used for dsa_area handles.  dsa_handle values can be shared with
  * other processes, so that they can attach to them.  This provides a way to
@@ -105,8 +118,7 @@ extern void dsa_unpin(dsa_area *area);
 extern void dsa_set_size_limit(dsa_area *area, Size limit);
 extern Size dsa_minimum_size(void);
 extern dsa_handle dsa_get_handle(dsa_area *area);
-extern dsa_pointer dsa_allocate(dsa_area *area, Size size);
-extern dsa_pointer dsa_allocate0(dsa_area *area, Size size);
+extern dsa_pointer dsa_allocate_extended(dsa_area *area, Size size, int flags);
 extern void dsa_free(dsa_area *area, dsa_pointer dp);
 extern void *dsa_get_address(dsa_area *area, dsa_pointer dp);
 extern void dsa_trim(dsa_area *area);

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d..e42d6ea 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -644,16 +644,32 @@ dsa_pin_mapping(dsa_area *area)
  * that can be passed to other processes, and converted to a local pointer
  * with dsa_get_address.  If no memory is available, returns
  * InvalidDsaPointer.
+ *
+ * 'flags' is a bitmap which should be constructed from the following values:
+ * DSA_ALLOC_HUGE to allow very large allocations, DSA_ALLOC_NO_OOM to return
+ * InvalidDsaPointer instead of raising an error on failure to allocate due to
+ * lack of memory (including failure due to exhausting a size limit imposed by
+ * calling dsa_set_size_limit) and DSA_ALLOC_ZERO to request that the memory
+ * be set to all zeroes.  These flags correspond to similarly named flags used
+ * by MemoryContextAllocExtended().  See also the macros dsa_allocate and
+ * dsa_allocate0 which expand to a call to this function with commonly used
+ * flags.
  */
 dsa_pointer
-dsa_allocate(dsa_area *area, Size size)
+dsa_allocate_extended(dsa_area *area, Size size, int flags)
 {
 	uint16		size_class;
 	dsa_pointer start_pointer;
 	dsa_segment_map *segment_map;
+	dsa_pointer result;
 
 	Assert(size > 0);
 
+	/* Sanity check on huge individual allocation size. */
+	if (((flags & DSA_ALLOC_HUGE) != 0 && !AllocHugeSizeIsValid(size)) ||
+		((flags & DSA_ALLOC_HUGE) == 0 && !AllocSizeIsValid(size)))
+		elog(ERROR, "invalid DSA memory alloc request size %zu", size);
+
 	/*
 	 * If bigger than the largest size class, just grab a run of pages from
 	 * the free page manager, instead of allocating an object from a pool.
@@ -684,6 +700,14 @@ dsa_allocate(dsa_area *area, Size size)
 			/* Can't make any more segments: game over. */
 			LWLockRelease(DSA_AREA_LOCK(area));
 			dsa_free(area, span_pointer);
+
+			/* Raise error unless asked not to. */
+			if ((flags & MCXT_ALLOC_NO_OOM) == 0)
+				ereport(ERROR,
+						(errcode(ERRCODE_OUT_OF_MEMORY),
+						 errmsg("out of memory"),
+						 errdetail("Failed on DSA request of size %zu.",
+								   size)));
 			return InvalidDsaPointer;
 		}
 
@@ -710,6 +734,10 @@ dsa_allocate(dsa_area *area, Size size)
 		segment_map->pagemap[first_page] = span_pointer;
 		LWLockRelease(DSA_SCLASS_LOCK(area, DSA_SCLASS_SPAN_LARGE));
 
+		/* Zero-initialize the memory if requested. */
+		if ((flags & DSA_ALLOC_ZERO) != 0)
+			memset(dsa_get_address(area, start_pointer), 0, size);
+
 		return start_pointer;
 	}
 
@@ -748,27 +776,28 @@ dsa_allocate(dsa_area *area, Size size)
 	Assert(size <= dsa_size_classes[size_class]);
 	Assert(size_class == 0 || size > dsa_size_classes[size_class - 1]);
 
-	/*
-	 * Attempt to allocate an object from the appropriate pool.  This might
-	 * return InvalidDsaPointer if there's no space available.
-	 */
-	return alloc_object(area, size_class);
-}
+	/* Attempt to allocate an object from the appropriate pool. */
+	result = alloc_object(area, size_class);
 
-/*
- * As dsa_allocate, but zeroes the allocated memory.
- */
-dsa_pointer
-dsa_allocate0(dsa_area *area, Size size)
-{
-	dsa_pointer dp;
-	char	   *object;
+	/* Check for failure to allocate. */
+	if (!DsaPointerIsValid(result))
+	{
+		/* Raise error unless asked not to. */
+		if ((flags & DSA_ALLOC_NO_OOM) == 0)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_OUT_OF_MEMORY),
+					 errmsg("out of memory"),
+					 errdetail("Failed on DSA request of size %zu.", size)));
+		}
+		return InvalidDsaPointer;
+	}
 
-	dp = dsa_allocate(area, size);
-	object = dsa_get_address(area, dp);
-	memset(object, 0, size);
+	/* Zero-initialize the memory if requested. */
+	if ((flags & DSA_ALLOC_ZERO) != 0)
+		memset(dsa_get_address(area, result), 0, size);
 
-	return dp;
+	return result;
 }
 
 /*
diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 3fd9bbd..c4b122e 100644
--- a/src/include/utils/dsa.h
+++ b/src/include/utils/dsa.h
@@ -71,12 +71,25 @@ typedef pg_atomic_uint64 dsa_pointer_atomic;
 #define DSA_POINTER_FORMAT "%016" INT64_MODIFIER "x"
 #endif
 
+/* Flags for dsa_allocate_extended. */
+#define DSA_ALLOC_HUGE		0x01	/* allow huge allocation (> 1 GB) */
+#define DSA_ALLOC_NO_OOM	0x02	/* no failure if out-of-memory */
+#define DSA_ALLOC_ZERO		0x04	/* zero allocated memory */
+
 /* A sentinel value for dsa_pointer used to indicate failure to allocate. */
 #define InvalidDsaPointer ((dsa_pointer) 0)
 
 /* Check if a dsa_pointer value is valid. */
 #define DsaPointerIsValid(x) ((x) != InvalidDsaPointer)
 
+/* Allocate uninitialized memory with error on out-of-memory. */
+#define dsa_allocate(area, size) \
+	dsa_allocate_extended(area, size, 0)
+
+/* Allocate zero-initialized memory with error on out-of-memory. */
+#define dsa_allocate0(area, size) \
+	dsa_allocate_extended(area, size, DSA_ALLOC_ZERO)
+
 /*
  * The type used for dsa_area handles.  dsa_handle values can be shared with
  * other processes, so that they can attach to them.  This provides a way to
@@ -105,8 +118,7 @@ extern void dsa_unpin(dsa_area *area);
 extern void dsa_set_size_limit(dsa_area *area, Size limit);
 extern Size dsa_minimum_size(void);
 extern dsa_handle dsa_get_handle(dsa_area *area);
-extern dsa_pointer dsa_allocate(dsa_area *area, Size size);
-extern dsa_pointer dsa_allocate0(dsa_area *area, Size size);
+extern dsa_pointer dsa_allocate_extended(dsa_area *area, Size size, int flags);
 extern void dsa_free(dsa_area *area, dsa_pointer dp);
 extern void *dsa_get_address(dsa_area *area, dsa_pointer dp);
 extern void dsa_trim(dsa_area *area);