wasm/alloc: async dealloc large chunks

We need to memset(0) memory for wasmtime, however if the user has
configured large chunks of memory, then this can cause us to go over the
task budget. In an effort to prevent that, we deallocate asynchronously.

Since allocation is now asynchronous (due to a pending zero operation)
we need a way to fit that into wasmtime's custom memory allocator API.
In order to do that we use a thread local variable as a side channel to
pass the allocated buffer into the VM, since wasmtime's APIs don't give
us the ability to pass any data into the allocation request.

Signed-off-by: Tyler Rockwood <[email protected]>

Author: rockwotj

src/v/wasm/allocator.cc

@@ -10,13 +10,20 @@
  */
 #include "wasm/allocator.h"
 
+#include "ssx/future-util.h"
 #include "vassert.h"
 #include "vlog.h"
 #include "wasm/logger.h"
 
 #include <seastar/core/align.hh>
 #include <seastar/core/aligned_buffer.hh>
+#include <seastar/core/future.hh>
+#include <seastar/core/loop.hh>
+#include <seastar/core/lowres_clock.hh>
 #include <seastar/core/print.hh>
+#include <seastar/core/when_all.hh>
+#include <seastar/coroutine/maybe_yield.hh>
+#include <seastar/util/later.hh>
 
 #include <sys/mman.h>
 
@@ -26,36 +33,56 @@
 
 namespace wasm {
 
-heap_allocator::heap_allocator(config c) {
+heap_allocator::heap_allocator(config c)
+  : _memset_chunk_size(c.memset_chunk_size) {
     size_t page_size = ::getpagesize();
-    _max_size = ss::align_up(c.heap_memory_size, page_size);
+    _size = ss::align_up(c.heap_memory_size, page_size);
     for (size_t i = 0; i < c.num_heaps; ++i) {
-        auto buffer = ss::allocate_aligned_buffer<uint8_t>(
-          _max_size, page_size);
-        std::memset(buffer.get(), 0, _max_size);
-        _memory_pool.emplace_back(std::move(buffer), _max_size);
+        auto buffer = ss::allocate_aligned_buffer<uint8_t>(_size, page_size);
+        _memory_pool.push_back(
+          async_zero_memory({std::move(buffer), _size}, _size));
     }
 }
 
-std::optional<heap_memory> heap_allocator::allocate(request req) {
+ss::future<> heap_allocator::stop() {
+    auto pool = std::exchange(_memory_pool, {});
+    co_await ss::when_all_succeed(pool.begin(), pool.end());
+}
+
+ss::future<std::optional<heap_memory>> heap_allocator::allocate(request req) {
     if (_memory_pool.empty()) {
-        return std::nullopt;
+        co_return std::nullopt;
     }
-    size_t size = _memory_pool.front().size;
-    if (size < req.minimum || size > req.maximum) {
-        return std::nullopt;
+    if (_size < req.minimum || _size > req.maximum) {
+        co_return std::nullopt;
     }
-    heap_memory front = std::move(_memory_pool.front());
+    ss::future<heap_memory> front = std::move(_memory_pool.front());
     _memory_pool.pop_front();
-    return front;
+    co_return co_await std::move(front);
 }
 
 void heap_allocator::deallocate(heap_memory m, size_t used_amount) {
-    std::memset(m.data.get(), 0, used_amount);
-    _memory_pool.push_back(std::move(m));
+    _memory_pool.push_back(async_zero_memory(std::move(m), used_amount));
+}
+
+ss::future<heap_memory>
+heap_allocator::async_zero_memory(heap_memory m, size_t used_amount) {
+    uint8_t* data = m.data.get();
+    size_t remaining = used_amount;
+    while (true) {
+        if (remaining <= _memset_chunk_size) {
+            std::memset(data, 0, remaining);
+            co_return m;
+        }
+        std::memset(data, 0, _memset_chunk_size);
+        // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
+        data += _memset_chunk_size;
+        remaining -= _memset_chunk_size;
+        co_await ss::coroutine::maybe_yield();
+    }
 }
 
-size_t heap_allocator::max_size() const { return _max_size; }
+size_t heap_allocator::max_size() const { return _size; }
 
 stack_memory::stack_memory(stack_bounds bounds, allocated_memory data)
   : _bounds(bounds)

src/v/wasm/allocator.h

@@ -14,7 +14,9 @@
 
 #include <seastar/core/aligned_buffer.hh>
 #include <seastar/core/chunked_fifo.hh>
-#include <seastar/util/optimized_optional.hh>
+#include <seastar/core/condition-variable.hh>
+#include <seastar/core/future.hh>
+#include <seastar/core/gate.hh>
 
 #include <absl/container/btree_set.h>
 
@@ -58,10 +60,18 @@ class heap_allocator {
         size_t heap_memory_size;
         // The total number of heaps allocated per core.
         size_t num_heaps;
+        // The amount of memory we zero out at once.
+        size_t memset_chunk_size;
     };
 
     explicit heap_allocator(config);
 
+    /**
+     * Stop this allocator, waiting for any asynchronous zero'ing of bytes to
+     * finish.
+     */
+    ss::future<> stop();
+
     /**
      * A request of heap memory based on the following bounds.
      */
@@ -71,18 +81,22 @@ class heap_allocator {
     };
 
     /**
-     * Allocate heap memory by taking a memory instance from the pool.
+     * Allocate heap memory by taking a memory instance from the pool, returns
+     * std::nullopt if the memory request cannot be fulfilled either because the
+     * request does not fix within our bounds or because all memory is currently
+     * allocated.
      *
      * Memory returned from this method will be zero-filled.
      */
-    std::optional<heap_memory> allocate(request);
+    ss::future<std::optional<heap_memory>> allocate(request);
 
     /**
      * Deallocate heap memory by returing a memory instance to the pool.
      *
      * used_amount is to zero out the used memory from the heap, this is
      * required so that if only a portion of a large memory space was used we
-     * don't have to touch all the bytes.
+     * don't have to touch all the bytes. If the used_amount is over some
+     * threshold then we make zero'ing out the bytes an asynchronous task.
      */
     void deallocate(heap_memory, size_t used_amount);
 
@@ -92,10 +106,13 @@ class heap_allocator {
     size_t max_size() const;
 
 private:
-    size_t _max_size;
+    ss::future<heap_memory> async_zero_memory(heap_memory, size_t used_amount);
+
+    size_t _memset_chunk_size;
+    size_t _size;
     // We expect this list to be small, so override the chunk to be smaller too.
     static constexpr size_t items_per_chunk = 16;
-    ss::chunked_fifo<heap_memory, items_per_chunk> _memory_pool;
+    ss::chunked_fifo<ss::future<heap_memory>, items_per_chunk> _memory_pool;
 };
 
 /**

src/v/wasm/tests/wasm_allocator_test.cc

@@ -10,8 +10,12 @@
  */
 
 #include "gmock/gmock.h"
+#include "units.h"
 #include "wasm/allocator.h"
 
+#include <seastar/core/reactor.hh>
+#include <seastar/core/when_all.hh>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
@@ -22,16 +26,22 @@
 
 namespace wasm {
 
+constexpr static auto default_memset_chunk_size = 10_MiB;
+
 using ::testing::Optional;
 
 TEST(HeapAllocatorParamsTest, SizeIsAligned) {
     size_t page_size = ::getpagesize();
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size + 3,
       .num_heaps = 1,
+      .memset_chunk_size = default_memset_chunk_size,
     });
-    auto mem = allocator.allocate(
-      {.minimum = 0, .maximum = std::numeric_limits<size_t>::max()});
+    auto mem = allocator
+                 .allocate(
+                   {.minimum = 0,
+                    .maximum = std::numeric_limits<size_t>::max()})
+                 .get();
     ASSERT_TRUE(mem.has_value());
     EXPECT_EQ(mem->size, page_size * 2);
 }
@@ -41,8 +51,10 @@ TEST(HeapAllocatorTest, CanAllocateOne) {
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size,
       .num_heaps = 1,
+      .memset_chunk_size = default_memset_chunk_size,
     });
-    auto mem = allocator.allocate({.minimum = page_size, .maximum = page_size});
+    auto mem
+      = allocator.allocate({.minimum = page_size, .maximum = page_size}).get();
     ASSERT_TRUE(mem.has_value());
     EXPECT_EQ(mem->size, page_size);
 }
@@ -52,14 +64,17 @@ TEST(HeapAllocatorTest, MustAllocateWithinBounds) {
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size,
       .num_heaps = 1,
+      .memset_chunk_size = default_memset_chunk_size,
     });
     // minimum too large
