ntex_net/

lib.rs

//! Utility for async runtime abstraction
#![deny(clippy::pedantic)]
#![allow(
    clippy::clone_on_copy,
    clippy::cast_possible_truncation,
    clippy::missing_fields_in_debug,
    clippy::must_use_candidate,
    clippy::missing_errors_doc,
    clippy::missing_panics_doc
)]
use std::{any::Any, io, net, net::SocketAddr, panic};

use ntex_io::Io;
use ntex_rt::{BlockFuture, Driver, Runner};
use ntex_service::cfg::SharedCfg;

pub mod channel;
pub mod connect;

#[cfg(unix)]
pub mod polling;

#[cfg(target_os = "linux")]
pub mod uring;

#[cfg(windows)]
pub mod iocp;

#[cfg(any(unix, windows))]
mod helpers;

#[cfg(feature = "tokio")]
pub mod tokio;

#[cfg(feature = "compio")]
pub mod compio;

#[allow(clippy::wrong_self_convention)]
pub trait Reactor: Driver {
    fn tcp_connect(&self, addr: net::SocketAddr, cfg: SharedCfg) -> channel::Receiver<Io>;

    fn unix_connect(
        &self,
        addr: std::path::PathBuf,
        cfg: SharedCfg,
    ) -> channel::Receiver<Io>;

    /// Convert std `TcpStream` to `Io`
    fn from_tcp_stream(&self, stream: net::TcpStream, cfg: SharedCfg) -> io::Result<Io>;

    #[cfg(unix)]
    /// Convert std `UnixStream` to `Io`
    fn from_unix_stream(
        &self,
        _: std::os::unix::net::UnixStream,
        _: SharedCfg,
    ) -> io::Result<Io>;
}

#[inline]
/// Opens a TCP connection to a remote host.
pub async fn tcp_connect(addr: SocketAddr, cfg: SharedCfg) -> io::Result<Io> {
    with_current(|driver| driver.tcp_connect(addr, cfg)).await
}

#[inline]
/// Opens a unix stream connection.
pub async fn unix_connect<'a, P>(addr: P, cfg: SharedCfg) -> io::Result<Io>
where
    P: AsRef<std::path::Path> + 'a,
{
    with_current(|driver| driver.unix_connect(addr.as_ref().into(), cfg)).await
}

#[inline]
/// Convert std `TcpStream` to `TcpStream`
pub fn from_tcp_stream(stream: net::TcpStream, cfg: SharedCfg) -> io::Result<Io> {
    with_current(|driver| driver.from_tcp_stream(stream, cfg))
}

#[cfg(unix)]
#[inline]
/// Convert std `UnixStream` to `UnixStream`
pub fn from_unix_stream(
    stream: std::os::unix::net::UnixStream,
    cfg: SharedCfg,
) -> io::Result<Io> {
    with_current(|driver| driver.from_unix_stream(stream, cfg))
}

fn with_current<T, F: FnOnce(&dyn Reactor) -> T>(f: F) -> T {
    #[cold]
    fn not_in_ntex_driver() -> ! {
        panic!("not in a ntex driver")
    }

    if CURRENT_DRIVER.is_set() {
        CURRENT_DRIVER.with(|d| f(&**d))
    } else {
        not_in_ntex_driver()
    }
}

#[allow(clippy::borrowed_box)]
/// Sets the current reactor and runs the provided closure.
pub fn with_reactor<R, F: FnOnce() -> R>(r: &Box<dyn Reactor>, f: F) -> R {
    #[cold]
    fn reactor_is_set() -> ! {
        panic!("reactor is already set");
    }

    if CURRENT_DRIVER.is_set() {
        reactor_is_set()
    } else {
        CURRENT_DRIVER.set(r, f)
    }
}

scoped_tls::scoped_thread_local!(static CURRENT_DRIVER: Box<dyn Reactor>);

/// The default runtime.
///
/// Automatically selects the runtime implementation based on the
/// configured features and the platform on which it runs.
#[derive(Copy, Clone, Debug)]
pub struct DefaultRuntime;

impl Runner for DefaultRuntime {
    #[allow(unused_variables, clippy::too_many_lines)]
    fn block_on(&self, fut: BlockFuture) -> Result<(), Box<dyn Any + Send>> {
        #[cfg(feature = "tokio")]
        {
            let driver: Box<dyn Reactor> = Box::new(self::tokio::Reactor);

            with_reactor(&driver, || crate::tokio::block_on(fut));
            Ok(())
        }

        #[cfg(all(feature = "compio", not(feature = "tokio")))]
        {
            let driver: Box<dyn Reactor> = Box::new(self::compio::Reactor);

            with_reactor(&driver, || crate::compio::block_on(fut));
            Ok(())
        }

        #[cfg(all(windows, not(feature = "tokio"), not(feature = "compio")))]
        {
            let driver: Box<dyn Reactor> =
                Box::new(crate::iocp::Reactor::new().expect("Cannot construct driver"));

            with_reactor(&driver, || {
                panic::catch_unwind(panic::AssertUnwindSafe(|| {
                    let rt = ntex_rt::Runtime::new(driver.handle());
                    rt.block_on(fut, &*driver);
                }))
            })
        }

        #[cfg(all(unix, not(feature = "tokio"), not(feature = "compio")))]
        {
            #[cfg(feature = "neon-polling")]
            {
                let driver: Box<dyn Reactor> = Box::new(
                    crate::polling::Reactor::new()
                        .expect("Cannot construct polling reactor"),
                );

                with_reactor(&driver, || {
                    panic::catch_unwind(panic::AssertUnwindSafe(|| {
                        let rt = ntex_rt::Runtime::new(driver.handle());
                        rt.block_on(fut, &*driver);
                    }))
                })
            }

            #[cfg(all(target_os = "linux", feature = "neon-uring"))]
            {
                let driver: Box<dyn Reactor> = Box::new(
                    crate::uring::Reactor::new(2048)
                        .expect("Cannot construct io-uring reactor"),
                );

                with_reactor(&driver, || {
                    panic::catch_unwind(panic::AssertUnwindSafe(|| {
                        let rt = ntex_rt::Runtime::new(driver.handle());
                        rt.block_on(fut, &*driver);
                    }))
                })
            }

            #[cfg(all(not(feature = "neon-uring"), not(feature = "neon-polling")))]
            {
                #[cfg(target_os = "linux")]
                let driver: Box<dyn Reactor> =
                    if let Ok(reactor) = crate::uring::Reactor::new(2048) {
                        Box::new(reactor)
                    } else {
                        Box::new(
                            crate::polling::Reactor::new()
                                .expect("Cannot construct io-uring reactor"),
                        )
                    };

                #[cfg(not(target_os = "linux"))]
                let driver: Box<dyn Reactor> = Box::new(
                    crate::polling::Reactor::new()
                        .expect("Cannot construct polling reactor"),
                );

                with_reactor(&driver, || {
                    panic::catch_unwind(panic::AssertUnwindSafe(|| {
                        let rt = ntex_rt::Runtime::new(driver.handle());
                        rt.block_on(fut, &*driver);
                    }))
                })
            }
        }
    }
}