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/callback.hpp"
24 : : #include "wirestead/diagnostics/runtime_stats_counter.hpp"
25 : : #include "wirestead/interface/channel.hpp"
26 : : #include "wirestead/transport/base/bp_utils.hpp"
27 : :
28 : : // Shared "pre-check -> route to tx/pending -> report backpressure -> drain on
29 : : // terminate" state machine, factored out of ~6 independently-maintained
30 : : // per-transport copies (jwsung91/wirestead#434). Deliberately non-owning/
31 : : // ref-based rather than template-owning the queues themselves: buffer
32 : : // element types genuinely differ across transports (UDP's TxItem carries a
33 : : // destination endpoint that stream transports have no analogue for), so
34 : : // this header only decides what a transport's own enqueue/write-completion
35 : : // code should do, and calls back into transport-supplied hooks for the
36 : : // pieces that must stay transport-specific (moving a buffer element from
37 : : // pending_ into tx_, kicking off a write).
38 : : namespace wirestead {
39 : : namespace transport {
40 : : namespace queue_util {
41 : :
42 : : // The subset of a transport's Impl state this state machine needs to read
43 : : // and mutate. All fields are references into the transport's own members -
44 : : // this struct never owns anything and is cheap to construct per call.
45 : : struct BackpressureFields {
46 : : std::atomic<size_t>& queue_bytes;
47 : : std::atomic<size_t>& pending_bytes;
48 : : std::atomic<bool>& backpressure_active;
49 : : size_t bp_high;
50 : : size_t bp_low;
51 : : size_t bp_limit;
52 : : ::wirestead::base::constants::BackpressureStrategy strategy;
53 : : std::mutex* reservation_mutex = nullptr;
54 : : };
55 : :
56 : : enum class EnqueueDecision {
57 : : Immediate, // push the buffer onto tx_ and proceed to write it now
58 : : Pending, // Backpressure already active: preserve accepted work in pending_
59 : : Rejected, // Defensive hard-limit guard; fixed-limit reserved admission prevents this
60 : : };
61 : :
62 : : // Acceptance already reserved hard-limit capacity. Strategy determines caller
63 : : // admission (try/refuse versus wait), never removal of accepted queue entries.
64 : : // Keep the projection/observer parameters for existing queue element adapters.
65 : : template <typename Deque, typename Project = IdentityProjection, typename OnDrop = IgnoreDroppedBuffer>
66 : 17265 : inline EnqueueDecision decide_enqueue(BackpressureFields& f, size_t added, Deque&, DropAccounting& dropped_out,
67 : : Project = Project{}, OnDrop = OnDrop{}) {
68 : 17265 : dropped_out = {};
69 : 17265 : const size_t queued = f.queue_bytes.load(std::memory_order_relaxed);
70 : 17265 : const size_t pending = f.pending_bytes.load(std::memory_order_relaxed);
71 [ + - + - : 17265 : if (queued > f.bp_limit || pending > f.bp_limit - queued || added > f.bp_limit - queued - pending)
+ + ]
72 : 1 : return EnqueueDecision::Rejected;
73 [ + + ]: 17264 : return f.backpressure_active.load(std::memory_order_relaxed) ? EnqueueDecision::Pending : EnqueueDecision::Immediate;
74 : : }
75 : :
76 : : // Runs the ON / OFF-with-reflush / re-ARM state machine and fires on_bp
77 : : // accordingly, mirroring UdpChannel::report_backpressure() (the
78 : : // best-behaved existing copy). `flush_pending_into_tx` is called exactly
79 : : // once, only on the OFF transition, and must move every element out of the
80 : : // transport's own pending_ deque into tx_ (preserving whatever
81 : : // transport-specific fields those elements carry) and return the number of
82 : : // bytes moved; `kick_write` is called if the OFF transition leaves anything
83 : : // in tx_ and nothing is currently being written.
84 : : template <typename FlushFn, typename KickFn>
85 : 286065 : inline void report_backpressure(BackpressureFields& f, size_t queued_bytes,
86 : : const interface::Channel::OnBackpressure& on_bp,
87 : : diagnostics::RuntimeStatsCounters& stats, FlushFn&& flush_pending_into_tx,
88 : : KickFn&& kick_write) {
89 [ + + + + : 286065 : if (!f.backpressure_active.load(std::memory_order_relaxed) && queued_bytes >= f.bp_high) {
+ + ]
90 : 655 : f.backpressure_active.store(true, std::memory_order_relaxed);
91 : 655 : stats.record_backpressure_event();
92 [ + + ]: 655 : if (on_bp) {
93 : 567 : diagnostics::invoke_callback("transport", "on_backpressure", on_bp, queued_bytes);
94 : : }
95 : 655 : return;
96 : : }
97 : :
98 [ + + + + : 285410 : if (f.backpressure_active.load(std::memory_order_relaxed) && queued_bytes <= f.bp_low) {
+ + ]
99 : : {
100 : : // Admission must not observe the gap between removing pending bytes and
101 : : // adding them to tx. Callbacks and kick_write run after releasing this lock.
102 : 333 : std::unique_lock<std::mutex> reservation;
103 [ + + + - ]: 333 : if (f.reservation_mutex) reservation = std::unique_lock<std::mutex>(*f.reservation_mutex);
104 : 333 : const size_t moved = flush_pending_into_tx();
105 : 333 : f.queue_bytes.fetch_add(moved, std::memory_order_relaxed);
106 : 333 : }
107 : 333 : f.backpressure_active.store(false, std::memory_order_relaxed);
108 : 333 : stats.record_backpressure_event();
109 [ + + ]: 333 : if (on_bp) {
110 : 274 : diagnostics::invoke_callback("transport", "on_backpressure", on_bp,
111 : : queued_bytes); // fire OFF with pre-flush queue size, matching every existing copy
112 : : }
113 : :
114 : 333 : const size_t post_flush = f.queue_bytes.load(std::memory_order_relaxed);
115 [ + + ]: 333 : if (post_flush >= f.bp_high) {
116 : 52 : f.backpressure_active.store(true, std::memory_order_relaxed);
117 : 52 : stats.record_backpressure_event();
118 [ + + ]: 52 : if (on_bp) {
119 : 21 : diagnostics::invoke_callback("transport", "on_backpressure", on_bp, post_flush);
120 : : }
121 : : }
122 : 333 : kick_write();
123 : : }
124 : : }
125 : :
126 : : // Drops all queued and pending writes and clears backpressure, firing on_bp
127 : : // once (with a final queued-bytes value of 0) if it was active - and doing
128 : : // nothing at all if it wasn't. This is the *terminal* drain path (stop/error),
129 : : // deliberately distinct from report_backpressure()'s normal-operation flush:
130 : : // report_backpressure() moves pending_ back into tx_ and can immediately
131 : : // re-arm backpressure_active_, which would strand a Reliable-mode sender
132 : : // blocked on a condition variable forever once nothing will ever call
133 : : // do_write() again (jwsung91/wirestead#427, #452). `clear_queues` must clear
134 : : // every transport-specific queue (tx_, pending_, and their byte counters).
135 : : template <typename ClearFn>
136 : 2032 : inline void drain_and_clear_backpressure(BackpressureFields& f, const interface::Channel::OnBackpressure& on_bp,
137 : : ClearFn&& clear_queues) {
138 : 2032 : clear_queues();
139 : 2032 : const bool had_backpressure = f.backpressure_active.load(std::memory_order_relaxed);
140 : 2032 : f.backpressure_active.store(false, std::memory_order_relaxed);
141 [ + + ]: 2032 : if (!had_backpressure) return;
142 [ + + ]: 170 : if (on_bp) {
143 : 79 : diagnostics::invoke_callback("transport", "on_backpressure", on_bp, 0);
144 : : }
145 : : }
146 : :
147 : : } // namespace queue_util
148 : : } // namespace transport
149 : : } // namespace wirestead
|
