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 <cstddef>
21 : :
22 : : #include "wirestead/base/constants.hpp"
23 : : #include "wirestead/diagnostics/runtime_stats_counter.hpp"
24 : : #include "wirestead/interface/channel.hpp"
25 : : #include "wirestead/transport/base/bp_utils.hpp"
26 : :
27 : : // Shared "pre-check -> route to tx/pending -> report backpressure -> drain on
28 : : // terminate" state machine, factored out of ~6 independently-maintained
29 : : // per-transport copies (jwsung91/wirestead#434). Deliberately non-owning/
30 : : // ref-based rather than template-owning the queues themselves: buffer
31 : : // element types genuinely differ across transports (UDP's TxItem carries a
32 : : // destination endpoint that stream transports have no analogue for), so
33 : : // this header only decides what a transport's own enqueue/write-completion
34 : : // code should do, and calls back into transport-supplied hooks for the
35 : : // pieces that must stay transport-specific (moving a buffer element from
36 : : // pending_ into tx_, kicking off a write).
37 : : namespace wirestead {
38 : : namespace transport {
39 : : namespace queue_util {
40 : :
41 : : // The subset of a transport's Impl state this state machine needs to read
42 : : // and mutate. All fields are references into the transport's own members -
43 : : // this struct never owns anything and is cheap to construct per call.
44 : : struct BackpressureFields {
45 : : std::atomic<size_t>& queue_bytes;
46 : : std::atomic<size_t>& pending_bytes;
47 : : std::atomic<bool>& backpressure_active;
48 : : size_t bp_high;
49 : : size_t bp_low;
50 : : size_t bp_limit;
51 : : ::wirestead::base::constants::BackpressureStrategy strategy;
52 : : };
53 : :
54 : : enum class EnqueueDecision {
55 : : Immediate, // push the buffer onto tx_ and proceed to write it now
56 : : Pending, // Reliable + backpressure already active: route to pending_ instead
57 : : Rejected, // over bp_limit_ even after BestEffort trimming; caller must drop it
58 : : };
59 : :
60 : : // Decides how a new buffer of `added` bytes should be routed, mirroring the
61 : : // branch structure every transport's enqueue function hand-copied. Does NOT
62 : : // push `added` into any queue itself - the caller pushes its own
63 : : // buffer/destination element into tx_ or pending_ based on the returned
64 : : // decision, and is responsible for incrementing the corresponding byte
65 : : // counter (queue_bytes/pending_bytes) to match. For BestEffort, trims `tx`
66 : : // via the existing maybe_flush_for_keep_latest() before returning
67 : : // Immediate, recording anything dropped into `dropped_out`. `project`
68 : : // defaults to identity (tx's elements are the BufferVariant directly);
69 : : // UDP's tx_ holds TxItem{BufferVariant, destination} instead and supplies a
70 : : // projection extracting `.buffer`.
71 : : template <typename Deque, typename Project = IdentityProjection>
72 : 12364 : inline EnqueueDecision decide_enqueue(BackpressureFields& f, size_t added, Deque& tx, DropAccounting& dropped_out,
73 : : Project project = Project{}) {
74 : : using Strategy = ::wirestead::base::constants::BackpressureStrategy;
75 : :
76 [ + + + + : 12364 : if (f.strategy == Strategy::Reliable && f.backpressure_active.load(std::memory_order_relaxed)) {
+ + ]
77 : 20986 : if (f.queue_bytes.load(std::memory_order_relaxed) + f.pending_bytes.load(std::memory_order_relaxed) + added >
78 : 10493 : f.bp_limit) {
79 : 1 : return EnqueueDecision::Rejected;
80 : : }
81 : 10492 : return EnqueueDecision::Pending;
82 : : }
83 : :
84 [ + + + + : 1885 : if (f.strategy == Strategy::BestEffort && (f.backpressure_active.load(std::memory_order_relaxed) ||
+ + ]
85 [ + + ]: 28 : f.queue_bytes.load(std::memory_order_relaxed) + added > f.bp_high)) {
86 : : dropped_out =
87 : 23 : maybe_flush_for_keep_latest(f.strategy, added, f.bp_high, tx, f.queue_bytes, f.backpressure_active, project);
88 : : }
89 : :
90 [ - + ]: 3742 : if (f.queue_bytes.load(std::memory_order_relaxed) + added > f.bp_limit) {
91 : 0 : return EnqueueDecision::Rejected;
92 : : }
93 : 1871 : return EnqueueDecision::Immediate;
94 : : }
95 : :
96 : : // Runs the ON / OFF-with-reflush / re-ARM state machine and fires on_bp
97 : : // accordingly, mirroring UdpChannel::report_backpressure() (the
98 : : // best-behaved existing copy). `flush_pending_into_tx` is called exactly
99 : : // once, only on the OFF transition, and must move every element out of the
100 : : // transport's own pending_ deque into tx_ (preserving whatever
101 : : // transport-specific fields those elements carry) and return the number of
102 : : // bytes moved; `kick_write` is called if the OFF transition leaves anything
103 : : // in tx_ and nothing is currently being written.
104 : : template <typename FlushFn, typename KickFn>
105 : 346241 : inline void report_backpressure(BackpressureFields& f, size_t queued_bytes,
106 : : const interface::Channel::OnBackpressure& on_bp,
107 : : diagnostics::RuntimeStatsCounters& stats, FlushFn&& flush_pending_into_tx,
108 : : KickFn&& kick_write) {
109 [ + + + + : 346241 : if (!f.backpressure_active.load(std::memory_order_relaxed) && queued_bytes >= f.bp_high) {
+ + ]
110 : 78 : f.backpressure_active.store(true, std::memory_order_relaxed);
111 : 78 : stats.record_backpressure_event();
112 [ + + ]: 78 : if (on_bp) {
113 : : try {
114 : 74 : on_bp(queued_bytes);
115 : 2 : } catch (...) {
116 : : }
117 : : }
118 : 78 : return;
119 : : }
120 : :
121 [ + + + + : 346163 : if (f.backpressure_active.load(std::memory_order_relaxed) && queued_bytes <= f.bp_low) {
+ + ]
122 : 51 : const size_t moved = flush_pending_into_tx();
123 : 51 : f.queue_bytes.fetch_add(moved, std::memory_order_relaxed);
124 : 51 : f.backpressure_active.store(false, std::memory_order_relaxed);
125 : 51 : stats.record_backpressure_event();
126 [ + + ]: 51 : if (on_bp) {
127 : : try {
128 : 47 : on_bp(queued_bytes);
129 : 0 : } catch (...) {
130 : : } // fire OFF with pre-flush queue size, matching every existing copy
131 : : }
132 : :
133 : 51 : const size_t post_flush = f.queue_bytes.load(std::memory_order_relaxed);
134 [ + + ]: 51 : if (post_flush >= f.bp_high) {
135 : 7 : f.backpressure_active.store(true, std::memory_order_relaxed);
136 : 7 : stats.record_backpressure_event();
137 [ + + ]: 7 : if (on_bp) {
138 : : try {
139 : 4 : on_bp(post_flush);
140 : 0 : } catch (...) {
141 : : }
142 : : }
143 : : }
144 : 51 : kick_write();
145 : : }
146 : : }
147 : :
148 : : // Drops all queued and pending writes and clears backpressure, firing on_bp
149 : : // once (with a final queued-bytes value of 0) if it was active - and doing
150 : : // nothing at all if it wasn't. This is the *terminal* drain path (stop/error),
151 : : // deliberately distinct from report_backpressure()'s normal-operation flush:
152 : : // report_backpressure() moves pending_ back into tx_ and can immediately
153 : : // re-arm backpressure_active_, which would strand a Reliable-mode sender
154 : : // blocked on a condition variable forever once nothing will ever call
155 : : // do_write() again (jwsung91/wirestead#427, #452). `clear_queues` must clear
156 : : // every transport-specific queue (tx_, pending_, and their byte counters).
157 : : template <typename ClearFn>
158 : 393 : inline void drain_and_clear_backpressure(BackpressureFields& f, const interface::Channel::OnBackpressure& on_bp,
159 : : ClearFn&& clear_queues) {
160 : 393 : clear_queues();
161 : 393 : const bool had_backpressure = f.backpressure_active.load(std::memory_order_relaxed);
162 : 393 : f.backpressure_active.store(false, std::memory_order_relaxed);
163 [ + + ]: 393 : if (!had_backpressure) return;
164 [ + + ]: 9 : if (on_bp) {
165 : : try {
166 : 3 : on_bp(0);
167 : 0 : } catch (...) {
168 : : }
169 : : }
170 : : }
171 : :
172 : : } // namespace queue_util
173 : : } // namespace transport
174 : : } // namespace wirestead
|