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, 14 insertions, 2 deletions

diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 3fd9bbd3a50..c4b122e34c9 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);