reentrant_allocator

Summary:
[Folly] `reentrant_allocator`, a multi-thread-safe and async-signal-safe allocator.

Based on `mmap` for dynamic allocation. While, technically, `mmap` is not documented to be async-signal-safe, in practice it is so at least on Linux. Take advantage.

Details:
* Large allocations are handled directly by `mmap` and deallocations by `munmap`, where large-size is 2^12 (1 page) by default; they are not tracked; they are aligned only to page boundaries.
* Small allocations are handled by an `mmap`-backed refcounted arena list. Arena sections are block-size bytes less overhead by default, where block-size is 2^16 (16 pages) by default; they are allocated on demand; and they are linked together to allow for `munmap` when the allocator refcounted arena list is destroyed.

Reviewed By: nbronson, luciang

Differential Revision: D19222635

fbshipit-source-id: adf30580c1bf3dd7f8dab13b1d4ac1e10b2f5c52

CMakeLists.txt

@@ -749,6 +749,7 @@ if (BUILD_TESTS)
 
     DIRECTORY memory/test/
       TEST arena_test SOURCES ArenaTest.cpp
+      TEST mmap_allocator_test SOURCES MmapAllocatorTest.cpp
       TEST thread_cached_arena_test WINDOWS_DISABLED
         SOURCES ThreadCachedArenaTest.cpp
       TEST mallctl_helper_test SOURCES MallctlHelperTest.cpp

folly/memory/ReentrantAllocator.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/memory/ReentrantAllocator.h>
+
+#include <new>
+#include <utility>
+
+#include <folly/lang/Bits.h>
+#include <folly/lang/SafeAssert.h>
+#include <folly/portability/SysMman.h>
+
+namespace folly {
+
+namespace {
+
+max_align_t dummy; // return value for zero-sized allocations
+
+void* reentrant_allocate(std::size_t const n) noexcept {
+  FOLLY_SAFE_CHECK(n, "zero-sized");
+  auto const prot = PROT_READ | PROT_WRITE;
+  auto const flags = MAP_ANONYMOUS | MAP_PRIVATE;
+  auto const addr = ::mmap(nullptr, n, prot, flags, 0, 0);
+  FOLLY_SAFE_PCHECK(addr != MAP_FAILED, "mmap failed");
+  return addr;
+}
+
+void reentrant_deallocate(void* const p, std::size_t const n) noexcept {
+  FOLLY_SAFE_CHECK(p, "null-pointer");
+  FOLLY_SAFE_CHECK(n, "zero-sized");
+  auto const err = ::munmap(p, n);
+  FOLLY_SAFE_PCHECK(!err, "munmap failed");
+}
+
+} // namespace
+
+namespace detail {
+
+reentrant_allocator_base::reentrant_allocator_base(
+    reentrant_allocator_options const& options) noexcept {
+  meta_ = static_cast<meta_t*>(reentrant_allocate(sizeof(meta_t)));
+  ::new (meta_) meta_t(options);
+}
+
+reentrant_allocator_base::reentrant_allocator_base(
+    reentrant_allocator_base const& that) noexcept {
+  meta_ = that.meta_;
+  meta_->refs.fetch_add(1, std::memory_order_relaxed);
+}
+
+reentrant_allocator_base& reentrant_allocator_base::operator=(
+    reentrant_allocator_base const& that) noexcept {
+  if (this != &that) {
+    if (meta_->refs.fetch_sub(1, std::memory_order_acq_rel) - 1 == 0) {
+      obliterate();
+    }
+    meta_ = that.meta_;
+    meta_->refs.fetch_add(1, std::memory_order_relaxed);
+  }
+  return *this;
+}
+
+reentrant_allocator_base::~reentrant_allocator_base() {
+  if (meta_->refs.fetch_sub(1, std::memory_order_acq_rel) - 1 == 0) {
+    obliterate();
+  }
+}
+
+void* reentrant_allocator_base::allocate(
+    std::size_t const n,
+    std::size_t const a) noexcept {
+  if (!n) {
+    return &dummy;
+  }
+  //  large requests are handled directly
+  if (n >= meta_->large_size) {
+    return reentrant_allocate(n);
+  }
+  auto const block_size = meta_->block_size;
+  //  small requests are handled from the shared arena list:
+  //  * if the list is empty or the list head has insufficient space, c/x a new
+  //    list head, starting over on failure
+  //  * then c/x the list head size to the new size, starting over on failure
+  while (true) {
+    //  load head - non-const because used in c/x below
+    auto head = meta_->head.load(std::memory_order_acquire);
+    //  load size - non-const because used in c/x below
+    //  size is where the prev allocation ends, if any
+    auto size = head //
+        ? head->size.load(std::memory_order_acquire)
+        : block_size;
+    //  offset is where the next allocation starts, and is aligned as a
+    auto const offset = (size + a - 1) & ~(a - 1);
+    //  if insufficient space in current segment or no current segment at all
+    if (offset + n > block_size || !head) {
+      //  mmap a new segment and try to c/x it in to be the segment list head
+      auto const newhead = static_cast<node_t*>(reentrant_allocate(block_size));
+      ::new (newhead) node_t(head);
+      auto const exchanged = meta_->head.compare_exchange_weak(
+          head, newhead, std::memory_order_release, std::memory_order_relaxed);
+      if (!exchanged) {
+        //  lost the race - munmap the new segment and start over
+        reentrant_deallocate(newhead, block_size);
+        continue;
+      }
+      head = newhead;
+    }
+    //  compute the new size and try to c/x it in to be the head segment size
+    auto const newsize = offset + n;
+    auto const exchanged = head->size.compare_exchange_weak(
+        size, newsize, std::memory_order_release, std::memory_order_relaxed);
+    if (!exchanged) {
+      //  lost the race - start over
+      continue;
+    }
+    return reinterpret_cast<char*>(head) + offset;
+  }
+}
+
+void reentrant_allocator_base::deallocate(
+    void* const p,
+    std::size_t const n) noexcept {
+  if (p == &dummy) {
+    FOLLY_SAFE_CHECK(n == 0, "unexpected non-zero size");
+    return;
+  }
+  if (!n || !p) {
+    return;
+  }
+  //  large requests are handled directly
+  if (n >= meta_->large_size) {
+    reentrant_deallocate(p, n);
+    return;
+  }
+  //  small requests are deferred to allocator destruction, so no-op here
+}
+
+void reentrant_allocator_base::obliterate() noexcept {
+  auto head = meta_->head.load(std::memory_order_acquire);
+  while (head != nullptr) {
+    auto const prev = std::exchange(head, head->next);
+    reentrant_deallocate(prev, meta_->block_size);
+  }
+  reentrant_deallocate(meta_, sizeof(meta_));
+  meta_ = nullptr;
+}
+
+} // namespace detail
+
+} // namespace folly

