@@ -642,18 +642,39 @@ dsa_pin_mapping(dsa_area *area)
642642/*
643643 * Allocate memory in this storage area. The return value is a dsa_pointer
644644 * that can be passed to other processes, and converted to a local pointer
645- * with dsa_get_address. If no memory is available, returns
646- * InvalidDsaPointer.
645+ * with dsa_get_address. 'flags' is a bitmap which should be constructed
646+ * from the following values:
647+ *
648+ * DSA_ALLOC_HUGE allows allocations >= 1GB. Otherwise, such allocations
649+ * will result in an ERROR.
650+ *
651+ * DSA_ALLOC_NO_OOM causes this function to return InvalidDsaPointer when
652+ * no memory is available or a size limit establed by set_dsa_size_limit
653+ * would be exceeded. Otherwise, such allocations will result in an ERROR.
654+ *
655+ * DSA_ALLOC_ZERO causes the allocated memory to be zeroed. Otherwise, the
656+ * contents of newly-allocated memory are indeterminate.
657+ *
658+ * These flags correspond to similarly named flags used by
659+ * MemoryContextAllocExtended(). See also the macros dsa_allocate and
660+ * dsa_allocate0 which expand to a call to this function with commonly used
661+ * flags.
647662 */
648663dsa_pointer
649- dsa_allocate (dsa_area * area ,Size size )
664+ dsa_allocate_extended (dsa_area * area ,Size size , int flags )
650665{
651666uint16 size_class ;
652667dsa_pointer start_pointer ;
653668dsa_segment_map * segment_map ;
669+ dsa_pointer result ;
654670
655671Assert (size > 0 );
656672
673+ /* Sanity check on huge individual allocation size. */
674+ if (((flags & DSA_ALLOC_HUGE )!= 0 && !AllocHugeSizeIsValid (size ))||
675+ ((flags & DSA_ALLOC_HUGE )== 0 && !AllocSizeIsValid (size )))
676+ elog (ERROR ,"invalid DSA memory alloc request size %zu" ,size );
677+
657678/*
658679 * If bigger than the largest size class, just grab a run of pages from
659680 * the free page manager, instead of allocating an object from a pool.
@@ -684,6 +705,14 @@ dsa_allocate(dsa_area *area, Size size)
684705/* Can't make any more segments: game over. */
685706LWLockRelease (DSA_AREA_LOCK (area ));
686707dsa_free (area ,span_pointer );
708+
709+ /* Raise error unless asked not to. */
710+ if ((flags & MCXT_ALLOC_NO_OOM )== 0 )
711+ ereport (ERROR ,
712+ (errcode (ERRCODE_OUT_OF_MEMORY ),
713+ errmsg ("out of memory" ),
714+ errdetail ("Failed on DSA request of size %zu." ,
715+ size )));
687716return InvalidDsaPointer ;
688717}
689718
@@ -710,6 +739,10 @@ dsa_allocate(dsa_area *area, Size size)
710739segment_map -> pagemap [first_page ]= span_pointer ;
711740LWLockRelease (DSA_SCLASS_LOCK (area ,DSA_SCLASS_SPAN_LARGE ));
712741
742+ /* Zero-initialize the memory if requested. */
743+ if ((flags & DSA_ALLOC_ZERO )!= 0 )
744+ memset (dsa_get_address (area ,start_pointer ),0 ,size );
745+
713746return start_pointer ;
714747}
715748
@@ -748,27 +781,28 @@ dsa_allocate(dsa_area *area, Size size)
748781Assert (size <=dsa_size_classes [size_class ]);
749782Assert (size_class == 0 || size > dsa_size_classes [size_class - 1 ]);
750783
751- /*
752- * Attempt to allocate an object from the appropriate pool. This might
753- * return InvalidDsaPointer if there's no space available.
754- */
755- return alloc_object (area ,size_class );
756- }
784+ /* Attempt to allocate an object from the appropriate pool. */
785+ result = alloc_object (area ,size_class );
757786
758- /*
759- * As dsa_allocate, but zeroes the allocated memory.
760- */
761- dsa_pointer
762- dsa_allocate0 (dsa_area * area ,Size size )
763- {
764- dsa_pointer dp ;
765- char * object ;
787+ /* Check for failure to allocate. */
788+ if (!DsaPointerIsValid (result ))
789+ {
790+ /* Raise error unless asked not to. */
791+ if ((flags & DSA_ALLOC_NO_OOM )== 0 )
792+ {
793+ ereport (ERROR ,
794+ (errcode (ERRCODE_OUT_OF_MEMORY ),
795+ errmsg ("out of memory" ),
796+ errdetail ("Failed on DSA request of size %zu." ,size )));
797+ }
798+ return InvalidDsaPointer ;
799+ }
766800
767- dp = dsa_allocate ( area , size );
768- object = dsa_get_address ( area , dp );
769- memset (object ,0 ,size );
801+ /* Zero-initialize the memory if requested. */
802+ if (( flags & DSA_ALLOC_ZERO ) != 0 )
803+ memset (dsa_get_address ( area , result ) ,0 ,size );
770804
771- return dp ;
805+ return result ;
772806}
773807
774808/*