Merge pull request #15101 from rockwotj/async-memset

wasm: asynchronously zero deallocated heaps

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/CMakeLists.txt

@@ -112,8 +112,6 @@ rp_test(
   LABELS wasm
 )
 
-# TODO(rockwood): Enable on CI when we determine why the module
-# OOMs in this environment
 rp_test(
   BENCHMARK_TEST
   BINARY_NAME wasm_transform
@@ -129,5 +127,4 @@ rp_test(
     "${TESTDATA_DIR}/identity.wasm"
   LABELS
     wasm 
-    disable_on_ci
 )

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()));
 }