1 | #pragma once |
2 | |
3 | #include <boost/noncopyable.hpp> |
4 | |
5 | #include <Common/ProfileEvents.h> |
6 | #include <Common/Allocator.h> |
7 | |
8 | #include <Common/Exception.h> |
9 | #include <Core/Defines.h> |
10 | |
11 | |
12 | namespace ProfileEvents |
13 | { |
14 | extern const Event IOBufferAllocs; |
15 | extern const Event IOBufferAllocBytes; |
16 | } |
17 | |
18 | |
19 | namespace DB |
20 | { |
21 | |
22 | |
23 | /** Replacement for std::vector<char> to use in buffers. |
24 | * Differs in that is doesn't do unneeded memset. (And also tries to do as little as possible.) |
25 | * Also allows to allocate aligned piece of memory (to use with O_DIRECT, for example). |
26 | */ |
27 | template <typename Allocator = Allocator<false>> |
28 | struct Memory : boost::noncopyable, Allocator |
29 | { |
30 | /// Padding is needed to allow usage of 'memcpySmallAllowReadWriteOverflow15' function with this buffer. |
31 | static constexpr size_t pad_right = 15; |
32 | |
33 | size_t m_capacity = 0; /// With padding. |
34 | size_t m_size = 0; |
35 | char * m_data = nullptr; |
36 | size_t alignment = 0; |
37 | |
38 | Memory() {} |
39 | |
40 | /// If alignment != 0, then allocate memory aligned to specified value. |
41 | Memory(size_t size_, size_t alignment_ = 0) : m_capacity(size_), m_size(m_capacity), alignment(alignment_) |
42 | { |
43 | alloc(); |
44 | } |
45 | |
46 | ~Memory() |
47 | { |
48 | dealloc(); |
49 | } |
50 | |
51 | Memory(Memory && rhs) noexcept |
52 | { |
53 | *this = std::move(rhs); |
54 | } |
55 | |
56 | Memory & operator=(Memory && rhs) noexcept |
57 | { |
58 | std::swap(m_capacity, rhs.m_capacity); |
59 | std::swap(m_size, rhs.m_size); |
60 | std::swap(m_data, rhs.m_data); |
61 | std::swap(alignment, rhs.alignment); |
62 | |
63 | return *this; |
64 | } |
65 | |
66 | size_t size() const { return m_size; } |
67 | const char & operator[](size_t i) const { return m_data[i]; } |
68 | char & operator[](size_t i) { return m_data[i]; } |
69 | const char * data() const { return m_data; } |
70 | char * data() { return m_data; } |
71 | |
72 | void resize(size_t new_size) |
73 | { |
74 | if (0 == m_capacity) |
75 | { |
76 | m_size = new_size; |
77 | m_capacity = new_size; |
78 | alloc(); |
79 | } |
80 | else if (new_size <= m_capacity - pad_right) |
81 | { |
82 | m_size = new_size; |
83 | return; |
84 | } |
85 | else |
86 | { |
87 | size_t new_capacity = align(new_size + pad_right, alignment); |
88 | m_data = static_cast<char *>(Allocator::realloc(m_data, m_capacity, new_capacity, alignment)); |
89 | m_capacity = new_capacity; |
90 | m_size = m_capacity - pad_right; |
91 | } |
92 | } |
93 | |
94 | private: |
95 | static size_t align(const size_t value, const size_t alignment) |
96 | { |
97 | if (!alignment) |
98 | return value; |
99 | |
100 | return (value + alignment - 1) / alignment * alignment; |
101 | } |
102 | |
103 | void alloc() |
104 | { |
105 | if (!m_capacity) |
106 | { |
107 | m_data = nullptr; |
108 | return; |
109 | } |
110 | |
111 | size_t padded_capacity = m_capacity + pad_right; |
112 | |
113 | ProfileEvents::increment(ProfileEvents::IOBufferAllocs); |
114 | ProfileEvents::increment(ProfileEvents::IOBufferAllocBytes, padded_capacity); |
115 | |
116 | size_t new_capacity = align(padded_capacity, alignment); |
117 | m_data = static_cast<char *>(Allocator::alloc(new_capacity, alignment)); |
118 | m_capacity = new_capacity; |
119 | m_size = m_capacity - pad_right; |
120 | } |
121 | |
122 | void dealloc() |
123 | { |
124 | if (!m_data) |
125 | return; |
126 | |
127 | Allocator::free(m_data, m_capacity); |
128 | m_data = nullptr; /// To avoid double free if next alloc will throw an exception. |
129 | } |
130 | }; |
131 | |
132 | |
133 | /** Buffer that could own its working memory. |
134 | * Template parameter: ReadBuffer or WriteBuffer |
135 | */ |
136 | template <typename Base> |
137 | class BufferWithOwnMemory : public Base |
138 | { |
139 | protected: |
140 | Memory<> memory; |
141 | public: |
142 | /// If non-nullptr 'existing_memory' is passed, then buffer will not create its own memory and will use existing_memory without ownership. |
143 | BufferWithOwnMemory(size_t size = DBMS_DEFAULT_BUFFER_SIZE, char * existing_memory = nullptr, size_t alignment = 0) |
144 | : Base(nullptr, 0), memory(existing_memory ? 0 : size, alignment) |
145 | { |
146 | Base::set(existing_memory ? existing_memory : memory.data(), size); |
147 | Base::padded = !existing_memory; |
148 | } |
149 | }; |
150 | |
151 | |
152 | } |
153 | |