Optimize MPT node allocation with custom allocator and SharedPtr

category/core/mem/allocators.hpp

@@ -292,6 +292,100 @@ namespace allocators
             throw;
         }
     }
+
+    /**************************************************************************/
+    //! \brief A STL allocator for use with `std::allocate_shared` that
+    //! allocates extra bytes beyond sizeof(T) for trailing variable-length
+    //! data.
+    //!
+    //! This allocator is designed for types with flexible array members or
+    //! trailing data. When used with `std::allocate_shared`, it ensures the
+    //! control block and object with trailing data are allocated together.
+    //!
+    //! \tparam T The type to allocate
+    template <typename T>
+    struct variable_size_allocator
+    {
+        using value_type = T;
+        using size_type = std::size_t;
+        using difference_type = std::ptrdiff_t;
+
+        //! \brief Construct allocator with total storage size
+        //! \param storage_bytes Total bytes needed (sizeof(T) + trailing data)
+        //!
+        //! The extra_bytes_ member stores the additional bytes beyond sizeof(T)
+        //! needed for trailing data (path, value, child data, etc.)
+        explicit variable_size_allocator(size_t storage_bytes) noexcept
+            : extra_bytes_(storage_bytes - sizeof(T))
+        {
+            MONAD_ASSERT(storage_bytes >= sizeof(T));
+        }
+
+        //! \brief Rebind constructor for allocator conversion
+        template <typename U>
+        // NOLINTNEXTLINE(google-explicit-constructor)
+        variable_size_allocator(
+            variable_size_allocator<U> const &other) noexcept
+            : extra_bytes_(other.extra_bytes_)
+        {
+        }
+
+        //! \brief Allocate memory for n objects of type T plus extra bytes
+        //! \param n Number of objects to allocate (must be 1)
+        //!
+        //! For std::allocate_shared:
+        //! - If T is the object type: allocates sizeof(T) + extra_bytes
+        //! - If T is control block: allocates sizeof(control_block) +
+        //! extra_bytes
+        //!   (control block already includes sizeof(object), so this gives
+        //!   control block + object trailing data)
+        [[nodiscard]] T *allocate(size_type n)
+        {
+            MONAD_ASSERT(n == 1);
+            size_t const bytes = sizeof(T) + extra_bytes_;
+
+            if constexpr (alignof(T) > alignof(max_align_t)) {
+                return reinterpret_cast<T *>(
+                    ::operator new(bytes, std::align_val_t{alignof(T)}));
+            }
+            return reinterpret_cast<T *>(::operator new(bytes));
+        }
+
+        //! \brief Deallocate memory
+        void deallocate(T *p, size_type) noexcept
+        {
+            if constexpr (alignof(T) > alignof(max_align_t)) {
+                ::operator delete(p, std::align_val_t{alignof(T)});
+            }
+            else {
+                ::operator delete(p);
+            }
+        }
+
+        //! \brief Rebind to allocate different types
+        template <typename U>
+        struct rebind
+        {
+            using other = variable_size_allocator<U>;
+        };
+
+        bool operator==(variable_size_allocator const &other) const noexcept
+        {
+            return extra_bytes_ == other.extra_bytes_;
+        }
+
+        bool operator!=(variable_size_allocator const &other) const noexcept
+        {
+            return !(*this == other);
+        }
+
+        template <typename U>
+        friend struct variable_size_allocator;
+
+    private:
+        //! Extra bytes beyond sizeof(T) for trailing data
+        size_t extra_bytes_;
+    };
 }
 
 MONAD_NAMESPACE_END

category/mpt/copy_trie.cpp

@@ -33,9 +33,9 @@
 
 MONAD_MPT_NAMESPACE_BEGIN
 