folly/memory/ReentrantAllocator.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+
+namespace folly {
+
+class reentrant_allocator_options {
+ public:
+  //  block_size_lg
+  //
+  //  The log2 of the block size, which is the size of the blocks from which
+  //  small allocations are returned.
+  std::size_t block_size_lg() const noexcept {
+    return block_size_lg_;
+  }
+  reentrant_allocator_options& block_size_lg(std::size_t const value) noexcept {
+    block_size_lg_ = value;
+    return *this;
+  }
+
+  //  large_size_lg
+  //
+  //  The log2 of the large size, which is the size starting at which
+  //  allocations are considered large and are returned directly from mmap.
+  std::size_t large_size_lg() const noexcept {
+    return large_size_lg_;
+  }
+  reentrant_allocator_options& large_size_lg(std::size_t const value) noexcept {
+    large_size_lg_ = value;
+    return *this;
+  }
+
+ private:
+  std::size_t block_size_lg_ = 16;
+  std::size_t large_size_lg_ = 12;
+};
+
+namespace detail {
+
+class reentrant_allocator_base {
+ public:
+  explicit reentrant_allocator_base(
+      reentrant_allocator_options const& options) noexcept;
+  reentrant_allocator_base(reentrant_allocator_base const& that) noexcept;
+  reentrant_allocator_base& operator=(
+      reentrant_allocator_base const& that) noexcept;
+  ~reentrant_allocator_base();
+
+  void* allocate(std::size_t n, std::size_t a) noexcept;
+  void deallocate(void* p, std::size_t n) noexcept;
+
+  friend bool operator==(
+      reentrant_allocator_base const& a,
+      reentrant_allocator_base const& b) noexcept {
+    return a.meta_ == b.meta_;
+  }
+  friend bool operator!=(
+      reentrant_allocator_base const& a,
+      reentrant_allocator_base const& b) noexcept {
+    return a.meta_ != b.meta_;
+  }
+
+ private:
+  //  For small sizes, maintain a shared list of segments. Segments are each
+  //  allocated via mmap, chained together into a list, and collectively
+  //  refcounted. When the last copy of the allocator is destroyed, segments
+  //  are deallocated all at once via munmap. Node is the header data
+  //  structure prefixing each segment while Meta is the data structure
+  //  representing shared ownership of the segment list. Serve allocations
+  //  from the head segment if it exists and has space, otherwise mmap and
+  //  chain a new segment and serve allocations from it. Serve deallocations
+  //  by doing nothing at all.
+  struct node_t {
+    node_t* next = nullptr;
+    std::atomic<std::size_t> size{sizeof(node_t)};
+    explicit node_t(node_t* next_) noexcept : next{next_} {}
+  };
+
+  //  The shared state which all copies of the allocator share.
+  struct meta_t {
+    //  Small allocations are served from block-sized segments.
+    std::size_t const block_size = 0;
+    //  Large allocations are served directly.
+    std::size_t const large_size = 0;
+    //  The refcount is atomic to permit some copies of the allocator to be
+    //  destroyed concurrently with uses of other copies of the allocator.
+    //  This lets an allocator be copied in a signal handler and the copy
+    //  be destroyed outside the signal handler.
+    std::atomic<std::size_t> refs{1};
+    //  The segment list head. All small allocations happen via the head node
+    //  if possible, or via a new head node otherwise.
+    std::atomic<node_t*> head{nullptr};
+
+    explicit meta_t(reentrant_allocator_options const& options) noexcept
+        : block_size{std::size_t(1) << options.block_size_lg()},
+          large_size{std::size_t(1) << options.large_size_lg()} {}
+  };
+
+  //  Deduplicates code between dtor and copy-assignment.
+  void obliterate() noexcept;
+
+  //  The allocator has all state in the shared state, keeping only a pointer.
+  meta_t* meta_{nullptr};
+};
+
+} // namespace detail
+
+//  reentrant_allocator
+//
+//  A reentrant mmap-based allocator.
+//
+//  Safety:
+//  * multi-thread-safe
+//  * async-signal-safe
+//
+//  The basic approach is in two parts:
+//  * For large sizes, serve allocations and deallocations directly via mmap and
+//    munmap and without any extra tracking.
+//  * For small sizes, serve allocations from a refcounted shared list of
+//    segments and defer deallocations to amortize calls to mmap and munmap - in
+//    other words, a shared arena list.
+//
+//  Large allocations are aligned to page boundaries, even if the type's natural
+//  alignment is larger.
+//
+//  Assumptions:
+//  * The mmap and munmap libc functions are async-signal-safe in practice even
+//    though POSIX does not require them to be.
+//  * The instances of std::atomic over size_t and pointer types are lock-free
+//    and operations on them are async-signal-safe.
+template <typename T>
+class reentrant_allocator : private detail::reentrant_allocator_base {
+ private:
+  template <typename>
+  friend class reentrant_allocator;
+
+  using base = detail::reentrant_allocator_base;
+
+ public:
+  using value_type = T;
+
+  using base::base;
+
+  template <typename U>
+  /* implicit */ reentrant_allocator(
+      reentrant_allocator<U> const& that) noexcept
+      : base{that} {}
+
+  T* allocate(std::size_t n) {
+    return static_cast<T*>(base::allocate(n * sizeof(T), alignof(T)));
+  }
+  void deallocate(T* p, std::size_t n) {
+    base::deallocate(p, n * sizeof(T));
+  }
+
+  template <typename A, typename B>
+  friend bool operator==(
+      reentrant_allocator<A> const& a,
+      reentrant_allocator<B> const& b) noexcept;
+  template <typename A, typename B>
+  friend bool operator!=(
+      reentrant_allocator<A> const& a,
+      reentrant_allocator<B> const& b) noexcept;
+};
+
+template <typename A, typename B>
+bool operator==(
+    reentrant_allocator<A> const& a,
+    reentrant_allocator<B> const& b) noexcept {
+  using base = detail::reentrant_allocator_base;
+  return static_cast<base const&>(a) == static_cast<base const&>(b);
+}
+template <typename A, typename B>
+bool operator!=(
+    reentrant_allocator<A> const& a,
+    reentrant_allocator<B> const& b) noexcept {
+  using base = detail::reentrant_allocator_base;
+  return static_cast<base const&>(a) != static_cast<base const&>(b);
+}
+
+} // namespace folly

