Add TSharedMutableRef::AllocateViaMmap

commit_hash:1666d5c27b2dfe54460efdc686d98ad955b0c4f1

5 files changed, 220 insertions, 1 deletions

diff --git a/library/cpp/yt/memory/ref-inl.h b/library/cpp/yt/memory/ref-inl.h
index 34a0e365c19..632a98f130d 100644
--- a/library/cpp/yt/memory/ref-inl.h
+++ b/library/cpp/yt/memory/ref-inl.h
@@ -226,6 +226,11 @@ Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t s
     return AllocatePageAligned<TDefaultSharedBlobTag>(size, options);
 }
 
+Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options)
+{
+    return AllocateViaMmap<TDefaultSharedBlobTag>(size, options);
+}
+
 template <class TTag>
 Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::MakeCopy(TRef ref)
 {
@@ -260,6 +265,13 @@ Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t s
     return AllocatePageAligned(size, options, GetRefCountedTypeCookie<TTag>());
 }
 
+template <class TTag>
+Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options)
+{
+    static_assert(IsEmptyClass<TTag>());
+    return AllocateViaMmap(size, options, GetRefCountedTypeCookie<TTag>());
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 Y_FORCE_INLINE size_t GetByteSize(TRef ref)
diff --git a/library/cpp/yt/memory/ref.cpp b/library/cpp/yt/memory/ref.cpp
index 780f098233c..8601c2ba3e1 100644
--- a/library/cpp/yt/memory/ref.cpp
+++ b/library/cpp/yt/memory/ref.cpp
@@ -3,15 +3,28 @@
 #include "blob.h"
 #include "poison.h"
 
+#include <library/cpp/yt/exception/exception.h>
+
 #include <library/cpp/yt/malloc/malloc.h>
 
 #include <library/cpp/yt/misc/port.h>
 
 #include <library/cpp/yt/string/format.h>
 
+#include <util/generic/size_literals.h>
+
+#include <util/string/printf.h>
+
 #include <util/system/info.h>
 #include <util/system/align.h>
 
+#ifdef _linux_
+#include <errno.h>
+#include <string.h>
+
+#include <sys/mman.h>
+#endif
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -259,6 +272,71 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+#ifdef _linux_
+
+class TMmapAllocationHolder
+    : public TAllocationHolderBase<TMmapAllocationHolder>
+{
+public:
+    TMmapAllocationHolder(
+        size_t size,
+        TSharedMutableRefAllocateViaMmapOptions options,
+        TRefCountedTypeCookie cookie)
+    {
+        // Round the mapping up to the huge page size when huge pages are
+        // requested so that the whole region can be backed by them.
+        auto alignment = options.UseThp ? TransparentPageSize : GetPageSize();
+        MappedSize_ = AlignUp(size, alignment);
+        auto* ptr = ::mmap(
+            nullptr,
+            MappedSize_,
+            PROT_READ | PROT_WRITE,
+            MAP_PRIVATE | MAP_ANONYMOUS,
+            -1,
+            0);
+        if (ptr == MAP_FAILED) {
+            throw TSimpleException(Sprintf("Failed to mmap %" PRISZT " bytes: %s", MappedSize_, strerror(errno)));
+        }
+        Begin_ = static_cast<char*>(ptr);
+#ifdef MADV_HUGEPAGE
+        if (options.UseThp) {
+            // Hint before the first touch so that faults are served by huge
+            // pages directly rather than relying on khugepaged to collapse
+            // already-populated small pages.
+            YT_VERIFY(::madvise(Begin_, MappedSize_, MADV_HUGEPAGE) == 0);
+        }
+#endif
+        Initialize(size, {.InitializeStorage = options.InitializeStorage}, cookie);
+    }
+
+    ~TMmapAllocationHolder()
+    {
+        Finalize();
+        YT_VERIFY(::munmap(Begin_, MappedSize_) == 0);
+    }
+
+    char* GetBegin()
+    {
+        return Begin_;
+    }
+
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return MappedSize_;
+    }
+
+private:
+    static constexpr size_t TransparentPageSize = 2_MB;
+
+    char* Begin_;
+    size_t MappedSize_;
+};
+
+#endif
+
+////////////////////////////////////////////////////////////////////////////////
+
 TRef TRef::FromBlob(const TBlob& blob)
 {
     return TRef(blob.Begin(), blob.Size());
@@ -364,6 +442,21 @@ TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t size, TSharedMut
     return TSharedMutableRef(ref, std::move(holder));
 }
 