-    auto mem = allocator.allocate(
-      {.minimum = page_size * 2, .maximum = page_size * 3});
+    auto mem = allocator
+                 .allocate({.minimum = page_size * 2, .maximum = page_size * 3})
+                 .get();
     EXPECT_FALSE(mem.has_value());
     // maximum too small
-    mem = allocator.allocate(
-      {.minimum = page_size / 2, .maximum = page_size - 1});
+    mem = allocator
+            .allocate({.minimum = page_size / 2, .maximum = page_size - 1})
+            .get();
     EXPECT_FALSE(mem.has_value());
 }
 
@@ -68,10 +83,13 @@ TEST(HeapAllocatorTest, Exhaustion) {
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size,
       .num_heaps = 1,
+      .memset_chunk_size = default_memset_chunk_size,
     });
-    auto mem = allocator.allocate({.minimum = page_size, .maximum = page_size});
+    auto mem
+      = allocator.allocate({.minimum = page_size, .maximum = page_size}).get();
     EXPECT_TRUE(mem.has_value());
-    mem = allocator.allocate({.minimum = page_size, .maximum = page_size});
+    mem
+      = allocator.allocate({.minimum = page_size, .maximum = page_size}).get();
     EXPECT_FALSE(mem.has_value());
 }
 
@@ -80,25 +98,90 @@ TEST(HeapAllocatorTest, CanReturnMemoryToThePool) {
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size,
       .num_heaps = 3,
+      .memset_chunk_size = default_memset_chunk_size,
     });
     heap_allocator::request req{.minimum = page_size, .maximum = page_size};
     std::vector<heap_memory> allocated;
     for (int i = 0; i < 3; ++i) {
-        auto mem = allocator.allocate(req);
+        auto mem = allocator.allocate(req).get();
         ASSERT_TRUE(mem.has_value());
         allocated.push_back(std::move(*mem));
     }
