Merge pull request #15063 from rockwotj/memset

Author: rockwotj

src/v/wasm/allocator.cc

@@ -20,6 +20,7 @@
 
 #include <sys/mman.h>
 
+#include <span>
 #include <stdexcept>
 #include <unistd.h>
 
@@ -31,6 +32,7 @@ heap_allocator::heap_allocator(config c) {
     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);
     }
 }
@@ -48,7 +50,8 @@ std::optional<heap_memory> heap_allocator::allocate(request req) {
     return front;
 }
 
-void heap_allocator::deallocate(heap_memory m) {
+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));
 }
 
@@ -97,6 +100,7 @@ stack_memory stack_allocator::allocate(size_t size) {
     if (_tracking_enabled) {
         _live_stacks.emplace(mem.bounds());
     }
+    std::memset(mem.bounds().bottom, 0, mem.size());
     return mem;
 }
 

src/v/wasm/allocator.h

@@ -72,13 +72,19 @@ class heap_allocator {
 
     /**
      * Allocate heap memory by taking a memory instance from the pool.
+     *
+     * Memory returned from this method will be zero-filled.
      */
     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.
      */
-    void deallocate(heap_memory);
+    void deallocate(heap_memory, size_t used_amount);
 
     /**
      * The maximum size of a heap_memory that can be allocated.
@@ -143,7 +149,8 @@ class stack_memory {
 // Stack memory must be page-aligned and a multiple of page size.
 // Some architectures require this alignment for stacks. Additionally, we
 // always allocate a guard page at the bottom, which is unprotected when memory
-// is "deallocated".
+// is "deallocated". Lastly, memory returned from this allocator is always
+// zero-filled as assumed by the Wasm VM.
 //
 // The allocator also supports the ability to query if the current stack being
 // used has been allocated from this allocator and returns the bounds. This is

src/v/wasm/tests/wasm_allocator_test.cc

@@ -16,11 +16,14 @@
 #include <gtest/gtest.h>
 
 #include <limits>
+#include <span>
 #include <stdexcept>
 #include <unistd.h>
 
 namespace wasm {
 
+using ::testing::Optional;
+
 TEST(HeapAllocatorParamsTest, SizeIsAligned) {
     size_t page_size = ::getpagesize();
     heap_allocator allocator(heap_allocator::config{
@@ -89,13 +92,41 @@ TEST(HeapAllocatorTest, CanReturnMemoryToThePool) {
     EXPECT_FALSE(mem.has_value());
     mem = std::move(allocated.back());
     allocated.pop_back();
-    allocator.deallocate(std::move(*mem));
+    allocator.deallocate(std::move(*mem), /*used_amount=*/0);
     mem = allocator.allocate(req);
     EXPECT_TRUE(mem.has_value());
     mem = allocator.allocate(req);
     EXPECT_FALSE(mem.has_value());
 }
 