folly/memory/test/ReentrantAllocatorTest.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/memory/ReentrantAllocator.h>
+
+#include <memory>
+#include <thread>
+#include <tuple>
+#include <vector>
+
+#include <folly/Utility.h>
+#include <folly/portability/GTest.h>
+
+class ReentrantAllocatorTest : public testing::Test {};
+
+TEST_F(ReentrantAllocatorTest, equality) {
+  folly::reentrant_allocator<void> a{folly::reentrant_allocator_options{}};
+  folly::reentrant_allocator<void> b{folly::reentrant_allocator_options{}};
+  EXPECT_TRUE(a == a);
+  EXPECT_FALSE(a == b);
+  EXPECT_FALSE(a != a);
+  EXPECT_TRUE(a != b);
+
+  folly::reentrant_allocator<int> a2{a};
+  folly::reentrant_allocator<int> b2{b};
+  EXPECT_TRUE(a2 == a);
+  EXPECT_FALSE(a2 == b);
+  EXPECT_TRUE(a2 == a2);
+  EXPECT_FALSE(a2 == b2);
+  EXPECT_FALSE(a2 != a);
+  EXPECT_TRUE(a2 != b);
+  EXPECT_FALSE(a2 != a2);
+  EXPECT_TRUE(a2 != b2);
+}
+
+TEST_F(ReentrantAllocatorTest, small) {
+  folly::reentrant_allocator<void> alloc{folly::reentrant_allocator_options{}};
+  std::vector<size_t, folly::reentrant_allocator<size_t>> vec{alloc};
+  for (auto i = 0u; i < (1u << 16); ++i) {
+    vec.push_back(i);
+  }
+  for (auto i = 0u; i < (1u << 16); ++i) {
+    EXPECT_EQ(i, vec.at(i));
+  }
+}
+
+TEST_F(ReentrantAllocatorTest, large) {
+  constexpr size_t const large_size_lg = 6;
+  struct alignas(1u << large_size_lg) type : std::tuple<size_t> {};
+  auto const opts = folly::reentrant_allocator_options{} //
+                        .large_size_lg(large_size_lg);
+  folly::reentrant_allocator<void> alloc{opts};
+  std::vector<type, folly::reentrant_allocator<type>> vec{alloc};
+  for (auto i = 0u; i < (1u << 16); ++i) {
+    vec.push_back({i});
+  }
+  for (auto i = 0u; i < (1u << 16); ++i) {
+    EXPECT_EQ(i, std::get<0>(vec.at(i)));
+  }
+}
+
+TEST_F(ReentrantAllocatorTest, zero) {
+  folly::reentrant_allocator<int> a{folly::reentrant_allocator_options{}};
+  a.deallocate(nullptr, 0);
+  auto ptr = a.allocate(0);
+  EXPECT_NE(nullptr, ptr);
+  a.deallocate(ptr, 0);
+}
+
+TEST_F(ReentrantAllocatorTest, self_assignment) {
+  folly::reentrant_allocator<int> a{folly::reentrant_allocator_options{}};
+  auto& i = *a.allocate(1);
+  ::new (&i) int(7);
+  EXPECT_EQ(7, i);
+  a = folly::as_const(a);
+  EXPECT_EQ(7, i);
+  a.deallocate(&i, 1);
+}
+
+TEST_F(ReentrantAllocatorTest, stress) {
+  struct alignas(256) big {};
+  folly::reentrant_allocator<void> a{folly::reentrant_allocator_options{}};
+  std::vector<std::thread> threads{4};
+  std::atomic<bool> done{false};
+  for (auto& th : threads) {
+    th = std::thread([&done, a] {
+      while (!done.load(std::memory_order_relaxed)) {
+        std::allocate_shared<big>(a);
+      }
+    });
+  }
+  /* sleep override */ std::this_thread::sleep_for(std::chrono::seconds(1));
+  done.store(true, std::memory_order_relaxed);
+  for (auto& th : threads) {
+    th.join();
+  }
+}