-    auto mem = allocator.allocate(req);
+    auto mem = allocator.allocate(req).get();
     EXPECT_FALSE(mem.has_value());
     mem = std::move(allocated.back());
     allocated.pop_back();
     allocator.deallocate(std::move(*mem), /*used_amount=*/0);
-    mem = allocator.allocate(req);
+    mem = allocator.allocate(req).get();
     EXPECT_TRUE(mem.has_value());
-    mem = allocator.allocate(req);
+    mem = allocator.allocate(req).get();
     EXPECT_FALSE(mem.has_value());
 }
 
+// We want to test a specific scenario where the deallocation happens
+// asynchronously, however, release mode continuations can be "inlined" into the
+// current executing task, so we can't enforce the scenario we want to test, so
+// this test only runs in debug mode, which forces as many scheduling points as
+// possible and the zeroing task always happens asynchronously.
+#ifndef NDEBUG
+
+using ::testing::_;
+
+TEST(HeapAllocatorTest, AsyncDeallocationOnlyOneAwakened) {
+    size_t page_size = ::getpagesize();
+    // force deallocations to be asynchronous.
+    size_t test_chunk_size = page_size / 4;
+    heap_allocator allocator(heap_allocator::config{
+      .heap_memory_size = page_size,
+      .num_heaps = 2,
+      .memset_chunk_size = test_chunk_size,
+    });
+    heap_allocator::request req{.minimum = page_size, .maximum = page_size};
+    // Start on deallocation in the background.
+    allocator.deallocate(allocator.allocate(req).get().value(), page_size);
+    // Start another deallocation in the background so there is no memory left.
+    allocator.deallocate(allocator.allocate(req).get().value(), page_size);
+    auto waiter1 = allocator.allocate(req);
+    auto waiter2 = allocator.allocate(req);
+    // There should not be memory available, so both our requests for
+    // memory should be waiting.
+    EXPECT_FALSE(waiter1.available());
+    EXPECT_FALSE(waiter2.available());
+    // waiter1 should be notified first there is memory available, but waiter2
+    // should not be notified yet as there is a deallocation in flight.
+    // waiter2 may or may not be ready depending on task execution order (which
+    // is randomized in debug mode), so we cannot assert that it's not
+    // completed.
+    EXPECT_THAT(waiter1.get(), Optional(_));
+    // The second waiter completes and we can allocate memory
+    EXPECT_THAT(waiter2.get(), Optional(_));
+}
+
+TEST(HeapAllocatorTest, AsyncDeallocationNotEnoughMemory) {
+    size_t page_size = ::getpagesize();
+    // force deallocations to be asynchronous.
+    size_t test_chunk_size = page_size / 4;
+    heap_allocator allocator(heap_allocator::config{
+      .heap_memory_size = page_size,
+      .num_heaps = 1,
+      .memset_chunk_size = test_chunk_size,
+    });
+    heap_allocator::request req{.minimum = page_size, .maximum = page_size};
+    allocator.deallocate(allocator.allocate(req).get().value(), page_size);
+
+    // The first request succeeded and is waiting for the deallocation to finish
+    auto waiter1 = allocator.allocate(req);
+    EXPECT_FALSE(waiter1.available());
+
+    // Future requests fail immediately, no memory is available
+    auto waiter2 = allocator.allocate(req);
+    EXPECT_TRUE(waiter2.available());
+
+    EXPECT_THAT(waiter1.get(), Optional(_));
+    EXPECT_EQ(waiter2.get(), std::nullopt);
+}
+#endif
+
 MATCHER(HeapIsZeroed, "is zeroed") {
     std::span<uint8_t> d = {arg.data.get(), arg.size};
     for (auto e : d) {
@@ -114,15 +197,16 @@ TEST(HeapAllocatorTest, MemoryIsZeroFilled) {
     heap_allocator allocator(heap_allocator::config{
       .heap_memory_size = page_size,
       .num_heaps = 1,
+      .memset_chunk_size = default_memset_chunk_size,
     });
     heap_allocator::request req{.minimum = page_size, .maximum = page_size};
-    auto allocated = allocator.allocate(req);
+    auto allocated = allocator.allocate(req).get();
     ASSERT_TRUE(allocated.has_value());
     EXPECT_THAT(allocated, Optional(HeapIsZeroed()));
     std::fill_n(allocated->data.get(), 4, 1);
     allocator.deallocate(*std::move(allocated), 4);
 
-    allocated = allocator.allocate(req);
+    allocated = allocator.allocate(req).get();
     ASSERT_TRUE(allocated.has_value());
     EXPECT_THAT(allocated, Optional(HeapIsZeroed()));
 }