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 : : #pragma once
18 : :
19 : : #include <array>
20 : : #include <atomic>
21 : : #include <cstddef>
22 : : #include <memory>
23 : : #include <mutex>
24 : : #include <string>
25 : : #include <unordered_map>
26 : : #include <vector>
27 : :
28 : : #include "wirestead/base/visibility.hpp"
29 : :
30 : : namespace wirestead {
31 : : namespace memory {
32 : :
33 : : /**
34 : : * @brief Selective simplified memory pool with optimized performance
35 : : *
36 : : * Core design principles:
37 : : * - Small pools: Lock-based (fast allocation, low overhead)
38 : : * - Large pools: Lock-free (high concurrency)
39 : : * - Memory alignment: 64-byte alignment for buffers >= 4KB
40 : : * - Minimal statistics: Basic stats only to minimize overhead
41 : : */
42 : : class WIRESTEAD_API MemoryPool {
43 : : public:
44 : : // Basic statistics
45 : : struct PoolStats {
46 : : size_t total_allocations{0};
47 : : size_t pool_hits{0};
48 : : };
49 : :
50 : : // Basic health metrics
51 : : struct HealthMetrics {
52 : : double hit_rate{0.0};
53 : : };
54 : :
55 : : // Predefined buffer sizes for common use cases
56 : : enum class BufferSize : size_t {
57 : : SMALL = 1024, // 1KB - small messages
58 : : MEDIUM = 4096, // 4KB - typical network packets
59 : : LARGE = 16384, // 16KB - large data transfers
60 : : XLARGE = 65536 // 64KB - bulk operations
61 : : };
62 : :
63 : : // initial_pool_size buffers are created up front, split evenly across the
64 : : // four size buckets. It defaults to 0 because the buckets run 1 KiB to
65 : : // 64 KiB: prefilling 400 would reserve 8.3 MiB before the first acquire().
66 : : explicit MemoryPool(size_t initial_pool_size = 0, size_t max_pool_size = 2000);
67 : 494 : ~MemoryPool() = default;
68 : :
69 : : // Non-copyable, non-movable
70 : : MemoryPool(const MemoryPool&) = delete;
71 : : MemoryPool& operator=(const MemoryPool&) = delete;
72 : : MemoryPool(MemoryPool&&) = delete;
73 : : MemoryPool& operator=(MemoryPool&&) = delete;
74 : :
75 : : std::unique_ptr<uint8_t[]> acquire(size_t size);
76 : : std::unique_ptr<uint8_t[]> acquire(BufferSize buffer_size);
77 : : void release(std::unique_ptr<uint8_t[]> buffer, size_t size);
78 : : PoolStats stats() const;
79 : : double hit_rate() const;
80 : : // Returns (bytes currently checked out, peak bytes ever checked out
81 : : // concurrently). Unlike a monotonically-increasing allocation counter,
82 : : // both reflect actual acquire()/release() traffic - the first value goes
83 : : // down as buffers are released (#451).
84 : : std::pair<size_t, size_t> memory_usage() const;
85 : : HealthMetrics health_metrics() const;
86 : :
87 : : private:
88 : : // Simple pool bucket
89 : : struct PoolBucket {
90 : : std::vector<std::unique_ptr<uint8_t[]>> buffers_;
91 : : mutable std::mutex mutex_;
92 : : size_t size_;
93 : : size_t capacity_;
94 : :
95 : 2180 : PoolBucket() : size_{0}, capacity_{0} {}
96 : : PoolBucket(PoolBucket&& other) noexcept;
97 : : PoolBucket& operator=(PoolBucket&& other) noexcept;
98 : : PoolBucket(const PoolBucket&) = delete;
99 : : PoolBucket& operator=(const PoolBucket&) = delete;
100 : : };
101 : :
102 : : std::array<PoolBucket, 4> buckets_; // For SMALL, MEDIUM, LARGE, XLARGE
103 : :
104 : : // Internal statistics (atomic for thread safety)
105 : : std::atomic<size_t> total_allocations_{0};
106 : : std::atomic<size_t> pool_hits_{0};
107 : : // #451: real, fluctuating usage tracking for memory_usage() - incremented
108 : : // in acquire(), decremented in release(), unlike total_allocations_ above
109 : : // which only ever grows.
110 : : std::atomic<size_t> outstanding_bytes_{0};
111 : : std::atomic<size_t> peak_bytes_{0};
112 : :
113 : : // Helper functions
114 : : PoolBucket& bucket(size_t size);
115 : : size_t bucket_index(size_t size) const;
116 : : void track_acquire(size_t size);
117 : : void track_release(size_t size);
118 : :
119 : : // Allocation functions
120 : : std::unique_ptr<uint8_t[]> acquire_from_bucket(PoolBucket& bucket);
121 : : std::unique_ptr<uint8_t[]> create_buffer(size_t size);
122 : :
123 : : // Release functions
124 : : void release_to_bucket(PoolBucket& bucket, std::unique_ptr<uint8_t[]> buffer);
125 : :
126 : : // Utility functions
127 : : void validate_size(size_t size) const;
128 : : };
129 : :
130 : : /**
131 : : * @brief Global memory pool instance
132 : : */
133 : : class WIRESTEAD_API GlobalMemoryPool {
134 : : public:
135 : : static MemoryPool& instance();
136 : :
137 : : // Both factories tune capacity, not prefill: they raise max_pool_size so more
138 : : // released buffers stay resident under concurrency. Their prefill arguments
139 : : // (800 and 1200) were written while initial_pool_size was discarded and so
140 : : // allocated nothing; #575 made the parameter real, which would have turned
141 : : // them into roughly 17 MB and 26 MB allocated up front. Both start empty and
142 : : // fill as buffers are released, matching the constructor's default.
143 : :
144 : : // Factory method to create optimized memory pool
145 : : static std::unique_ptr<MemoryPool> create_optimized() {
146 : : return std::make_unique<MemoryPool>(0, 4000); // Optimized default sizes
147 : : }
148 : :
149 : : // Factory method to create size-optimized memory pool
150 : : static std::unique_ptr<MemoryPool> create_size_optimized() {
151 : : return std::make_unique<MemoryPool>(0, 6000); // Even larger for better concurrency
152 : : }
153 : :
154 : : // Non-copyable, non-movable
155 : : GlobalMemoryPool() = delete;
156 : : GlobalMemoryPool(const GlobalMemoryPool&) = delete;
157 : : GlobalMemoryPool& operator=(const GlobalMemoryPool&) = delete;
158 : : };
159 : :
160 : : /**
161 : : * @brief RAII wrapper for memory pool buffers with enhanced safety
162 : : */
163 : : class WIRESTEAD_API PooledBuffer {
164 : : public:
165 : : explicit PooledBuffer(size_t size);
166 : : explicit PooledBuffer(MemoryPool::BufferSize buffer_size);
167 : : // #443: draw from a specific pool (e.g. a per-channel instance) instead of
168 : : // the process-wide GlobalMemoryPool singleton, to avoid cross-channel
169 : : // contention on the singleton's bucket mutexes.
170 : : PooledBuffer(size_t size, MemoryPool& pool);
171 : : PooledBuffer(MemoryPool::BufferSize buffer_size, MemoryPool& pool);
172 : : ~PooledBuffer();
173 : :
174 : : // Non-copyable, movable
175 : : PooledBuffer(const PooledBuffer&) = delete;
176 : : PooledBuffer& operator=(const PooledBuffer&) = delete;
177 : : PooledBuffer(PooledBuffer&& other) noexcept;
178 : : PooledBuffer& operator=(PooledBuffer&& other) noexcept;
179 : :
180 : : // Safe access methods
181 : : uint8_t* data() const;
182 : : size_t size() const;
183 : : bool valid() const;
184 : :
185 : : // Safe array access with bounds checking
186 : : uint8_t& operator[](size_t index);
187 : : const uint8_t& operator[](size_t index) const;
188 : :
189 : : // Safe array access with bounds checking
190 : : uint8_t& at(size_t index);
191 : : const uint8_t& at(size_t index) const;
192 : :
193 : : // Explicit conversion methods (no implicit conversion)
194 : : uint8_t* get() const { return data(); }
195 : : explicit operator bool() const { return valid(); }
196 : :
197 : : private:
198 : : std::unique_ptr<uint8_t[]> buffer_;
199 : : size_t size_;
200 : : MemoryPool* pool_;
201 : :
202 : : // Helper for bounds checking
203 : : void check_bounds(size_t index) const;
204 : : };
205 : :
206 : : } // namespace memory
207 : : } // namespace wirestead
|