+MATCHER(HeapIsZeroed, "is zeroed") {
+    std::span<uint8_t> d = {arg.data.get(), arg.size};
+    for (auto e : d) {
+        if (e != 0) {
+            return false;
+        }
+    }
+    return true;
+}
+
+TEST(HeapAllocatorTest, MemoryIsZeroFilled) {
+    size_t page_size = ::getpagesize();
+    heap_allocator allocator(heap_allocator::config{
+      .heap_memory_size = page_size,
+      .num_heaps = 1,
+    });
+    heap_allocator::request req{.minimum = page_size, .maximum = page_size};
+    auto allocated = allocator.allocate(req);
+    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);
+    ASSERT_TRUE(allocated.has_value());
+    EXPECT_THAT(allocated, Optional(HeapIsZeroed()));
+}
+
 TEST(StackAllocatorParamsTest, TrackingCanBeEnabled) {
     stack_allocator allocator(stack_allocator::config{
       .tracking_enabled = true,
@@ -117,8 +148,6 @@ TEST(StackAllocatorTest, CanAllocateOne) {
     EXPECT_EQ(stack.bounds().top - stack.bounds().bottom, page_size * 4);
 }
 
-using ::testing::Optional;
-
 TEST(StackAllocatorTest, CanLookupMemory) {
     stack_allocator allocator(stack_allocator::config{
       .tracking_enabled = true,
@@ -164,4 +193,27 @@ TEST(StackAllocatorTest, CanReturnMemoryToThePool) {
     EXPECT_EQ(stack.bounds(), bounds);
 }
 
+MATCHER(StackIsZeroed, "is zeroed") {
+    std::span<uint8_t> d = {arg.bounds().bottom, arg.size()};
+    for (auto e : d) {
+        if (e != 0) {
+            return false;
+        }
+    }
+    return true;
+}
+
+TEST(StackAllocatorTest, MemoryIsZeroFilled) {
+    stack_allocator allocator(stack_allocator::config{
+      .tracking_enabled = true,
+    });
+    size_t page_size = ::getpagesize();
+    auto stack = allocator.allocate(page_size * 4);
+    EXPECT_THAT(stack, StackIsZeroed());
+    std::fill_n(stack.bounds().bottom, 4, 1);
+    allocator.deallocate(std::move(stack));
+    stack = allocator.allocate(page_size * 4);
+    EXPECT_THAT(stack, StackIsZeroed());
+}
+
 } // namespace wasm

src/v/wasm/wasmtime.cc

@@ -1311,14 +1311,11 @@ wasmtime_error_t* wasmtime_runtime::allocate_stack_memory(
     auto stack = _stack_allocator.local().allocate(size);
     struct vm_stack {
         stack_memory underlying;
-        ss::noncopyable_function<void(stack_memory)> reclaimer;
+        wasm::stack_allocator* allocator;
     };
     memory_ret->env = new vm_stack{
       .underlying = std::move(stack),
-      .reclaimer =
-        [this](stack_memory mem) {
-            _stack_allocator.local().deallocate(std::move(mem));
-        },
+      .allocator = &_stack_allocator.local(),
     };
     memory_ret->get_stack_memory = [](void* env, size_t* len_ret) -> uint8_t* {
         auto* mem = static_cast<vm_stack*>(env);
@@ -1327,7 +1324,7 @@ wasmtime_error_t* wasmtime_runtime::allocate_stack_memory(
     };
     memory_ret->finalizer = [](void* env) {
         auto* mem = static_cast<vm_stack*>(env);
-        mem->reclaimer(std::move(mem->underlying));
+        mem->allocator->deallocate(std::move(mem->underlying));
         // NOLINTNEXTLINE(cppcoreguidelines-owning-memory)
         delete mem;
     };
@@ -1368,37 +1365,41 @@ wasmtime_error_t* wasmtime_runtime::allocate_heap_memory(
     }
     struct linear_memory {
         heap_memory underlying;
-        ss::noncopyable_function<void(heap_memory)> reclaimer;
+        size_t used_memory;
+        wasm::heap_allocator* allocator;
     };
     memory_ret->env = new linear_memory{
       .underlying = *std::move(memory),
-      .reclaimer =
-        [this](heap_memory mem) {
-            _heap_allocator.local().deallocate(std::move(mem));
-        },
+      .used_memory = minimum,
+      .allocator = &_heap_allocator.local(),
     };
     memory_ret->finalizer = [](void* env) {
         auto* mem = static_cast<linear_memory*>(env);
-        mem->reclaimer(std::move(mem->underlying));
+        mem->allocator->deallocate(
+          std::move(mem->underlying), /*used_amount=*/mem->used_memory);
         // NOLINTNEXTLINE(cppcoreguidelines-owning-memory)
         delete mem;
     };
     memory_ret->get_memory =
       // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
       [](void* env, size_t* byte_size, size_t* max_byte_size) {
           auto* mem = static_cast<linear_memory*>(env);
-          *byte_size = mem->underlying.size;
+          *byte_size = mem->used_memory;
           *max_byte_size = mem->underlying.size;
           return mem->underlying.data.get();
       };
     memory_ret->grow_memory =
       [](void* env, size_t new_size) -> wasmtime_error_t* {
         auto* mem = static_cast<linear_memory*>(env);
-        if (new_size == mem->underlying.size) {
+        if (new_size <= mem->underlying.size) {
+            mem->used_memory = std::max(new_size, mem->used_memory);
             return nullptr;
         }
-        return wasmtime_error_new(
-          "linear memories are fixed size and ungrowable");
+        auto msg = ss::format(
+          "unable to grow memory past {} to {}",
+          human::bytes(double(mem->underlying.size)),
+          human::bytes(double(new_size)));
+        return wasmtime_error_new(msg.c_str());
     };
     return nullptr;
 }