dope_session/pool/client/

mod.rs

use std::collections::VecDeque;
use std::marker::PhantomData;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Instant;

use dope_core::driver::{Backend, DriverConfig, PoolDriver};
use dope_runtime::runtime::Runtime;

use crate::CodecLayer;
use crate::protocol::client::ClientProtocol;
use dope_core::profile::Profile;
use dope_transport::{Tcp, Transport};

pub mod connect;
pub mod connect_source;
pub mod recv;
pub mod state;
pub mod write;

use self::connect_source::{ConnectAction, ConnectSource};
use self::state::{CLIENT_IOV_CAP, ClientSlotState};
use super::Direction;
use super::Pool;
use super::core::PoolCore;
use super::slot::{ConnFlags, INGRESS_BUF_CAP};
use crate::SlotId;

pub type ClientSession<P> =
    crate::session::ClientSessionGeneric<P, INGRESS_BUF_CAP, CLIENT_IOV_CAP>;

pub struct Client<P: ClientProtocol, T: Transport, S: ConnectSource<T>> {
    pub(crate) protocol: P,
    pub(crate) connect_source: S,
    pub(crate) socket_opts: dope_transport::config::StreamSocketConfig,
    pub(crate) connecting_slots: VecDeque<u16>,
    transport: PhantomData<T>,
}

impl<P: ClientProtocol, T: Transport, S: ConnectSource<T>> Client<P, T, S> {
    pub fn new(protocol: P, connect_source: S) -> Self {
        Self {
            protocol,
            connect_source,
            socket_opts: dope_transport::config::StreamSocketConfig::default(),
            connecting_slots: VecDeque::with_capacity(16),
            transport: PhantomData,
        }
    }

    #[inline(always)]
    pub fn protocol(&self) -> &P {
        &self.protocol
    }

    #[inline(always)]
    pub fn protocol_mut(&mut self) -> &mut P {
        &mut self.protocol
    }

    #[inline(always)]
    pub fn connect_source(&self) -> &S {
        &self.connect_source
    }

    #[inline(always)]
    pub fn connect_source_mut(&mut self) -> &mut S {
        &mut self.connect_source
    }

    #[inline(always)]
    pub(crate) fn has_pending(&self) -> bool {
        !self.connecting_slots.is_empty()
    }
}

impl<P, T, S, B> Direction<B> for Client<P, T, S>
where
    P: ClientProtocol,
    T: Transport,
    T::Addr: Clone,
    S: ConnectSource<T>,
    B: PoolDriver,
{
    type Transport = T;
    type SlotUser = ClientSlotState<P, B>;
    type CodecLayer = P::CodecLayer;

    #[inline(always)]
    fn has_pending(&self) -> bool {
        Self::has_pending(self)
    }

    #[inline]
    fn tick(
        &mut self,
        core: &mut PoolCore<Self::Transport, Self::CodecLayer, Self::SlotUser, B>,
        driver: &mut B::Driver,
    ) {
        let now = Instant::now();
        let max_new_connects = core.free_conns.len();
        for _ in 0..max_new_connects {
            match self.connect_source.poll_connect(now) {
                ConnectAction::Connect { addr, tag } => {
                    if !self.spawn_connect::<B>(core, driver, addr, tag) {
                        self.connect_source.on_connect_outcome(tag, false, now);
                        break;
                    }
                }
                ConnectAction::Backoff { .. } | ConnectAction::Idle => break,
            }
        }
        core.flush_recv_rearm(driver);
    }
}

