Make dsa_allocate interface more like MemoryContextAlloc.

A new function dsa_allocate_extended now takes flags which indicate
that huge allocations should be permitted, that out-of-memory
conditions should not throw an error, and/or that the returned memory
should be zero-filled, just like MemoryContextAllocateExtended.

Commit 9acb85597f1223ac26a5b19a9345849c43d0ff54, which added
dsa_allocate0, was broken because it failed to account for the
possibility that dsa_allocate() might return InvalidDsaPointer.
This fixes that problem along the way.

Thomas Munro, with some comment changes by me.

Discussion: http://postgr.es/m/CA+Tgmobt7CcF_uQP2UQwWmu4K9qCHehMJP9_9m1urwP8hbOeHQ@mail.gmail.com

1 files changed, 55 insertions, 21 deletions

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d9a4e..49e68b4d08a 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -642,18 +642,39 @@ dsa_pin_mapping(dsa_area *area)
 /*
  * Allocate memory in this storage area.  The return value is a dsa_pointer
  * that can be passed to other processes, and converted to a local pointer
- * with dsa_get_address.  If no memory is available, returns
- * InvalidDsaPointer.
+ * with dsa_get_address.  'flags' is a bitmap which should be constructed
+ * from the following values:
+ *
+ * DSA_ALLOC_HUGE allows allocations >= 1GB.  Otherwise, such allocations
+ * will result in an ERROR.
+ *
+ * DSA_ALLOC_NO_OOM causes this function to return InvalidDsaPointer when
+ * no memory is available or a size limit establed by set_dsa_size_limit
+ * would be exceeded.  Otherwise, such allocations will result in an ERROR.
+ *
+ * DSA_ALLOC_ZERO causes the allocated memory to be zeroed.  Otherwise, the
+ * contents of newly-allocated memory are indeterminate.
+ *
+ * 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 +705,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 +739,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 +781,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;
 }
 
 /*