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 <algorithm>
20 : : #include <atomic>
21 : : #include <memory>
22 : : #include <stdexcept>
23 : : #include <utility>
24 : :
25 : : #include "wirestead/wrapper/receive_limits.hpp"
26 : :
27 : : namespace wirestead::wrapper::detail {
28 : : class ReceiveScope;
29 : : class ReceiveCharge;
30 : :
31 : : class ReceiveBudget : public std::enable_shared_from_this<ReceiveBudget> {
32 : : public:
33 : 1858 : explicit ReceiveBudget(ReceiveLimits limits) : limits_(limits) { limits_.validate(); }
34 : 1179 : const ReceiveLimits& limits() const { return limits_; }
35 : : std::shared_ptr<ReceiveScope> open_scope();
36 : 16298 : ReceiveMemoryStats stats() const {
37 : 16298 : const auto used = used_.load();
38 : 16298 : ReceiveMemoryStats result{used, std::max(used, peak_.load()), sessions_.load(), overflows_.load(),
39 : 65192 : overflow_bytes_.load()};
40 [ + + ]: 162980 : for (size_t i = 0; i < reasons_.size(); ++i) result.overflow_by_reason[i] = reasons_[i].load();
41 : 16298 : return result;
42 : : }
43 : 2670 : void reset_stats() {
44 : 5340 : peak_ = used_.load();
45 : 2670 : overflows_ = 0;
46 : 2670 : overflow_bytes_ = 0;
47 [ + + ]: 13350 : for (auto& reason : reasons_) reason = 0;
48 : 2670 : }
49 : 28 : void overflow(ReceiveOverflowReason reason, size_t bytes) {
50 : 28 : ++overflows_;
51 : 28 : overflow_bytes_ += bytes;
52 : 28 : ++reasons_[static_cast<size_t>(reason)];
53 : 28 : }
54 : :
55 : : private:
56 : : friend class ReceiveScope;
57 : 19753 : bool acquire(size_t bytes) {
58 : 39506 : auto used = used_.load();
59 : : do {
60 [ + + ]: 24206 : if (bytes > limits_.max_bytes - used) return false;
61 [ + + ]: 38120 : } while (!used_.compare_exchange_weak(used, used + bytes));
62 : 14607 : auto peak = peak_.load();
63 [ + + - + : 15165 : while (peak < used + bytes && !peak_.compare_exchange_weak(peak, used + bytes)) {
- + ]
64 : : }
65 : 14607 : return true;
66 : : }
67 : : ReceiveLimits limits_;
68 : : std::atomic<size_t> used_{0}, peak_{0}, sessions_{0};
69 : : std::atomic<uint64_t> overflows_{0}, overflow_bytes_{0};
70 : : std::array<std::atomic<uint64_t>, 4> reasons_{};
71 : : };
72 : :
73 : : class ReceiveScope : public std::enable_shared_from_this<ReceiveScope> {
74 : : friend class ReceiveBudget;
75 : : struct ConstructionKey {};
76 : :
77 : : public:
78 : 1959 : ReceiveScope(std::shared_ptr<ReceiveBudget> budget, ConstructionKey) : budget_(std::move(budget)) {}
79 : 1959 : ~ReceiveScope() { --budget_->sessions_; }
80 : : std::shared_ptr<ReceiveCharge> reserve(size_t bytes);
81 : 22 : void overflow(size_t bytes, ReceiveOverflowReason reason = ReceiveOverflowReason::ByteLimit) {
82 : 22 : ++overflows_;
83 : 22 : overflow_bytes_ += bytes;
84 : 22 : budget_->overflow(reason, bytes);
85 : 22 : }
86 : 4 : ReceiveMemoryStats stats() const {
87 : 4 : const auto used = used_.load();
88 : 16 : return {used, std::max(used, peak_.load()), 1, overflows_.load(), overflow_bytes_.load()};
89 : : }
90 : 1179 : const ReceiveLimits& limits() const { return budget_->limits(); }
91 : :
92 : : private:
93 : : friend class ReceiveCharge;
94 : 14609 : void release(size_t bytes) {
95 : 14609 : used_ -= bytes;
96 : 14609 : budget_->used_ -= bytes;
97 : 14609 : }
98 : : std::shared_ptr<ReceiveBudget> budget_;
99 : : std::atomic<size_t> used_{0}, peak_{0};
100 : : std::atomic<uint64_t> overflows_{0}, overflow_bytes_{0};
101 : : };
102 : :
103 : : // The charge follows retained library storage through moves, timer delivery and
104 : : // callback execution. Removing a session cannot release bytes still in flight.
105 : : class ReceiveCharge {
106 : : friend class ReceiveScope;
107 : : struct ConstructionKey {};
108 : :
109 : : public:
110 : 14607 : ReceiveCharge(std::shared_ptr<ReceiveScope> scope, size_t bytes, ConstructionKey)
111 : 14607 : : scope_(std::move(scope)), bytes_(bytes) {}
112 : 14607 : ~ReceiveCharge() { scope_->release(bytes_); }
113 : : ReceiveCharge(const ReceiveCharge&) = delete;
114 : : ReceiveCharge& operator=(const ReceiveCharge&) = delete;
115 : 4 : size_t size() const { return bytes_; }
116 : 3 : void merge(ReceiveCharge& other) {
117 [ + + ]: 3 : if (this == &other) return;
118 [ + + + - ]: 2 : if (scope_ != other.scope_) throw std::logic_error("receive charge scopes differ");
119 : 1 : bytes_ += std::exchange(other.bytes_, 0);
120 : : }
121 : 3 : void shrink(size_t bytes) {
122 [ + + + - ]: 3 : if (bytes > bytes_) throw std::logic_error("receive charge cannot grow without reservation");
123 : 2 : scope_->release(bytes_ - bytes);
124 : 2 : bytes_ = bytes;
125 : 2 : }
126 : :
127 : : private:
128 : : std::shared_ptr<ReceiveScope> scope_;
129 : : size_t bytes_;
130 : : };
131 : :
132 : 19753 : inline std::shared_ptr<ReceiveCharge> ReceiveScope::reserve(size_t bytes) {
133 [ + + ]: 19753 : if (!budget_->acquire(bytes)) return {};
134 : 14607 : const auto used = used_.fetch_add(bytes) + bytes;
135 : 14607 : auto peak = peak_.load();
136 [ + + - + : 15227 : while (peak < used && !peak_.compare_exchange_weak(peak, used)) {
- + ]
137 : : }
138 : : try {
139 : 14607 : return std::make_shared<ReceiveCharge>(shared_from_this(), bytes, ReceiveCharge::ConstructionKey{});
140 : 0 : } catch (...) {
141 : 0 : release(bytes);
142 : 0 : throw;
143 : 0 : }
144 : : }
145 : :
146 : 1963 : inline std::shared_ptr<ReceiveScope> ReceiveBudget::open_scope() {
147 : 3926 : auto count = sessions_.load();
148 : : do {
149 [ + + ]: 1963 : if (count >= limits_.max_sessions) {
150 : 4 : overflow(ReceiveOverflowReason::SessionLimit, 0);
151 : 4 : return {};
152 : : }
153 [ - + ]: 3918 : } while (!sessions_.compare_exchange_weak(count, count + 1));
154 : : try {
155 : 1959 : return std::make_shared<ReceiveScope>(shared_from_this(), ReceiveScope::ConstructionKey{});
156 : 0 : } catch (...) {
157 : 0 : --sessions_;
158 : 0 : throw;
159 : 0 : }
160 : : }
161 : : } // namespace wirestead::wrapper::detail
|
