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