impl<P, T, S, F, B> Pool<Client<P, T, S>, F, B>
where
    P: ClientProtocol,
    T: Transport,
    T::Addr: Clone,
    S: ConnectSource<T>,
    F: Profile,
    B: PoolDriver,
{
    pub fn new(
        direction: Client<P, T, S>,
        max_conn: usize,
        slot_id_base: usize,
    ) -> std::io::Result<Self> {
        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
        let core =
            PoolCore::<T, P::CodecLayer, ClientSlotState<P, B>, B>::new(max_conn, slot_id_base)?;
        Ok(Self {
            direction,
            core,
            profile: PhantomData,
        })
    }

    pub fn slot_id_range(&self) -> std::ops::Range<usize> {
        let base = self.core.slot_id_base;
        base..(base + self.core.slots.len())
    }

    pub fn conn_slot_count(&self) -> usize {
        self.core.slots.len()
    }

    pub fn direction_mut(&mut self) -> &mut Client<P, T, S> {
        &mut self.direction
    }

    pub fn live_conns(&self) -> impl Iterator<Item = SlotId> + '_ {
        self.core
            .slots
            .iter()
            .enumerate()
            .filter(|(_, slot)| slot.conn.flags.contains(ConnFlags::LIVE))
            .map(|(i, slot)| SlotId::from_parts(i as u32, slot.conn.generation))
    }

    pub fn is_live(&self, conn_id: SlotId) -> bool {
        self.core.resolve_live_slot(conn_id).is_ok()
    }

    pub fn session_mut(&mut self, conn_id: SlotId) -> Option<&mut ClientSession<P>> {
        let slot_idx = self.core.resolve_live_slot(conn_id).ok()?;
        Some(&mut self.core.slots[slot_idx].user.get_mut().session)
    }

    /// Out-of-callback access to both the session and the paired codec state.
    ///
    /// Note: mutations made through this accessor do **not** auto-poke the
    /// outbound path. After mutating codec state in a way that should produce
    /// wire bytes (e.g. triggering a TLS handshake), call
    /// [`Pool::poke_outbound`] to wake the slot.
    pub fn session_with_codec_mut(
        &mut self,
        conn_id: SlotId,
    ) -> Option<(
        &mut ClientSession<P>,
        &mut <P::CodecLayer as CodecLayer>::ConnState,
    )> {
        let slot_idx = self.core.resolve_live_slot(conn_id).ok()?;
        let entry = &mut self.core.slots[slot_idx];
        let codec_state = entry.codec_state.get_mut();
        let session = &mut entry.user.get_mut().session;
        Some((session, codec_state))
    }
}

#[derive(Clone, Debug, Default)]
pub struct TcpClientConfig {
    pub max_conn: usize,
    pub socket_opts: dope_transport::config::StreamSocketConfig,
    pub shutdown: Option<Arc<AtomicBool>>,
    pub cpu_id: Option<u16>,
}

impl<P, S, F, B> Pool<Client<P, Tcp, S>, F, B>
where
    P: ClientProtocol,
    S: ConnectSource<Tcp>,
    F: Profile,
    B: PoolDriver,
{
    pub fn build(
        protocol: P,
        connect_source: S,
        cfg: TcpClientConfig,
    ) -> std::io::Result<Runtime<Self>> {
        let shutdown = cfg.shutdown.clone();
        let max_conn = cfg.max_conn;
        let cpu_id = cfg.cpu_id;
        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);

        let driver_cfg = <<B as Backend>::Config as DriverConfig>::for_tcp_profile(
            F::RING_ENTRIES,
            max_conn,
            crate::pool::slot::INGRESS_BUF_CAP,
        )
        .with_cpu_id(cpu_id);
        let mut driver = <B as Backend>::new_driver(driver_cfg)?;
        let pool = Self::build_raw(protocol, connect_source, cfg, &mut driver, 0)?;
        Ok(Runtime::new(pool, driver).with_shutdown(shutdown))
    }

    pub fn build_raw(
        protocol: P,
        connect_source: S,
        cfg: TcpClientConfig,
        _driver: &mut B::Driver,
        slot_id_base: usize,
    ) -> std::io::Result<Self> {
        let TcpClientConfig {
            max_conn,
            socket_opts,
            shutdown: _,
            cpu_id: _,
        } = cfg;
        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);
        let mut direction = Client::new(protocol, connect_source);
        direction.socket_opts = socket_opts;
        Self::new(direction, max_conn, slot_id_base)
    }
}

impl<P, T, S, F, B> super::PoolTick<B> for Pool<Client<P, T, S>, F, B>
where
    P: ClientProtocol,
    T: Transport,
    T::Addr: Clone,
    S: ConnectSource<T>,
    F: Profile,
    B: PoolDriver,
{
    #[inline(always)]
    fn tick(&mut self, driver: &mut B::Driver) {
        self.direction.tick(&mut self.core, driver);
        self.poll_connect_requests(driver);
        let _ = self.pump_outbound(driver);
    }
}