[wasm] Implement cached PKU-protected assembler buffers

This adds the implementation of the {AssemblerBufferCache} class.
PKU-protected memory is allocated via the {WasmCodeManager}, which has
access to the actual protection key.

R=thibaudm@chromium.org

Bug: v8:12809
Change-Id: Id26abd6f98248d5c646ae337ccb903d3e168bed1
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3593137Reviewed-by: Thibaud Michaud <thibaudm@chromium.org>
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/main@{#80181}

src/base/address-region.h

@@ -74,14 +74,19 @@ class AddressRegion {
 };
 ASSERT_TRIVIALLY_COPYABLE(AddressRegion);
 
+// Construct an AddressRegion from a start pointer and a size.
+template <typename T>
+inline AddressRegion AddressRegionOf(T* ptr, size_t size) {
+  return AddressRegion{reinterpret_cast<AddressRegion::Address>(ptr),
+                       sizeof(T) * size};
+}
+
 // Construct an AddressRegion from anything providing a {data()} and {size()}
 // accessor.
-template <typename Container,
-          typename = decltype(std::declval<Container>().data()),
-          typename = decltype(std::declval<Container>().size())>
-inline constexpr AddressRegion AddressRegionOf(Container&& c) {
-  return AddressRegion{reinterpret_cast<AddressRegion::Address>(c.data()),
-                       sizeof(*c.data()) * c.size()};
+template <typename Container>
+inline auto AddressRegionOf(Container&& c)
+    -> decltype(AddressRegionOf(c.data(), c.size())) {
+  return AddressRegionOf(c.data(), c.size());
 }
 
 inline std::ostream& operator<<(std::ostream& out, AddressRegion region) {

src/wasm/assembler-buffer-cache.cc

@@ -4,14 +4,64 @@
 
 #include "src/wasm/assembler-buffer-cache.h"
 
+#include <algorithm>
+
+#include "src/base/bits.h"
 #include "src/codegen/assembler.h"
+#include "src/wasm/wasm-engine.h"
 
 namespace v8::internal::wasm {
 
+class CachedAssemblerBuffer final : public AssemblerBuffer {
+ public:
+  CachedAssemblerBuffer(AssemblerBufferCache* cache, base::AddressRegion region)
+      : cache_(cache), region_(region) {}
+
+  ~CachedAssemblerBuffer() override { cache_->Return(region_); }
+
+  uint8_t* start() const override {
+    return reinterpret_cast<uint8_t*>(region_.begin());
+  }
+
+  int size() const override { return static_cast<int>(region_.size()); }
+
+  std::unique_ptr<AssemblerBuffer> Grow(int new_size) override {
+    return cache_->GetAssemblerBuffer(new_size);
+  }
+
+ private:
+  AssemblerBufferCache* const cache_;
+  const base::AddressRegion region_;
+};
+
+AssemblerBufferCache::~AssemblerBufferCache() {
+  for (base::AddressRegion region : available_memory_.regions()) {
+    GetWasmCodeManager()->FreeAssemblerBufferSpace(region);
+  }
+}
+
 std::unique_ptr<AssemblerBuffer> AssemblerBufferCache::GetAssemblerBuffer(
     int size) {
-  // TODO(12809): Return PKU-protected buffers, and cache them.
-  return NewAssemblerBuffer(size);
+  DCHECK_LT(0, size);
+  base::AddressRegion region = available_memory_.Allocate(size);
+  if (region.is_empty()) {
+    static constexpr int kMinimumReservation = 64 * KB;
+    int minimum_allocation =
+        std::max(kMinimumReservation, std::max(total_allocated_ / 4, size));
+    base::AddressRegion new_space =
+        GetWasmCodeManager()->AllocateAssemblerBufferSpace(minimum_allocation);
+    available_memory_.Merge(new_space);
+    CHECK_GE(kMaxInt - total_allocated_, new_space.size());
+    total_allocated_ += new_space.size();
+
+    region = available_memory_.Allocate(size);
+    DCHECK(!region.is_empty());
+  }
+  return std::make_unique<CachedAssemblerBuffer>(this, region);
+}
+
+void AssemblerBufferCache::Return(base::AddressRegion region) {
+  available_memory_.Merge(region);
 }
 
 }  // namespace v8::internal::wasm

src/wasm/assembler-buffer-cache.h

@@ -11,22 +11,34 @@
 
 #include <memory>
 
+#include "src/wasm/wasm-code-manager.h"
+
 namespace v8::internal {
 class AssemblerBuffer;
 }
 
 namespace v8::internal::wasm {
 
+class CachedAssemblerBuffer;
+
 // Creating assembler buffers can be expensive, in particular if PKU is used,
 // which requires an {mmap} and {pkey_protect} system call for each new buffer.
 // Hence pool-allocate a larger memory region and reuse it if assembler buffers
 // are freed.
-// For now, this class only implements the interface without actually caching
-// anything.
-// TODO(12809): Actually cache the assembler buffers.
 class AssemblerBufferCache final {
  public:
+  ~AssemblerBufferCache();
   std::unique_ptr<AssemblerBuffer> GetAssemblerBuffer(int size);
+
+ private:
+  friend class CachedAssemblerBuffer;
+
+  // Called when CachedAssemblerBuffers get destroyed, to return memory to the
+  // cache.
+  void Return(base::AddressRegion);
+
+  DisjointAllocationPool available_memory_;
+  int total_allocated_ = 0;
 };
 
 }  // namespace v8::internal::wasm

src/wasm/memory-protection-key.cc

@@ -146,10 +146,13 @@ void FreeMemoryProtectionKey(int key) {
 
 int GetProtectionFromMemoryPermission(PageAllocator::Permission permission) {
 #if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
-  // Mappings for PKU are either RWX (on this level) or no access.
+  // Mappings for PKU are either RWX (for code), no access (for uncommitted
+  // memory), or RW (for assembler buffers).
   switch (permission) {
     case PageAllocator::kNoAccess:
       return PROT_NONE;
+    case PageAllocator::kReadWrite:
+      return PROT_READ | PROT_WRITE;
     case PageAllocator::kReadWriteExecute:
       return PROT_READ | PROT_WRITE | PROT_EXEC;
     default:

src/wasm/wasm-code-manager.cc

@@ -2182,6 +2182,49 @@ void WasmCodeManager::InitializeMemoryProtectionKeyPermissionsIfSupported()
   }
 }
 
+base::AddressRegion WasmCodeManager::AllocateAssemblerBufferSpace(int size) {
+#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+  if (MemoryProtectionKeysEnabled()) {
+    auto* page_allocator = GetPlatformPageAllocator();
+    size_t page_size = page_allocator->AllocatePageSize();
+    size = RoundUp(size, page_size);
+    void* mapped = AllocatePages(page_allocator, nullptr, size, page_size,
+                                 PageAllocator::kNoAccess);
+    if (V8_UNLIKELY(!mapped)) {
+      constexpr auto format = base::StaticCharVector(
+          "Cannot allocate %zu more bytes for assembler buffers");
+      constexpr int kMaxMessageLength =
+          format.size() - 3 + std::numeric_limits<size_t>::digits10;
+      base::EmbeddedVector<char, kMaxMessageLength + 1> message;
+      SNPrintF(message, format.begin(), size);
+      V8::FatalProcessOutOfMemory(nullptr, message.begin());
+      UNREACHABLE();
+    }
+    auto region =
+        base::AddressRegionOf(reinterpret_cast<uint8_t*>(mapped), size);
+    CHECK(SetPermissionsAndMemoryProtectionKey(page_allocator, region,
+                                               PageAllocator::kReadWrite,
+                                               memory_protection_key_));
+    return region;
+  }
+#endif  // defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+  DCHECK(!MemoryProtectionKeysEnabled());
+  return base::AddressRegionOf(new uint8_t[size], size);
+}
+
+void WasmCodeManager::FreeAssemblerBufferSpace(base::AddressRegion region) {
+#if defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+  if (MemoryProtectionKeysEnabled()) {
+    auto* page_allocator = GetPlatformPageAllocator();
+    FreePages(page_allocator, reinterpret_cast<void*>(region.begin()),
+              region.size());
+    return;
+  }
+#endif  // defined(V8_OS_LINUX) && defined(V8_HOST_ARCH_X64)
+  DCHECK(!MemoryProtectionKeysEnabled());
+  delete[] reinterpret_cast<uint8_t*>(region.begin());
+}
+
 std::shared_ptr<NativeModule> WasmCodeManager::NewNativeModule(
     Isolate* isolate, const WasmFeatures& enabled, size_t code_size_estimate,
     std::shared_ptr<const WasmModule> module) {

src/wasm/wasm-code-manager.h

@@ -140,12 +140,12 @@ class V8_EXPORT_PRIVATE DisjointAllocationPool final {
   // obtained from a previous Allocate. Returns the merged region.
   base::AddressRegion Merge(base::AddressRegion);
 
-  // Allocate a contiguous region of size {size}. Return an empty pool on
+  // Allocate a contiguous region of size {size}. Return an empty region on
   // failure.
   base::AddressRegion Allocate(size_t size);
 
   // Allocate a contiguous region of size {size} within {region}. Return an
-  // empty pool on failure.
+  // empty region on failure.
   base::AddressRegion AllocateInRegion(size_t size, base::AddressRegion);
 
   bool IsEmpty() const { return regions_.empty(); }
@@ -1069,6 +1069,12 @@ class V8_EXPORT_PRIVATE WasmCodeManager final {
   // if we have support.
   void InitializeMemoryProtectionKeyPermissionsIfSupported() const;
 
+  // Allocate new memory for assembler buffers, potentially protected by PKU.
+  base::AddressRegion AllocateAssemblerBufferSpace(int size);
+
+  // Free previously allocated space for assembler buffers.
+  void FreeAssemblerBufferSpace(base::AddressRegion region);
+
  private:
   friend class WasmCodeAllocator;
   friend class WasmEngine;