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
12namespace ProfileEvents
13{
14 extern const Event IOBufferAllocs;
15 extern const Event IOBufferAllocBytes;
16}
17
18
19namespace 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 */
27template <typename Allocator = Allocator<false>>
28struct 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
94private:
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 */
136template <typename Base>
137class BufferWithOwnMemory : public Base
138{
139protected:
140 Memory<> memory;
141public:
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