YT-25976: fix blocks alignment in io requests

commit_hash:785d73458fcb036340236b123b33a6f3b0b1c1ce

2 files changed, 51 insertions, 0 deletions

diff --git a/library/cpp/yt/memory/ref.cpp b/library/cpp/yt/memory/ref.cpp
index fc153bb02e3..5a6ca1e3373 100644
--- a/library/cpp/yt/memory/ref.cpp
+++ b/library/cpp/yt/memory/ref.cpp
@@ -171,6 +171,44 @@ public:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+class TCustomAlignedAllocationHolder
+    : public TAllocationHolderBase<TCustomAlignedAllocationHolder>
+{
+public:
+    TCustomAlignedAllocationHolder(
+        size_t size,
+        size_t alignment,
+        TSharedMutableRefAllocateOptions options,
+        TRefCountedTypeCookie cookie)
+        : Begin_(static_cast<char*>(::aligned_malloc(size, alignment)))
+        , Alignment_(alignment)
+    {
+        Initialize(size, options, cookie);
+    }
+
+    ~TCustomAlignedAllocationHolder()
+    {
+        ::free(Begin_);
+    }
+
+    char* GetBegin()
+    {
+        return Begin_;
+    }
+
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return AlignUp(Size_, Alignment_);
+    }
+
+private:
+    char* const Begin_;
+    size_t Alignment_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 class TPageAlignedAllocationHolder
     : public TAllocationHolderBase<TPageAlignedAllocationHolder>
 {
@@ -311,6 +349,13 @@ TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t size, TSharedMut
     return TSharedMutableRef(ref, std::move(holder));
 }
 
+TSharedMutableRef TSharedMutableRef::AllocateAligned(size_t size, size_t alignment, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie)
+{
+    auto holder = New<TCustomAlignedAllocationHolder>(size, alignment, options, tagCookie);
+    auto ref = holder->GetRef();
+    return TSharedMutableRef(ref, std::move(holder));
+}
+
 TSharedMutableRef TSharedMutableRef::FromBlob(TBlob&& blob)
 {
     auto ref = TMutableRef::FromBlob(blob);
diff --git a/library/cpp/yt/memory/ref.h b/library/cpp/yt/memory/ref.h
index a30fa2d3034..e9708afae4d 100644
--- a/library/cpp/yt/memory/ref.h
+++ b/library/cpp/yt/memory/ref.h
@@ -262,6 +262,12 @@ public:
     //! The memory is marked with a given tag.
     static TSharedMutableRef AllocatePageAligned(size_t size, TSharedMutableRefAllocateOptions 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.
+    //! Unlike AllocatePageAligned, this method also stores size_t inside holder.
+    static TSharedMutableRef AllocateAligned(size_t size, size_t alignment, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie tagCookie);
+
     //! Creates a TSharedMutableRef for the whole blob taking ownership of its content.
     static TSharedMutableRef FromBlob(TBlob&& blob);