Add TSharedMutableRef::AllocateViaMmap

commit_hash:1666d5c27b2dfe54460efdc686d98ad955b0c4f1

1 files changed, 93 insertions, 0 deletions

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);