aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/yt/memory/chunked_memory_pool.h
blob: a22e2ec27c4e4cbee37015ba331a10ad7814ce37 (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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
#pragma once

#include "public.h"
#include "ref.h"

#include <util/generic/size_literals.h>

namespace NYT {

////////////////////////////////////////////////////////////////////////////////

struct TDefaultChunkedMemoryPoolTag { };

// TAllocationHolder is polymorphic. So we cannot use TWithExtraSpace mixin
// because it needs the most derived type as a template argument and
// it would require GetExtraSpacePtr/GetRef methods to be virtual.

class TAllocationHolder
{
public:
    TAllocationHolder(TMutableRef ref, TRefCountedTypeCookie cookie);
    TAllocationHolder(const TAllocationHolder&) = delete;
    TAllocationHolder(TAllocationHolder&&) = default;
    virtual ~TAllocationHolder();

    void operator delete(void* ptr) noexcept;

    TMutableRef GetRef() const;

    template <class TDerived>
    static TDerived* Allocate(size_t size, TRefCountedTypeCookie cookie);

private:
    const TMutableRef Ref_;
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
    const TRefCountedTypeCookie Cookie_;
#endif
};

////////////////////////////////////////////////////////////////////////////////

struct IMemoryChunkProvider
    : public TRefCounted
{
    virtual std::unique_ptr<TAllocationHolder> Allocate(size_t size, TRefCountedTypeCookie cookie) = 0;
};

DEFINE_REFCOUNTED_TYPE(IMemoryChunkProvider)

const IMemoryChunkProviderPtr& GetDefaultMemoryChunkProvider();

////////////////////////////////////////////////////////////////////////////////

class TChunkedMemoryPool
    : private TNonCopyable
{
public:
    static constexpr size_t DefaultStartChunkSize = 4_KB;
    static constexpr size_t RegularChunkSize = 36_KB - 512;

    TChunkedMemoryPool(
        TRefCountedTypeCookie tagCookie,
        IMemoryChunkProviderPtr chunkProvider,
        size_t startChunkSize = DefaultStartChunkSize);

    explicit TChunkedMemoryPool(
        TRefCountedTypeCookie tagCookie,
        size_t startChunkSize = DefaultStartChunkSize);

    TChunkedMemoryPool();

    template <class TTag>
    explicit TChunkedMemoryPool(
        TTag,
        size_t startChunkSize = DefaultStartChunkSize);

    //! Allocates #sizes bytes without any alignment.
    char* AllocateUnaligned(size_t size);

    //! Allocates #size bytes aligned with 8-byte granularity.
    char* AllocateAligned(size_t size, int align = 8);

    //! Allocates #n uninitialized instances of #T.
    template <class T>
    T* AllocateUninitialized(int n, int align = alignof(T));

    //! Allocates space and copies #src inside it.
    template <class T>
    TMutableRange<T> Capture(TRange<T> src, int align = alignof(T));

    //! Frees memory range if possible: namely, if the free region is a suffix of last allocated region.
    void Free(char* from, char* to);

    //! Marks all previously allocated small chunks as free for subsequent allocations but
    //! does not deallocate them.
    //! Purges all large blocks.
    void Clear();

    //! Purges all allocated memory, including small chunks.
    void Purge();

    //! Returns the number of allocated bytes.
    size_t GetSize() const;

    //! Returns the number of reserved bytes.
    size_t GetCapacity() const;

    //! Returns the number of bytes that can be acquired in the current chunk
    //! without additional allocations.
    size_t GetCurrentChunkSpareSize() const;

    //! Moves all the allocated memory from other memory pool to the current one.
    //! The other pool becomes empty, like after Purge() call.
    void Absorb(TChunkedMemoryPool&& other);

private:
    const TRefCountedTypeCookie TagCookie_;
    // A common usecase is to construct TChunkedMemoryPool with the default
    // memory chunk provider. The latter is ref-counted and is shared between
    // a multitude of TChunkedMemoryPool instances. This could potentially
    // lead to a contention over IMemoryChunkProvider's ref-counter.
    // To circumvent this, we keep both an owning (#ChunkProviderHolder_) and
    // a non-owning (#ChunkProvider_) reference to the underlying provider.
    // In case of the default chunk provider, the owning reference is not used.
    const IMemoryChunkProviderPtr ChunkProviderHolder_;
    IMemoryChunkProvider* const ChunkProvider_;

    int NextChunkIndex_ = 0;
    size_t NextSmallSize_;

    size_t Size_ = 0;
    size_t Capacity_ = 0;

    // Chunk memory layout:
    //   |AAAA|....|UUUU|
    // Legend:
    //   A aligned allocations
    //   U unaligned allocations
    //   . free zone
    char* FreeZoneBegin_;
    char* FreeZoneEnd_;

    std::vector<std::unique_ptr<TAllocationHolder>> Chunks_;
    std::vector<std::unique_ptr<TAllocationHolder>> OtherBlocks_;

    void Initialize(size_t startChunkSize);

    char* AllocateUnalignedSlow(size_t size);
    char* AllocateAlignedSlow(size_t size, int align);
    char* AllocateSlowCore(size_t size);
};

////////////////////////////////////////////////////////////////////////////////

} // namespace NYT

#define CHUNKED_MEMORY_POOL_INL_H_
#include "chunked_memory_pool-inl.h"
#undef CHUNKED_MEMORY_POOL_INL_H_