liminal-server 0.25.4

//! The TCP transport: a non-blocking socket armed through beamr's readiness
//! facility, its inbound probe and read, and the TCP process's constructor.

use super::{
    ConnectionProcess, ConnectionTransport, InboundPending, READ_BUFFER_BYTES, ReadStatus,
};
use crate::ServerError;
use crate::server::connection::readiness_refusal::{ReadinessCall, readiness_refusal};
use crate::server::connection::supervisor::runtime::ConnectionRuntime;
use crate::server::mount::MountKind;
use beamr::native::native_process::NativeContext;
use beamr::scheduler::{Interest, ReadinessToken};
use liminal_protocol::wire::ConnectionIncarnation;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::os::fd::AsRawFd;
use std::sync::{Arc, Mutex};

/// The TCP transport: a non-blocking socket armed through beamr's `RawFd`
/// readiness facility.
#[derive(Debug)]
pub(in crate::server::connection) struct TcpTransport {
    stream: Option<TcpStream>,
    /// One registration for this connection's lifetime. The same token is copied
    /// into the host record so external death can deregister it.
    readiness_token: Option<ReadinessToken>,
}

impl ConnectionTransport for TcpTransport {
    const MOUNT: MountKind = MountKind::Tcp;

    fn is_connected(&self) -> bool {
        self.stream.is_some()
    }

    fn read_available(&mut self, buffer: &mut Vec<u8>) -> Result<ReadStatus, ServerError> {
        let Some(stream) = self.stream.as_mut() else {
            // Unreachable: the caller checks `is_connected` first. Kept total so
            // the read half binds without an unwrap.
            return Ok(ReadStatus::Closed);
        };
        read_available(stream, buffer)
    }

    fn sink(&mut self) -> Option<&mut dyn Write> {
        self.stream
            .as_mut()
            .map(|stream| stream as &mut dyn std::io::Write)
    }

    fn probe_inbound(&self) -> Result<bool, ServerError> {
        self.stream
            .as_ref()
            .map_or(Ok(true), InboundPending::inbound_pending)
    }

    fn arm_readiness(
        &mut self,
        pid: u64,
        ctx: &NativeContext<'_>,
        interest: Interest,
        runtime: &ConnectionRuntime,
    ) -> Result<(), ServerError> {
        let facility = ctx
            .readiness_facility()
            .ok_or_else(|| ServerError::ListenerAccept {
                message: "connection scheduler started without its required readiness service"
                    .to_owned(),
            })?;
        if let Some(token) = self.readiness_token {
            return facility
                .rearm(&token, interest)
                .map_err(|error| readiness_refusal(pid, ReadinessCall::Rearm, &error));
        }
        let stream = self
            .stream
            .as_ref()
            .ok_or_else(|| ServerError::ListenerAccept {
                message: "cannot register readiness for a missing connection stream".to_owned(),
            })?;
        let token = facility
            .register(stream.as_raw_fd(), interest, pid, runtime.ready_atom())
            .map_err(|error| readiness_refusal(pid, ReadinessCall::Register, &error))?;
        if let Err(error) = runtime.set_readiness_token_once(pid, token, stream.as_raw_fd()) {
            runtime.deregister_unpublished_readiness(token);
            return Err(error);
        }
        self.readiness_token = Some(token);
        Ok(())
    }

    /// Nothing: a socket is told by the readiness facility, armed each slice.
    fn install_wake(&mut self, _pid: u64, _runtime: &ConnectionRuntime) -> Result<(), ServerError> {
        Ok(())
    }

    fn release(&mut self) {
        self.stream.take();
    }

    #[cfg(test)]
    fn note_process_drop(&mut self, runtime: &ConnectionRuntime) {
        // External scheduler termination can remove the process-table entry while
        // the native handler is still executing. Tests that need the descriptor
        // allocator boundary must observe the stream's actual drop, not table
        // absence. Production keeps the ordinary field-drop path byte-for-byte.
        if let Some(stream) = self.stream.take() {
            let fd = stream.as_raw_fd();
            drop(stream);
            runtime.record_process_stream_drop(fd);
        }
    }
}

impl ConnectionProcess {
    pub(in crate::server::connection) fn from_holder(
        runtime: Arc<ConnectionRuntime>,
        peer_addr: Option<SocketAddr>,
        holder: &Arc<Mutex<Option<TcpStream>>>,
        connection_incarnation: Option<ConnectionIncarnation>,
    ) -> Self {
        // The `NativeHandlerFactory` is `Fn + Send + Sync`, so the accepted
        // `TcpStream` cannot be moved into the closure (a `Fn` captures by shared
        // reference and may be invoked more than once for restart). The shared
        // `Arc<Mutex<Option<TcpStream>>>` is the interior-mutability proxy that
        // lets the FIRST handler build take the stream out exactly once; the
        // Mutex is required by the `Sync` bound, not incidental.
        //
        // If the lock is poisoned the take silently yields `None`, and the
        // process would later stop with a bare crash and no root cause. Log the
        // poisoning clearly (with the peer address) so a missing-stream handoff
        // is diagnosable instead of a mystery crash.
        let stream = match holder.lock() {
            Ok(mut held) => held.take(),
            Err(poisoned) => {
                tracing::error!(
                    peer_addr = ?peer_addr,
                    error = %poisoned,
                    "connection stream handoff failed: stream holder mutex was poisoned; \
                     the connection process will start without a stream and stop immediately"
                );
                None
            }
        };
        Self::over_transport(
            runtime,
            peer_addr,
            TcpTransport {
                stream,
                readiness_token: None,
            },
            connection_incarnation,
        )
    }
}

impl InboundPending for TcpStream {
    fn inbound_pending(&self) -> Result<bool, ServerError> {
        let mut byte = [0_u8; 1];
        match self.peek(&mut byte) {
            Ok(_) => Ok(true),
            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => Ok(false),
            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => Ok(true),
            Err(error) => Err(ServerError::ListenerAccept {
                message: format!("connection readiness probe failed: {error}"),
            }),
        }
    }
}

fn read_available(stream: &mut TcpStream, buffer: &mut Vec<u8>) -> Result<ReadStatus, ServerError> {
    let mut chunk = [0_u8; READ_BUFFER_BYTES];
    match stream.read(&mut chunk) {
        Ok(0) => Ok(ReadStatus::Closed),
        Ok(bytes_read) => {
            // `unwrap_or(&[])` masks exactly one case: a `Read` implementation
            // reporting more bytes read than the buffer it was given can hold.
            // That is a broken `Read`, and with a correct one the slice is
            // structurally always in range. The masked outcome — silently
            // appending nothing instead of panicking on the out-of-range slice —
            // is preferred to taking a connection process down over a defect
            // that cannot originate in this crate.
            buffer.extend_from_slice(chunk.get(..bytes_read).unwrap_or(&[]));
            Ok(ReadStatus::Read)
        }
        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => Ok(ReadStatus::WouldBlock),
        Err(error) if error.kind() == std::io::ErrorKind::Interrupted => Ok(ReadStatus::WouldBlock),
        Err(error) => Err(ServerError::ListenerAccept {
            message: format!("failed to read connection stream: {error}"),
        }),
    }
}