Branch data Line data Source code
1 : : /*
2 : : * Copyright 2025 Jinwoo Sung
3 : : *
4 : : * Licensed under the Apache License, Version 2.0 (the "License");
5 : : * you may not use this file except in compliance with the License.
6 : : * You may obtain a copy of the License at
7 : : *
8 : : * http://www.apache.org/licenses/LICENSE-2.0
9 : : *
10 : : * Unless required by applicable law or agreed to in writing, software
11 : : * distributed under the License is distributed on an "AS IS" BASIS,
12 : : * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 : : * See the License for the specific language governing permissions and
14 : : * limitations under the License.
15 : : */
16 : :
17 : : #include "wirestead/memory/memory_pool.hpp"
18 : :
19 : : #include <algorithm>
20 : : #include <cstdlib>
21 : : #include <mutex>
22 : : #include <stdexcept>
23 : :
24 : : #include "wirestead/memory/memory_tracker.hpp"
25 : :
26 : : namespace wirestead {
27 : : namespace memory {
28 : :
29 : 52 : MemoryPool& GlobalMemoryPool::instance() {
30 : 52 : static MemoryPool* pool = new MemoryPool();
31 : 52 : return *pool;
32 : : }
33 : :
34 : : // ============================================================================
35 : : // SelectiveMemoryPool Implementation
36 : : // ============================================================================
37 : :
38 : 545 : MemoryPool::MemoryPool(size_t initial_pool_size, size_t max_pool_size) {
39 : : // Initialize 4 fixed-size pools
40 : : static constexpr std::array<size_t, 4> BUCKET_SIZES = {
41 : : static_cast<size_t>(BufferSize::SMALL), // 1KB
42 : : static_cast<size_t>(BufferSize::MEDIUM), // 4KB
43 : : static_cast<size_t>(BufferSize::LARGE), // 16KB
44 : : static_cast<size_t>(BufferSize::XLARGE) // 64KB
45 : : };
46 : :
47 [ + + ]: 5450 : for (size_t i = 0; i < buckets_.size(); ++i) {
48 : 2180 : buckets_[i].size_ = BUCKET_SIZES[i];
49 : 4360 : buckets_[i].capacity_ = max_pool_size / buckets_.size();
50 : 2180 : buckets_[i].buffers_.reserve(buckets_[i].capacity_);
51 : : }
52 : :
53 : : // Prefill, split evenly across the buckets. This used to be discarded, so
54 : : // the pool always started empty whatever the caller asked for. It defaults
55 : : // to 0 rather than to the old nominal 400: the buckets run 1 KiB to 64 KiB,
56 : : // so honouring 400 would have meant 8.3 MiB reserved before the first
57 : : // acquire() - a cost the previous behaviour never actually charged, and one
58 : : // that matters on the embedded targets this library is aimed at.
59 : 545 : const size_t per_bucket = initial_pool_size / buckets_.size();
60 [ + + ]: 2725 : for (auto& bucket : buckets_) {
61 : 2180 : const size_t count = std::min(per_bucket, bucket.capacity_);
62 [ + + ]: 3216 : for (size_t n = 0; n < count; ++n) {
63 : 1036 : bucket.buffers_.push_back(create_buffer(bucket.size_));
64 : : }
65 : : }
66 : 545 : }
67 : :
68 : 4886 : std::unique_ptr<uint8_t[]> MemoryPool::acquire(size_t size) {
69 : 4886 : validate_size(size);
70 : :
71 : : // Optimization: Bypass pool for large allocations > 64KB (XLARGE)
72 : : // This avoids mutex contention and fixes truncation bug for sizes > 64KB.
73 [ + + ]: 4881 : if (size > static_cast<size_t>(BufferSize::XLARGE)) {
74 : 102 : auto buffer = create_buffer(size);
75 : : // create_buffer throws on failure, so if we are here, we succeeded.
76 : 102 : total_allocations_.fetch_add(1, std::memory_order_relaxed);
77 : 102 : track_acquire(size);
78 : 102 : return buffer;
79 : 102 : }
80 : :
81 : 4779 : auto& bkt = bucket(size);
82 : 4779 : auto buffer = acquire_from_bucket(bkt);
83 [ + - ]: 4779 : if (buffer) track_acquire(bkt.size_);
84 : 4779 : return buffer;
85 : 4779 : }
86 : :
87 : 16 : std::unique_ptr<uint8_t[]> MemoryPool::acquire(BufferSize buffer_size) {
88 : 16 : return acquire(static_cast<size_t>(buffer_size));
89 : : }
90 : :
91 : 4865 : void MemoryPool::release(std::unique_ptr<uint8_t[]> buffer, size_t size) {
92 [ + + ]: 4865 : if (!buffer) return;
93 : :
94 : 4864 : validate_size(size);
95 : :
96 [ + + ]: 4864 : if (size > static_cast<size_t>(BufferSize::XLARGE)) {
97 : : MEMORY_TRACK_DEALLOCATION(buffer.get());
98 : 102 : track_release(size);
99 : 102 : return; // unique_ptr destructor handles cleanup
100 : : }
101 : :
102 : 4762 : auto& bkt = bucket(size);
103 : 4762 : track_release(bkt.size_);
104 : 4762 : release_to_bucket(bkt, std::move(buffer));
105 : : }
106 : :
107 : 57 : MemoryPool::PoolStats MemoryPool::stats() const {
108 : 57 : PoolStats stats;
109 : 57 : stats.total_allocations = total_allocations_.load(std::memory_order_relaxed);
110 : 57 : stats.pool_hits = pool_hits_.load(std::memory_order_relaxed);
111 : 57 : return stats;
112 : : }
113 : :
114 : 6 : double MemoryPool::hit_rate() const {
115 : 6 : size_t total = total_allocations_.load(std::memory_order_relaxed);
116 [ + + ]: 6 : if (total == 0) return 0.0;
117 : :
118 : 3 : size_t hits = pool_hits_.load(std::memory_order_relaxed);
119 : 3 : return static_cast<double>(hits) / static_cast<double>(total);
120 : : }
121 : :
122 : 6 : std::pair<size_t, size_t> MemoryPool::memory_usage() const {
123 : 6 : return std::make_pair(outstanding_bytes_.load(std::memory_order_relaxed),
124 : 18 : peak_bytes_.load(std::memory_order_relaxed));
125 : : }
126 : :
127 : 4881 : void MemoryPool::track_acquire(size_t size) {
128 : 4881 : size_t new_outstanding = outstanding_bytes_.fetch_add(size, std::memory_order_relaxed) + size;
129 : 4881 : size_t prev_peak = peak_bytes_.load(std::memory_order_relaxed);
130 [ + + - + ]: 6214 : while (new_outstanding > prev_peak &&
131 [ - + ]: 2666 : !peak_bytes_.compare_exchange_weak(prev_peak, new_outstanding, std::memory_order_relaxed)) {
132 : : }
133 : 4881 : }
134 : :
135 : 4864 : void MemoryPool::track_release(size_t size) { outstanding_bytes_.fetch_sub(size, std::memory_order_relaxed); }
136 : :
137 : 2 : MemoryPool::HealthMetrics MemoryPool::health_metrics() const {
138 : 2 : HealthMetrics metrics;
139 : 2 : metrics.hit_rate = hit_rate();
140 : 2 : return metrics;
141 : : }
142 : :
143 : : // ============================================================================
144 : : // Private helper functions
145 : : // ============================================================================
146 : :
147 : 9541 : MemoryPool::PoolBucket& MemoryPool::bucket(size_t size) { return buckets_[bucket_index(size)]; }
148 : :
149 : 9541 : size_t MemoryPool::bucket_index(size_t size) const {
150 : : // Optimized: Unrolled checks for faster lookup
151 [ + + ]: 9541 : if (size <= static_cast<size_t>(BufferSize::SMALL)) return 0;
152 [ + + ]: 2772 : if (size <= static_cast<size_t>(BufferSize::MEDIUM)) return 1;
153 [ + + ]: 306 : if (size <= static_cast<size_t>(BufferSize::LARGE)) return 2;
154 : 107 : return 3; // XLARGE
155 : : }
156 : :
157 : 4779 : std::unique_ptr<uint8_t[]> MemoryPool::acquire_from_bucket(PoolBucket& bucket) {
158 : 4779 : std::unique_ptr<uint8_t[]> buffer;
159 : :
160 : : {
161 : 4779 : std::lock_guard<std::mutex> lock(bucket.mutex_);
162 : :
163 : : // Get from stack
164 [ + + ]: 4779 : if (!bucket.buffers_.empty()) {
165 : 3458 : buffer = std::move(bucket.buffers_.back());
166 : 3458 : bucket.buffers_.pop_back();
167 : : }
168 : 4779 : }
169 : :
170 [ + + ]: 4779 : if (buffer) {
171 : 3458 : pool_hits_.fetch_add(1, std::memory_order_relaxed);
172 : 3458 : total_allocations_.fetch_add(1, std::memory_order_relaxed);
173 : 3458 : return buffer;
174 : : }
175 : :
176 : : // Create new buffer outside lock
177 : 1321 : buffer = create_buffer(bucket.size_);
178 [ + - ]: 1321 : if (buffer) {
179 : 1321 : total_allocations_.fetch_add(1, std::memory_order_relaxed);
180 : : }
181 : :
182 : 1321 : return buffer;
183 : 0 : }
184 : :
185 : 4762 : void MemoryPool::release_to_bucket(PoolBucket& bucket, std::unique_ptr<uint8_t[]> buffer) {
186 : : {
187 : 4762 : std::lock_guard<std::mutex> lock(bucket.mutex_);
188 : :
189 : : // Add buffer back to pool (stack)
190 [ + + ]: 4762 : if (bucket.buffers_.size() < bucket.capacity_) {
191 : 4260 : bucket.buffers_.push_back(std::move(buffer));
192 : 4260 : return;
193 : : }
194 : 4762 : }
195 : :
196 : : // If pool is full, discard buffer (auto-release)
197 : : // buffer is unique_ptr so it will be automatically deleted when out of scope
198 : : MEMORY_TRACK_DEALLOCATION(buffer.get());
199 : : }
200 : :
201 : 2459 : std::unique_ptr<uint8_t[]> MemoryPool::create_buffer(size_t size) {
202 : 2459 : uint8_t* raw_buffer = new (std::nothrow) uint8_t[size];
203 [ - + ]: 2459 : if (raw_buffer) {
204 : : MEMORY_TRACK_ALLOCATION(raw_buffer, size);
205 : : } else {
206 : 0 : throw std::bad_alloc(); // Or handle error appropriately
207 : : }
208 : 2459 : return std::unique_ptr<uint8_t[]>(raw_buffer);
209 : : }
210 : :
211 : 9750 : void MemoryPool::validate_size(size_t size) const {
212 [ + + + + ]: 9750 : if (size == 0 || size > 64 * 1024 * 1024) { // 64MB maximum
213 : 5 : throw std::invalid_argument("Invalid buffer size");
214 : : }
215 : 9745 : }
216 : :
217 : : // ============================================================================
218 : : // PoolBucket Move Constructor/Assignment Operator
219 : : // ============================================================================
220 : :
221 : 0 : MemoryPool::PoolBucket::PoolBucket(PoolBucket&& other) noexcept
222 : 0 : : buffers_(std::move(other.buffers_)), mutex_(), size_(other.size_), capacity_(other.capacity_) {
223 : 0 : other.size_ = 0;
224 : 0 : other.capacity_ = 0;
225 : 0 : }
226 : :
227 : 0 : MemoryPool::PoolBucket& MemoryPool::PoolBucket::operator=(PoolBucket&& other) noexcept {
228 [ # # ]: 0 : if (this != &other) {
229 : 0 : buffers_ = std::move(other.buffers_);
230 : 0 : size_ = other.size_;
231 : 0 : capacity_ = other.capacity_;
232 : :
233 : 0 : other.size_ = 0;
234 : 0 : other.capacity_ = 0;
235 : : }
236 : 0 : return *this;
237 : : }
238 : :
239 : : // ============================================================================
240 : : // PooledBuffer Implementation
241 : : // ============================================================================
242 : :
243 : 5 : PooledBuffer::PooledBuffer(size_t size) : size_(size), pool_(&GlobalMemoryPool::instance()) {
244 : 5 : buffer_ = pool_->acquire(size);
245 : 5 : }
246 : :
247 : 4 : PooledBuffer::PooledBuffer(MemoryPool::BufferSize buffer_size)
248 : 4 : : size_(static_cast<size_t>(buffer_size)), pool_(&GlobalMemoryPool::instance()) {
249 : 4 : buffer_ = pool_->acquire(size_);
250 : 4 : }
251 : :
252 : 1295 : PooledBuffer::PooledBuffer(size_t size, MemoryPool& pool) : size_(size), pool_(&pool) {
253 : 1295 : buffer_ = pool_->acquire(size);
254 : 1295 : }
255 : :
256 : 4 : PooledBuffer::PooledBuffer(MemoryPool::BufferSize buffer_size, MemoryPool& pool)
257 : 4 : : size_(static_cast<size_t>(buffer_size)), pool_(&pool) {
258 : 4 : buffer_ = pool_->acquire(size_);
259 : 4 : }
260 : :
261 : 10539 : PooledBuffer::~PooledBuffer() {
262 [ + + + - : 10539 : if (buffer_ && pool_) {
+ + ]
263 : 1305 : pool_->release(std::move(buffer_), size_);
264 : : }
265 : 10539 : }
266 : :
267 : 9233 : PooledBuffer::PooledBuffer(PooledBuffer&& other) noexcept
268 : 9233 : : buffer_(std::move(other.buffer_)), size_(other.size_), pool_(other.pool_) {
269 : 9233 : other.buffer_ = nullptr;
270 : 9233 : other.size_ = 0;
271 : 9233 : other.pool_ = nullptr;
272 : 9233 : }
273 : :
274 : 2 : PooledBuffer& PooledBuffer::operator=(PooledBuffer&& other) noexcept {
275 [ + + ]: 2 : if (this != &other) {
276 [ + - + - : 1 : if (buffer_ && pool_) {
+ - ]
277 : : try {
278 : 1 : pool_->release(std::move(buffer_), size_);
279 : 0 : } catch (...) {
280 : : // Ignore exception and continue (noexcept function)
281 : 0 : }
282 : : }
283 : :
284 : 1 : buffer_ = std::move(other.buffer_);
285 : 1 : size_ = other.size_;
286 : 1 : pool_ = other.pool_;
287 : :
288 : 1 : other.buffer_ = nullptr;
289 : 1 : other.size_ = 0;
290 : 1 : other.pool_ = nullptr;
291 : : }
292 : 2 : return *this;
293 : : }
294 : :
295 : 2573 : uint8_t* PooledBuffer::data() const { return buffer_.get(); }
296 : :
297 : 3845 : size_t PooledBuffer::size() const { return size_; }
298 : :
299 : 1311 : bool PooledBuffer::valid() const { return buffer_ != nullptr; }
300 : :
301 : 307 : uint8_t& PooledBuffer::operator[](size_t index) { return buffer_[index]; }
302 : :
303 : 100 : const uint8_t& PooledBuffer::operator[](size_t index) const { return buffer_[index]; }
304 : :
305 : 103 : uint8_t& PooledBuffer::at(size_t index) {
306 [ + - + + : 103 : if (!buffer_ || index >= size_) {
+ + ]
307 : 3 : throw std::out_of_range("Buffer index out of range");
308 : : }
309 : 100 : return buffer_[index];
310 : : }
311 : :
312 : 101 : const uint8_t& PooledBuffer::at(size_t index) const {
313 [ + - + + : 101 : if (!buffer_ || index >= size_) {
+ + ]
314 : 1 : throw std::out_of_range("Buffer index out of range");
315 : : }
316 : 100 : return buffer_[index];
317 : : }
318 : :
319 : 0 : void PooledBuffer::check_bounds(size_t index) const {
320 [ # # # # : 0 : if (!buffer_ || index >= size_) {
# # ]
321 : 0 : throw std::out_of_range("Buffer index out of range");
322 : : }
323 : 0 : }
324 : :
325 : : } // namespace memory
326 : : } // namespace wirestead
|
