aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/coroutine/engine/stack/stack_allocator.inl
blob: 5d2de52f0deebfc5306aeaf1f1c57fc3b2acc213 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#include "stack_guards.h"
#include "stack_pool.h"
#include "stack_utils.h"

#include <util/generic/hash.h>

#ifdef _linux_
#include <unistd.h>
#endif


namespace NCoro::NStack {

    template<typename TGuard>
    class TPoolAllocator final : public IAllocator {
    public:
        explicit TPoolAllocator(const TPoolAllocatorSettings& settings);

        TArrayRef<char> GetStackWorkspace(void* stack, size_t size) noexcept override {
            return Guard_.GetWorkspace(stack, size);
        }
        bool CheckStackOverflow(void* stack) const noexcept override {
            return Guard_.CheckOverflow(stack);
        }
        bool CheckStackOverride(void* stack, size_t size) const noexcept override {
            return Guard_.CheckOverride(stack, size);
        }

        TAllocatorStats GetStackStats() const noexcept override {
            TAllocatorStats stats;
            for (const auto& i : Pools_) {
                stats.ReleasedSize += i.second.GetReleasedSize();
                stats.NotReleasedSize += i.second.GetFullSize();
                stats.NumOfAllocated += i.second.GetNumOfAllocated();
            }
            return stats;
        }

    private: // methods
        NDetails::TStack DoAllocStack(size_t size, const char* name) override;
        void DoFreeStack(NDetails::TStack& stack) noexcept override;

    private: // data
        const TPoolAllocatorSettings PoolSettings_;
        const TGuard& Guard_;
        THashMap<size_t, TPool<TGuard>> Pools_; // key - stack size
    };

    template<typename TGuard>
    TPoolAllocator<TGuard>::TPoolAllocator(const TPoolAllocatorSettings& settings)
        : PoolSettings_(settings)
        , Guard_(GetGuard<TGuard>())
    {
#ifdef _linux_
        Y_ABORT_UNLESS(sysconf(_SC_PAGESIZE) == PageSize);
#endif
    }

    template<typename TGuard>
    NDetails::TStack TPoolAllocator<TGuard>::DoAllocStack(size_t alignedSize, const char* name) {
        Y_ASSERT(alignedSize > Guard_.GetSize());

        auto pool = Pools_.find(alignedSize);
        if (pool == Pools_.end()) {
            Y_ASSERT(Pools_.size() < 1000); // too many different sizes for coroutine stacks
            auto [newPool, success] = Pools_.emplace(alignedSize, TPool<TGuard>{alignedSize, PoolSettings_, Guard_});
            Y_ABORT_UNLESS(success, "Failed to add new coroutine pool");
            pool = newPool;
        }
        return pool->second.AllocStack(name);
    }

    template<typename TGuard>
    void TPoolAllocator<TGuard>::DoFreeStack(NDetails::TStack& stack) noexcept {
        auto pool = Pools_.find(stack.GetSize());
        Y_ABORT_UNLESS(pool != Pools_.end(), "Attempt to free stack from another allocator");
        pool->second.FreeStack(stack);
    }

    // ------------------------------------------------------------------------
    //
    template<typename TGuard>
    class TSimpleAllocator : public IAllocator {
    public:
        explicit TSimpleAllocator();

        TArrayRef<char> GetStackWorkspace(void* stack, size_t size) noexcept override {
            return Guard_.GetWorkspace(stack, size);
        }
        bool CheckStackOverflow(void* stack) const noexcept override {
            return Guard_.CheckOverflow(stack);
        }
        bool CheckStackOverride(void* stack, size_t size) const noexcept override {
            return Guard_.CheckOverride(stack, size);
        }

        TAllocatorStats GetStackStats() const noexcept override { return {}; } // not used for simple allocator

    private: // methods
        NDetails::TStack DoAllocStack(size_t size, const char* name) override;
        void DoFreeStack(NDetails::TStack& stack) noexcept override;

    private: // data
        const TGuard& Guard_;
    };


    template<typename TGuard>
    TSimpleAllocator<TGuard>::TSimpleAllocator()
        : Guard_(GetGuard<TGuard>())
    {}

    template<typename TGuard>
    NDetails::TStack TSimpleAllocator<TGuard>::DoAllocStack(size_t alignedSize, const char* name) {
        Y_ASSERT(alignedSize > Guard_.GetSize());

        char* rawPtr = nullptr;
        char* alignedPtr = nullptr; // with extra space for previous guard in this type of allocator

        Y_ABORT_UNLESS(GetAlignedMemory((alignedSize + Guard_.GetPageAlignedSize()) / PageSize, rawPtr, alignedPtr)); // + memory for previous guard
        char* alignedStackMemory = alignedPtr + Guard_.GetPageAlignedSize(); // after previous guard

        // Default allocator sets both guards, because it doesn't have memory chunk with previous stack and guard on it
        Guard_.Protect((void*)alignedPtr, Guard_.GetPageAlignedSize(), false); // first guard should be before stack memory
        Guard_.Protect(alignedStackMemory, alignedSize, true); // second guard is placed on stack memory

        return NDetails::TStack{rawPtr, alignedStackMemory, alignedSize, name};
    }

    template<typename TGuard>
    void TSimpleAllocator<TGuard>::DoFreeStack(NDetails::TStack& stack) noexcept {
        Guard_.RemoveProtection(stack.GetAlignedMemory() - Guard_.GetPageAlignedSize(), Guard_.GetPageAlignedSize());
        Guard_.RemoveProtection(stack.GetAlignedMemory(), stack.GetSize());

        free(stack.GetRawMemory());
        stack.Reset();
    }
}