+TSharedMutableRef TSharedMutableRef::AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options, TRefCountedTypeCookie tagCookie)
+{
+#ifdef _linux_
+    auto holder = New<TMmapAllocationHolder>(size, options, tagCookie);
+    auto ref = holder->GetRef();
+    return TSharedMutableRef(ref, std::move(holder));
+#else
+    // No mmap/huge page support; fall back to a plain page-aligned allocation.
+    return AllocatePageAligned(
+        size,
+        TSharedMutableRefAllocateOptions{.InitializeStorage = options.InitializeStorage},
+        tagCookie);
+#endif
+}
+
 TSharedMutableRef TSharedMutableRef::AllocateAligned(size_t size, size_t alignment, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie)
 {
     auto holder = New<TCustomAlignedAllocationHolder>(size, alignment, options, tagCookie);
diff --git a/library/cpp/yt/memory/ref.h b/library/cpp/yt/memory/ref.h
index 028afffd354..a6c6c4f2260 100644
--- a/library/cpp/yt/memory/ref.h
+++ b/library/cpp/yt/memory/ref.h
@@ -195,13 +195,25 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-//! Various options for allocating TSharedMutableRef.
+//! Various options for allocating TSharedMutableRef via malloc.
 struct TSharedMutableRefAllocateOptions
 {
+    //! If true, the storage is zero-initialized; otherwise its contents are undefined.
     bool InitializeStorage = true;
+    //! If true, the allocation is extended to the usable size reported by the allocator.
     bool ExtendToUsableSize = false;
 };
 
+//! Various options for allocating TSharedMutableRef via mmap.
+struct TSharedMutableRefAllocateViaMmapOptions
+{
+    //! See TSharedMutableRefAllocateOptions::InitializeStorage.
+    bool InitializeStorage = true;
+    //! Hints the kernel to back the region with transparent huge pages.
+    //! Best-effort; ignored on non-Linux systems.
+    bool UseThp = false;
+};
+
 //! A reference to a mutable range of memory with shared ownership.
 //! Use with caution :)
 class TSharedMutableRef
@@ -262,6 +274,21 @@ public:
     //! The memory is marked with a given tag.
     static TSharedMutableRef AllocatePageAligned(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie);
 
+    //! Allocates a new shared block of memory backed by an anonymous mmap region.
+    //! Optionally requests transparent huge pages; see #UseThp.
+    //! Falls back to AllocatePageAligned on non-Linux systems.
+    //! The memory is marked with a given tag.
+    template <class TTag>
+    static TSharedMutableRef AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options = {});
+
+    //! Allocates a new shared block of memory backed by an anonymous mmap region.
+    //! The memory is marked with TDefaultSharedBlobTag.
+    static TSharedMutableRef AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options = {});
+
+    //! Allocates a new shared block of memory backed by an anonymous mmap region.
+    //! The memory is marked with a given tag.
+    static TSharedMutableRef AllocateViaMmap(size_t size, TSharedMutableRefAllocateViaMmapOptions options, TRefCountedTypeCookie tagCookie);
+
     //! Allocates a new aligned shared block of memory.
     //! #size must be divisible by alignment size.
     //! The memory is marked with a given tag.
diff --git a/library/cpp/yt/memory/unittests/ref_ut.cpp b/library/cpp/yt/memory/unittests/ref_ut.cpp
index 8ecd57e15d2..734fc030d29 100644
--- a/library/cpp/yt/memory/unittests/ref_ut.cpp
+++ b/library/cpp/yt/memory/unittests/ref_ut.cpp
@@ -4,6 +4,8 @@
 
 #include <library/cpp/yt/memory/ref.h>
 
