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

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;
 }
 
 /*

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);