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 : : #include "wirestead/transport/tcp_server/tcp_server_session.hpp"
18 : :
19 : : #include <cstring>
20 : : #include <iostream>
21 : :
22 : : #include "wirestead/diagnostics/callback.hpp"
23 : : #include "wirestead/memory/memory_pool.hpp"
24 : : #include "wirestead/transport/base/bp_utils.hpp"
25 : : #include "wirestead/transport/base/stop_test_hook.hpp"
26 : : #include "wirestead/transport/tcp_server/boost_tcp_socket.hpp"
27 : :
28 : : namespace wirestead {
29 : : namespace transport {
30 : :
31 : 316 : TcpServerSession::TcpServerSession(net::io_context& ioc, tcp::socket sock, size_t backpressure_threshold,
32 : : int idle_timeout_ms, base::constants::BackpressureStrategy strategy,
33 : 316 : bool enable_memory_pool, size_t read_buffer_size)
34 : 316 : : ioc_(ioc),
35 : 316 : strand_(ioc.get_executor()),
36 : 316 : idle_timer_(ioc),
37 : 316 : socket_(std::make_unique<BoostTcpSocket>(std::move(sock))),
38 : 316 : enable_memory_pool_(enable_memory_pool),
39 : 316 : writing_(false),
40 : 316 : queue_bytes_(0),
41 : 316 : bp_strategy_(strategy),
42 : 316 : bp_high_(backpressure_threshold),
43 : 316 : idle_timeout_ms_(idle_timeout_ms),
44 : 316 : alive_(false),
45 : 1264 : cleanup_done_(false) {
46 : 316 : bp_limit_ = std::min(std::max(bp_high_ * 4, base::constants::DEFAULT_BACKPRESSURE_THRESHOLD),
47 : : base::constants::MAX_BUFFER_SIZE);
48 [ + - ]: 316 : bp_low_ = bp_high_ > 1 ? bp_high_ / 2 : bp_high_;
49 [ - + ]: 316 : if (bp_low_ == 0) bp_low_ = 1;
50 : 316 : rx_.resize(
51 : 316 : std::clamp(read_buffer_size, base::constants::MIN_READ_BUFFER_SIZE, base::constants::MAX_READ_BUFFER_SIZE));
52 : 316 : }
53 : :
54 : 206 : TcpServerSession::TcpServerSession(net::io_context& ioc, std::unique_ptr<interface::TcpSocketInterface> socket,
55 : : size_t backpressure_threshold, int idle_timeout_ms,
56 : : base::constants::BackpressureStrategy strategy, bool enable_memory_pool,
57 : 206 : size_t read_buffer_size)
58 : 206 : : ioc_(ioc),
59 : 206 : strand_(ioc.get_executor()),
60 : 206 : idle_timer_(ioc),
61 : 206 : socket_(std::move(socket)),
62 : 206 : enable_memory_pool_(enable_memory_pool),
63 : 206 : writing_(false),
64 : 206 : queue_bytes_(0),
65 : 206 : bp_strategy_(strategy),
66 : 206 : bp_high_(backpressure_threshold),
67 : 206 : idle_timeout_ms_(idle_timeout_ms),
68 : 206 : alive_(false),
69 : 824 : cleanup_done_(false) {
70 : 206 : bp_limit_ = std::min(std::max(bp_high_ * 4, base::constants::DEFAULT_BACKPRESSURE_THRESHOLD),
71 : : base::constants::MAX_BUFFER_SIZE);
72 [ + - ]: 206 : bp_low_ = bp_high_ > 1 ? bp_high_ / 2 : bp_high_;
73 [ - + ]: 206 : if (bp_low_ == 0) bp_low_ = 1;
74 : 206 : rx_.resize(
75 : 206 : std::clamp(read_buffer_size, base::constants::MIN_READ_BUFFER_SIZE, base::constants::MAX_READ_BUFFER_SIZE));
76 : 206 : }
77 : :
78 : 174 : void TcpServerSession::start() { start_with_notification({}); }
79 : :
80 : 522 : void TcpServerSession::start_with_notification(std::function<void()> notify) {
81 [ - + ]: 522 : if (alive_.exchange(true)) return;
82 : 522 : auto self = shared_from_this();
83 : 522 : net::post(strand_, [self, notify = std::move(notify)] {
84 [ + + - + : 522 : if (self->closing_ || !self->alive_) return;
+ + ]
85 : 490 : self->reset_idle_timer();
86 : : // No-op on a plain socket, the TLS handshake on an encrypted one. Reading
87 : : // before it completes would hand the session ciphertext, so the first read
88 : : // waits on it - and a failed handshake closes rather than reads.
89 : 490 : self->socket_->async_handshake([self](const boost::system::error_code& ec) {
90 : : // Plain sockets may complete inline. Queue read startup so connection
91 : : // notification finishes first, just as it does for an async TLS handshake.
92 : 490 : net::post(self->strand_, [self, ec] {
93 [ + + - + : 490 : if (self->closing_ || !self->alive_) return;
+ + ]
94 [ + + ]: 471 : if (ec) {
95 : 30 : WIRESTEAD_LOG_WARNING("tcp_server_session", "handshake", "Handshake failed: " + ec.message());
96 : 30 : self->do_close();
97 : 30 : return;
98 : : }
99 : 441 : self->reset_idle_timer();
100 : 441 : self->start_read();
101 : : });
102 : 490 : });
103 : 490 : diagnostics::invoke_callback("tcp_server_session", "on_connect", notify);
104 : : });
105 : 522 : }
106 : :
107 : 894 : bool TcpServerSession::async_write_copy(memory::ConstByteSpan data) {
108 : 894 : const auto result = write_copy(data);
109 [ + + + - ]: 894 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
110 : 1788 : return result.accepted();
111 : : }
112 : :
113 : 1115 : wrapper::SendResult TcpServerSession::write_copy(memory::ConstByteSpan data) {
114 : 1115 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
115 [ + + + - ]: 1115 : if (auto hook = detail::g_tcp_session_write_admission_hook.load()) hook();
116 [ + + + + : 1115 : if (!alive_ || closing_) {
+ + ]
117 : 30 : stats_.record_failed_send();
118 : 30 : return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
119 : : } // Don't queue writes if session is not alive
120 : :
121 : 1085 : size_t size = data.size();
122 [ + + ]: 1085 : if (size == 0) {
123 : 6 : stats_.record_failed_send();
124 : 6 : return wrapper::SendResult::reject(wrapper::SendRejection::InvalidArgument);
125 : : }
126 [ + + ]: 1079 : if (size > base::constants::MAX_BUFFER_SIZE) {
127 : 1 : WIRESTEAD_LOG_ERROR("tcp_server_session", "write", "Write size exceeds maximum allowed");
128 : 1 : stats_.record_failed_send();
129 : 1 : return wrapper::SendResult::reject(wrapper::SendRejection::TooLarge);
130 : : }
131 : :
132 : : // Use memory pool for better performance (only for reasonable sizes)
133 [ + + + + ]: 1078 : if (size <= base::constants::LARGE_BUFFER_THRESHOLD && enable_memory_pool_) { // Only use pool for buffers <= 64KB
134 : 577 : memory::PooledBuffer pooled_buffer(size, pool_);
135 [ + - + - ]: 577 : if (pooled_buffer.valid()) {
136 : : // Copy data to pooled buffer safely
137 : 577 : base::safe_memory::safe_memcpy(pooled_buffer.data(), data.data(), size);
138 [ + - + + ]: 577 : if (!queue_util::try_reserve_limit_bytes(write_reserve_mtx_, queue_bytes_, pending_bytes_, inflight_bytes_, size,
139 : : bp_limit_)) {
140 : 25 : stats_.record_failed_send();
141 : 25 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
142 : : }
143 : 552 : Ledger::Admission admission(send_accounting_, size);
144 : 552 : stats_.record_accepted(size);
145 : 552 : net::post(strand_,
146 : 552 : [self = shared_from_this(), buf = std::move(pooled_buffer), request = admission.request()]() mutable {
147 : 552 : const auto added = buf.size();
148 [ + + + + : 552 : if (!self->alive_ || self->closing_) { // Double-check in case session was closed
+ + ]
149 : 5 : queue_util::release_reserved_limit_bytes(self->write_reserve_mtx_, self->inflight_bytes_, added);
150 : 5 : self->stats_.record_failed_send();
151 : 5 : return;
152 : : }
153 : 547 : self->route_enqueued_buffer(TrackedBuffer{BufferVariant{std::move(buf)}, request}, added);
154 : : });
155 : 552 : admission.commit();
156 : 552 : return wrapper::SendResult::accept();
157 : 552 : }
158 : 577 : }
159 : :
160 : : // Fallback to regular allocation for large buffers or pool exhaustion
161 [ + - + + ]: 501 : if (!queue_util::try_reserve_limit_bytes(write_reserve_mtx_, queue_bytes_, pending_bytes_, inflight_bytes_, size,
162 : : bp_limit_)) {
163 : 54 : stats_.record_failed_send();
164 : 54 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
165 : : }
166 : 447 : std::vector<uint8_t> fallback(data.begin(), data.end());
167 : 447 : Ledger::Admission admission(send_accounting_, size);
168 : 447 : stats_.record_accepted(size);
169 : :
170 : 447 : net::post(strand_,
171 : 447 : [self = shared_from_this(), buf = std::move(fallback), size, request = admission.request()]() mutable {
172 [ + + + + : 447 : if (!self->alive_ || self->closing_) { // Double-check in case session was closed
+ + ]
173 : 5 : queue_util::release_reserved_limit_bytes(self->write_reserve_mtx_, self->inflight_bytes_, size);
174 : 5 : self->stats_.record_failed_send();
175 : 5 : return;
176 : : }
177 : 442 : self->route_enqueued_buffer(TrackedBuffer{BufferVariant{std::move(buf)}, request}, size);
178 : : });
179 : 447 : admission.commit();
180 : 447 : return wrapper::SendResult::accept();
181 : 1115 : }
182 : :
183 : 24531 : bool TcpServerSession::async_write_move(std::vector<uint8_t>&& data) {
184 : 24531 : const auto result = write_move(std::move(data));
185 [ + + + - ]: 24531 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
186 : 49062 : return result.accepted();
187 : : }
188 : :
189 : 24531 : wrapper::SendResult TcpServerSession::write_move(std::vector<uint8_t>&& data) {
190 : 24531 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
191 [ + + + - ]: 24531 : if (auto hook = detail::g_tcp_session_write_admission_hook.load()) hook();
192 [ + + + + : 24531 : if (!alive_ || closing_) {
+ + ]
193 : 14 : stats_.record_failed_send();
194 : 14 : return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
195 : : }
196 : 24517 : const auto added = data.size();
197 [ + + ]: 24517 : if (added == 0) {
198 : 5 : stats_.record_failed_send();
199 : 5 : return wrapper::SendResult::reject(wrapper::SendRejection::InvalidArgument);
200 : : }
201 [ + + ]: 24512 : if (added > base::constants::MAX_BUFFER_SIZE) {
202 : 1 : WIRESTEAD_LOG_ERROR("tcp_server_session", "write", "Write size exceeds maximum allowed");
203 : 1 : stats_.record_failed_send();
204 : 1 : return wrapper::SendResult::reject(wrapper::SendRejection::TooLarge);
205 : : }
206 [ + - + + ]: 24511 : if (!queue_util::try_reserve_limit_bytes(write_reserve_mtx_, queue_bytes_, pending_bytes_, inflight_bytes_, added,
207 : : bp_limit_)) {
208 : 13522 : stats_.record_failed_send();
209 : 13522 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
210 : : }
211 : 10989 : Ledger::Admission admission(send_accounting_, added);
212 : 10989 : stats_.record_accepted(added);
213 : 10989 : net::post(strand_,
214 : 10989 : [self = shared_from_this(), buf = std::move(data), added, request = admission.request()]() mutable {
215 [ + + + + : 10989 : if (!self->alive_ || self->closing_) {
+ + ]
216 : 19 : queue_util::release_reserved_limit_bytes(self->write_reserve_mtx_, self->inflight_bytes_, added);
217 : 19 : self->stats_.record_failed_send();
218 : 19 : return;
219 : : }
220 : 10970 : self->route_enqueued_buffer(TrackedBuffer{BufferVariant{std::move(buf)}, request}, added);
221 : : });
222 : 10989 : admission.commit();
223 : 10989 : return wrapper::SendResult::accept();
224 : 24531 : }
225 : :
226 : 458 : bool TcpServerSession::async_write_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
227 : 458 : const auto result = write_shared(std::move(data));
228 [ + + + - ]: 458 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
229 : 916 : return result.accepted();
230 : : }
231 : :
232 : 458 : wrapper::SendResult TcpServerSession::write_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
233 : 458 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
234 [ + + + - ]: 458 : if (auto hook = detail::g_tcp_session_write_admission_hook.load()) hook();
235 [ + + + + : 458 : if (!alive_ || closing_) {
+ + ]
236 : 8 : stats_.record_failed_send();
237 : 8 : return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
238 : : }
239 [ + + + + : 450 : if (!data || data->empty()) {
+ + ]
240 : 6 : stats_.record_failed_send();
241 : 6 : return wrapper::SendResult::reject(wrapper::SendRejection::InvalidArgument);
242 : : }
243 : 444 : const auto added = data->size();
244 [ + + ]: 444 : if (added > base::constants::MAX_BUFFER_SIZE) {
245 : 1 : WIRESTEAD_LOG_ERROR("tcp_server_session", "write", "Write size exceeds maximum allowed");
246 : 1 : stats_.record_failed_send();
247 : 1 : return wrapper::SendResult::reject(wrapper::SendRejection::TooLarge);
248 : : }
249 [ + - + + ]: 443 : if (!queue_util::try_reserve_limit_bytes(write_reserve_mtx_, queue_bytes_, pending_bytes_, inflight_bytes_, added,
250 : : bp_limit_)) {
251 : 6 : stats_.record_failed_send();
252 : 6 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
253 : : }
254 : 437 : Ledger::Admission admission(send_accounting_, added);
255 : 437 : stats_.record_accepted(added);
256 : 437 : net::post(strand_,
257 : 437 : [self = shared_from_this(), buf = std::move(data), added, request = admission.request()]() mutable {
258 [ + + + + : 437 : if (!self->alive_ || self->closing_) {
+ + ]
259 : 7 : queue_util::release_reserved_limit_bytes(self->write_reserve_mtx_, self->inflight_bytes_, added);
260 : 7 : self->stats_.record_failed_send();
261 : 7 : return;
262 : : }
263 : 430 : self->route_enqueued_buffer(TrackedBuffer{BufferVariant{std::move(buf)}, request}, added);
264 : : });
265 : 437 : admission.commit();
266 : 437 : return wrapper::SendResult::accept();
267 : 458 : }
268 : :
269 : 457 : bool TcpServerSession::async_try_write_copy(memory::ConstByteSpan data) {
270 : 457 : const auto result = try_write_copy(data);
271 [ + + + - ]: 457 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
272 : 914 : return result.accepted();
273 : : }
274 : :
275 : 515 : wrapper::SendResult TcpServerSession::try_write_copy(memory::ConstByteSpan data) {
276 [ + + + + : 515 : if (data.empty() || data.size() > base::constants::MAX_BUFFER_SIZE) {
+ + ]
277 : 6 : stats_.record_failed_send();
278 [ + + ]: 6 : return wrapper::SendResult::reject(data.empty() ? wrapper::SendRejection::InvalidArgument
279 : 6 : : wrapper::SendRejection::TooLarge);
280 : : }
281 : 1018 : return try_write_move(std::vector<uint8_t>(data.begin(), data.end()));
282 : : }
283 : :
284 : 471 : bool TcpServerSession::async_try_write_move(std::vector<uint8_t>&& data) {
285 : 471 : const auto result = try_write_move(std::move(data));
286 [ + + + - ]: 471 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
287 : 942 : return result.accepted();
288 : : }
289 : :
290 : 980 : wrapper::SendResult TcpServerSession::try_write_move(std::vector<uint8_t>&& data) {
291 : 980 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
292 [ + + + - ]: 980 : if (auto hook = detail::g_tcp_session_write_admission_hook.load()) hook();
293 [ + + + + : 980 : if (!alive_ || closing_) {
+ + ]
294 : 16 : stats_.record_failed_send();
295 : 16 : return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
296 : : }
297 : 964 : const auto added = data.size();
298 [ + + + + ]: 964 : if (added == 0 || added > base::constants::MAX_BUFFER_SIZE) {
299 : 6 : stats_.record_failed_send();
300 [ + + ]: 6 : return wrapper::SendResult::reject(added == 0 ? wrapper::SendRejection::InvalidArgument
301 : 6 : : wrapper::SendRejection::TooLarge);
302 : : }
303 : 22 : const auto reject_for_pressure = [this, added]() {
304 [ + + ]: 22 : if (bp_strategy_ == base::constants::BackpressureStrategy::BestEffort) {
305 : 14 : stats_.record_dropped(1, added);
306 : : } else {
307 : 8 : stats_.record_failed_send();
308 : : }
309 : 980 : };
310 [ + + + + : 1894 : if (backpressure_active_.load() || queue_bytes_ + added > bp_high_ ||
+ + ]
311 [ - + ]: 936 : queue_bytes_ + pending_bytes_ + added > bp_limit_) {
312 : 22 : reject_for_pressure();
313 : 22 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
314 : : }
315 [ - + ]: 936 : if (!queue_util::try_reserve_write_bytes(write_reserve_mtx_, inflight_bytes_, queue_bytes_, pending_bytes_,
316 : 936 : backpressure_active_, added, bp_high_, bp_limit_)) {
317 : 0 : reject_for_pressure();
318 : 0 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
319 : : }
320 : 936 : Ledger::Admission admission(send_accounting_, added);
321 : 936 : stats_.record_accepted(added);
322 : :
323 : 936 : net::post(strand_,
324 : 936 : [self = shared_from_this(), buf = std::move(data), added, request = admission.request()]() mutable {
325 [ + + + + : 936 : if (!self->alive_ || self->closing_) {
+ + ]
326 : 26 : queue_util::release_reserved_write_bytes(self->queue_bytes_, added);
327 : 26 : self->stats_.record_failed_send();
328 : 26 : return;
329 : : }
330 : :
331 : 910 : self->tx_.push_back(TrackedBuffer{BufferVariant{std::move(buf)}, request});
332 : 910 : self->observe_queue();
333 : 910 : self->report_backpressure(self->queue_bytes_);
334 [ + + ]: 910 : if (!self->writing_) self->do_write();
335 : : });
336 : 936 : admission.commit();
337 : 936 : return wrapper::SendResult::accept();
338 : 980 : }
339 : :
340 : 459 : bool TcpServerSession::async_try_write_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
341 : 459 : const auto result = try_write_shared(std::move(data));
342 [ + + + - ]: 459 : if (auto hook = detail::g_tcp_session_write_result_hook.load()) hook(result);
343 : 918 : return result.accepted();
344 : : }
345 : :
346 : 1824 : wrapper::SendResult TcpServerSession::try_write_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
347 : 1824 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
348 [ + + + - ]: 1824 : if (auto hook = detail::g_tcp_session_write_admission_hook.load()) hook();
349 [ + + + + : 1824 : if (!alive_ || closing_) {
+ + ]
350 : 10 : stats_.record_failed_send();
351 : 10 : return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
352 : : }
353 [ + - + + : 1814 : if (!data || data->empty()) {
+ + ]
354 : 5 : stats_.record_failed_send();
355 : 5 : return wrapper::SendResult::reject(wrapper::SendRejection::InvalidArgument);
356 : : }
357 : 1809 : const auto added = data->size();
358 [ + + ]: 1809 : if (added > base::constants::MAX_BUFFER_SIZE) {
359 : 1 : stats_.record_failed_send();
360 : 1 : return wrapper::SendResult::reject(wrapper::SendRejection::TooLarge);
361 : : }
362 : 229 : const auto reject_for_pressure = [this, added]() {
363 [ + + ]: 229 : if (bp_strategy_ == base::constants::BackpressureStrategy::BestEffort) {
364 : 7 : stats_.record_dropped(1, added);
365 : : } else {
366 : 222 : stats_.record_failed_send();
367 : : }
368 : 2037 : };
369 [ + + + + : 3387 : if (backpressure_active_.load() || queue_bytes_ + added > bp_high_ ||
+ + ]
370 [ - + ]: 1579 : queue_bytes_ + pending_bytes_ + added > bp_limit_) {
371 : 229 : reject_for_pressure();
372 : 229 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
373 : : }
374 [ - + ]: 1579 : if (!queue_util::try_reserve_write_bytes(write_reserve_mtx_, inflight_bytes_, queue_bytes_, pending_bytes_,
375 : 1579 : backpressure_active_, added, bp_high_, bp_limit_)) {
376 : 0 : reject_for_pressure();
377 : 0 : return wrapper::SendResult::reject(wrapper::SendRejection::WouldBlock);
378 : : }
379 : 1579 : Ledger::Admission admission(send_accounting_, added);
380 : 1579 : stats_.record_accepted(added);
381 : :
382 : 1579 : net::post(strand_,
383 : 1579 : [self = shared_from_this(), buf = std::move(data), added, request = admission.request()]() mutable {
384 [ + + + + : 1579 : if (!self->alive_ || self->closing_) {
+ + ]
385 : 17 : queue_util::release_reserved_write_bytes(self->queue_bytes_, added);
386 : 17 : self->stats_.record_failed_send();
387 : 17 : return;
388 : : }
389 : :
390 : 1562 : self->tx_.push_back(TrackedBuffer{BufferVariant{std::move(buf)}, request});
391 : 1562 : self->observe_queue();
392 : 1562 : self->report_backpressure(self->queue_bytes_);
393 [ + + ]: 1562 : if (!self->writing_) self->do_write();
394 : : });
395 : 1579 : admission.commit();
396 : 1579 : return wrapper::SendResult::accept();
397 : 1824 : }
398 : :
399 : 353 : void TcpServerSession::on_bytes(OnBytes cb) {
400 : 353 : auto self = shared_from_this();
401 : 353 : net::dispatch(strand_, [self, cb = std::move(cb)]() mutable {
402 [ + - - + : 353 : if (self->closing_.load() || self->cleanup_done_.load()) return;
- + ]
403 : 353 : self->on_bytes_ = std::move(cb);
404 : : });
405 : 353 : }
406 : 350 : void TcpServerSession::on_backpressure(OnBackpressure cb) {
407 : 350 : auto self = shared_from_this();
408 : 350 : net::dispatch(strand_, [self, cb = std::move(cb)]() mutable {
409 [ + - - + : 350 : if (self->closing_.load() || self->cleanup_done_.load()) return;
- + ]
410 : 350 : self->on_bp_ = std::move(cb);
411 : : });
412 : 350 : }
413 : 355 : void TcpServerSession::on_close(OnClose cb) {
414 : 355 : auto self = shared_from_this();
415 : 355 : net::dispatch(strand_, [self, cb = std::move(cb)]() mutable {
416 [ + - - + : 355 : if (self->closing_.load() || self->cleanup_done_.load()) return;
- + ]
417 : 355 : self->on_close_ = std::move(cb);
418 : : });
419 : 355 : }
420 : :
421 : 298 : std::optional<wrapper::SendResult> TcpServerSession::poll_write_wait() const {
422 : 298 : std::lock_guard<std::mutex> lock(submission_mtx_);
423 [ + + ]: 298 : if (wait_ended_by_) return wrapper::SendResult::reject(*wait_ended_by_);
424 [ + - - + : 256 : if (!alive_ || closing_) return wrapper::SendResult::reject(wrapper::SendRejection::NotReady);
- + ]
425 [ + + ]: 256 : if (!backpressure_active_) return wrapper::SendResult::accept();
426 : 156 : return std::nullopt;
427 : 298 : }
428 : :
429 : 136 : void TcpServerSession::request_stop() {
430 : 136 : std::lock_guard<std::mutex> lock(submission_mtx_);
431 : 136 : send_accounting_.end(Ledger::Cause::ExplicitStop);
432 : 136 : closing_ = true;
433 [ + + ]: 136 : if (!wait_ended_by_) wait_ended_by_ = wrapper::SendRejection::CancelledWhileWaiting;
434 : 136 : }
435 : :
436 : 130 : void TcpServerSession::cancel_write_wait() {
437 : 130 : std::lock_guard<std::mutex> lock(submission_mtx_);
438 [ + + ]: 130 : if (!wait_ended_by_) wait_ended_by_ = wrapper::SendRejection::CancelledWhileWaiting;
439 : 130 : }
440 : :
441 : 2299 : bool TcpServerSession::alive() const { return alive_.load(); }
442 : :
443 : 1120 : wrapper::RuntimeStats TcpServerSession::stats() const {
444 : : auto snapshot =
445 : 3360 : stats_.snapshot(queue_bytes_.load(std::memory_order_relaxed), pending_bytes_.load(std::memory_order_relaxed),
446 : 1120 : backpressure_active_.load(std::memory_order_relaxed));
447 : 1120 : snapshot.send_accounting = send_accounting_.snapshot();
448 : 1120 : return snapshot;
449 : : }
450 : :
451 : 26 : void TcpServerSession::reset_stats() {
452 : 26 : std::lock_guard<std::mutex> lock(submission_mtx_);
453 : 26 : send_accounting_.reset();
454 : 78 : stats_.reset(queue_bytes_.load(std::memory_order_relaxed) + pending_bytes_.load(std::memory_order_relaxed));
455 : 26 : }
456 : :
457 : 229 : void TcpServerSession::stop() { async_stop({}); }
458 : :
459 : 368 : void TcpServerSession::async_stop(std::function<void()> completion) {
460 : 368 : std::lock_guard<std::mutex> admission_lock(submission_mtx_);
461 : 368 : send_accounting_.end(Ledger::Cause::ExplicitStop);
462 : 368 : closing_.store(true);
463 [ + + ]: 368 : if (!wait_ended_by_) wait_ended_by_ = wrapper::SendRejection::NotReady;
464 : 368 : auto self = shared_from_this();
465 : 368 : net::post(strand_, [self, completion = std::move(completion)] {
466 : 358 : self->on_bytes_ = nullptr;
467 : 358 : self->on_bp_ = nullptr;
468 : 358 : self->on_close_ = nullptr;
469 : 358 : self->idle_timer_.cancel();
470 : 358 : self->do_close();
471 [ + + ]: 358 : if (completion) completion();
472 : 358 : });
473 : 368 : }
474 : :
475 : 2 : void TcpServerSession::cancel() {
476 : 2 : auto self = shared_from_this();
477 : 2 : net::dispatch(strand_, [self] {
478 : 2 : self->idle_timer_.cancel();
479 : 2 : boost::system::error_code ec;
480 : : // Cancelling the socket via close() causes ongoing operations to complete with operation_aborted.
481 : : // Unlike stop(), this does NOT set closing_ flag immediately, allowing the
482 : : // error handler to run normally and trigger do_close() via the error path.
483 [ + - ]: 2 : if (self->socket_) {
484 : 2 : self->socket_->close(ec);
485 : : }
486 : 2 : });
487 : 2 : }
488 : :
489 : 3354 : void TcpServerSession::start_read() {
490 [ + + - + : 3354 : if (closing_ || !alive_) return;
+ + ]
491 : 3346 : auto self = shared_from_this();
492 : 6692 : socket_->async_read_some(net::buffer(rx_.data(), rx_.size()),
493 : 6692 : [self](const boost::system::error_code& ec, std::size_t n) {
494 : 3336 : net::dispatch(self->strand_, [self, ec, n] {
495 [ + + - + : 3336 : if (self->closing_ || !self->alive_) return;
+ + ]
496 [ + + ]: 3127 : if (ec) {
497 : 214 : self->do_close();
498 : 214 : return;
499 : : }
500 : 2913 : self->reset_idle_timer();
501 [ + - ]: 2913 : if (n > 0) self->stats_.record_received(n);
502 : 2913 : diagnostics::invoke_callback("tcp_server_session", "on_bytes", self->on_bytes_,
503 : 2913 : memory::ConstByteSpan(self->rx_.data(), n));
504 : 2913 : self->start_read();
505 : : });
506 : 3336 : });
507 : 3346 : }
508 : :
509 : 950 : void TcpServerSession::do_write() {
510 : : // Serialize the actual handoff with stop, not only the enqueue handler.
511 : 950 : std::unique_lock<std::mutex> lock(submission_mtx_);
512 [ + + + - : 950 : if (closing_ || !alive_ || tx_.empty() || writing_) return;
+ + - + +
+ ]
513 : 730 : writing_ = true;
514 : 730 : const size_t bytes_to_write = queue_util::take_gather_batch(tx_, current_write_batch_, current_write_views_, payload);
515 : : // Admission still fences stop/loss; amortize the ledger lock over the batch.
516 : 9298 : send_accounting_.begin_batch(current_write_batch_, [](const auto& item) { return item.request; });
517 : 730 : auto self = shared_from_this();
518 : : // Initiation may write inline; never hold admission across the syscall.
519 : : // The batch members are strand-owned and the completion posts back to it.
520 : 730 : lock.unlock();
521 : : try {
522 : 751 : socket_->async_write(current_write_views_, [self, bytes_to_write](const boost::system::error_code& ec, size_t n) {
523 : : // The interface erases associated executors. Post explicitly, including
524 : : // for endpoints that complete inline during initiation.
525 : 709 : net::post(self->strand_, [self, bytes_to_write, ec, n] {
526 [ + + + + ]: 709 : const bool failed = ec || n != bytes_to_write;
527 : : {
528 : 709 : std::lock_guard<std::mutex> completion_lock(self->submission_mtx_);
529 [ + + - + : 709 : if (self->closing_ || !self->alive_) {
+ + ]
530 : 96 : self->current_write_batch_.clear();
531 : 96 : return;
532 : : }
533 : 613 : self->send_accounting_.complete_batch(
534 : 4778 : self->current_write_batch_, std::min(n, bytes_to_write), [](const auto& item) { return item.request; },
535 : 4165 : [](const auto& item) {
536 : 8330 : return std::visit([](const auto& b) { return queue_util::variant_buffer_size(b); }, item.buffer);
537 : : });
538 : 613 : self->current_write_batch_.clear();
539 [ + + ]: 613 : if (failed) {
540 : 28 : self->send_accounting_.end(Ledger::Cause::ConnectionLoss);
541 : 28 : self->closing_ = true;
542 [ + - ]: 28 : if (!self->wait_ended_by_) self->wait_ended_by_ = wrapper::SendRejection::NotReady;
543 : : }
544 : 709 : }
545 [ + + ]: 613 : if (failed) {
546 : 28 : self->do_close();
547 : 28 : return;
548 : : }
549 : 585 : queue_util::release_reserved_write_bytes(self->queue_bytes_, bytes_to_write);
550 : 585 : self->stats_.record_sent(n);
551 : : // Keep writing_ set while callbacks move pending data or enqueue writes.
552 : 585 : self->report_backpressure(self->queue_bytes_);
553 : 585 : self->writing_ = false;
554 [ + - - + : 585 : if (self->closing_ || !self->alive_) return;
- + ]
555 : 585 : self->reset_idle_timer();
556 : 585 : self->do_write();
557 : : });
558 : 709 : });
559 : 7 : } catch (...) {
560 : 7 : lock.lock();
561 : 7 : send_accounting_.end(Ledger::Cause::ConnectionLoss);
562 : 7 : closing_ = true;
563 [ + - ]: 7 : if (!wait_ended_by_) wait_ended_by_ = wrapper::SendRejection::NotReady;
564 : 7 : lock.unlock();
565 : 7 : do_close();
566 : 7 : }
567 : 950 : }
568 : :
569 : 641 : void TcpServerSession::do_close() {
570 [ + + ]: 641 : if (cleanup_done_.exchange(true)) return; // Ensures cleanup runs only once
571 : :
572 : : {
573 : 512 : std::lock_guard<std::mutex> lock(submission_mtx_);
574 : 512 : send_accounting_.end(Ledger::Cause::ConnectionLoss);
575 [ + + ]: 512 : if (!wait_ended_by_) wait_ended_by_ = wrapper::SendRejection::NotReady;
576 : 512 : alive_.store(false);
577 : 512 : closing_.store(true);
578 : 512 : }
579 : :
580 : : // Safely invoke on_close callback
581 : 512 : auto close_cb = std::move(on_close_);
582 : :
583 : 512 : WIRESTEAD_LOG_INFO("tcp_server_session", "disconnect", "Client disconnected");
584 : 512 : boost::system::error_code ec;
585 : 512 : socket_->shutdown(tcp::socket::shutdown_both, ec);
586 : 512 : socket_->close(ec);
587 : :
588 : : // Drain queued/pending writes and unconditionally clear backpressure,
589 : : // notifying any waiter directly - shares UdpChannel's terminal-drain
590 : : // helper (#434). Must run before on_bp_ is cleared below: otherwise a
591 : : // Reliable-mode caller blocked in send_to_blocking() for this client
592 : : // would never be woken up when the client disconnects via a read error
593 : : // or idle timeout (jwsung91/wirestead#452).
594 : : {
595 : 512 : auto f = bp_fields();
596 : 512 : queue_util::drain_and_clear_backpressure(f, on_bp_, [&]() {
597 : 512 : tx_.clear();
598 : 512 : queue_bytes_ = 0;
599 : 512 : pending_.clear();
600 : 512 : pending_bytes_ = 0;
601 : 512 : });
602 : : }
603 : :
604 : : // Clear all callbacks to prevent any further invocations
605 : 512 : on_bytes_ = nullptr;
606 : 512 : on_bp_ = nullptr;
607 : 512 : on_close_ = nullptr;
608 : 512 : idle_timer_.cancel();
609 : :
610 [ + + ]: 512 : if (close_cb) {
611 : : try {
612 : 232 : close_cb();
613 : 0 : } catch (const std::exception& e) {
614 : 0 : WIRESTEAD_LOG_ERROR("tcp_server_session", "on_close", "Exception in on_close callback: " + std::string(e.what()));
615 : 0 : } catch (...) {
616 : 0 : WIRESTEAD_LOG_ERROR("tcp_server_session", "on_close", "Unknown exception in on_close callback");
617 : 0 : }
618 : : }
619 : 512 : }
620 : :
621 : 17844 : queue_util::BackpressureFields TcpServerSession::bp_fields() {
622 : 17844 : return queue_util::BackpressureFields{queue_bytes_, pending_bytes_, backpressure_active_, bp_high_,
623 : 17844 : bp_low_, bp_limit_, bp_strategy_, &write_reserve_mtx_};
624 : : }
625 : :
626 : 12389 : void TcpServerSession::route_enqueued_buffer(TrackedBuffer&& buf, size_t added) {
627 : 12389 : auto f = bp_fields();
628 : 12389 : queue_util::DropAccounting dropped;
629 : 12389 : auto decision = queue_util::decide_enqueue(f, added, tx_, dropped, payload, [this](const TrackedBuffer& item) {
630 : : send_accounting_.discard(item.request, Ledger::Cause::QueuePressure);
631 : : });
632 [ - + - - ]: 12389 : if (dropped.any()) stats_.record_dropped(dropped.messages, dropped.bytes);
633 : :
634 [ - + ]: 12389 : if (decision == queue_util::EnqueueDecision::Rejected) {
635 : 0 : send_accounting_.discard(buf.request, Ledger::Cause::QueuePressure);
636 : 0 : WIRESTEAD_LOG_ERROR("tcp_server_session", "write", "Queue limit exceeded, dropping message");
637 : : // #448: record as dropped so it's reflected in RuntimeStats instead of
638 : : // silently vanishing after being counted as accepted.
639 : 0 : stats_.record_dropped(1, added);
640 : 0 : queue_util::release_reserved_limit_bytes(write_reserve_mtx_, inflight_bytes_, added);
641 : 0 : report_backpressure(queue_bytes_ + added);
642 : 10502 : return;
643 : : }
644 [ + + ]: 12389 : if (decision == queue_util::EnqueueDecision::Pending) {
645 : 10502 : queue_util::commit_reserved_limit_bytes(write_reserve_mtx_, pending_bytes_, inflight_bytes_, added);
646 : 10502 : pending_.emplace_back(std::move(buf));
647 : 10502 : observe_queue();
648 : 10502 : return;
649 : : }
650 : 1887 : queue_util::commit_reserved_limit_bytes(write_reserve_mtx_, queue_bytes_, inflight_bytes_, added);
651 : 1887 : tx_.emplace_back(std::move(buf));
652 : 1887 : observe_queue();
653 : 1887 : report_backpressure(queue_bytes_);
654 [ + + + - ]: 1887 : if (!writing_) do_write();
655 : : }
656 : :
657 : 19902 : void TcpServerSession::observe_queue() {
658 : 59706 : stats_.observe_queue(queue_bytes_.load(std::memory_order_relaxed) + pending_bytes_.load(std::memory_order_relaxed));
659 : 19902 : }
660 : :
661 : 4944 : void TcpServerSession::report_backpressure(size_t queued_bytes) {
662 [ + + - + : 4944 : if (closing_ || !alive_) return;
+ + ]
663 : 4943 : observe_queue();
664 : 4943 : auto f = bp_fields();
665 : 4943 : queue_util::report_backpressure(
666 : 4943 : f, queued_bytes, on_bp_, stats_,
667 : 0 : [&]() -> size_t {
668 : 98 : const size_t moved = pending_bytes_.exchange(0);
669 [ + + ]: 102 : while (!pending_.empty()) {
670 : 4 : tx_.emplace_back(std::move(pending_.front()));
671 : 4 : pending_.pop_front();
672 : : }
673 : 98 : return moved;
674 : : },
675 : 4943 : [&]() {
676 : 98 : observe_queue();
677 [ - + ]: 98 : if (!writing_) do_write();
678 : 98 : });
679 : : }
680 : :
681 : 4429 : void TcpServerSession::reset_idle_timer() {
682 [ + + ]: 4429 : if (idle_timeout_ms_ <= 0) return;
683 : :
684 : : // Cancel any existing timer
685 : 3 : idle_timer_.cancel();
686 : :
687 : : // Reset timer
688 : 3 : idle_timer_.expires_after(std::chrono::milliseconds(idle_timeout_ms_));
689 : :
690 : 3 : auto self = shared_from_this();
691 : 3 : idle_timer_.async_wait(net::bind_executor(strand_, [self](const boost::system::error_code& ec) {
692 [ + + ]: 3 : if (ec == boost::asio::error::operation_aborted) return;
693 [ + - - + : 1 : if (!self->alive_ || self->closing_) return;
- + ]
694 : :
695 : 1 : WIRESTEAD_LOG_WARNING("tcp_server_session", "timeout", "Connection idle timeout expired, closing session");
696 : 1 : self->do_close();
697 : : }));
698 : 3 : }
699 : :
700 : : } // namespace transport
701 : : } // namespace wirestead
|
