Line data Source code
1 : #include "connection_tracker.hpp"
2 : #include "../../util/logger.hpp"
3 : #include <boost/asio/post.hpp>
4 : #include <utility>
5 :
6 : namespace thinger::http {
7 :
8 : // tracked_connection
9 :
10 28098 : tracked_connection::~tracked_connection() {
11 28098 : untrack();
12 28098 : }
13 :
14 28092 : void tracked_connection::track(std::shared_ptr<connection_tracker> tracker, std::weak_ptr<void> self,
15 : boost::asio::io_context& io_context, connection_type type) {
16 28092 : tracker_ = tracker;
17 28092 : self_ = std::move(self);
18 28092 : io_context_ = &io_context;
19 28092 : type_ = type;
20 28092 : tracker->add(*this);
21 28092 : }
22 :
23 42403 : void tracked_connection::untrack() {
24 : // (a tracker gone no longer lists it)
25 42403 : auto tracker = tracker_.lock();
26 42403 : tracker_.reset();
27 42403 : if (tracker) tracker->remove(*this);
28 42403 : }
29 :
30 7191 : void tracked_connection::set_type(connection_type type) {
31 7191 : if (auto tracker = tracker_.lock()) tracker->retype(*this, type);
32 7191 : else type_ = type;
33 7191 : }
34 :
35 : // connection_tracker
36 :
37 28092 : size_t connection_tracker::shard_of(const boost::asio::io_context& io_context) {
38 : // The connections of an io_context share a shard, so it is not contended by the others
39 28092 : auto address = reinterpret_cast<uintptr_t>(&io_context);
40 28092 : return ((address >> 6) * 0x9E3779B97F4A7C15ull) >> 60;
41 : }
42 :
43 28092 : void connection_tracker::add(tracked_connection& connection) {
44 28092 : connection.shard_ = static_cast<uint8_t>(shard_of(*connection.io_context_));
45 28092 : auto& shard = shards_[connection.shard_];
46 28092 : std::scoped_lock lock(shard.mutex);
47 28092 : connection.next_ = shard.head;
48 28092 : if (shard.head) shard.head->prev_ = &connection;
49 28092 : shard.head = &connection;
50 28092 : connection.linked_ = true;
51 28092 : ++shard.active[static_cast<size_t>(connection.type_)];
52 28092 : if (connection.type_ == connection_type::http) ++shard.total_connections;
53 28092 : }
54 :
55 28091 : void connection_tracker::remove(tracked_connection& connection) {
56 : {
57 28091 : auto& shard = shards_[connection.shard_];
58 28091 : std::scoped_lock lock(shard.mutex);
59 28091 : if (!connection.linked_) return;
60 28091 : if (connection.prev_) connection.prev_->next_ = connection.next_;
61 21948 : else shard.head = connection.next_;
62 28091 : if (connection.next_) connection.next_->prev_ = connection.prev_;
63 28091 : connection.prev_ = connection.next_ = nullptr;
64 28091 : connection.linked_ = false;
65 28091 : --shard.active[static_cast<size_t>(connection.type_)];
66 28091 : }
67 : // Only while a shutdown waits for the last one. Set before notify_when_empty() counted the
68 : // connections: if that count saw this one, this sees it set (the lock of the shard orders
69 : // them), and the last one to leave tells it
70 28091 : if (watching_.load(std::memory_order_acquire)) [[unlikely]] check_empty();
71 : }
72 :
73 7191 : void connection_tracker::retype(tracked_connection& connection, connection_type type) {
74 7191 : auto& shard = shards_[connection.shard_];
75 7191 : std::scoped_lock lock(shard.mutex);
76 7191 : if (connection.linked_) {
77 7191 : --shard.active[static_cast<size_t>(connection.type_)];
78 7191 : ++shard.active[static_cast<size_t>(type)];
79 : }
80 7191 : connection.type_ = type;
81 7191 : }
82 :
83 7906 : server_stats connection_tracker::stats() const {
84 7906 : server_stats stats;
85 134402 : for (const auto& shard : shards_) {
86 126496 : std::scoped_lock lock(shard.mutex);
87 126496 : stats.http_connections += shard.active[static_cast<size_t>(connection_type::http)];
88 126496 : stats.websocket_connections += shard.active[static_cast<size_t>(connection_type::websocket)];
89 126496 : stats.sse_connections += shard.active[static_cast<size_t>(connection_type::sse)];
90 126496 : stats.tunnel_connections += shard.active[static_cast<size_t>(connection_type::tunnel)];
91 126496 : stats.taken_over_connections += shard.active[static_cast<size_t>(connection_type::taken_over)];
92 126496 : stats.total_connections += shard.total_connections;
93 126496 : }
94 7906 : return stats;
95 : }
96 :
97 9871 : size_t connection_tracker::active() const {
98 9871 : size_t active = 0;
99 167807 : for (const auto& shard : shards_) {
100 157936 : std::scoped_lock lock(shard.mutex);
101 947616 : for (auto count : shard.active) active += count;
102 157936 : }
103 9871 : return active;
104 : }
105 :
106 4464 : void connection_tracker::notify_when_empty(std::function<void()> empty) {
107 : {
108 4464 : std::scoped_lock lock(empty_mutex_);
109 4464 : on_empty_ = std::move(empty);
110 4464 : watching_.store(static_cast<bool>(on_empty_), std::memory_order_release);
111 4464 : }
112 : // None left already (or the last left before it was set)
113 4464 : check_empty();
114 4464 : }
115 :
116 4710 : void connection_tracker::check_empty() {
117 4710 : std::function<void()> empty;
118 : {
119 : // (held while counting: told once, by whoever counts none first)
120 4710 : std::scoped_lock lock(empty_mutex_);
121 4710 : if (!on_empty_ || active() > 0) return;
122 228 : empty = std::exchange(on_empty_, nullptr);
123 228 : watching_.store(false, std::memory_order_release);
124 4710 : }
125 228 : empty();
126 4710 : }
127 :
128 25665 : std::vector<std::pair<std::shared_ptr<void>, tracked_connection*>> connection_tracker::snapshot() const {
129 25665 : std::vector<std::pair<std::shared_ptr<void>, tracked_connection*>> connections;
130 436305 : for (const auto& shard : shards_) {
131 410640 : std::scoped_lock lock(shard.mutex);
132 415633 : for (auto* connection = shard.head; connection; connection = connection->next_) {
133 : // (one being destroyed waits for the lock to leave: it cannot be locked anymore)
134 4993 : if (auto owner = connection->self_.lock()) connections.emplace_back(std::move(owner), connection);
135 : }
136 410640 : }
137 25665 : return connections;
138 0 : }
139 :
140 3252 : void connection_tracker::drain_all() {
141 4161 : for (auto& [owner, connection] : snapshot()) {
142 909 : boost::asio::post(*connection->io_context_, [owner = std::move(owner), connection = connection] {
143 900 : connection->shutdown_drain();
144 900 : });
145 3252 : }
146 3252 : }
147 :
148 22413 : void connection_tracker::cancel_and_close_all(std::function<void()> done) {
149 22413 : auto connections = snapshot();
150 22413 : if (connections.empty()) {
151 20372 : if (done) done();
152 20372 : return;
153 : }
154 : struct pending_closes {
155 : std::atomic<size_t> left;
156 : std::function<void()> done;
157 : };
158 2041 : auto pending = std::make_shared<pending_closes>(connections.size(), std::move(done));
159 5823 : for (auto& [owner, connection] : connections) {
160 3782 : auto* io_context = connection->io_context_;
161 3782 : boost::asio::post(*io_context, [io_context, owner = std::move(owner), connection = connection, pending]() mutable {
162 : // In one call: the handlers interrupted resume once closed. Posted apart, the
163 : // first could run, and a handler resume and write, before the second was posted
164 602 : connection->shutdown_cancel();
165 602 : connection->shutdown_close();
166 : // Closed once what was posted before the close ran too: the handlers interrupted
167 : // resume (on a closed connection) before the server may stop
168 602 : boost::asio::post(*io_context, [owner = std::move(owner), pending] {
169 936 : if (pending->left.fetch_sub(1, std::memory_order_acq_rel) == 1 && pending->done) pending->done();
170 468 : });
171 602 : });
172 : }
173 22413 : }
174 :
175 : // taken_over_tracking
176 :
177 : namespace detail {
178 :
179 4306 : std::shared_ptr<taken_over_tracking> taken_over_tracking::track(std::shared_ptr<connection_tracker> tracker,
180 : std::shared_ptr<asio::socket> socket) {
181 4306 : auto& io_context = socket->get_io_context();
182 4306 : std::shared_ptr<taken_over_tracking> tracking(new taken_over_tracking(std::move(socket)));
183 4306 : tracking->tracked_connection::track(std::move(tracker), tracking, io_context, connection_type::taken_over);
184 4306 : return tracking;
185 0 : }
186 :
187 4306 : taken_over_tracking::~taken_over_tracking() = default;
188 :
189 4306 : upgraded_connection taken_over_tracking::hand_out(std::string buffered) {
190 : // The socket handed out shares the ownership of this tracking: the connection counts while
191 : // the new owner keeps it, and the socket never knows of the server
192 4306 : auto self = shared_from_this();
193 8612 : return upgraded_connection{std::shared_ptr<asio::socket>(self, socket_.get()), std::move(buffered),
194 8612 : server_scope(self)};
195 4306 : }
196 :
197 4306 : void taken_over_tracking::handed_over() {
198 4306 : auto tracker = this->tracker();
199 4306 : if (!tracker) return;
200 4214 : switch (tracker->current_phase()) {
201 4120 : case connection_tracker::phase::serving: break;
202 9 : case connection_tracker::phase::draining: shutdown_drain(); break;
203 85 : case connection_tracker::phase::closing: shutdown_close(); break;
204 : }
205 4306 : }
206 :
207 3702 : void taken_over_tracking::on_shutdown(sync_callback<void()> callback) {
208 : // Closed, or told already: never told (again)
209 3702 : if (!tracked() || notified_) return;
210 3702 : if (!draining_) {
211 3684 : on_shutdown_ = std::move(callback);
212 3684 : return;
213 : }
214 : // Set once the drain started (e.g. by a coroutine handler, after the takeover): told now,
215 : // posted, unless closed meanwhile. Not kept here, as it may own the socket
216 18 : notified_ = true;
217 18 : boost::asio::post(socket_->get_io_context(), [weak = weak_from_this(), callback = std::move(callback)] {
218 18 : auto self = weak.lock();
219 18 : if (self && self->tracked()) notify(callback);
220 18 : });
221 : }
222 :
223 3666 : void taken_over_tracking::closed() {
224 3666 : on_shutdown_ = nullptr;
225 3666 : untrack();
226 3666 : }
227 :
228 165 : void taken_over_tracking::notify(const sync_callback<void()>& callback) {
229 : try {
230 165 : callback();
231 0 : } catch (const std::exception& e) {
232 0 : LOG_ERROR("error notifying a connection of the server shutdown: {}", e.what());
233 0 : }
234 165 : }
235 :
236 228 : void taken_over_tracking::shutdown_drain() {
237 300 : if (draining_ || !tracked()) return;
238 219 : draining_ = true;
239 : // Called once: released before, as it may own the socket (which owns this). Without one
240 : // yet, the owner is told once it sets it (on_shutdown())
241 219 : auto callback = std::move(on_shutdown_);
242 219 : on_shutdown_ = nullptr;
243 219 : if (!callback) return;
244 147 : notified_ = true;
245 147 : notify(callback);
246 219 : }
247 :
248 194 : void taken_over_tracking::shutdown_close() {
249 194 : if (!tracked()) return;
250 185 : on_shutdown_ = nullptr;
251 185 : socket_->close();
252 185 : untrack();
253 : }
254 :
255 : }
256 :
257 : // server_scope
258 :
259 7230 : void server_scope::set_type(connection_type type) const {
260 7230 : if (auto tracking = tracking_.lock()) tracking->set_type(type);
261 7230 : }
262 :
263 3744 : void server_scope::on_shutdown(sync_callback<void()> callback) const {
264 3744 : if (auto tracking = tracking_.lock()) tracking->on_shutdown(std::move(callback));
265 3744 : }
266 :
267 3732 : void server_scope::closed() const {
268 3732 : if (auto tracking = tracking_.lock()) tracking->closed();
269 3732 : }
270 :
271 : }
|
