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 <array>
20 : : #include <atomic>
21 : : #include <boost/asio.hpp>
22 : : #include <cstdint>
23 : : #include <deque>
24 : : #include <functional>
25 : : #include <memory>
26 : : #include <mutex>
27 : : #include <optional>
28 : : #include <variant>
29 : : #include <vector>
30 : :
31 : : #include "wirestead/base/constants.hpp"
32 : : #include "wirestead/base/platform.hpp"
33 : : #include "wirestead/base/visibility.hpp"
34 : : #include "wirestead/diagnostics/error_handler.hpp"
35 : : #include "wirestead/diagnostics/logger.hpp"
36 : : #include "wirestead/diagnostics/runtime_stats_counter.hpp"
37 : : #include "wirestead/interface/channel.hpp"
38 : : #include "wirestead/interface/itcp_socket.hpp"
39 : : #include "wirestead/memory/memory_pool.hpp"
40 : : #include "wirestead/memory/safe_span.hpp"
41 : : #include "wirestead/transport/base/bp_state_machine.hpp"
42 : :
43 : : namespace wirestead {
44 : : namespace transport {
45 : :
46 : : namespace net = boost::asio;
47 : :
48 : : using base::LinkState;
49 : : using interface::TcpSocketInterface;
50 : : using tcp = net::ip::tcp;
51 : :
52 : : class WIRESTEAD_API TcpServerSession : public std::enable_shared_from_this<TcpServerSession> {
53 : : public:
54 : : using OnBytes = interface::Channel::OnBytes;
55 : : using OnBackpressure = interface::Channel::OnBackpressure;
56 : : using OnClose = std::function<void()>;
57 : : using BufferVariant =
58 : : std::variant<memory::PooledBuffer, std::vector<uint8_t>, std::shared_ptr<const std::vector<uint8_t>>>;
59 : :
60 : : TcpServerSession(net::io_context& ioc, tcp::socket sock,
61 : : size_t backpressure_threshold = base::constants::DEFAULT_BACKPRESSURE_THRESHOLD,
62 : : int idle_timeout_ms = 0,
63 : : base::constants::BackpressureStrategy strategy = base::constants::BackpressureStrategy::Reliable,
64 : : bool enable_memory_pool = true, size_t read_buffer_size = base::constants::DEFAULT_READ_BUFFER_SIZE);
65 : : // Constructor for testing with dependency injection
66 : : TcpServerSession(net::io_context& ioc, std::unique_ptr<interface::TcpSocketInterface> socket,
67 : : size_t backpressure_threshold = base::constants::DEFAULT_BACKPRESSURE_THRESHOLD,
68 : : int idle_timeout_ms = 0,
69 : : base::constants::BackpressureStrategy strategy = base::constants::BackpressureStrategy::Reliable,
70 : : bool enable_memory_pool = true, size_t read_buffer_size = base::constants::DEFAULT_READ_BUFFER_SIZE);
71 : :
72 : : void start();
73 : : bool async_write_copy(memory::ConstByteSpan data);
74 : : bool async_write_move(std::vector<uint8_t>&& data);
75 : : bool async_write_shared(std::shared_ptr<const std::vector<uint8_t>> data);
76 : : bool async_try_write_copy(memory::ConstByteSpan data);
77 : : bool async_try_write_move(std::vector<uint8_t>&& data);
78 : : bool async_try_write_shared(std::shared_ptr<const std::vector<uint8_t>> data);
79 : : void on_bytes(OnBytes cb);
80 : : void on_backpressure(OnBackpressure cb);
81 : : void on_close(OnClose cb);
82 : : bool alive() const;
83 : 12 : bool is_backpressure_active() const { return backpressure_active_.load(); }
84 : : wrapper::RuntimeStats stats() const;
85 : : void reset_stats();
86 : : void stop();
87 : : void cancel();
88 : :
89 : : private:
90 : : void start_read();
91 : : void do_write();
92 : : void do_close();
93 : : void report_backpressure(size_t queued_bytes);
94 : : void reset_idle_timer();
95 : : void observe_queue();
96 : : // Shared decide_enqueue()/route dispatch used by all 3 async_write_* variants (#434).
97 : : void route_enqueued_buffer(BufferVariant&& buf, size_t added);
98 : : queue_util::BackpressureFields bp_fields();
99 : :
100 : : private:
101 : : net::io_context& ioc_;
102 : : net::strand<net::io_context::executor_type> strand_;
103 : : net::steady_timer idle_timer_;
104 : : std::unique_ptr<interface::TcpSocketInterface> socket_;
105 : : // #443: per-session pool instead of the process-wide GlobalMemoryPool
106 : : // singleton - avoids cross-channel contention on the singleton's bucket
107 : : // mutexes. Capacity is much smaller than the old shared default since
108 : : // it's no longer amortized across every connected client in the process.
109 : : // Prefill stays 0. This literal was written while MemoryPool discarded
110 : : // initial_pool_size, so 50 allocated nothing; #575 made the parameter real
111 : : // and turned it into ~1 MiB eagerly allocated per accepted session, on the
112 : : // accept handler itself. The pool fills as buffers are released.
113 : : memory::MemoryPool pool_{0, 200};
114 : : bool enable_memory_pool_ = true;
115 : : // Sized from the server's read_buffer_size rather than being a fixed
116 : : // std::array: this buffer exists per connection, so a server trades memory
117 : : // against read completions here with max_connections as the multiplier.
118 : : std::vector<uint8_t> rx_;
119 : : std::deque<BufferVariant> tx_;
120 : : std::deque<BufferVariant> pending_;
121 : : std::atomic<size_t> pending_bytes_{0};
122 : : // Buffers handed to the in-flight gather write; `current_write_views_`
123 : : // points into the batch, so neither is touched while a write is in flight.
124 : : std::vector<BufferVariant> current_write_batch_;
125 : : std::vector<net::const_buffer> current_write_views_;
126 : : bool writing_ = false;
127 : : std::atomic<size_t> queue_bytes_{0};
128 : : // Bytes accepted by a plain async_write_* call but not yet routed onto the
129 : : // strand - reserved via try_reserve_limit_bytes() to close the
130 : : // accept-then-drop race (jwsung91/wirestead#517). inflight_bytes_ mutations
131 : : // and the queue_bytes_/pending_bytes_ increments that promote a
132 : : // reservation both go through write_reserve_mtx_ - see bp_utils.hpp.
133 : : std::atomic<size_t> inflight_bytes_{0};
134 : : std::mutex write_reserve_mtx_;
135 : : base::constants::BackpressureStrategy bp_strategy_{base::constants::BackpressureStrategy::Reliable};
136 : : size_t bp_high_; // Configurable backpressure threshold
137 : : size_t bp_limit_; // Hard cap for queued bytes
138 : : size_t bp_low_; // Backpressure relief threshold
139 : : std::atomic<bool> backpressure_active_{false};
140 : : diagnostics::RuntimeStatsCounters stats_;
141 : : int idle_timeout_ms_ = 0;
142 : :
143 : : OnBytes on_bytes_;
144 : : OnBackpressure on_bp_;
145 : : OnClose on_close_;
146 : : std::atomic<bool> alive_{false};
147 : : std::atomic<bool> closing_{false};
148 : : std::atomic<bool> cleanup_done_{false};
149 : : };
150 : : } // namespace transport
151 : : } // namespace wirestead
|
