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/wrapper/tcp_client/tcp_client.hpp"
18 : :
19 : : #include <atomic>
20 : : #include <boost/asio/executor_work_guard.hpp>
21 : : #include <boost/asio/io_context.hpp>
22 : : #include <boost/asio/steady_timer.hpp>
23 : : #include <chrono>
24 : : #include <mutex>
25 : : #include <shared_mutex>
26 : : #include <stdexcept>
27 : : #include <stop_token>
28 : : #include <thread>
29 : : #include <vector>
30 : :
31 : : #include "wirestead/base/common.hpp"
32 : : #include "wirestead/base/constants.hpp"
33 : : #include "wirestead/concurrency/io_thread_hook.hpp"
34 : : #include "wirestead/config/tcp_client_config.hpp"
35 : : #include "wirestead/diagnostics/error_mapping.hpp"
36 : : #include "wirestead/factory/channel_factory.hpp"
37 : : #include "wirestead/transport/tcp_client/tcp_client.hpp"
38 : : #include "wirestead/wrapper/callback_guard.hpp"
39 : : #include "wirestead/wrapper/error_context_builder.hpp"
40 : :
41 : : namespace wirestead {
42 : : namespace wrapper {
43 : :
44 : : struct TcpClient::Impl : public std::enable_shared_from_this<Impl> {
45 : : mutable std::shared_mutex mutex_;
46 : : std::mutex bp_mutex_;
47 : : std::condition_variable bp_cv_;
48 : : std::string host_;
49 : : uint16_t port_;
50 : : std::shared_ptr<interface::Channel> channel_;
51 : : std::shared_ptr<boost::asio::io_context> external_ioc_;
52 : : std::atomic<bool> use_external_context_{false};
53 : : std::atomic<bool> manage_external_context_{false};
54 : : std::jthread external_thread_;
55 : : std::unique_ptr<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> work_guard_;
56 : :
57 : : std::vector<std::promise<bool>> pending_promises_;
58 : : std::atomic<bool> started_{false};
59 : 128 : std::shared_ptr<bool> alive_marker_{std::make_shared<bool>(true)};
60 : :
61 : : // Shared snapshots: the io thread copies one out per received chunk, and a
62 : : // std::function copy allocates whenever the user handler outgrows its
63 : : // small-object buffer. See interface::SharedCallback.
64 : : interface::SharedCallback<MessageHandler> data_handler_;
65 : : interface::SharedCallback<BatchMessageHandler> data_batch_handler_;
66 : : ConnectionHandler connect_handler_{nullptr};
67 : : ConnectionHandler disconnect_handler_{nullptr};
68 : : ErrorHandler error_handler_{nullptr};
69 : : std::function<void(size_t)> bp_handler_{nullptr};
70 : : interface::SharedCallback<MessageHandler> message_handler_;
71 : : interface::SharedCallback<BatchMessageHandler> message_batch_handler_;
72 : :
73 : : std::shared_ptr<framer::IFramer> framer_{nullptr};
74 : :
75 : : // Batching logic
76 : : std::vector<MessageContext> data_batch_queue_;
77 : : std::vector<MessageContext> message_batch_queue_;
78 : : std::unique_ptr<boost::asio::steady_timer> batch_timer_;
79 : : size_t max_batch_size_ = 100;
80 : 128 : std::chrono::milliseconds max_batch_latency_{1};
81 : :
82 : : std::atomic<bool> auto_start_ = false;
83 : : std::chrono::milliseconds retry_interval_{base::constants::DEFAULT_RETRY_INTERVAL_MS};
84 : : int max_retries_ = base::constants::DEFAULT_MAX_RETRIES;
85 : : std::chrono::milliseconds connection_timeout_{base::constants::DEFAULT_CONNECTION_TIMEOUT_MS};
86 : : bool tls_enabled_{false};
87 : : std::string tls_ca_file_;
88 : : std::chrono::milliseconds idle_timeout_{base::constants::DEFAULT_IDLE_TIMEOUT_MS};
89 : : IdleTimeoutAction idle_timeout_action_{IdleTimeoutAction::Reconnect};
90 : : size_t backpressure_threshold_{base::constants::DEFAULT_BACKPRESSURE_THRESHOLD};
91 : : // Atomic rather than mutex-guarded: read from the send()/send_line() fast
92 : : // path on arbitrary caller threads while the setter can be called
93 : : // concurrently from any other thread (#436).
94 : : std::atomic<base::constants::BackpressureStrategy> backpressure_strategy_{
95 : : base::constants::BackpressureStrategy::Reliable};
96 : : bool tcp_no_delay_ = base::constants::DEFAULT_TCP_NO_DELAY;
97 : : bool keep_alive_ = base::constants::DEFAULT_KEEP_ALIVE;
98 : : size_t send_buffer_size_ = 0;
99 : : size_t receive_buffer_size_ = 0;
100 : : size_t read_buffer_size_ = base::constants::DEFAULT_READ_BUFFER_SIZE;
101 : :
102 : 101 : Impl(const std::string& host, uint16_t port) : host_(host), port_(port), started_(false) {}
103 : :
104 : 8 : Impl(const std::string& host, uint16_t port, std::shared_ptr<boost::asio::io_context> external_ioc)
105 : 8 : : host_(host),
106 : 8 : port_(port),
107 : 8 : external_ioc_(std::move(external_ioc)),
108 : 8 : use_external_context_(external_ioc_ != nullptr),
109 : 8 : manage_external_context_(false),
110 : 24 : started_(false) {}
111 : :
112 : 19 : explicit Impl(std::shared_ptr<interface::Channel> channel)
113 : 57 : : host_(""), port_(0), channel_(std::move(channel)), started_(false) {
114 : : // #450: setup_internal_handlers() captures weak_from_this() - calling it
115 : : // from inside this constructor would capture an empty weak_ptr, since
116 : : // enable_shared_from_this isn't wired up until make_shared() finishes
117 : : // constructing the object. Deferred to TcpClient's own constructor,
118 : : // which runs after impl_ is a fully-formed shared_ptr<Impl>.
119 : 19 : }
120 : :
121 : 128 : ~Impl() {
122 : : try {
123 : 128 : stop();
124 : 0 : } catch (...) {
125 : 0 : }
126 : 128 : }
127 : :
128 : 327 : void fulfill_all_locked(bool value) {
129 [ + + ]: 426 : for (auto& p : pending_promises_) {
130 : : try {
131 : 99 : p.set_value(value);
132 : 0 : } catch (...) {
133 : 0 : }
134 : : }
135 : 327 : pending_promises_.clear();
136 : 327 : }
137 : :
138 : 1 : void flush_batches() {
139 : 1 : std::unique_lock<std::shared_mutex> lock(mutex_);
140 [ + - ]: 1 : if (!data_batch_queue_.empty()) {
141 : 1 : auto handler = data_batch_handler_;
142 : 1 : auto batch = std::move(data_batch_queue_);
143 : 1 : data_batch_queue_.clear();
144 [ + - ]: 1 : if (handler) {
145 : 1 : lock.unlock();
146 : 1 : detail::invoke_user_callback("tcp_client", "on_data_batch", handler, batch);
147 : 1 : lock.lock();
148 : : }
149 : 1 : }
150 [ + - ]: 1 : if (!message_batch_queue_.empty()) {
151 : 1 : auto handler = message_batch_handler_;
152 : 1 : auto batch = std::move(message_batch_queue_);
153 : 1 : message_batch_queue_.clear();
154 [ + - ]: 1 : if (handler) {
155 : 1 : lock.unlock();
156 : 1 : detail::invoke_user_callback("tcp_client", "on_message_batch", handler, batch);
157 : 1 : lock.lock();
158 : : }
159 : 1 : }
160 [ + - ]: 1 : if (batch_timer_) {
161 : 1 : batch_timer_->cancel();
162 : : }
163 : 1 : }
164 : :
165 : 5 : void schedule_batch_timer() {
166 [ - + ]: 5 : if (!batch_timer_) return;
167 : 5 : batch_timer_->expires_after(max_batch_latency_);
168 : 10 : batch_timer_->async_wait([this, weak_impl = weak_from_this(),
169 : 5 : weak_alive = std::weak_ptr<bool>(alive_marker_)](const boost::system::error_code& ec) {
170 [ + + ]: 2 : if (ec) return;
171 : : // #450: keep Impl alive for the duration of this callback - on an
172 : : // externally-owned io_context, stop() doesn't join/wait for
173 : : // in-flight handlers, so a bare `this` could otherwise dangle.
174 : 1 : auto impl_keepalive = weak_impl.lock();
175 [ - + ]: 1 : if (!impl_keepalive) return;
176 : 1 : auto alive = weak_alive.lock();
177 [ - + ]: 1 : if (!alive) return;
178 : 1 : flush_batches();
179 : 1 : });
180 : : }
181 : :
182 : 104 : std::future<bool> start() {
183 : 104 : std::unique_lock<std::shared_mutex> lock(mutex_);
184 [ + + + - : 104 : if (channel_ && channel_->is_connected()) {
+ + + + ]
185 : 5 : started_.store(true);
186 : 5 : std::promise<bool> p;
187 : 5 : p.set_value(true);
188 : 5 : return p.get_future();
189 : 5 : }
190 : 99 : std::promise<bool> p;
191 : 99 : auto f = p.get_future();
192 : 99 : pending_promises_.push_back(std::move(p));
193 [ - + ]: 99 : if (started_.load()) return f;
194 : :
195 [ + + ]: 99 : if (!alive_marker_) {
196 : 2 : alive_marker_ = std::make_shared<bool>(true);
197 : : }
198 : :
199 [ + + ]: 99 : if (!channel_) {
200 : 92 : config::TcpClientConfig config;
201 : 92 : config.host = host_;
202 : 92 : config.port = port_;
203 : 92 : config.retry_interval_ms = static_cast<unsigned int>(retry_interval_.count());
204 : 92 : config.max_retries = max_retries_;
205 : 92 : config.connection_timeout_ms = static_cast<unsigned>(connection_timeout_.count());
206 : 92 : config.tls_enabled = tls_enabled_;
207 : 92 : config.tls_ca_file = tls_ca_file_;
208 : 92 : config.idle_timeout_ms = static_cast<unsigned>(idle_timeout_.count());
209 : 92 : config.idle_timeout_action = idle_timeout_action_;
210 : 92 : config.backpressure_threshold = backpressure_threshold_;
211 : 92 : config.backpressure_strategy = backpressure_strategy_;
212 : 92 : config.tcp_no_delay = tcp_no_delay_;
213 : 92 : config.keep_alive = keep_alive_;
214 : 92 : config.send_buffer_size = send_buffer_size_;
215 : 92 : config.receive_buffer_size = receive_buffer_size_;
216 : 92 : config.read_buffer_size = read_buffer_size_;
217 : 92 : channel_ = factory::ChannelFactory::create(config, external_ioc_);
218 : 92 : setup_internal_handlers();
219 : 92 : }
220 : :
221 : 99 : started_.store(true);
222 : 99 : channel_->start();
223 [ + + + + : 99 : if (use_external_context_.load() && manage_external_context_.load() && !external_thread_.joinable()) {
+ - + + ]
224 [ + - + - : 3 : if (external_ioc_ && external_ioc_->stopped()) {
- + - + ]
225 : 0 : external_ioc_->restart();
226 : : }
227 : 6 : work_guard_ = std::make_unique<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>>(
228 : 9 : external_ioc_->get_executor());
229 : 6 : external_thread_ = std::jthread([ioc = external_ioc_](std::stop_token st) {
230 : 3 : wirestead::concurrency::run_io_thread_init();
231 : : try {
232 : 6 : std::stop_callback cb(st, [ioc] { ioc->stop(); });
233 : 3 : ioc->run();
234 : 3 : } catch (...) {
235 : 0 : }
236 : 6 : });
237 : : }
238 : 99 : return f;
239 : 104 : }
240 : :
241 : 223 : void stop() {
242 : 223 : bool should_join = false;
243 : : {
244 : 223 : std::unique_lock<std::shared_mutex> lock(mutex_);
245 [ + + ]: 223 : if (!started_.load()) {
246 : 120 : fulfill_all_locked(false);
247 : 120 : return;
248 : : }
249 : 103 : started_.store(false);
250 : 103 : bp_cv_.notify_all();
251 : 103 : alive_marker_.reset();
252 [ + - ]: 103 : if (batch_timer_) {
253 : 103 : batch_timer_->cancel();
254 : 103 : batch_timer_.reset();
255 : : }
256 [ + - ]: 103 : if (channel_) {
257 : 103 : auto ch = channel_;
258 : 103 : lock.unlock();
259 : 103 : ch->stop();
260 : 103 : lock.lock();
261 [ + - ]: 103 : if (channel_ == ch) {
262 : 103 : channel_->on_bytes(nullptr);
263 : 103 : channel_->on_state(nullptr);
264 : 103 : channel_->on_backpressure(nullptr);
265 : : }
266 : 103 : }
267 [ + + + + : 103 : if (use_external_context_.load() && manage_external_context_.load()) {
+ + ]
268 [ + - ]: 3 : if (work_guard_) work_guard_.reset();
269 [ + - + - ]: 3 : if (external_ioc_) external_ioc_->stop();
270 : 3 : should_join = true;
271 : : }
272 : 103 : fulfill_all_locked(false);
273 : 223 : }
274 [ + + + - : 103 : if (should_join && external_thread_.joinable()) {
+ + ]
275 : : try {
276 [ + - ]: 3 : if (std::this_thread::get_id() != external_thread_.get_id()) {
277 : 3 : external_thread_.request_stop();
278 : 3 : external_thread_.join();
279 : : } else {
280 : 0 : external_thread_.detach();
281 : : }
282 : 0 : } catch (...) {
283 : 0 : }
284 : : }
285 : 103 : std::unique_lock<std::shared_mutex> lock(mutex_);
286 : 103 : channel_.reset();
287 [ + + + - ]: 103 : if (framer_) framer_->reset();
288 : 103 : }
289 : :
290 : 2011 : bool try_send(std::string_view data) {
291 : 2011 : std::shared_lock<std::shared_mutex> lock(mutex_);
292 [ + - + - : 2011 : if (channel_ && channel_->is_connected()) {
+ + + + ]
293 : 2009 : auto binary_view = base::safe_convert::string_to_bytes(data);
294 : 2009 : return channel_->async_try_write_copy(memory::ConstByteSpan(binary_view.first, binary_view.second));
295 : : }
296 : 2 : return false;
297 : 2011 : }
298 : :
299 : 3 : bool try_send_move(std::vector<uint8_t>&& data) {
300 : 3 : std::shared_lock<std::shared_mutex> lock(mutex_);
301 [ + - + - : 3 : if (channel_ && channel_->is_connected()) {
+ - + - ]
302 : 3 : return channel_->async_try_write_move(std::move(data));
303 : : }
304 : 0 : return false;
305 : 3 : }
306 : :
307 : 5 : bool try_send_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
308 [ + + + + : 5 : if (!data || data->empty()) return false;
+ + ]
309 : 3 : std::shared_lock<std::shared_mutex> lock(mutex_);
310 [ + - + - : 3 : if (channel_ && channel_->is_connected()) {
+ - + - ]
311 : 3 : return channel_->async_try_write_shared(std::move(data));
312 : : }
313 : 0 : return false;
314 : 3 : }
315 : :
316 : 144 : bool send(std::string_view data) {
317 [ + + ]: 144 : if (backpressure_strategy_ == base::constants::BackpressureStrategy::Reliable) return send_blocking(data);
318 : 2 : return try_send(data);
319 : : }
320 : :
321 : : // channel_->on_backpressure() calls bp_cv_.notify_all() from the transport's io_context
322 : : // thread without holding bp_mutex_ (backpressure_active_ is a plain atomic on the transport
323 : : // side, not guarded by bp_mutex_ at all). That makes a classic lost-wakeup race possible: a
324 : : // waiter can check the predicate, find it still blocking, and be in the process of
325 : : // registering to wait when the notify fires - in the rare case that race is lost, an
326 : : // unbounded wait() would block forever. Poll with a bounded timeout instead so a missed
327 : : // notify only costs a short delay rather than a permanent hang (see #427, #431).
328 : : //
329 : : // Returns false instead of waiting if called from the channel's own io
330 : : // thread while backpressure is active - e.g. a blocking send() called
331 : : // from inside an on_data/on_message callback. Clearing backpressure
332 : : // requires that same io thread to make progress, so blocking here would
333 : : // deadlock forever rather than eventually clear (#449).
334 : 1754 : bool wait_for_backpressure_clear(std::unique_lock<std::mutex>& bp_lock) {
335 : 1758 : auto predicate = [this] {
336 : 1758 : std::shared_lock<std::shared_mutex> lock(mutex_);
337 [ + - + - : 3516 : return !started_.load() || !channel_ || !channel_->is_backpressure_active();
+ - + + ]
338 : 1758 : };
339 [ + - + + ]: 1754 : if (predicate()) return true;
340 [ + + ]: 2 : if (detail::in_data_callback()) return false;
341 [ + - + + ]: 2 : while (!bp_cv_.wait_for(bp_lock, std::chrono::milliseconds(50), predicate)) {
342 : : }
343 : 1 : return true;
344 : : }
345 : :
346 : : // #509: wait_for_backpressure_clear()'s condition (is_backpressure_active(),
347 : : // tied to bp_high) and the transport's own hard queue-byte cap
348 : : // (bp_limit = bp_high * 4, rechecked inside each async_write_*) are
349 : : // different thresholds observed at different times - something else can
350 : : // refill the queue in the narrow window between the wait exiting and this
351 : : // write's own cap check, rejecting a write on an otherwise perfectly
352 : : // healthy channel. Retry a bounded number of times rather than giving up
353 : : // after one attempt, since the failure is almost always transient, not
354 : : // permanent. Bounded (rather than unbounded like the wait loop itself) so
355 : : // a payload that can never fit under a very small configured
356 : : // backpressure_threshold still fails in bounded time instead of hanging.
357 : : static constexpr int kMaxBlockingSendAttempts = 5;
358 : :
359 : 514 : bool send_move(std::vector<uint8_t>&& data) {
360 [ + + ]: 514 : if (backpressure_strategy_ == base::constants::BackpressureStrategy::Reliable) {
361 [ + - ]: 513 : for (int attempt = 0; attempt < kMaxBlockingSendAttempts; ++attempt) {
362 : 513 : std::unique_lock<std::mutex> bp_lock(bp_mutex_);
363 [ + - - + ]: 513 : if (!wait_for_backpressure_clear(bp_lock)) return false;
364 : 513 : bp_lock.unlock();
365 : 513 : std::shared_lock<std::shared_mutex> lock(mutex_);
366 [ + - + - : 513 : if (!started_.load() || !channel_ || !channel_->is_connected()) return false;
+ - - + -
+ ]
367 : : // async_write_move only actually moves from `data` on success (see
368 : : // TcpClient::async_write_move), so retrying with the same `data`
369 : : // after a `false` return is safe.
370 [ + - + - ]: 513 : if (channel_->async_write_move(std::move(data))) return true;
371 : 1026 : }
372 : 0 : return false;
373 : : }
374 : 1 : return try_send_move(std::move(data));
375 : : }
376 : :
377 : 2 : bool send_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
378 [ + - - + : 2 : if (!data || data->empty()) return false;
- + ]
379 [ + + ]: 2 : if (backpressure_strategy_ == base::constants::BackpressureStrategy::Reliable) {
380 [ + - ]: 1 : for (int attempt = 0; attempt < kMaxBlockingSendAttempts; ++attempt) {
381 : 1 : std::unique_lock<std::mutex> bp_lock(bp_mutex_);
382 [ + - - + ]: 1 : if (!wait_for_backpressure_clear(bp_lock)) return false;
383 : 1 : bp_lock.unlock();
384 : 1 : std::shared_lock<std::shared_mutex> lock(mutex_);
385 [ + - + - : 1 : if (!started_.load() || !channel_ || !channel_->is_connected()) return false;
+ - - + -
+ ]
386 [ + - + - ]: 1 : if (channel_->async_write_shared(data)) return true;
387 : 2 : }
388 : 0 : return false;
389 : : }
390 : 1 : return try_send_shared(std::move(data));
391 : : }
392 : :
393 : 1002 : bool send_line(std::string_view line) {
394 [ + + ]: 1002 : if (backpressure_strategy_ == base::constants::BackpressureStrategy::Reliable) return send_line_blocking(line);
395 : 1 : return try_send_line(line);
396 : : }
397 : :
398 : 6 : bool try_send_line(std::string_view line) { return try_send(std::string(line) + "\n"); }
399 : :
400 : 1239 : bool send_blocking(std::string_view data) {
401 : 1239 : auto binary_view = base::safe_convert::string_to_bytes(data);
402 : 1239 : memory::ConstByteSpan span(binary_view.first, binary_view.second);
403 [ + - ]: 1240 : for (int attempt = 0; attempt < kMaxBlockingSendAttempts; ++attempt) {
404 : 1240 : std::unique_lock<std::mutex> bp_lock(bp_mutex_);
405 [ + - + + ]: 1240 : if (!wait_for_backpressure_clear(bp_lock)) return false;
406 : 1239 : bp_lock.unlock();
407 : 1239 : std::shared_lock<std::shared_mutex> lock(mutex_);
408 [ + - - + : 1239 : if (!started_.load() || !channel_) return false;
- + ]
409 [ + - + + ]: 1239 : if (channel_->async_write_copy(span)) return true;
410 : 2478 : }
411 : 0 : return false;
412 : : }
413 : :
414 : 3003 : bool send_line_blocking(std::string_view line) { return send_blocking(std::string(line) + "\n"); }
415 : :
416 : 3056 : bool connected() const {
417 : 3056 : std::shared_lock<std::shared_mutex> lock(mutex_);
418 [ + + + - : 6112 : return channel_ && channel_->is_connected();
+ + ]
419 : 3056 : }
420 : :
421 : 8 : RuntimeStats stats() const {
422 : 8 : std::shared_lock<std::shared_mutex> lock(mutex_);
423 [ + - + - ]: 16 : return channel_ ? channel_->stats() : RuntimeStats{};
424 : 8 : }
425 : :
426 : 1 : void reset_stats() {
427 : 1 : std::shared_lock<std::shared_mutex> lock(mutex_);
428 [ + - + - ]: 1 : if (channel_) channel_->reset_stats();
429 : 1 : }
430 : :
431 : 111 : void setup_internal_handlers() {
432 [ - + ]: 111 : if (!channel_) return;
433 : :
434 : 111 : batch_timer_ = std::make_unique<boost::asio::steady_timer>(channel_->get_executor());
435 : :
436 : 111 : std::weak_ptr<bool> weak_alive = alive_marker_;
437 : 111 : std::weak_ptr<Impl> weak_impl = weak_from_this();
438 : :
439 : 111 : channel_->on_bytes([this, weak_impl, weak_alive](memory::ConstByteSpan data) {
440 : : // #450: keep Impl alive for the duration of this callback - on an
441 : : // externally-owned io_context, stop() doesn't join/wait for in-flight
442 : : // handlers, so a bare `this` could otherwise dangle.
443 : 125 : auto impl_keepalive = weak_impl.lock();
444 [ - + ]: 125 : if (!impl_keepalive) return;
445 : 125 : auto alive = weak_alive.lock();
446 [ - + ]: 125 : if (!alive) return;
447 : :
448 : : // #449: everything below (data_handler_/data_batch_handler_, and
449 : : // transitively message_handler_/message_batch_handler_ via
450 : : // framer_to_push->push_bytes() below) runs synchronously on this io
451 : : // thread. Mark it so a blocking send() called from within one of
452 : : // these callbacks fails fast instead of deadlocking waiting for this
453 : : // same thread to clear backpressure.
454 : 125 : detail::CallbackGuard callback_guard;
455 : :
456 : : // #441: snapshot the handler/framer pointers under a shared_lock (not
457 : : // unique_lock) - this is a pure read, matching try_send's locking
458 : : // level so it no longer blocks concurrent sends even briefly.
459 : : bool batch_mode;
460 : 125 : interface::SharedCallback<MessageHandler> handler;
461 : 125 : std::shared_ptr<framer::IFramer> framer_to_push;
462 : : {
463 : 125 : std::shared_lock<std::shared_mutex> lock(mutex_);
464 : 125 : batch_mode = static_cast<bool>(data_batch_handler_);
465 : 125 : handler = data_handler_;
466 : 125 : framer_to_push = framer_;
467 : 125 : }
468 : :
469 [ + + ]: 125 : if (batch_mode) {
470 : : // #441: build the copy before taking the exclusive lock, so the
471 : : // lock is only held for the queue mutation itself, not the
472 : : // allocation.
473 : 4 : MessageContext ctx(0, memory::SafeDataBuffer(data));
474 : 4 : interface::SharedCallback<BatchMessageHandler> flush_handler;
475 : 4 : std::vector<MessageContext> batch;
476 : : {
477 : 4 : std::unique_lock<std::shared_mutex> lock(mutex_);
478 : 4 : data_batch_queue_.emplace_back(std::move(ctx));
479 [ + + ]: 4 : if (data_batch_queue_.size() >= max_batch_size_) {
480 : 1 : flush_handler = data_batch_handler_;
481 : 1 : batch = std::move(data_batch_queue_);
482 : 1 : data_batch_queue_.clear();
483 [ + - ]: 3 : } else if (data_batch_queue_.size() == 1) {
484 : 3 : schedule_batch_timer();
485 : : }
486 : 4 : }
487 : 4 : detail::invoke_user_callback("tcp_client", "on_data_batch", flush_handler, batch);
488 : 4 : } else {
489 : 242 : detail::invoke_user_callback("tcp_client", "on_data", handler, MessageContext(0, data));
490 : : }
491 : :
492 [ + + + - ]: 125 : if (framer_to_push) framer_to_push->push_bytes(data);
493 : 125 : });
494 : :
495 : 111 : channel_->on_state([this, weak_impl, weak_alive](base::LinkState state) {
496 : 278 : auto impl_keepalive = weak_impl.lock();
497 [ - + ]: 278 : if (!impl_keepalive) return;
498 : 278 : auto alive = weak_alive.lock();
499 [ - + ]: 278 : if (!alive) return;
500 : 278 : ConnectionHandler connect_handler;
501 : 278 : ConnectionHandler disconnect_handler;
502 : 278 : ErrorHandler error_handler;
503 : 278 : std::shared_ptr<interface::Channel> channel_snapshot;
504 : :
505 [ + + ]: 278 : if (state == base::LinkState::Connected) {
506 : : {
507 : 97 : std::unique_lock<std::shared_mutex> lock(mutex_);
508 : 97 : fulfill_all_locked(true);
509 : 97 : connect_handler = connect_handler_;
510 : 97 : }
511 : 194 : detail::invoke_user_callback("tcp_client", "on_connect", connect_handler, ConnectionContext(0));
512 [ + + + + ]: 181 : } else if (state == base::LinkState::Closed || state == base::LinkState::Error) {
513 : : {
514 : 7 : std::unique_lock<std::shared_mutex> lock(mutex_);
515 : 7 : fulfill_all_locked(false);
516 [ + + ]: 7 : if (state == base::LinkState::Closed) {
517 : 3 : disconnect_handler = disconnect_handler_;
518 : : } else {
519 : 4 : error_handler = error_handler_;
520 : 4 : channel_snapshot = channel_;
521 : : }
522 : 7 : }
523 [ + + ]: 7 : if (state == base::LinkState::Closed) {
524 : 6 : detail::invoke_user_callback("tcp_client", "on_disconnect", disconnect_handler, ConnectionContext(0));
525 [ + - ]: 4 : } else if (state == base::LinkState::Error) {
526 [ + - ]: 4 : detail::invoke_user_callback("tcp_client", "on_error", error_handler,
527 : 4 : channel_snapshot
528 [ + - + - : 8 : ? detail::build_error_context(*channel_snapshot, "Connection error")
- - ]
529 : : : ErrorContext(ErrorCode::IoError, "Connection error"));
530 : : }
531 : : }
532 : 278 : });
533 : :
534 : 111 : channel_->on_backpressure([this, weak_impl, weak_alive](size_t queued) {
535 : 3 : bp_cv_.notify_all();
536 : 3 : auto impl_keepalive = weak_impl.lock();
537 [ - + ]: 3 : if (!impl_keepalive) return;
538 : 3 : auto alive = weak_alive.lock();
539 [ - + ]: 3 : if (!alive) return;
540 : 3 : std::function<void(size_t)> handler;
541 : : {
542 : 3 : std::shared_lock<std::shared_mutex> lock(mutex_);
543 : 3 : handler = bp_handler_;
544 : 3 : }
545 : 3 : detail::invoke_user_callback("tcp_client", "on_backpressure", handler, queued);
546 : 3 : });
547 : 111 : }
548 : :
549 : 4 : void attach_framer_callback() {
550 [ - + ]: 4 : if (!framer_) return;
551 : 4 : framer_->on_message([this](memory::ConstByteSpan msg) {
552 : : // #441: snapshot under a shared_lock (pure read), build the copy
553 : : // before taking the exclusive lock for queue mutation.
554 : : bool batch_mode;
555 : 3 : interface::SharedCallback<MessageHandler> handler;
556 : : {
557 : 3 : std::shared_lock<std::shared_mutex> lock(mutex_);
558 : 3 : batch_mode = static_cast<bool>(message_batch_handler_);
559 : 3 : handler = message_handler_;
560 : 3 : }
561 : :
562 [ + - ]: 3 : if (batch_mode) {
563 : 3 : MessageContext ctx(0, memory::SafeDataBuffer(msg));
564 : 3 : interface::SharedCallback<BatchMessageHandler> flush_handler;
565 : 3 : std::vector<MessageContext> batch;
566 : : {
567 : 3 : std::unique_lock<std::shared_mutex> lock(mutex_);
568 : 3 : message_batch_queue_.emplace_back(std::move(ctx));
569 [ + + ]: 3 : if (message_batch_queue_.size() >= max_batch_size_) {
570 : 1 : flush_handler = message_batch_handler_;
571 : 1 : batch = std::move(message_batch_queue_);
572 : 1 : message_batch_queue_.clear();
573 [ + - ]: 2 : } else if (message_batch_queue_.size() == 1) {
574 : 2 : schedule_batch_timer();
575 : : }
576 : 3 : }
577 : 3 : detail::invoke_user_callback("tcp_client", "on_message_batch", flush_handler, batch);
578 : 3 : return;
579 : 3 : }
580 : :
581 : 0 : detail::invoke_user_callback("tcp_client", "on_message", handler, MessageContext(0, msg));
582 : 3 : });
583 : : }
584 : :
585 : 3 : void set_framer(std::unique_ptr<framer::IFramer> framer) {
586 : 3 : std::unique_lock<std::shared_mutex> lock(mutex_);
587 : 3 : framer_ = std::shared_ptr<framer::IFramer>(std::move(framer));
588 [ + - + - : 3 : if (framer_ && (message_handler_ || message_batch_handler_)) attach_framer_callback();
- + - + -
- ]
589 : 3 : }
590 : :
591 : 1 : void on_message(MessageHandler handler) {
592 : 1 : std::unique_lock<std::shared_mutex> lock(mutex_);
593 : 1 : message_handler_ = interface::share_callback(std::move(handler));
594 [ + - + - ]: 1 : if (framer_) attach_framer_callback();
595 : 1 : }
596 : :
597 : 3 : void on_message_batch(BatchMessageHandler handler) {
598 : 3 : std::unique_lock<std::shared_mutex> lock(mutex_);
599 : 3 : message_batch_handler_ = interface::share_callback(std::move(handler));
600 [ + - + - ]: 3 : if (framer_) attach_framer_callback();
601 : 3 : }
602 : : };
603 : :
604 : 101 : TcpClient::TcpClient(const std::string& h, uint16_t p) : impl_(std::make_shared<Impl>(h, p)) {}
605 : 8 : TcpClient::TcpClient(const std::string& h, uint16_t p, std::shared_ptr<boost::asio::io_context> ioc)
606 : 8 : : impl_(std::make_shared<Impl>(h, p, ioc)) {}
607 : 19 : TcpClient::TcpClient(std::shared_ptr<interface::Channel> ch) : impl_(std::make_shared<Impl>(ch)) {
608 : 19 : impl_->setup_internal_handlers();
609 : 19 : }
610 : 220 : TcpClient::~TcpClient() = default;
611 : :
612 : 0 : TcpClient::TcpClient(TcpClient&&) noexcept = default;
613 : 0 : TcpClient& TcpClient::operator=(TcpClient&&) noexcept = default;
614 : :
615 : 104 : std::future<bool> TcpClient::start() { return impl_->start(); }
616 : 95 : void TcpClient::stop() { impl_->stop(); }
617 : 144 : bool TcpClient::send(std::string_view data) { return impl_->send(data); }
618 : 2007 : bool TcpClient::try_send(std::string_view data) { return impl_->try_send(data); }
619 : 1002 : bool TcpClient::send_line(std::string_view line) { return impl_->send_line(line); }
620 : 1 : bool TcpClient::try_send_line(std::string_view line) { return impl_->try_send_line(line); }
621 : 514 : bool TcpClient::send_move(std::vector<uint8_t>&& data) { return impl_->send_move(std::move(data)); }
622 : 2 : bool TcpClient::try_send_move(std::vector<uint8_t>&& data) { return impl_->try_send_move(std::move(data)); }
623 : 2 : bool TcpClient::send_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
624 : 2 : return impl_->send_shared(std::move(data));
625 : : }
626 : 4 : bool TcpClient::try_send_shared(std::shared_ptr<const std::vector<uint8_t>> data) {
627 : 4 : return impl_->try_send_shared(std::move(data));
628 : : }
629 : 96 : bool TcpClient::send_blocking(std::string_view data) { return impl_->send_blocking(data); }
630 : 0 : bool TcpClient::send_line_blocking(std::string_view line) { return impl_->send_line_blocking(line); }
631 : 3056 : bool TcpClient::connected() const { return get_impl()->connected(); }
632 : 8 : RuntimeStats TcpClient::stats() const { return get_impl()->stats(); }
633 : 1 : void TcpClient::reset_stats() { impl_->reset_stats(); }
634 : :
635 : 94 : TcpClient& TcpClient::on_data(MessageHandler h) {
636 : 94 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
637 : 94 : impl_->data_handler_ = interface::share_callback(std::move(h));
638 : 94 : return *this;
639 : 94 : }
640 : 3 : TcpClient& TcpClient::on_data_batch(BatchMessageHandler h) {
641 : 3 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
642 : 3 : impl_->data_batch_handler_ = interface::share_callback(std::move(h));
643 : 3 : return *this;
644 : 3 : }
645 : 17 : TcpClient& TcpClient::on_connect(ConnectionHandler h) {
646 : 17 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
647 : 17 : impl_->connect_handler_ = std::move(h);
648 : 17 : return *this;
649 : 17 : }
650 : 6 : TcpClient& TcpClient::on_disconnect(ConnectionHandler h) {
651 : 6 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
652 : 6 : impl_->disconnect_handler_ = std::move(h);
653 : 6 : return *this;
654 : 6 : }
655 : 88 : TcpClient& TcpClient::on_error(ErrorHandler h) {
656 : 88 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
657 : 88 : impl_->error_handler_ = std::move(h);
658 : 88 : return *this;
659 : 88 : }
660 : :
661 : 2 : TcpClient& TcpClient::on_backpressure(std::function<void(size_t)> h) {
662 : 2 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
663 : 2 : impl_->bp_handler_ = std::move(h);
664 : 2 : return *this;
665 : 2 : }
666 : :
667 : 3 : TcpClient& TcpClient::framer(std::unique_ptr<framer::IFramer> f) {
668 : 3 : impl_->set_framer(std::move(f));
669 : 3 : return *this;
670 : : }
671 : 1 : TcpClient& TcpClient::on_message(MessageHandler h) {
672 : 1 : impl_->on_message(std::move(h));
673 : 1 : return *this;
674 : : }
675 : 3 : TcpClient& TcpClient::on_message_batch(BatchMessageHandler h) {
676 : 3 : impl_->on_message_batch(std::move(h));
677 : 3 : return *this;
678 : : }
679 : :
680 : 8 : TcpClient& TcpClient::auto_start(bool m) {
681 : 8 : impl_->auto_start_.store(m);
682 [ + - + - : 8 : if (impl_->auto_start_.load() && !impl_->started_.load()) start();
+ - + - ]
683 : 8 : return *this;
684 : : }
685 : :
686 : 2 : TcpClient& TcpClient::batch_size(size_t size) {
687 : 2 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
688 : 2 : impl_->max_batch_size_ = size;
689 : 2 : return *this;
690 : 2 : }
691 : :
692 : 2 : TcpClient& TcpClient::batch_latency(std::chrono::milliseconds latency) {
693 : 2 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
694 : 2 : impl_->max_batch_latency_ = latency;
695 : 2 : return *this;
696 : 2 : }
697 : :
698 : 5 : TcpClient& TcpClient::retry_interval(std::chrono::milliseconds i) {
699 : 5 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
700 : 5 : impl_->retry_interval_ = i;
701 [ + + ]: 5 : if (impl_->channel_) {
702 : 1 : auto transport_client = std::dynamic_pointer_cast<transport::TcpClient>(impl_->channel_);
703 [ + - + - ]: 1 : if (transport_client) transport_client->set_retry_interval(static_cast<unsigned int>(i.count()));
704 : 1 : }
705 : 5 : return *this;
706 : 5 : }
707 : :
708 : 12 : TcpClient& TcpClient::max_retries(int m) {
709 : 12 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
710 : 12 : impl_->max_retries_ = m;
711 [ + + ]: 12 : if (impl_->channel_) {
712 : 1 : auto transport = std::dynamic_pointer_cast<transport::TcpClient>(impl_->channel_);
713 [ + - + - ]: 1 : if (transport) transport->set_max_retries(m);
714 : 1 : }
715 : 12 : return *this;
716 : 12 : }
717 : 2 : TcpClient& TcpClient::tls(const std::string& ca_file) {
718 : 2 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
719 : 2 : impl_->tls_enabled_ = true;
720 : 2 : impl_->tls_ca_file_ = ca_file;
721 : 2 : return *this;
722 : 2 : }
723 : :
724 : 4 : TcpClient& TcpClient::connection_timeout(std::chrono::milliseconds t) {
725 : 4 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
726 : 4 : impl_->connection_timeout_ = t;
727 [ + + ]: 4 : if (impl_->channel_) {
728 : 1 : auto transport = std::dynamic_pointer_cast<transport::TcpClient>(impl_->channel_);
729 [ + - + - ]: 1 : if (transport) transport->set_connection_timeout(static_cast<unsigned>(t.count()));
730 : 1 : }
731 : 4 : return *this;
732 : 4 : }
733 : :
734 : 5 : TcpClient& TcpClient::idle_timeout(std::chrono::milliseconds t) {
735 : 5 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
736 : 5 : impl_->idle_timeout_ = t;
737 [ + + ]: 5 : if (impl_->channel_) {
738 : 1 : auto transport = std::dynamic_pointer_cast<transport::TcpClient>(impl_->channel_);
739 [ + - + - ]: 1 : if (transport) transport->set_idle_timeout(static_cast<unsigned>(t.count()));
740 : 1 : }
741 : 5 : return *this;
742 : 5 : }
743 : :
744 : 4 : TcpClient& TcpClient::idle_timeout_action(IdleTimeoutAction action) {
745 : 4 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
746 : 4 : impl_->idle_timeout_action_ = action;
747 [ + + ]: 4 : if (impl_->channel_) {
748 : 1 : auto transport = std::dynamic_pointer_cast<transport::TcpClient>(impl_->channel_);
749 [ + - + - ]: 1 : if (transport) transport->set_idle_timeout_action(action);
750 : 1 : }
751 : 4 : return *this;
752 : 4 : }
753 : :
754 : 6 : TcpClient& TcpClient::manage_external_context(bool m) {
755 : 6 : impl_->manage_external_context_.store(m);
756 : 6 : return *this;
757 : : }
758 : 94 : TcpClient& TcpClient::backpressure_threshold(size_t threshold) {
759 : 94 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
760 : 94 : impl_->backpressure_threshold_ = threshold;
761 : 94 : return *this;
762 : 94 : }
763 : 8 : TcpClient& TcpClient::backpressure_strategy(base::constants::BackpressureStrategy strategy) {
764 : 8 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
765 : 8 : impl_->backpressure_strategy_ = strategy;
766 [ + + ]: 8 : if (impl_->channel_) {
767 : 4 : auto* tc = dynamic_cast<transport::TcpClient*>(impl_->channel_.get());
768 [ - + - - ]: 4 : if (tc) tc->set_backpressure_strategy(strategy);
769 : : }
770 : 8 : return *this;
771 : 8 : }
772 : :
773 : 3 : size_t TcpClient::backpressure_threshold() const {
774 : 3 : std::shared_lock<std::shared_mutex> lock(impl_->mutex_);
775 : 6 : return impl_->backpressure_threshold_;
776 : 3 : }
777 : :
778 : 3 : base::constants::BackpressureStrategy TcpClient::backpressure_strategy() const {
779 : 3 : std::shared_lock<std::shared_mutex> lock(impl_->mutex_);
780 : 6 : return impl_->backpressure_strategy_;
781 : 3 : }
782 : :
783 : 1 : TcpClient& TcpClient::tcp_no_delay(bool enable) {
784 : 1 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
785 : 1 : impl_->tcp_no_delay_ = enable;
786 : 1 : return *this;
787 : 1 : }
788 : :
789 : 1 : TcpClient& TcpClient::keep_alive(bool enable) {
790 : 1 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
791 : 1 : impl_->keep_alive_ = enable;
792 : 1 : return *this;
793 : 1 : }
794 : :
795 : 1 : TcpClient& TcpClient::send_buffer_size(size_t bytes) {
796 : 1 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
797 : 1 : impl_->send_buffer_size_ = bytes;
798 : 1 : return *this;
799 : 1 : }
800 : :
801 : 1 : TcpClient& TcpClient::receive_buffer_size(size_t bytes) {
802 : 1 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
803 : 1 : impl_->receive_buffer_size_ = bytes;
804 : 1 : return *this;
805 : 1 : }
806 : :
807 : 1 : TcpClient& TcpClient::read_buffer_size(size_t bytes) {
808 : 1 : std::unique_lock<std::shared_mutex> lock(impl_->mutex_);
809 : 1 : impl_->read_buffer_size_ = bytes;
810 : 1 : return *this;
811 : 1 : }
812 : :
813 : : } // namespace wrapper
814 : : } // namespace wirestead
|