MallocAllocator: add Options struct and malloc-based contiguous allocations (facebookincubator#16695)

Summary:
X-link: facebookincubator/nimble#562


MallocAllocator currently uses mmap/munmap for contiguous allocations despite
its name suggesting it delegates everything to malloc. The mmap/munmap syscall
pair on every contiguous alloc/free cycle adds overhead.

This diff:
1. Adds an `Options` struct to `MallocAllocator` (mirroring `MmapAllocator::Options`)
   for cleaner configuration.
2. Adds a `mallocContiguousEnabled` option that, when true, uses
   `aligned_alloc`/`free` for contiguous allocations instead of `mmap`/`munmap`.
3. Wires the option through `MemoryManager::Options`.
4. Updates `CachedBufferedInputTest` to use the new `Options` struct.

The option defaults to false to preserve existing behavior. When enabled,
contiguous allocations use `aligned_alloc(kPageSize, maxBytes)` which
immediately commits physical memory (vs mmap's lazy page faulting), but this
is acceptable for MallocAllocator's use case.

`growContiguous` requires no changes since both mmap and malloc allocate
`maxPages` upfront, so growth stays within already-allocated memory.

Reviewed By: duxiao1212, xiaoxmeng

Differential Revision: D95875673

Author: tanjialiang

velox/common/memory/MallocAllocator.cpp

@@ -21,10 +21,11 @@
 #include <sys/mman.h>
 
 namespace facebook::velox::memory {
-MallocAllocator::MallocAllocator(size_t capacity, uint32_t reservationByteLimit)
+MallocAllocator::MallocAllocator(const Options& options)
     : kind_(MemoryAllocator::Kind::kMalloc),
-      capacity_(capacity),
-      reservationByteLimit_(reservationByteLimit),
+      mallocContiguousEnabled_(options.mallocContiguousEnabled),
+      capacity_(options.capacity),
+      reservationByteLimit_(options.reservationByteLimit),
       reserveFunc_(
           [this](uint32_t& counter, uint32_t increment, std::mutex& lock) {
             return incrementUsageWithReservationFunc(counter, increment, lock);
@@ -145,11 +146,7 @@ bool MallocAllocator::allocateContiguousImpl(
   }
   auto numContiguousCollateralPages = allocation.numPages();
   if (numContiguousCollateralPages > 0) {
-    useHugePages(allocation, false);
-    if (::munmap(allocation.data(), allocation.maxSize()) < 0) {
-      VELOX_MEM_LOG(ERROR) << "munmap got " << folly::errnoStr(errno) << "for "
-                           << allocation.data() << ", " << allocation.size();
-    }
+    dispatchFreeContiguous(allocation);
     numMapped_.fetch_sub(numContiguousCollateralPages);
     numAllocated_.fetch_sub(numContiguousCollateralPages);
     numExternalMapped_.fetch_sub(numContiguousCollateralPages);
@@ -175,19 +172,23 @@ bool MallocAllocator::allocateContiguousImpl(
   numAllocated_.fetch_add(numPages);
   numMapped_.fetch_add(numPages);
   numExternalMapped_.fetch_add(numPages);
-  void* data = ::mmap(
-      nullptr,
-      AllocationTraits::pageBytes(maxPages),
-      PROT_READ | PROT_WRITE,
-      MAP_PRIVATE | MAP_ANONYMOUS,
-      -1,
-      0);
-  // TODO: add handling of MAP_FAILED.
-  allocation.set(
-      data,
-      AllocationTraits::pageBytes(numPages),
-      AllocationTraits::pageBytes(maxPages));
-  useHugePages(allocation, true);
+  const auto maxBytes = AllocationTraits::pageBytes(maxPages);
+  void* data = dispatchAllocateContiguous(maxBytes);
+  if (FOLLY_UNLIKELY(data == nullptr)) {
+    numAllocated_.fetch_sub(numPages);
+    numMapped_.fetch_sub(numPages);
+    numExternalMapped_.fetch_sub(numPages);
+    decrementUsage(static_cast<int64_t>(AllocationTraits::pageBytes(numPages)));
+    const auto errorMsg = fmt::format(
+        "Failed to allocate {} of contiguous memory", succinctBytes(maxBytes));
+    VELOX_MEM_LOG(WARNING) << errorMsg;
+    setAllocatorFailureMessage(errorMsg);
+    return false;
+  }
+  allocation.set(data, AllocationTraits::pageBytes(numPages), maxBytes);
+  if (!mallocContiguousEnabled_) {
+    useHugePages(allocation, true);
+  }
   return true;
 }
 
@@ -222,21 +223,46 @@ void MallocAllocator::freeContiguousImpl(ContiguousAllocation& allocation) {
   if (allocation.empty()) {
     return;
   }
-  useHugePages(allocation, false);
   const auto bytes = allocation.size();
   const auto numPages = allocation.numPages();
-  if (::munmap(allocation.data(), allocation.maxSize()) < 0) {
-    VELOX_MEM_LOG(ERROR) << "Error for munmap(" << allocation.data() << ", "
-                         << succinctBytes(bytes) << "): '"
-                         << folly::errnoStr(errno) << "'";
-  }
+  dispatchFreeContiguous(allocation);
   numMapped_.fetch_sub(numPages);
   numAllocated_.fetch_sub(numPages);
   numExternalMapped_.fetch_sub(numPages);
   decrementUsage(bytes);
   allocation.clear();
 }
 
+void* MallocAllocator::dispatchAllocateContiguous(size_t maxBytes) {
+  if (testingHasInjectedFailure(InjectedFailure::kAllocate)) {
+    return nullptr;
+  }
+  if (mallocContiguousEnabled_) {
+    return ::aligned_alloc(AllocationTraits::kPageSize, maxBytes);
+  }
+  // TODO: add handling of MAP_FAILED.
+  return ::mmap(
+      nullptr,
+      maxBytes,
+      PROT_READ | PROT_WRITE,
+      MAP_PRIVATE | MAP_ANONYMOUS,
+      -1,
+      0);
+}
+
+void MallocAllocator::dispatchFreeContiguous(ContiguousAllocation& allocation) {
+  if (mallocContiguousEnabled_) {
+    ::free(allocation.data());
+  } else {
+    useHugePages(allocation, false);
+    if (::munmap(allocation.data(), allocation.maxSize()) < 0) {
+      VELOX_MEM_LOG(ERROR) << "Error for munmap(" << allocation.data() << ", "
+                           << succinctBytes(allocation.size()) << "): '"
+                           << folly::errnoStr(errno) << "'";
+    }
+  }
+}
+
 bool MallocAllocator::growContiguousWithoutRetry(
     MachinePageCount increment,
     ContiguousAllocation& allocation) {

velox/common/memory/MallocAllocator.h

@@ -26,7 +26,19 @@ namespace facebook::velox::memory {
 /// The implementation of MemoryAllocator using malloc.
 class MallocAllocator : public MemoryAllocator {
  public:
-  MallocAllocator(size_t capacity, uint32_t reservationByteLimit);
+  struct Options {
+    /// Capacity in bytes, default unlimited.
+    size_t capacity{kMaxMemory};
+
+    /// Allocation size threshold below which allocations use sharded local
+    /// counters instead of updating the global counter. Default 1MB.
+    uint32_t reservationByteLimit{1 << 20};
+
+    /// If true, use malloc for contiguous allocations instead of mmap/munmap.
+    bool mallocContiguousEnabled{false};
+  };
+
+  explicit MallocAllocator(const Options& options);
 
   ~MallocAllocator() override;
 
@@ -102,6 +114,15 @@ class MallocAllocator : public MemoryAllocator {
       ContiguousAllocation& allocation,
       MachinePageCount maxPages);
 
+  // Allocates 'maxBytes' of contiguous memory using malloc or mmap depending
+  // on 'mallocContiguousEnabled_'. Returns the allocated pointer, or nullptr
+  // on failure.
+  void* dispatchAllocateContiguous(size_t maxBytes);
+
+  // Frees contiguous memory previously allocated by
+  // dispatchAllocateContiguous.
+  void dispatchFreeContiguous(ContiguousAllocation& allocation);
+
   void freeContiguousImpl(ContiguousAllocation& allocation);
 
   void* allocateBytesWithoutRetry(uint64_t bytes, uint16_t alignment) override;
@@ -217,6 +238,9 @@ class MallocAllocator : public MemoryAllocator {
 
   const Kind kind_;
 
+  // If true, use malloc for contiguous allocations instead of mmap/munmap.
+  const bool mallocContiguousEnabled_;
+
   // Capacity in bytes. Total allocation byte is not allowed to exceed this
   // value.
   const size_t capacity_;

velox/common/memory/Memory.cpp

@@ -54,9 +54,12 @@ std::shared_ptr<MemoryAllocator> createAllocator(
     mmapOptions.mmapArenaCapacityRatio = options.mmapArenaCapacityRatio;
     return std::make_shared<MmapAllocator>(mmapOptions);
   } else {
-    return std::make_shared<MallocAllocator>(
-        options.allocatorCapacity,
-        options.allocationSizeThresholdWithReservation);
+    MallocAllocator::Options mallocOptions;
+    mallocOptions.capacity = options.allocatorCapacity;
+    mallocOptions.reservationByteLimit =
+        options.allocationSizeThresholdWithReservation;
+    mallocOptions.mallocContiguousEnabled = options.mallocContiguousEnabled;
+    return std::make_shared<MallocAllocator>(mallocOptions);
   }
 }
 

velox/common/memory/Memory.h

@@ -139,6 +139,12 @@ class MemoryManager {
     /// NOTE: this only applies for MallocAllocator.
     uint32_t allocationSizeThresholdWithReservation{1 << 20};
 
+    /// If true, MallocAllocator uses malloc for contiguous allocations instead
+    /// of mmap/munmap.
+    ///
+    /// NOTE: this only applies for MallocAllocator.
+    bool mallocContiguousEnabled{false};
+
     /// ================== 'MemoryArbitrator' settings =================
 
     /// Memory capacity available for query/task memory pools. This capacity