ipcz: Introduce BlockAllocator

BlockAllocator manages a region of memory to dynamically allocate and
free blocks of a smaller fixed size within the region.

This only lands the implementation and unit tests for now, but
BlockAllocator will provide the basis for dynamic allocation of shared
memory fragments on each NodeLink.

Split out from the ipcz uber-CL: https://crrev.com/c/3570271

Bug: 1299283
Change-Id: I3a582364eca449e468512f0dfab75a9d9dd4c26f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3630827
Commit-Queue: Ken Rockot <rockot@google.com>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1002879}

third_party/ipcz/src/BUILD.gn

@@ -161,6 +161,7 @@ ipcz_source_set("impl") {
   visibility = [ ":*" ]
   public = [
     "ipcz/api_object.h",
+    "ipcz/block_allocator.h",
     "ipcz/driver_memory.h",
     "ipcz/driver_memory_mapping.h",
     "ipcz/driver_object.h",
@@ -184,6 +185,7 @@ ipcz_source_set("impl") {
   ]
   sources = [
     "ipcz/api_object.cc",
+    "ipcz/block_allocator.cc",
     "ipcz/driver_memory.cc",
     "ipcz/driver_memory_mapping.cc",
     "ipcz/driver_object.cc",
@@ -252,6 +254,7 @@ ipcz_source_set("ipcz_tests_sources") {
   sources = [
     "api_test.cc",
     "connect_test.cc",
+    "ipcz/block_allocator_test.cc",
     "ipcz/driver_memory_test.cc",
     "ipcz/driver_object_test.cc",
     "ipcz/driver_transport_test.cc",

third_party/ipcz/src/ipcz/block_allocator.cc

@@ -0,0 +1,203 @@
+// Copyright 2022 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "ipcz/block_allocator.h"
+
+#include <array>
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <thread>
+
+#include "ipcz/ipcz.h"
+#include "third_party/abseil-cpp/absl/base/macros.h"
+#include "util/safe_math.h"
+
+namespace ipcz {
+
+namespace {
+
+constexpr int16_t kFrontBlockIndex = 0;
+
+// Helper for legibility of relative/absolute index conversions.
+struct ForBaseIndex {
+  // Constructs a helper to compute absolute or relative indexes based around
+  // the successor to the block at `index`. See documentation on BlockHeader's
+  // `next` field.
+  explicit constexpr ForBaseIndex(int16_t index) : index_(index) {}
+
+  // Returns an index equivalent to `absolute_index`, but relative to the
+  // successor of the block at `index_`.
+  constexpr int16_t GetRelativeFromAbsoluteIndex(int16_t absolute_index) const {
+    // NOTE: We intentionally ignore overflow. Absolute indices are always
+    // validated before use anyway, and it would be redundant in some cases to
+    // validate these inputs.
+    return absolute_index - index_ - 1;
+  }
+
+  // Returns an absolute index which is equivalent to `relative_index` if taken
+  // as relative to the successor of the block at `index_`.
+  constexpr int16_t GetAbsoluteFromRelativeIndex(int16_t relative_index) const {
+    // NOTE: We intentionally ignore overflow. The returned index is only stored
+    // and maybe eventually used to reconstruct an absolute index. Absolute
+    // indices are always validated before use, and it would be redundant in
+    // some cases to validate these inputs.
+    return index_ + relative_index + 1;
+  }
+
+ private:
+  const int16_t index_;
+};
+
+}  // namespace
+
+BlockAllocator::BlockAllocator() = default;
+BlockAllocator::BlockAllocator(absl::Span<uint8_t> region, uint32_t block_size)
+    : region_(region), block_size_(block_size) {
+  // Require 8-byte alignment of the region and of block sizes, to ensure that
+  // each BlockHeader is itself 8-byte aligned. Also a non-zero block size is
+  // obviously a requirement.
+  ABSL_ASSERT((reinterpret_cast<uintptr_t>(region_.data()) & 7) == 0);
+  ABSL_ASSERT(block_size > 0);
+  ABSL_ASSERT((block_size & 7) == 0);
+
+  // BlockHeader uses a signed 16-bit index to reference other blocks, and block
+  // 0 must be able to reference any block; so the total number of blocks must
+  // not exceed the max value of an int16_t.
+  num_blocks_ = checked_cast<int16_t>(region.size() / block_size);
+  ABSL_ASSERT(num_blocks_ > 0);
+}
+
+BlockAllocator::BlockAllocator(const BlockAllocator&) = default;
+
+BlockAllocator& BlockAllocator::operator=(const BlockAllocator&) = default;
+
+BlockAllocator::~BlockAllocator() = default;
+
+void BlockAllocator::InitializeRegion() const {
+  // By zeroing the entire region, every block effectively points to its
+  // immediate successor as the next free block. See comments on the `next`
+  // field if BlockHeader.
+  memset(region_.data(), 0, region_.size());
+
+  // Ensure that the last block points back to the unallocable first block,
+  // indicating the end of the free-list.
+  free_block_at(last_block_index()).SetNextFreeBlock(kFrontBlockIndex);
+}
+
+void* BlockAllocator::Alloc() const {
+  BlockHeader front =
+      block_header_at(kFrontBlockIndex).load(std::memory_order_relaxed);
+  for (;;) {
+    const int16_t first_free_block_index =
+        ForBaseIndex(kFrontBlockIndex).GetAbsoluteFromRelativeIndex(front.next);
+    if (first_free_block_index == kFrontBlockIndex ||
+        !is_index_valid(first_free_block_index)) {
+      // Note that the front block can never be allocated, so if that's where
+      // the head of the free-list points then the free-list is empty. If it
+      // otherwise points to an out-of-range index, the allocator is in an
+      // invalid state. In either case, we fail the allocation.
+      return nullptr;
+    }
+
+    // Extract the index of the *next* free block from the header of the first.
+    FreeBlock first_free_block = free_block_at(first_free_block_index);
+    BlockHeader first_free_block_header =
+        first_free_block.header().load(std::memory_order_acquire);
+    const int16_t next_free_block_index =
+        ForBaseIndex(first_free_block_index)
+            .GetAbsoluteFromRelativeIndex(first_free_block_header.next);
+    if (!is_index_valid(next_free_block_index)) {
+      // Invalid block header, so we cannot proceed.
+      return nullptr;
+    }
+
+    if (TryUpdateFrontHeader(front, next_free_block_index)) {
+      // If we successfully update the front block's header to point at the
+      // second free block, we have effectively allocated the first free block
+      // by removing it from the free-list. This means we're done and the
+      // allocator will not touch the contents of this block (including header)
+      // until it's freed.
+      return first_free_block.address();
+    }
+
+    // Another thread must have modified the front block header since we fetched
+    // it above. `front` now has a newly updated copy, so we loop around again
+    // to retry allocation.
+  }
+}
+
+bool BlockAllocator::Free(void* ptr) const {
+  // Derive a block index from the given address, relative to the start of this
+  // allocator's managed region.
+  const int16_t new_free_index =
+      (reinterpret_cast<uint8_t*>(ptr) - region_.data()) / block_size_;
+  if (new_free_index == kFrontBlockIndex || !is_index_valid(new_free_index)) {
+    // The first block cannot be freed, and obviously neither can any block out
+    // of range for this allocator.
+    return false;
+  }
+
+  FreeBlock free_block = free_block_at(new_free_index);
+  BlockHeader front =
+      block_header_at(kFrontBlockIndex).load(std::memory_order_relaxed);
+  do {
+    const int16_t first_free_index =
+        ForBaseIndex(kFrontBlockIndex).GetAbsoluteFromRelativeIndex(front.next);
+    if (!is_index_valid(first_free_index)) {
+      // The front block header is in an invalid state, so we cannot proceed.
+      return false;
+    }
+
+    // The application calling Free() implies that it's done using the block.
+    // The allocator is therefore free to overwrite the contents with a new
+    // BlockHeader. Write a header which points to the current head of the
+    // free-list.
+    free_block.SetNextFreeBlock(first_free_index);
+
+    // And now try to update the front block so that this newly freed block
+    // becomes the new head of the free-list. Upon success, the block is
+    // effectively freed. Upon failure, `front` will have an updated copy of
+    // the front block header, so we can loop around and try to insert the freed
+    // block again.
+  } while (!TryUpdateFrontHeader(front, new_free_index));
+
+  return true;
+}
+
+bool BlockAllocator::TryUpdateFrontHeader(BlockHeader& last_known_header,
+                                          int16_t first_free_block) const {
+  // Note that `version` overflow is acceptable here. The version is only used
+  // as a tag to protect against the ABA problem. Because all alloc and free
+  // operations are completed via an atomic exchange of the front block header,
+  // and because they must always increment the header version at the same time,
+  // we effectively avoid one operation trampling the result of another.
+  const uint16_t new_version = last_known_header.version + 1;
+  const int16_t relative_next =
+      ForBaseIndex(kFrontBlockIndex)
+          .GetRelativeFromAbsoluteIndex(first_free_block);
+
+  // A weak compare/exchange is used since in practice this will always be
+  // called within a tight retry loop.
+  return block_header_at(kFrontBlockIndex)
+      .compare_exchange_weak(
+          last_known_header, {.version = new_version, .next = relative_next},
+          std::memory_order_release, std::memory_order_relaxed);
+}
+
+BlockAllocator::FreeBlock::FreeBlock(int16_t index, AtomicBlockHeader& header)
+    : index_(index), header_(header) {
+  ABSL_ASSERT(index > 0);
+}
+
+void BlockAllocator::FreeBlock::SetNextFreeBlock(int16_t next_free_block) {
+  const int16_t relative_next =
+      ForBaseIndex(index_).GetRelativeFromAbsoluteIndex(next_free_block);
+  header_.store({.version = 0, .next = relative_next},
+                std::memory_order_release);
+}
+
+}  // namespace ipcz

third_party/ipcz/src/ipcz/block_allocator.h

@@ -0,0 +1,170 @@
+// Copyright 2022 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef IPCZ_SRC_IPCZ_BLOCK_ALLOCATOR_H_
+#define IPCZ_SRC_IPCZ_BLOCK_ALLOCATOR_H_
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+
+#include "ipcz/ipcz.h"
+#include "third_party/abseil-cpp/absl/base/macros.h"
+#include "third_party/abseil-cpp/absl/types/span.h"
+
+namespace ipcz {
+
+// BlockAllocator manages a region of memory, dividing it into dynamically
+// allocable blocks of a smaller fixed size. Allocation prioritizes reuse of the
+// most recently freed blocks.
+//
+// This is a thread-safe, lock-free implementation which doesn't store heap
+// pointers within the managed region. Multiple BlockAllocators may therefore
+// cooperatively manage the same region of memory for the same block size across
+// different threads and processes.
+class BlockAllocator {
+ public:
+  // Constructs an empty BlockAllocator with no memory to manage and an
+  // unspecified block size. This cannot be used to allocate blocks.
+  BlockAllocator();
+
+  // Constructs a BlockAllocator to manage the memory within `region`,
+  // allocating blocks of `block_size` bytes. Note that this DOES NOT initialize
+  // the region. Before any BlockAllocators can allocate blocks from `region`,
+  // InitializeRegion() must be called once by a single BlockAllocator managing
+  // this region for the same block size.
+  BlockAllocator(absl::Span<uint8_t> region, uint32_t block_size);
+
+  BlockAllocator(const BlockAllocator&);
+  BlockAllocator& operator=(const BlockAllocator&);
+  ~BlockAllocator();
+
+  const absl::Span<uint8_t>& region() const { return region_; }
+
+  size_t block_size() const { return block_size_; }
+
+  size_t capacity() const {
+    // Note that the first block cannot be allocated, so real capacity is one
+    // less than the total number of blocks which fit within the region.
+    ABSL_ASSERT(num_blocks_ > 0);
+    return num_blocks_ - 1;
+  }
+
+  // Performs a one-time initialization of the memory region managed by this
+  // allocator. Many allocators may operate on the same region, but only one
+  // must initialize that region, and it must do so before any of them can
+  // allocate blocks.
+  void InitializeRegion() const;
+
+  // Allocates a block from the allocator's managed region of memory and returns
+  // a pointer to its base address, where `block_size()` contiguous bytes are
+  // then owned by the caller. May return null if out of blocks.
+  void* Alloc() const;
+
+  // Frees a block back to the allocator, given its base address as returned by
+  // a prior call to Alloc(). Returns true on success, or false on failure.
+  // Failure implies that `ptr` was not a valid block to free.
+  bool Free(void* ptr) const;
+
+ private:
+  // For a region of N bytes with a block size B, BlockAllocator divides the
+  // region into `N/B` contiguous blocks with this BlockHeader structure at the
+  // beginning of each. These headers form a singly-linked free-list of
+  // available blocks. Block 0 can never be allocated and is used exclusively
+  // for its header to point to the first free block. Each free block
+  // references the next free block in the list, and the last free block
+  // references back to block 0 to terminate the list.
+  //
+  // Once a block is allocated, its entire span of B bytes -- including the
+  // header space -- will remain untouched by BlockAllocator and is available
+  // for application use until freed. When a block is freed, its header is
+  // restored.
+  struct IPCZ_ALIGN(4) BlockHeader {
+    // Only meaningful within the first block.
+    //
+    // This field is incremented any time the block's header changes, and it's
+    // used to resolve races between concurrent Alloc() or Free() operations
+    // modifying the head of the free-list.
+    //
+    // For blocks other than the first block, this is always zero.
+    uint16_t version;
+
+    // A relative index to the next free block in the list. Note that this is
+    // not relative to the index of the header's own block, but rather to the
+    // index of the block which physically follows it within the region (this
+    // block's "successor".) For example if this header belongs to block 3,
+    // the value of `next` is relative to index 4.
+    //
+    // This scheme is chosen so that a region can be initialized efficiently by
+    // zeroing it out and then updating only the last block's header to
+    // terminate the list. That way block 0 begins by pointing to block 1 as
+    // the first free block, block 1 points to block 2 as the next free block,
+    // and so on.
+    int16_t next;
+  };
+  static_assert(sizeof(BlockHeader) == 4, "Invalid BlockHeader size");
+
+  using AtomicBlockHeader = std::atomic<BlockHeader>;
+  static_assert(AtomicBlockHeader::is_always_lock_free,
+                "ipcz requires lock-free atomics");
+
+  // Helper class which tracks an absolute block index, as well as a reference
+  // to that block's atomic header within the managed region.
+  class FreeBlock {
+   public:
+    FreeBlock(int16_t index, AtomicBlockHeader& header);
+
+    AtomicBlockHeader& header() { return header_; }
+
+    // Returns the address of the start of this block.
+    void* address() const { return &header_; }
+
+    // Atomically updates the header for this free block, implying that this
+    // block will become the new head of the allocator's free-list.
+    //
+    // `next_free_block` is the absolute index of the previous head of the
+    // free-list, which this block will now reference as the next free block.
+    void SetNextFreeBlock(int16_t next_free_block);
+
+   private:
+    const int16_t index_;
+    AtomicBlockHeader& header_;
+  };
+
+  // Attempts to update the front block's header to point to `first_free_block`
+  // as the free-list's new head block. `last_known_header` is a reference to
+  // a copy of the most recently known header value from the first block. This
+  // call can only succeed if that value can be atomically swapped out for a
+  // new value.
+  //
+  // On failure, `last_known_header` is updated to reflect the value of the
+  // current front block's header.
+  bool TryUpdateFrontHeader(BlockHeader& last_known_header,
+                            int16_t first_free_block) const;
+
+  int16_t last_block_index() const { return num_blocks_ - 1; }
+
+  bool is_index_valid(int16_t index) const {
+    return index >= 0 && index <= last_block_index();
+  }
+
+  // Returns a reference to the AtomicBlockHeader for the block at `index`.
+  AtomicBlockHeader& block_header_at(int16_t index) const {
+    ABSL_ASSERT(is_index_valid(index));
+    return *reinterpret_cast<AtomicBlockHeader*>(&region_[block_size_ * index]);
+  }
+
+  // Returns a FreeBlock corresponding to the block at `index`.
+  FreeBlock free_block_at(int16_t index) const {
+    return FreeBlock(index, block_header_at(index));
+  }
+
+  absl::Span<uint8_t> region_;
+  uint32_t block_size_ = 0;
+  int16_t num_blocks_ = 0;
+};
+
+}  // namespace ipcz
+
+#endif  // IPCZ_SRC_IPCZ_BLOCK_ALLOCATOR_H_