-Node::UniquePtr create_node_add_new_branch(
+Node::SharedPtr create_node_add_new_branch(
     UpdateAuxImpl &aux, Node *const node, unsigned char const new_branch,
-    Node::UniquePtr new_child, uint64_t const new_version,
+    Node::SharedPtr new_child, uint64_t const new_version,
     std::optional<byte_string_view> opt_value)
 {
     uint16_t const mask =
@@ -80,9 +80,9 @@ Node::UniquePtr create_node_add_new_branch(
         static_cast<int64_t>(new_version));
 }
 
-Node::UniquePtr create_node_with_two_children(
+Node::SharedPtr create_node_with_two_children(
     UpdateAuxImpl &aux, NibblesView const path, unsigned char const branch0,
-    Node::UniquePtr child0, unsigned char const branch1, Node::UniquePtr child1,
+    Node::SharedPtr child0, unsigned char const branch1, Node::SharedPtr child1,
     uint64_t const new_version, std::optional<byte_string_view> opt_value)
 {
     // mismatch: split node's path: turn node to a branch node with two
@@ -158,7 +158,7 @@ Node::SharedPtr copy_trie_impl(
     Node *parent = nullptr;
     unsigned char branch = INVALID_BRANCH;
     Node::SharedPtr node = dest_root;
-    Node::UniquePtr new_node{};
+    Node::SharedPtr new_node{};
     unsigned prefix_index = 0;
     unsigned node_prefix_index = 0;
 
@@ -216,7 +216,7 @@ Node::SharedPtr copy_trie_impl(
         if (node->mask & (1u << nibble)) {
             auto const index = node->to_child_index(nibble);
             if (node->next(index) == nullptr) {
-                Node::UniquePtr next_node_ondisk =
+                auto next_node_ondisk =
                     read_node_blocking(aux, node->fnext(index), dest_version);
                 MONAD_ASSERT(next_node_ondisk != nullptr);
                 node->set_next(index, std::move(next_node_ondisk));

category/mpt/deserialize_node_from_receiver_result.hpp

@@ -54,12 +54,12 @@ namespace detail
     }
 
     template <class NodeType, class ResultType>
-    inline NodeType::UniquePtr deserialize_node_from_receiver_result(
+    inline NodeType::SharedPtr deserialize_node_from_receiver_result(
         ResultType buffer_, uint16_t buffer_off,
         MONAD_ASYNC_NAMESPACE::erased_connected_operation *io_state)
     {
         MONAD_ASSERT(buffer_);
-        typename NodeType::UniquePtr node;
+        typename NodeType::SharedPtr node;
         if constexpr (std::is_same_v<
                           std::decay_t<ResultType>,
                           typename monad::async::read_single_buffer_sender::

category/mpt/find.cpp

@@ -52,7 +52,7 @@ find_cursor_result_type find_blocking(
                 MONAD_ASSERT(aux.is_on_disk());
                 auto g2(g.upgrade());
                 if (g2.upgrade_was_atomic() || !node->next(idx)) {
-                    Node::UniquePtr next_node_ondisk =
+                    auto next_node_ondisk =
                         read_node_blocking(aux, node->fnext(idx), version);
                     if (!next_node_ondisk) {
                         return {

category/mpt/node.cpp

@@ -456,7 +456,7 @@ void ChildData::erase()
 }
 
 void ChildData::finalize(
-    Node::UniquePtr node, Compute &compute, bool const cache)
+    Node::SharedPtr node, Compute &compute, bool const cache)
 {
     MONAD_DEBUG_ASSERT(is_valid());
     ptr = std::move(node);
@@ -488,13 +488,13 @@ void ChildData::copy_old_child(Node *const old, unsigned const i)
     MONAD_DEBUG_ASSERT(is_valid());
 }
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     Node &from, NibblesView const path,
     std::optional<byte_string_view> const value, int64_t const version)
 {
     auto const value_size =
         value.transform(&byte_string_view::size).value_or(0);
-    auto node = Node::make(
+    auto node = Node::make_shared(
         calculate_node_size(
             from.number_of_children(),
             from.child_data_len(),
@@ -530,17 +530,11 @@ Node::UniquePtr make_node(
     return node;
 }
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     uint16_t const mask, std::span<ChildData> const children,
     NibblesView const path, std::optional<byte_string_view> const value,
     size_t const data_size, int64_t const version)
 {
-    MONAD_DEBUG_ASSERT(data_size <= KECCAK256_SIZE);
-    if (value.has_value()) {
-        MONAD_DEBUG_ASSERT(
-            value->size() <=
-            std::numeric_limits<decltype(Node::value_len)>::max());
-    }
     for (size_t i = 0; i < 16; ++i) {
         MONAD_DEBUG_ASSERT(
             !std::ranges::contains(children, i, &ChildData::branch) ||
@@ -558,7 +552,7 @@ Node::UniquePtr make_node(
         }
     }
 
-    auto node = Node::make(
+    auto node = Node::make_shared(
         calculate_node_size(
             number_of_children,
             total_child_data_size,
@@ -596,7 +590,7 @@ Node::UniquePtr make_node(
     return node;
 }
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     uint16_t const mask, std::span<ChildData> const children,
     NibblesView const path, std::optional<byte_string_view> const value,
     byte_string_view const data, int64_t const version)
@@ -608,7 +602,7 @@ Node::UniquePtr make_node(
 
 // all children's offset are set before creating parent
 // create node with at least one child
-Node::UniquePtr create_node_with_children(
+Node::SharedPtr create_node_with_children(
     Compute &comp, uint16_t const mask, std::span<ChildData> const children,
     NibblesView const path, std::optional<byte_string_view> const value,
     int64_t const version)

category/mpt/node.hpp

@@ -158,7 +158,7 @@ class NodeBase
         max_disk_size + max_number_of_children * KECCAK256_SIZE;
 
     template <node_type DestNodeType, node_type SrcNodeType>
-    friend DestNodeType::UniquePtr copy_node(SrcNodeType const *const from);
+    friend DestNodeType::SharedPtr copy_node(SrcNodeType const *const from);
 
     /* 16-bit mask for children */
     uint16_t mask{0};
@@ -327,6 +327,16 @@ class Node final : public NodeBase
             std::forward<Args>(args)...);
     }
 
+    template <class... Args>
+    static SharedPtr make_shared(size_t bytes, Args &&...args)
+    {
+        allocators::variable_size_allocator<Node> alloc(bytes);
+        return std::allocate_shared<Node>(
+            alloc,
+            prevent_public_construction_tag{},
+            std::forward<Args>(args)...);
+    }
+
     SharedPtr *child_ptr(unsigned index) noexcept;
     SharedPtr const *child_ptr(unsigned index) const noexcept;
 
@@ -351,6 +361,7 @@ class CacheNode final : public NodeBase
     using Deleter = allocators::unique_ptr_aliasing_allocator_deleter<
         &allocators::aliasing_allocator_pair<CacheNode>>;
     using UniquePtr = std::unique_ptr<CacheNode, Deleter>;
+    using SharedPtr = std::shared_ptr<CacheNode>;
 
     CacheNode(prevent_public_construction_tag)
         : NodeBase()
@@ -368,6 +379,16 @@ class CacheNode final : public NodeBase
             std::forward<Args>(args)...);
     }
 
+    template <class... Args>
+    static SharedPtr make_shared(size_t bytes, Args &&...args)
+    {
+        allocators::variable_size_allocator<CacheNode> alloc(bytes);
+        return std::allocate_shared<CacheNode>(
+            alloc,
+            prevent_public_construction_tag{},
+            std::forward<Args>(args)...);
+    }
+
     void *next(size_t const index) const noexcept;
     void set_next(unsigned const index, void *const ptr) noexcept;
 
@@ -398,7 +419,7 @@ struct ChildData
 
     bool is_valid() const;
     void erase();
-    void finalize(Node::UniquePtr, Compute &, bool cache);
+    void finalize(Node::SharedPtr, Compute &, bool cache);
     void copy_old_child(Node *old, unsigned i);
 };
 
@@ -431,21 +452,21 @@ constexpr size_t MAX_VALUE_LEN_OF_LEAF =
         0 /* number_of_children */, 0 /* child_data_size */, 0 /* value_size */,
         KECCAK256_SIZE /* path_size */, KECCAK256_SIZE /* data_size*/);
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     Node &from, NibblesView path, std::optional<byte_string_view> value,
     int64_t version);
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     uint16_t mask, std::span<ChildData>, NibblesView path,
     std::optional<byte_string_view> value, size_t data_size, int64_t version);
 
-Node::UniquePtr make_node(
+Node::SharedPtr make_node(
     uint16_t mask, std::span<ChildData>, NibblesView path,
     std::optional<byte_string_view> value, byte_string_view data,
     int64_t version);
 
 // create node: either branch/extension, with or without leaf
-Node::UniquePtr create_node_with_children(
+Node::SharedPtr create_node_with_children(
     Compute &, uint16_t mask, std::span<ChildData> children, NibblesView path,
     std::optional<byte_string_view> value, int64_t version);
 
@@ -454,7 +475,7 @@ void serialize_node_to_buffer(
     uint32_t disk_size, unsigned offset = 0);
 
 template <node_type NodeType>
-inline NodeType::UniquePtr
+inline NodeType::SharedPtr
 deserialize_node_from_buffer(unsigned char const *read_pos, size_t max_bytes)
 {
     for (size_t n = 0; n < max_bytes; n += 64) {
@@ -473,7 +494,7 @@ deserialize_node_from_buffer(unsigned char const *read_pos, size_t max_bytes)
     if constexpr (std::same_as<NodeType, Node>) {
         auto const alloc_size = round_up_align<3>(base_size) +
                                 number_of_children * sizeof(Node::SharedPtr);
-        auto node = NodeType::make(alloc_size);
+        auto node = NodeType::make_shared(alloc_size);
         std::copy_n(read_pos, base_size, (unsigned char *)node.get());
         for (unsigned i = 0; i < node->number_of_children(); ++i) {
             new (node->child_ptr(i)) Node::SharedPtr();
@@ -484,7 +505,7 @@ deserialize_node_from_buffer(unsigned char const *read_pos, size_t max_bytes)
     else {
         auto const alloc_size = round_up_align<3>(base_size) +
                                 number_of_children * sizeof(NodeType *);
-        auto node = NodeType::make(alloc_size);
+        auto node = NodeType::make_shared(alloc_size);
         std::copy_n(read_pos, base_size, (unsigned char *)node.get());
         std::memset(
             node->next_data_aligned(),
@@ -496,7 +517,7 @@ deserialize_node_from_buffer(unsigned char const *read_pos, size_t max_bytes)
 }
 
 template <node_type DestNodeType, node_type SrcNodeType>
-DestNodeType::UniquePtr copy_node(SrcNodeType const *const from)
+DestNodeType::SharedPtr copy_node(SrcNodeType const *const from)
 {
     auto const number_of_children = from->number_of_children();
     auto const base_size = static_cast<unsigned>(
@@ -505,7 +526,7 @@ DestNodeType::UniquePtr copy_node(SrcNodeType const *const from)
     if constexpr (std::same_as<DestNodeType, Node>) {
         auto const alloc_size = round_up_align<3>(base_size) +
                                 number_of_children * sizeof(Node::SharedPtr);
-        auto node_copy = DestNodeType::make(alloc_size);
+        auto node_copy = DestNodeType::make_shared(alloc_size);
         std::copy_n(
             (unsigned char *)from, base_size, (unsigned char *)node_copy.get());
         for (unsigned i = 0; i < number_of_children; ++i) {
@@ -516,7 +537,7 @@ DestNodeType::UniquePtr copy_node(SrcNodeType const *const from)
     else {
         auto const next_ptrs_size = number_of_children * sizeof(void *);
         auto const alloc_size = round_up_align<3>(base_size) + next_ptrs_size;
-        auto node_copy = DestNodeType::make(alloc_size);
+        auto node_copy = DestNodeType::make_shared(alloc_size);
         std::copy_n(
             (unsigned char *)from, base_size, (unsigned char *)node_copy.get());
         // reset all in memory children

category/mpt/read_node_blocking.cpp

@@ -29,7 +29,7 @@
 
 MONAD_MPT_NAMESPACE_BEGIN
 
-Node::UniquePtr read_node_blocking(
+Node::SharedPtr read_node_blocking(
     UpdateAuxImpl const &aux, chunk_offset_t const node_offset,
     uint64_t const version)
 {
@@ -69,7 +69,7 @@ Node::UniquePtr read_node_blocking(
     return aux.version_is_valid_ondisk(version)
                ? deserialize_node_from_buffer<Node>(
                      buffer + buffer_off, size_t(bytes_read) - buffer_off)
-               : Node::UniquePtr{};
+               : Node::SharedPtr{};
 }
 
 MONAD_MPT_NAMESPACE_END