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 <atomic>
20 : : #include <boost/asio/buffer.hpp>
21 : : #include <cstddef>
22 : : #include <memory>
23 : : #include <mutex>
24 : : #include <type_traits>
25 : : #include <variant>
26 : : #include <vector>
27 : :
28 : : #include "wirestead/base/constants.hpp"
29 : :
30 : : namespace wirestead {
31 : : namespace transport {
32 : : namespace queue_util {
33 : :
34 : : struct DropAccounting {
35 : : size_t messages = 0;
36 : : size_t bytes = 0;
37 : :
38 [ + - - + ]: 17263 : bool any() const { return messages > 0 || bytes > 0; }
39 : : };
40 : :
41 : 272785 : inline bool try_reserve_write_bytes(std::atomic<size_t>& queue_bytes, const std::atomic<size_t>& pending_bytes,
42 : : const std::atomic<bool>& backpressure_active, size_t bytes, size_t bp_high,
43 : : size_t bp_limit) {
44 [ + - + + : 272785 : if (bytes == 0 || bytes > bp_limit || backpressure_active.load(std::memory_order_relaxed)) return false;
- + + + ]
45 : :
46 : 545566 : size_t current = queue_bytes.load(std::memory_order_relaxed);
47 : : for (;;) {
48 : 272824 : const size_t pending = pending_bytes.load(std::memory_order_relaxed);
49 [ + - - + ]: 272824 : if (current > bp_high || bytes > bp_high - current) return false;
50 [ + - + - : 272824 : if (current > bp_limit || pending > bp_limit - current || bytes > bp_limit - current - pending) return false;
- + ]
51 : :
52 [ + + ]: 545648 : if (queue_bytes.compare_exchange_weak(current, current + bytes, std::memory_order_acq_rel,
53 : : std::memory_order_relaxed)) {
54 : 272783 : return true;
55 : : }
56 : 41 : }
57 : : }
58 : :
59 : : // Nonblocking admission shares the plain-write reservation lock and includes
60 : : // accepted posts that have not reached the executor yet.
61 : 272785 : inline bool try_reserve_write_bytes(std::mutex& mutex, const std::atomic<size_t>& inflight_bytes,
62 : : std::atomic<size_t>& queue_bytes, const std::atomic<size_t>& pending_bytes,
63 : : const std::atomic<bool>& backpressure_active, size_t bytes, size_t bp_high,
64 : : size_t bp_limit) {
65 : 272785 : std::lock_guard<std::mutex> lock(mutex);
66 : 272785 : const size_t inflight = inflight_bytes.load(std::memory_order_relaxed);
67 [ - + ]: 272785 : if (inflight > bp_limit) return false;
68 : 272785 : return try_reserve_write_bytes(queue_bytes, pending_bytes, backpressure_active, bytes, bp_high, bp_limit - inflight);
69 : 272785 : }
70 : :
71 : 7625 : inline void release_reserved_write_bytes(std::atomic<size_t>& queue_bytes, size_t bytes) {
72 : 15250 : size_t current = queue_bytes.load(std::memory_order_relaxed);
73 : : for (;;) {
74 [ + + ]: 7640 : const size_t next = current > bytes ? current - bytes : 0;
75 [ + + ]: 15280 : if (queue_bytes.compare_exchange_weak(current, next, std::memory_order_acq_rel, std::memory_order_relaxed)) return;
76 : 15 : }
77 : : }
78 : :
79 : : // Reserves `added` bytes against bp_limit for the plain (blocking-capable)
80 : : // async_write_* entry points, accounting for bytes already
81 : : // reserved-but-not-yet-routed via `inflight_bytes` in addition to
82 : : // `queue_bytes`/`pending_bytes`. Closes an accept-then-drop race where the
83 : : // caller-thread precheck used to only read queue_bytes+pending_bytes
84 : : // non-atomically without reserving space: concurrent callers could all pass
85 : : // the check, then get rejected once actually routed onto the strand, where
86 : : // the real combined total exceeded bp_limit (jwsung91/wirestead#517).
87 : : //
88 : : // Guarded by `mtx` rather than a lock-free CAS on `inflight_bytes` alone:
89 : : // queue_bytes/pending_bytes/inflight_bytes are three independently-atomic
90 : : // counters, so no CAS retry loop touching only one of them can make a
91 : : // check spanning all three race-free - the strand's commit_reserved_limit_bytes()
92 : : // promoting inflight_bytes into queue_bytes/pending_bytes is itself two
93 : : // separate atomic ops, and a concurrent reservation's queue_bytes/pending_bytes
94 : : // reads can land in the gap between them while its CAS against a stale
95 : : // inflight_bytes snapshot still spuriously succeeds. This was caught by a
96 : : // CI run reproducing a 1-in-~10000 dropped message under this exact
97 : : // window. `mtx` must be the same mutex passed to commit_reserved_limit_bytes()
98 : : // for a given transport instance. Deliberately omits the
99 : : // bp_high/backpressure_active gate that try_reserve_write_bytes() applies
100 : : // above - a plain write is allowed to land in pending_ while backpressure
101 : : // is active, unlike the non-blocking try_write path.
102 : 32386 : inline bool try_reserve_limit_bytes(std::mutex& mtx, const std::atomic<size_t>& queue_bytes,
103 : : const std::atomic<size_t>& pending_bytes, std::atomic<size_t>& inflight_bytes,
104 : : size_t added, size_t bp_limit) {
105 [ + - + + ]: 32386 : if (added == 0 || added > bp_limit) return false;
106 : :
107 : 32378 : std::lock_guard<std::mutex> lock(mtx);
108 : 32378 : const size_t queue = queue_bytes.load(std::memory_order_relaxed);
109 : 32378 : const size_t pending = pending_bytes.load(std::memory_order_relaxed);
110 : 32378 : const size_t inflight = inflight_bytes.load(std::memory_order_relaxed);
111 [ + - + - : 32378 : if (queue > bp_limit || pending > bp_limit - queue || inflight > bp_limit - queue - pending ||
+ - ]
112 [ + + ]: 32378 : added > bp_limit - queue - pending - inflight) {
113 : 13802 : return false;
114 : : }
115 : 18576 : inflight_bytes.fetch_add(added, std::memory_order_relaxed);
116 : 18576 : return true;
117 : 32378 : }
118 : :
119 : : // Releases a reservation made by try_reserve_limit_bytes(), for the case
120 : : // where the message never gets promoted into queue_bytes_/pending_bytes_
121 : : // (route_enqueued_buffer's Rejected branch - expected to be effectively
122 : : // unreachable given the reservation above, but kept as a safety net).
123 : 1316 : inline void release_reserved_limit_bytes(std::mutex& mtx, std::atomic<size_t>& inflight_bytes, size_t bytes) {
124 : 1316 : std::lock_guard<std::mutex> lock(mtx);
125 : 1316 : const size_t current = inflight_bytes.load(std::memory_order_relaxed);
126 [ + + ]: 1316 : inflight_bytes.store(current > bytes ? current - bytes : 0, std::memory_order_relaxed);
127 : 1316 : }
128 : :
129 : : // Promotes a reservation made by try_reserve_limit_bytes() into `counter`
130 : : // (queue_bytes_ or pending_bytes_) once route_enqueued_buffer() has decided
131 : : // where the message lands. Must run under the same `mtx` as
132 : : // try_reserve_limit_bytes() and perform both the increment and the
133 : : // inflight_bytes release as one critical section - doing them as two
134 : : // separately-locked steps would reopen the exact gap described above.
135 : 17259 : inline void commit_reserved_limit_bytes(std::mutex& mtx, std::atomic<size_t>& counter,
136 : : std::atomic<size_t>& inflight_bytes, size_t bytes) {
137 : 17259 : std::lock_guard<std::mutex> lock(mtx);
138 : 17259 : counter.fetch_add(bytes, std::memory_order_relaxed);
139 : 17259 : const size_t current = inflight_bytes.load(std::memory_order_relaxed);
140 [ + + ]: 17259 : inflight_bytes.store(current > bytes ? current - bytes : 0, std::memory_order_relaxed);
141 : 17259 : }
142 : :
143 : : // Legacy locked increment utility. Built-in plain writes now always reserve
144 : : // capacity and use commit_reserved_limit_bytes() instead.
145 : : inline void commit_unreserved_limit_bytes(std::mutex& mtx, std::atomic<size_t>& counter, size_t bytes) {
146 : : std::lock_guard<std::mutex> lock(mtx);
147 : : counter.fetch_add(bytes, std::memory_order_relaxed);
148 : : }
149 : :
150 : : // Returns the byte size of a buffer held in a transport BufferVariant alternative.
151 : : // shared_ptr<const vector<uint8_t>> goes through ->size(); everything else via .size().
152 : : template <typename T>
153 : 24695 : inline size_t variant_buffer_size(const T& buf) {
154 : : if constexpr (std::is_same_v<T, std::shared_ptr<const std::vector<uint8_t>>>) {
155 [ + + ]: 6281 : return buf ? buf->size() : 0;
156 : : } else {
157 : 18414 : return buf.size();
158 : : }
159 : : }
160 : :
161 : : // ---------------------------------------------------------------------------
162 : : // Gather writes
163 : : // ---------------------------------------------------------------------------
164 : : //
165 : : // Stream transports used to hand exactly one queued buffer to async_write(),
166 : : // so a backlog of N queued messages cost N send syscalls even though they were
167 : : // all ready at once. Draining several into one scatter-gather write collapses
168 : : // those into one.
169 : : //
170 : : // Caps on how much of tx_ a single gather write may take. These bound the
171 : : // active batch and the work re-queued when a write fails.
172 : : // 16, not an arbitrary round number: asio fills at most 16 buffers per
173 : : // prepared_buffers (detail/consuming_buffers.hpp, max_buffers), which is also
174 : : // the iovec count a single sendmsg gets. Staying at or under it keeps each
175 : : // gather write inside one prepare. Larger batches were measured to corrupt the
176 : : // stream - at 64 buffers, queued messages came out of order on the wire while
177 : : // byte counts still matched, so only an ordering check catches it. Anything
178 : : // above this needs that verified again, not assumed.
179 : : inline constexpr size_t kMaxGatherBuffers = 16;
180 : : inline constexpr size_t kMaxGatherBytes = 256 * 1024;
181 : :
182 : : // Returns a const_buffer over a BufferVariant alternative.
183 : : template <typename T>
184 : 12268 : inline ::boost::asio::const_buffer variant_const_buffer(const T& buf) {
185 : : if constexpr (std::is_same_v<T, std::shared_ptr<const std::vector<uint8_t>>>) {
186 [ + - ]: 3130 : return buf ? ::boost::asio::const_buffer(buf->data(), buf->size()) : ::boost::asio::const_buffer();
187 : : } else {
188 : 9138 : return ::boost::asio::const_buffer(buf.data(), buf.size());
189 : : }
190 : : }
191 : : // Identity for plain buffers; tracked/endpoint queue items project their payload.
192 : : struct IdentityProjection {
193 : : template <typename T>
194 : 6 : constexpr T& operator()(T& x) const {
195 : 6 : return x;
196 : : }
197 : : };
198 : :
199 : : struct IgnoreDroppedBuffer {
200 : : template <typename T>
201 : 6 : void operator()(const T&) const {}
202 : : };
203 : :
204 : : // Moves buffers from the front of `tx` into `batch` (which is cleared first)
205 : : // up to the caps above, filling `views` with the matching const_buffers.
206 : : // Returns the total byte count, which is what the caller must subtract from
207 : : // queue_bytes_ once the write completes.
208 : : //
209 : : // Must run on the strand, with no write already in flight: `views` describes
210 : : // memory owned by `batch`, so neither may be touched until the write completes.
211 : : //
212 : : // `views` is handed to asio as an ordinary ConstBufferSequence, which asio
213 : : // copies into the composed operation. An earlier version passed a non-owning
214 : : // view to dodge that copy; asio's partial-write bookkeeping does not survive
215 : : // an aliasing sequence, and it silently duplicated or stalled queued buffers.
216 : : // Project extracts the payload while batch retains the complete queue item.
217 : : // One small copy per gather write is the right trade - it replaces N syscalls,
218 : : // and with several messages per write it is fewer allocations than before.
219 : : template <typename Deque, typename Batch, typename Project = IdentityProjection>
220 : 3500 : inline size_t take_gather_batch(Deque& tx, Batch& batch, std::vector<::boost::asio::const_buffer>& views,
221 : : Project project = Project{}) {
222 : 3500 : batch.clear();
223 : 3500 : views.clear();
224 : 3500 : size_t total = 0;
225 [ + + + + : 15768 : while (!tx.empty() && batch.size() < kMaxGatherBuffers && total < kMaxGatherBytes) {
+ - + + ]
226 : 24536 : const size_t n = std::visit([](const auto& b) { return variant_buffer_size(b); }, project(tx.front()));
227 : 12268 : batch.push_back(std::move(tx.front()));
228 : 12268 : tx.pop_front();
229 : 12268 : total += n;
230 : : }
231 : 3500 : views.reserve(batch.size());
232 [ + + ]: 15768 : for (const auto& b : batch) {
233 : 24536 : views.push_back(std::visit([](const auto& x) { return variant_const_buffer(x); }, project(b)));
234 : : }
235 : 3500 : return total;
236 : : }
237 : :
238 : : // Returns a failed batch to the front of `tx` in its original order, so a
239 : : // retry/teardown sees the queue exactly as it was.
240 : : template <typename Deque, typename Batch>
241 : 106 : inline void return_gather_batch(Deque& tx, Batch& batch) {
242 [ + + ]: 268 : for (auto it = batch.rbegin(); it != batch.rend(); ++it) {
243 : 162 : tx.push_front(std::move(*it));
244 : : }
245 : 106 : batch.clear();
246 : 106 : }
247 : :
248 : : // BestEffort queue-trimming shared by all stream transports (TCP client/server, UDS client/server)
249 : : // and, via a projection, UDP.
250 : : // Must be called on the strand immediately before enqueueing a new buffer of `added` bytes.
251 : : //
252 : : // No-op for Reliable strategy.
253 : : // For BestEffort:
254 : : // added >= bp_high → drop entire tx_ (full keep-latest replacement).
255 : : // otherwise → pop oldest tx_ entries until queue_bytes + added <= bp_high.
256 : : // on_drop observes each removed item before destruction and must not throw.
257 : : template <typename Deque, typename Project = IdentityProjection, typename OnDrop = IgnoreDroppedBuffer>
258 : 5 : inline DropAccounting maybe_flush_for_keep_latest(::wirestead::base::constants::BackpressureStrategy bp_strategy,
259 : : size_t added, size_t bp_high, Deque& tx,
260 : : std::atomic<size_t>& queue_bytes,
261 : : const std::atomic<bool>& backpressure_active,
262 : : Project project = Project{}, OnDrop on_drop = OnDrop{}) {
263 : 5 : DropAccounting dropped;
264 [ + + ]: 5 : if (bp_strategy != ::wirestead::base::constants::BackpressureStrategy::BestEffort) return dropped;
265 : :
266 [ + + ]: 4 : if (added >= bp_high) {
267 : 1 : size_t removed_bytes = 0;
268 [ + + ]: 3 : for (auto& buf : tx) {
269 : 2 : on_drop(buf);
270 : 4 : removed_bytes += std::visit([](const auto& b) { return variant_buffer_size(b); }, project(buf));
271 : : }
272 : 1 : dropped.messages = tx.size();
273 : 1 : dropped.bytes = removed_bytes;
274 : 1 : tx.clear();
275 : 1 : const size_t qb = queue_bytes.load(std::memory_order_relaxed);
276 [ - + ]: 1 : queue_bytes.store(qb > removed_bytes ? qb - removed_bytes : 0, std::memory_order_relaxed);
277 : 1 : return dropped;
278 : : }
279 : :
280 [ + - + + ]: 6 : if (backpressure_active.load(std::memory_order_relaxed) ||
281 [ + + ]: 6 : queue_bytes.load(std::memory_order_relaxed) + added > bp_high) {
282 [ + + ]: 6 : while (!tx.empty()) {
283 : 5 : const size_t qb = queue_bytes.load(std::memory_order_relaxed);
284 [ + + ]: 5 : if (qb + added <= bp_high) break;
285 : 8 : const size_t oldest = std::visit([](const auto& b) { return variant_buffer_size(b); }, project(tx.front()));
286 [ + - ]: 4 : queue_bytes.store(qb > oldest ? qb - oldest : 0, std::memory_order_relaxed);
287 : 4 : on_drop(tx.front());
288 : 4 : tx.pop_front();
289 : 4 : ++dropped.messages;
290 : 4 : dropped.bytes += oldest;
291 : : }
292 : : }
293 : 3 : return dropped;
294 : : }
295 : :
296 : : } // namespace queue_util
297 : : } // namespace transport
298 : : } // namespace wirestead
|
