intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

6 files changed, 493 insertions, 0 deletions

diff --git a/library/cpp/coroutine/engine/stack/ut/stack_allocator_ut.cpp b/library/cpp/coroutine/engine/stack/ut/stack_allocator_ut.cpp
new file mode 100644
index 0000000000..a7283d44a3
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/stack_allocator_ut.cpp
@@ -0,0 +1,115 @@
+#include <library/cpp/coroutine/engine/stack/stack_allocator.h>
+#include <library/cpp/coroutine/engine/stack/stack_common.h>
+#include <library/cpp/testing/gtest/gtest.h>
+
+
+using namespace testing;
+
+namespace NCoro::NStack::Tests {
+
+    enum class EAllocator {
+        Pool,    // allocates page-size aligned stacks from pools
+        Simple  // uses malloc/free for each stack
+    };
+
+    class TAllocatorParamFixture : public TestWithParam< std::tuple<EGuard, EAllocator> > {
+    protected: // methods
+    void SetUp() override {
+        EGuard guardType;
+        EAllocator allocType;
+        std::tie(guardType, allocType) = GetParam();
+
+        TMaybe<TPoolAllocatorSettings> poolSettings;
+        if (allocType == EAllocator::Pool) {
+            poolSettings = TPoolAllocatorSettings{};
+        }
+
+        Allocator_ = GetAllocator(poolSettings, guardType);
+    }
+
+    protected: // data
+        THolder<IAllocator> Allocator_;
+    };
+
+
+    TEST_P(TAllocatorParamFixture, StackAllocationAndRelease) {
+        uint64_t stackSize = PageSize * 12;
+        auto stack = Allocator_->AllocStack(stackSize, "test_stack");
+#if defined(_san_enabled_) || !defined(NDEBUG)
+        stackSize *= DebugOrSanStackMultiplier;
+#endif
+
+        // Correct stack should have
+        EXPECT_EQ(stack.GetSize(), stackSize); // predefined size
+        EXPECT_FALSE((uint64_t)stack.GetAlignedMemory() & PageSizeMask); // aligned pointer
+        // Writable workspace
+        auto workspace = Allocator_->GetStackWorkspace(stack.GetAlignedMemory(), stack.GetSize());
+        for (uint64_t i = 0; i < workspace.size(); i += 512) {
+            workspace[i] = 42;
+        }
+        EXPECT_TRUE(Allocator_->CheckStackOverflow(stack.GetAlignedMemory()));
+        EXPECT_TRUE(Allocator_->CheckStackOverride(stack.GetAlignedMemory(), stack.GetSize()));
+
+        Allocator_->FreeStack(stack);
+        EXPECT_FALSE(stack.GetRawMemory());
+    }
+
+    INSTANTIATE_TEST_SUITE_P(AllocatorTestParams, TAllocatorParamFixture,
+            Combine(Values(EGuard::Canary, EGuard::Page), Values(EAllocator::Pool, EAllocator::Simple)));
+
+
+    // ------------------------------------------------------------------------
+    // Test that allocated stack has guards
+    //
+    template<class AllocatorType>
+    THolder<IAllocator> GetAllocator(EGuard guardType);
+
+    struct TPoolTag {};
+    struct TSimpleTag {};
+
+    template<>
+    THolder<IAllocator> GetAllocator<TPoolTag>(EGuard guardType) {
+        TMaybe<TPoolAllocatorSettings> poolSettings = TPoolAllocatorSettings{};
+        return GetAllocator(poolSettings, guardType);
+    }
+
+    template<>
+    THolder<IAllocator> GetAllocator<TSimpleTag>(EGuard guardType) {
+        TMaybe<TPoolAllocatorSettings> poolSettings;
+        return GetAllocator(poolSettings, guardType);
+    }
+
+
+    template <class AllocatorType>
+    class TAllocatorFixture : public Test {
+    protected:
+        TAllocatorFixture()
+            : Allocator_(GetAllocator<AllocatorType>(EGuard::Page))
+        {}
+
+        const uint64_t StackSize_ = PageSize * 2;
+        THolder<IAllocator> Allocator_;
+    };
+
+    typedef Types<TPoolTag, TSimpleTag> Implementations;
+    TYPED_TEST_SUITE(TAllocatorFixture, Implementations);
+
+    TYPED_TEST(TAllocatorFixture, StackOverflow) {
+        ASSERT_DEATH({
+            auto stack = this->Allocator_->AllocStack(this->StackSize_, "test_stack");
+
+            // Overwrite previous guard, crash is here
+            *(stack.GetAlignedMemory() - 1) = 42;
+        }, "");
+    }
+
+    TYPED_TEST(TAllocatorFixture, StackOverride) {
+        ASSERT_DEATH({
+            auto stack = this->Allocator_->AllocStack(this->StackSize_, "test_stack");
+
+            // Overwrite guard, crash is here
+            *(stack.GetAlignedMemory() + stack.GetSize() - 1) = 42;
+        }, "");
+    }
+
+}
diff --git a/library/cpp/coroutine/engine/stack/ut/stack_guards_ut.cpp b/library/cpp/coroutine/engine/stack/ut/stack_guards_ut.cpp
new file mode 100644
index 0000000000..9da9a9b3d5
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/stack_guards_ut.cpp
@@ -0,0 +1,158 @@
+#include <library/cpp/coroutine/engine/stack/stack_common.h>
+#include <library/cpp/coroutine/engine/stack/stack_guards.h>
+#include <library/cpp/coroutine/engine/stack/stack_utils.h>
+#include <library/cpp/testing/gtest/gtest.h>
+
+
+using namespace testing;
+
+namespace NCoro::NStack::Tests {
+
+    template <class TGuard>
+    class TGuardFixture : public Test {
+    protected:
+        TGuardFixture() : Guard_(GetGuard<TGuard>()) {}
+
+        const TGuard& Guard_;
+    };
+
+    typedef Types<TCanaryGuard, TPageGuard> Implementations;
+    TYPED_TEST_SUITE(TGuardFixture, Implementations);
+
+    TYPED_TEST(TGuardFixture, GuardSize) {
+        const auto size = this->Guard_.GetSize();
+        EXPECT_GE(size, 64ul);
+        EXPECT_FALSE(size & 63ul); // check 64-byte alignment
+    }
+
+    TYPED_TEST(TGuardFixture, GuardAlignedSize) {
+        const auto size = this->Guard_.GetPageAlignedSize();
+        EXPECT_GE(size, PageSize);
+        EXPECT_FALSE(size & PageSizeMask); // check page-alignment
+    }
+
+    TYPED_TEST(TGuardFixture, StackWorkspace) {
+        for (uint64_t sizeInPages : {2, 5, 12}) {
+            char *rawPtr, *alignedPtr = nullptr;
+            ASSERT_TRUE(GetAlignedMemory(sizeInPages, rawPtr, alignedPtr));
+            auto workspace = this->Guard_.GetWorkspace(alignedPtr, sizeInPages * PageSize);
+            EXPECT_EQ(workspace.size(), sizeInPages * PageSize - this->Guard_.GetSize()) << " size in pages " << sizeInPages;
+
+            this->Guard_.Protect(alignedPtr, sizeInPages * PageSize, false);
+            workspace = this->Guard_.GetWorkspace(alignedPtr, sizeInPages * PageSize);
+            EXPECT_EQ(workspace.size(), sizeInPages * PageSize - this->Guard_.GetSize()) << " size in pages " << sizeInPages;
+
+            this->Guard_.RemoveProtection(alignedPtr, sizeInPages * PageSize);
+            workspace = this->Guard_.GetWorkspace(alignedPtr, sizeInPages * PageSize);
+            EXPECT_EQ(workspace.size(), sizeInPages * PageSize - this->Guard_.GetSize()) << " size in pages " << sizeInPages;
+
+            free(rawPtr);
+        }
+    }
+
+    TYPED_TEST(TGuardFixture, SetRemoveProtectionWorks) {
+        char *rawPtr, *alignedPtr = nullptr;
+        constexpr uint64_t sizeInPages = 4;
+        ASSERT_TRUE(GetAlignedMemory(sizeInPages + 1, rawPtr, alignedPtr));
+
+        this->Guard_.Protect(alignedPtr, PageSize, false); // set previous guard
+        alignedPtr += PageSize; // leave first page for previous guard
+        this->Guard_.Protect(alignedPtr, sizeInPages * PageSize, true);
+
+        EXPECT_TRUE(this->Guard_.CheckOverflow(alignedPtr));
+        EXPECT_TRUE(this->Guard_.CheckOverride(alignedPtr, sizeInPages * PageSize));
+
+        this->Guard_.RemoveProtection(alignedPtr, sizeInPages * PageSize);
+        this->Guard_.RemoveProtection(alignedPtr - PageSize, PageSize); // remove previous guard
+
+        free(rawPtr);
+    }
+
+    TEST(StackGuardTest, CanaryGuardTestOverflow) {
+        const auto& guard = GetGuard<TCanaryGuard>();
+
+        char *rawPtr, *alignedPtr = nullptr;
+        constexpr uint64_t sizeInPages = 4;
+        ASSERT_TRUE(GetAlignedMemory(sizeInPages + 1, rawPtr, alignedPtr));
+        guard.Protect(alignedPtr, PageSize, false); // set previous guard
+        alignedPtr += PageSize; // leave first page for previous guard
+        guard.Protect(alignedPtr, sizeInPages * PageSize, true);
+
+        EXPECT_TRUE(guard.CheckOverflow(alignedPtr));
+        EXPECT_TRUE(guard.CheckOverride(alignedPtr, sizeInPages * PageSize));
+
+        // Overwrite previous guard
+        *(alignedPtr - 1) = 42;
+
+        EXPECT_FALSE(guard.CheckOverflow(alignedPtr));
+
+        free(rawPtr);
+    }
+
+    TEST(StackGuardTest, CanaryGuardTestOverride) {
+        const auto& guard = GetGuard<TCanaryGuard>();
+
+        char *rawPtr, *alignedPtr = nullptr;
+        constexpr uint64_t sizeInPages = 4;
+        ASSERT_TRUE(GetAlignedMemory(sizeInPages + 1, rawPtr, alignedPtr));
+        guard.Protect(alignedPtr, PageSize, false); // set previous guard
+        alignedPtr += PageSize; // leave first page for previous guard
+        guard.Protect(alignedPtr, sizeInPages * PageSize, true);
+
+        EXPECT_TRUE(guard.CheckOverflow(alignedPtr));
+        EXPECT_TRUE(guard.CheckOverride(alignedPtr, sizeInPages * PageSize));
+
+        // Overwrite guard
+        *(alignedPtr + sizeInPages * PageSize - 1) = 42;
+
+        EXPECT_FALSE(guard.CheckOverride(alignedPtr, sizeInPages * PageSize));
+
+        free(rawPtr);
+    }
+
+    TEST(StackGuardDeathTest, PageGuardTestOverflow) {
+        ASSERT_DEATH({
+            const auto &guard = GetGuard<TPageGuard>();
+
+            char* rawPtr = nullptr;
+            char* alignedPtr = nullptr;
+            constexpr uint64_t sizeInPages = 4;
+            ASSERT_TRUE(GetAlignedMemory(sizeInPages + 1, rawPtr, alignedPtr));
+
+            guard.Protect(alignedPtr, PageSize, false); // set previous guard
+            alignedPtr += PageSize; // leave first page for previous guard
+            guard.Protect(alignedPtr, sizeInPages * PageSize, true);
+
+            // Overwrite previous guard, crash is here
+            *(alignedPtr - 1) = 42;
+
+            guard.RemoveProtection(alignedPtr, sizeInPages * PageSize);
+            guard.RemoveProtection(alignedPtr - PageSize, PageSize); // remove previous guard
+
+            free(rawPtr);
+        }, "");
+    }
+
+    TEST(StackGuardDeathTest, PageGuardTestOverride) {
+        ASSERT_DEATH({
+            const auto &guard = GetGuard<TPageGuard>();
+
+            char* rawPtr = nullptr;
+            char* alignedPtr = nullptr;
+            constexpr uint64_t sizeInPages = 4;
+            ASSERT_TRUE(GetAlignedMemory(sizeInPages + 1, rawPtr, alignedPtr));
+            guard.Protect(alignedPtr, PageSize, false); // set previous guard
+            alignedPtr += PageSize; // leave first page for previous guard
+            guard.Protect(alignedPtr, sizeInPages * PageSize, true);
+
+            // Overwrite guard, crash is here
+            *(alignedPtr + sizeInPages * PageSize - 1) = 42;
+
+            guard.RemoveProtection(alignedPtr, sizeInPages * PageSize);
+            guard.RemoveProtection(alignedPtr - PageSize, PageSize); // remove previous guard
+
+            free(rawPtr);
+        }, "");
+    }
+
+}
diff --git a/library/cpp/coroutine/engine/stack/ut/stack_pool_ut.cpp b/library/cpp/coroutine/engine/stack/ut/stack_pool_ut.cpp
new file mode 100644
index 0000000000..9e3e5e7117
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/stack_pool_ut.cpp
@@ -0,0 +1,70 @@
+#include <library/cpp/coroutine/engine/stack/stack_common.h>
+#include <library/cpp/coroutine/engine/stack/stack_guards.h>
+#include <library/cpp/coroutine/engine/stack/stack_pool.h>
+#include <library/cpp/testing/gtest/gtest.h>
+
+
+using namespace testing;
+
+namespace NCoro::NStack::Tests {
+
+    template <class TGuard>
+    class TPoolFixture : public Test {
+    protected:
+        TPoolFixture() : Guard_(GetGuard<TGuard>()), Pool_(StackSize_, TPoolAllocatorSettings{1, 1, 8, 32}, Guard_) {}
+
+        const uint64_t StackSize_ = PageSize * 4;
+        const TGuard& Guard_;
+        TPool<TGuard> Pool_;
+    };
+
+    typedef Types<TCanaryGuard, TPageGuard> Implementations;
+    TYPED_TEST_SUITE(TPoolFixture, Implementations);
+
+    TYPED_TEST(TPoolFixture, AllocAndFreeStack) {
+        auto stack = this->Pool_.AllocStack("test_stack");
+        this->Pool_.FreeStack(stack);
+        EXPECT_FALSE(stack.GetRawMemory());
+    }
+
+    TYPED_TEST(TPoolFixture, FreedStackReused) {
+        auto stack = this->Pool_.AllocStack("test_stack");
+        auto rawPtr = stack.GetRawMemory();
+        auto alignedPtr = stack.GetAlignedMemory();
+
+        this->Pool_.FreeStack(stack);
+        EXPECT_FALSE(stack.GetRawMemory());
+
+        auto stack2 = this->Pool_.AllocStack("test_stack");
+        EXPECT_EQ(rawPtr, stack2.GetRawMemory());
+        EXPECT_EQ(alignedPtr, stack2.GetAlignedMemory());
+
+        this->Pool_.FreeStack(stack2);
+        EXPECT_FALSE(stack2.GetRawMemory());
+    }
+
+    TYPED_TEST(TPoolFixture, MruFreedStackReused) {
+        auto stack = this->Pool_.AllocStack("test_stack");
+        auto rawPtr = stack.GetRawMemory();
+        auto alignedPtr = stack.GetAlignedMemory();
+        auto stack2 = this->Pool_.AllocStack("test_stack");
+        auto stack3 = this->Pool_.AllocStack("test_stack");
+
+        this->Pool_.FreeStack(stack2);
+        EXPECT_FALSE(stack2.GetRawMemory());
+
+        this->Pool_.FreeStack(stack);
+        EXPECT_FALSE(stack.GetRawMemory());
+
+        auto stack4 = this->Pool_.AllocStack("test_stack");
+        EXPECT_EQ(rawPtr, stack4.GetRawMemory());
+        EXPECT_EQ(alignedPtr, stack4.GetAlignedMemory());
+
+        this->Pool_.FreeStack(stack3);
+        EXPECT_FALSE(stack.GetRawMemory());
+
+        this->Pool_.FreeStack(stack4);
+        EXPECT_FALSE(stack4.GetRawMemory());
+    }
+
+}
diff --git a/library/cpp/coroutine/engine/stack/ut/stack_ut.cpp b/library/cpp/coroutine/engine/stack/ut/stack_ut.cpp
new file mode 100644
index 0000000000..31f8ad6b61
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/stack_ut.cpp
@@ -0,0 +1,60 @@
+#include <library/cpp/coroutine/engine/stack/stack.h>
+#include <library/cpp/coroutine/engine/stack/stack_common.h>
+#include <library/cpp/coroutine/engine/stack/stack_guards.h>
+#include <library/cpp/coroutine/engine/stack/stack_utils.h>
+#include <library/cpp/testing/gtest/gtest.h>
+
+
+using namespace testing;
+
+namespace NCoro::NStack::Tests {
+
+    constexpr uint64_t StackSizeInPages = 4;
+
+    template <class TGuard>
+    class TStackFixture : public Test {
+    protected: // methods
+        TStackFixture()
+            : Guard_(GetGuard<TGuard>())
+            , StackSize_(StackSizeInPages * PageSize)
+        {}
+
+        void SetUp() override {
+            ASSERT_TRUE(GetAlignedMemory(StackSizeInPages, RawMemory_, AlignedMemory_));
+            Stack_ = MakeHolder<NDetails::TStack>(RawMemory_, AlignedMemory_, StackSize_, "test_stack");
+            Guard_.Protect(AlignedMemory_, StackSize_, false);
+        }
+
+        void TearDown() override {
+            Guard_.RemoveProtection(AlignedMemory_, StackSize_);
+            free(Stack_->GetRawMemory());
+            Stack_->Reset();
+            EXPECT_EQ(Stack_->GetRawMemory(), nullptr);
+        }
+
+    protected: // data
+        const TGuard& Guard_;
+        const uint64_t StackSize_ = 0;
+        char* RawMemory_ = nullptr;
+        char* AlignedMemory_ = nullptr;
+        THolder<NDetails::TStack> Stack_;
+    };
+
+    typedef Types<TCanaryGuard, TPageGuard> Implementations;
+    TYPED_TEST_SUITE(TStackFixture, Implementations);
+
+    TYPED_TEST(TStackFixture, PointersAndSize) {
+        EXPECT_EQ(this->Stack_->GetRawMemory(), this->RawMemory_);
+        EXPECT_EQ(this->Stack_->GetAlignedMemory(), this->AlignedMemory_);
+        EXPECT_EQ(this->Stack_->GetSize(), this->StackSize_);
+    }
+
+    TYPED_TEST(TStackFixture, WriteStack) {
+        auto workspace = this->Guard_.GetWorkspace(this->Stack_->GetAlignedMemory(), this->Stack_->GetSize());
+        for (uint64_t i = 0; i < workspace.size(); i += 512) {
+            workspace[i] = 42;
+        }
+        EXPECT_TRUE(this->Guard_.CheckOverride(this->Stack_->GetAlignedMemory(), this->Stack_->GetSize()));
+    }
+
+}
diff --git a/library/cpp/coroutine/engine/stack/ut/stack_utils_ut.cpp b/library/cpp/coroutine/engine/stack/ut/stack_utils_ut.cpp
new file mode 100644
index 0000000000..dc0593dcf2
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/stack_utils_ut.cpp
@@ -0,0 +1,73 @@
+#include <library/cpp/coroutine/engine/stack/stack_common.h>
+#include <library/cpp/coroutine/engine/stack/stack_utils.h>
+#include <library/cpp/testing/gtest/gtest.h>
+
+
+using namespace testing;
+
+namespace NCoro::NStack::Tests {
+
+    TEST(StackUtilsTest, Allocation) {
+        char *rawPtr, *alignedPtr = nullptr;
+        for (uint64_t i : {1, 2, 3, 4, 11}) {
+            EXPECT_TRUE(GetAlignedMemory(i, rawPtr, alignedPtr));
+            EXPECT_TRUE(rawPtr);
+            EXPECT_TRUE(alignedPtr);
+            EXPECT_FALSE((uint64_t)alignedPtr & PageSizeMask);
+            free(rawPtr);
+        }
+    }
+
+#if !defined(_san_enabled_) && defined(_linux_)
+
+    TEST(StackUtilsTest, RssReleaseOnePage) {
+        char *rawPtr, *alignedPtr = nullptr;
+        for (uint64_t i : {1, 2, 8}) {
+            EXPECT_TRUE(GetAlignedMemory(i, rawPtr, alignedPtr));
+            EXPECT_TRUE(rawPtr);
+            EXPECT_TRUE(alignedPtr);
+            EXPECT_FALSE((uint64_t)alignedPtr & PageSizeMask);
+
+            ReleaseRss(alignedPtr, i); // allocator can provide reused memory with RSS memory on it
+            EXPECT_EQ(CountMapped(alignedPtr, i), 0ul); // no RSS memory allocated
+
+            *(alignedPtr + (i - 1) * PageSize) = 42; // map RSS memory
+            EXPECT_EQ(CountMapped(alignedPtr, i), 1ul);
+
+            ReleaseRss(alignedPtr, i);
+            EXPECT_EQ(CountMapped(alignedPtr, i), 0ul) << "number of pages " << i; // no RSS memory allocated
+
+            free(rawPtr);
+        }
+    }
+
+    TEST(StackUtilsTest, RssReleaseSeveralPages) {
+        char *rawPtr, *alignedPtr = nullptr;
+
+        for (uint64_t i : {1, 2, 5, 8}) {
+            EXPECT_TRUE(GetAlignedMemory(i, rawPtr, alignedPtr));
+            EXPECT_TRUE(rawPtr);
+            EXPECT_TRUE(alignedPtr);
+            EXPECT_FALSE((uint64_t)alignedPtr & PageSizeMask);
+
+            ReleaseRss(alignedPtr, i); // allocator can provide reused memory with RSS memory on it
+            EXPECT_EQ(CountMapped(alignedPtr, i), 0ul); // no RSS memory allocated
+
+            for (uint64_t page = 0; page < i; ++page) {
+                *(alignedPtr + page * PageSize) = 42; // map RSS memory
+                EXPECT_EQ(CountMapped(alignedPtr, page + 1), page + 1);
+            }
+
+            const uint64_t pagesToKeep = (i > 2) ? 2 : i;
+
+            ReleaseRss(alignedPtr, i - pagesToKeep);
+            EXPECT_EQ(CountMapped(alignedPtr, i), pagesToKeep) << "number of pages " << i; // no RSS memory allocated
+
+            free(rawPtr);
+        }
+    }
+
+#endif
+
+}
+
diff --git a/library/cpp/coroutine/engine/stack/ut/ya.make b/library/cpp/coroutine/engine/stack/ut/ya.make
new file mode 100644
index 0000000000..65c5af9b7f
--- /dev/null
+++ b/library/cpp/coroutine/engine/stack/ut/ya.make
@@ -0,0 +1,17 @@
+GTEST()
+
+OWNER(g:balancer)
+
+SRCS(
+    stack_allocator_ut.cpp
+    stack_guards_ut.cpp
+    stack_pool_ut.cpp
+    stack_ut.cpp
+    stack_utils_ut.cpp
+)
+
+PEERDIR(
+    library/cpp/coroutine/engine
+)
+
+END()
\ No newline at end of file