dope_session/pool/client/
mod.rs

1use std::collections::VecDeque;
2use std::marker::PhantomData;
3use std::sync::Arc;
4use std::sync::atomic::AtomicBool;
5use std::time::Instant;
6
7use dope_core::driver::{Backend, DriverConfig, PoolDriver};
8use dope_runtime::runtime::Runtime;
9
10use crate::Codec;
11use crate::protocol::client::ClientProtocol;
12use dope_core::profile::Profile;
13use dope_transport::{Tcp, Transport};
14
15pub mod connect;
16pub mod connect_source;
17pub mod recv;
18pub mod state;
19pub mod write;
20
21use self::connect_source::{ConnectAction, ConnectSource};
22use self::state::{CLIENT_IOV_CAP, ClientSlotState};
23use super::Direction;
24use super::Pool;
25use super::core::PoolCore;
26use super::slot::{ConnFlags, INGRESS_BUF_CAP};
27use crate::SlotId;
28
29pub type ClientSession<P> =
30    crate::session::ClientSessionGeneric<P, INGRESS_BUF_CAP, CLIENT_IOV_CAP>;
31
32pub struct Client<P: ClientProtocol, T: Transport, S: ConnectSource<T>> {
33    pub(crate) protocol: P,
34    pub(crate) connect_source: S,
35    pub(crate) socket_opts: dope_transport::config::StreamSocketConfig,
36    pub(crate) connecting_slots: VecDeque<u16>,
37    transport: PhantomData<T>,
38}
39
40impl<P: ClientProtocol, T: Transport, S: ConnectSource<T>> Client<P, T, S> {
41    pub fn new(protocol: P, connect_source: S) -> Self {
42        Self {
43            protocol,
44            connect_source,
45            socket_opts: dope_transport::config::StreamSocketConfig::default(),
46            connecting_slots: VecDeque::with_capacity(16),
47            transport: PhantomData,
48        }
49    }
50
51    #[inline(always)]
52    pub fn protocol(&self) -> &P {
53        &self.protocol
54    }
55
56    #[inline(always)]
57    pub fn protocol_mut(&mut self) -> &mut P {
58        &mut self.protocol
59    }
60
61    #[inline(always)]
62    pub fn connect_source(&self) -> &S {
63        &self.connect_source
64    }
65
66    #[inline(always)]
67    pub fn connect_source_mut(&mut self) -> &mut S {
68        &mut self.connect_source
69    }
70
71    #[inline(always)]
72    pub(crate) fn has_pending(&self) -> bool {
73        !self.connecting_slots.is_empty()
74    }
75}
76
77impl<P, T, S, B> Direction<B> for Client<P, T, S>
78where
79    P: ClientProtocol,
80    T: Transport,
81    T::Addr: Clone,
82    S: ConnectSource<T>,
83    B: PoolDriver,
84{
85    type Transport = T;
86    type SlotUser = ClientSlotState<P, B>;
87
88    #[inline(always)]
89    fn has_pending(&self) -> bool {
90        Self::has_pending(self)
91    }
92
93    #[inline]
94    fn tick<C: Codec>(
95        &mut self,
96        core: &mut PoolCore<Self::Transport, C, Self::SlotUser, B>,
97        driver: &mut B::Driver,
98    ) {
99        let now = Instant::now();
100        let max_new_connects = core.free_conns.len();
101        for _ in 0..max_new_connects {
102            match self.connect_source.poll_connect(now) {
103                ConnectAction::Connect { addr, tag } => {
104                    if !self.spawn_connect::<C, B>(core, driver, addr, tag) {
105                        self.connect_source.on_connect_outcome(tag, false, now);
106                        break;
107                    }
108                }
109                ConnectAction::Backoff { .. } | ConnectAction::Idle => break,
110            }
111        }
112        core.flush_recv_rearm(driver);
113    }
114}
115
116impl<P, T, S, C, F, B> Pool<Client<P, T, S>, C, F, B>
117where
118    P: ClientProtocol,
119    T: Transport,
120    T::Addr: Clone,
121    S: ConnectSource<T>,
122    C: Codec,
123    F: Profile,
124    B: PoolDriver,
125{
126    pub fn new(
127        direction: Client<P, T, S>,
128        max_conn: usize,
129        slot_id_base: usize,
130    ) -> std::io::Result<Self> {
131        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
132        let core = PoolCore::<T, C, ClientSlotState<P, B>, B>::new(max_conn, slot_id_base)?;
133        Ok(Self {
134            direction,
135            core,
136            profile: PhantomData,
137        })
138    }
139
140    pub fn slot_id_range(&self) -> std::ops::Range<usize> {
141        let base = self.core.slot_id_base;
142        base..(base + self.core.slots.len())
143    }
144
145    pub fn conn_slot_count(&self) -> usize {
146        self.core.slots.len()
147    }
148
149    pub fn direction_mut(&mut self) -> &mut Client<P, T, S> {
150        &mut self.direction
151    }
152
153    pub fn live_conns(&self) -> impl Iterator<Item = SlotId> + '_ {
154        self.core
155            .slots
156            .iter()
157            .enumerate()
158            .filter(|(_, slot)| slot.conn.flags.contains(ConnFlags::LIVE))
159            .map(|(i, slot)| SlotId::from_parts(i as u32, slot.conn.generation))
160    }
161
162    pub fn is_live(&self, conn_id: SlotId) -> bool {
163        self.core.resolve_live_slot(conn_id).is_ok()
164    }
165
166    pub fn session_mut(&mut self, conn_id: SlotId) -> Option<&mut ClientSession<P>> {
167        let slot_idx = self.core.resolve_live_slot(conn_id).ok()?;
168        Some(&mut self.core.slots[slot_idx].user.get_mut().session)
169    }
170}
171
172#[derive(Clone, Debug, Default)]
173pub struct TcpClientConfig {
174    pub max_conn: usize,
175    pub socket_opts: dope_transport::config::StreamSocketConfig,
176    pub shutdown: Option<Arc<AtomicBool>>,
177    pub cpu_id: Option<u16>,
178}
179
180impl<P, S, C, F, B> Pool<Client<P, Tcp, S>, C, F, B>
181where
182    P: ClientProtocol,
183    S: ConnectSource<Tcp>,
184    C: Codec,
185    F: Profile,
186    B: PoolDriver,
187{
188    pub fn build(
189        protocol: P,
190        connect_source: S,
191        cfg: TcpClientConfig,
192    ) -> std::io::Result<Runtime<Self>> {
193        let shutdown = cfg.shutdown.clone();
194        let max_conn = cfg.max_conn;
195        let cpu_id = cfg.cpu_id;
196        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
197
198        let driver_cfg = <<B as Backend>::Config as DriverConfig>::for_tcp_profile(
199            F::RING_ENTRIES,
200            max_conn,
201            crate::pool::slot::INGRESS_BUF_CAP,
202        )
203        .with_cpu_id(cpu_id);
204        let mut driver = <B as Backend>::new_driver(driver_cfg)?;
205        let pool = Self::build_raw(protocol, connect_source, cfg, &mut driver, 0)?;
206        Ok(Runtime::new(pool, driver).with_shutdown(shutdown))
207    }
208
209    pub fn build_raw(
210        protocol: P,
211        connect_source: S,
212        cfg: TcpClientConfig,
213        _driver: &mut B::Driver,
214        slot_id_base: usize,
215    ) -> std::io::Result<Self> {
216        let TcpClientConfig {
217            max_conn,
218            socket_opts,
219            shutdown: _,
220            cpu_id: _,
221        } = cfg;
222        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
223        let mut direction = Client::new(protocol, connect_source);
224        direction.socket_opts = socket_opts;
225        Self::new(direction, max_conn, slot_id_base)
226    }
227}
228
229impl<P, T, S, C, F, B> super::PoolTick<B> for Pool<Client<P, T, S>, C, F, B>
230where
231    P: ClientProtocol,
232    T: Transport,
233    T::Addr: Clone,
234    S: ConnectSource<T>,
235    C: Codec,
236    F: Profile,
237    B: PoolDriver,
238{
239    #[inline(always)]
240    fn tick(&mut self, driver: &mut B::Driver) {
241        self.direction.tick(&mut self.core, driver);
242        self.poll_connect_requests(driver);
243        let _ = self.pump_outbound(driver);
244    }
245}
dope_session/pool/client/mod.rs

dope_session/pool/client/
mod.rs