+#include <util/generic/size_literals.h>
+
 namespace NYT::NYson {
 namespace {
 
@@ -18,5 +20,88 @@ TEST(TSharedRefTest, Save)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+constexpr char Payload = 0x42;
+
+TEST(TSharedMutableRefTest, Allocate)
+{
+    auto ref = TSharedMutableRef::Allocate(64_KB);
+    EXPECT_EQ(ref.Size(), 64_KB);
+
+    // InitializeStorage defaults to true.
+    for (char ch : ref) {
+        EXPECT_EQ(ch, 0);
+    }
+
+    // The whole range must be writable.
+    ::memset(ref.Begin(), Payload, ref.Size());
+    EXPECT_EQ(ref[0], Payload);
+    EXPECT_EQ(ref[ref.Size() - 1], Payload);
+}
+
+TEST(TSharedMutableRefTest, AllocateNoInitialize)
+{
+    auto ref = TSharedMutableRef::Allocate(64_KB, {.InitializeStorage = false});
+    EXPECT_EQ(ref.Size(), 64_KB);
+    ::memset(ref.Begin(), Payload, ref.Size());
+}
+
+TEST(TSharedMutableRefTest, AllocateExtendToUsableSize)
+{
+    auto ref = TSharedMutableRef::Allocate(64_KB, {.ExtendToUsableSize = true});
+    EXPECT_GE(ref.Size(), 64_KB);
+
+    // InitializeStorage defaults to true.
+    for (char ch : ref) {
+        EXPECT_EQ(ch, 0);
+    }
+
+    ::memset(ref.Begin(), Payload, ref.Size());
+    EXPECT_EQ(ref[ref.Size() - 1], Payload);
+}
+
+TEST(TSharedMutableRefTest, AllocateNoInitializeExtendToUsableSize)
+{
+    auto ref = TSharedMutableRef::Allocate(
+        64_KB,
+        {.InitializeStorage = false, .ExtendToUsableSize = true});
+    EXPECT_GE(ref.Size(), 64_KB);
+    ::memset(ref.Begin(), Payload, ref.Size());
+}
+
+TEST(TSharedMutableRefTest, AllocateViaMmap)
+{
+    auto ref = TSharedMutableRef::AllocateViaMmap(64_KB);
+    EXPECT_EQ(ref.Size(), 64_KB);
+
+    // InitializeStorage defaults to true.
+    for (char ch : ref) {
+        EXPECT_EQ(ch, 0);
+    }
+
+    // The whole range must be writable.
+    ::memset(ref.Begin(), Payload, ref.Size());
+    EXPECT_EQ(ref[0], Payload);
+    EXPECT_EQ(ref[ref.Size() - 1], Payload);
+}
+
+TEST(TSharedMutableRefTest, AllocateViaMmapNoInitialize)
+{
+    auto ref = TSharedMutableRef::AllocateViaMmap(64_KB, {.InitializeStorage = false});
+    EXPECT_EQ(ref.Size(), 64_KB);
+    ::memset(ref.Begin(), Payload, ref.Size());
+}
+
+TEST(TSharedMutableRefTest, AllocateViaMmapThp)
+{
+    // Larger than a single huge page to exercise the rounded-up mapping.
+    constexpr size_t Size = 5_MB;
+    auto ref = TSharedMutableRef::AllocateViaMmap(Size, {.UseThp = true});
+    EXPECT_EQ(ref.Size(), Size);
+    ::memset(ref.Begin(), Payload, ref.Size());
+    EXPECT_EQ(ref[Size - 1], Payload);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace
 } // namespace NYT::NYson
diff --git a/library/cpp/yt/memory/ya.make b/library/cpp/yt/memory/ya.make
index 75904d1680e..37b05048b76 100644
--- a/library/cpp/yt/memory/ya.make
+++ b/library/cpp/yt/memory/ya.make
@@ -27,6 +27,7 @@ SRCS(
 PEERDIR(
     library/cpp/sanitizer/include
     library/cpp/yt/assert
+    library/cpp/yt/exception
     library/cpp/yt/misc
     library/cpp/yt/malloc
     library/cpp/yt/system
@@ -39,6 +40,7 @@ CHECK_DEPENDENT_DIRS(
     library
     util
     library/cpp/yt/assert
+    library/cpp/yt/exception
     library/cpp/yt/misc
     library/cpp/yt/malloc
 )