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

diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 3eb3d4d9a4e2b65f26925f2aa55d701c3022afd7..49e68b4d08a43a84070bbed42e25150cc815a3a0 100644 (file)
--- 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;
 }
 
 /*

diff --git a/src/include/utils/dsa.h b/src/include/utils/dsa.h
index 3fd9bbd3a5027a1ea490fd938a91d11276c76405..c4b122e34c9eb75f556cd377f93b0dc253cd84aa 100644 (file)
--- 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);