liminal-server 0.25.4

//! Per-connection process for the WebSocket sibling transport (R1.2/R1.3).
//!
//! One supervised beamr native process owns one upgraded
//! `tungstenite::WebSocket<TcpStream>`. Each scheduler slice services inbound
//! messages, the pending-reply table, the shared delivery pump, the Q-A
//! keepalive schedule, and the outbound drain — the SAME slice discipline,
//! semantic seam ([`apply_frame`](super::super::apply::apply_frame)), shared [`ConnectionRuntime`], and
//! close-cause/cleanup paths as the TCP [`super::super::process`] reference
//! implementation, with the byte-stream read/write halves replaced by the
//! one-binary-message-one-canonical-frame transport contract.
//!
//! RFC 6455 Ping/Pong and Close are transport control here: they are never
//! converted to liminal `Frame::Ping`/`Frame::Pong`/`Frame::Disconnect`, and
//! liminal control frames remain ordinary canonical binary-message payloads.

use std::net::{SocketAddr, TcpStream};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

#[path = "process/force_close.rs"]
mod force_close;

#[cfg(test)]
use inbound::pre_admission_close_frame;
use keepalive::KeepaliveSchedule;
#[path = "process/close.rs"]
mod close;
#[path = "process/control.rs"]
mod control;
#[path = "process/inbound.rs"]
mod inbound;
#[path = "process/keepalive.rs"]
mod keepalive;
#[path = "process/parking.rs"]
mod parking;

use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use beamr::scheduler::{Interest, ReadinessToken};

use tungstenite::protocol::WebSocket;

use liminal_protocol::wire::ConnectionIncarnation;

use super::super::delivery::service_subscriptions;
use super::super::outbound::{DrainOutcome, OutboundError};
use super::super::participant_delivery::{
    UNIT2_PUSH_SLICE_BUDGET, has_held_participant_head, service_participant_publications,
};
use super::super::state::ConnectionProcessState;
use super::super::supervisor::runtime::{ConnectionControl, ConnectionRuntime};
use super::AcceptorSettings;
use super::outbound::WebSocketOutbound;
use crate::ServerError;
use crate::server::connection::pre_admission::{PRE_ADMISSION_FRAME_BOUND, PreAdmissionOversize};
use crate::server::participant::ConnectionFateClass;

#[cfg(test)]
#[path = "process_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "process_peer_reset_tests.rs"]
mod peer_reset_tests;

#[cfg(test)]
#[path = "process_pre_admission_tests.rs"]
mod pre_admission_tests;

#[cfg(test)]
#[path = "process_readiness_stop_tests.rs"]
mod readiness_stop_tests;

/// Per-slice byte budget for the outbound drain, matching the TCP process's
/// read-buffer-sized drain budget so one connection cannot monopolize a
/// scheduler thread writing.
const DRAIN_SLICE_BYTES: usize = 8192;

/// Per-slice cap on complete inbound messages applied by one connection,
/// mirroring the delivery pump's slice budget rationale: bounds the work one
/// fast sender forces into a single slice so connections sharing the scheduler
/// stay fair. An exhausted budget re-queues the process (never parks it), so no
/// buffered message can be stranded behind a park.
const INBOUND_SLICE_MESSAGE_BUDGET: usize = 32;

/// Application stream id used for server-initiated push frames — the same
/// value the TCP process uses, so a push looks identical on both transports.
const PUSH_STREAM_ID: u32 = 1;

/// Whether a slice's socket/inbound servicing leaves the connection running or
/// has already resolved its lifecycle (the TCP process's `SliceStep` shape).
enum SliceStep {
    Continue,
    Stop(ExitReason),
}

#[derive(Debug)]
pub(in super::super) struct WebSocketConnectionProcess {
    runtime: Arc<ConnectionRuntime>,
    peer_addr: Option<SocketAddr>,
    socket: Option<WebSocket<TcpStream>>,
    state: ConnectionProcessState,
    outbound: WebSocketOutbound,
    readiness_token: Option<ReadinessToken>,
    /// The configured Q-A keepalive interval; `None` means pings are disabled.
    keepalive_interval: Option<Duration>,
    /// The live keepalive schedule, initialised on the first slice so that an
    /// interval overflowing the monotonic clock is a typed connection failure
    /// rather than a constructor panic.
    keepalive: Option<KeepaliveSchedule>,
    /// Set when the inbound slice budget was exhausted with messages possibly
    /// still buffered inside tungstenite; consulted by the final probe so the
    /// process re-queues instead of parking past buffered work.
    inbound_budget_exhausted: bool,
    /// The acceptor's pinned message bound, installed on the socket once the
    /// connection's `Connect` succeeds; until then the pre-admission bound is.
    admitted_bound: usize,
    /// When the listener accepted this connection and when this process was
    /// built, until its first outbound frame is flushed and logged as
    /// `websocket first frame sent` (the start-path row); `None` after.
    start: Option<(Instant, Instant)>,
    /// Whether [`Self::admitted_bound`] has replaced the pre-admission bound.
    admitted_bound_installed: bool,
}

impl WebSocketConnectionProcess {
    /// Builds the process from the spawn holder, taking the upgraded socket out
    /// exactly once — the same interior-mutability handoff the TCP process
    /// uses, with the same loud poisoned-holder diagnostics.
    pub(in super::super) fn from_holder(
        runtime: Arc<ConnectionRuntime>,
        peer_addr: Option<SocketAddr>,
        holder: &Arc<Mutex<Option<WebSocket<TcpStream>>>>,
        connection_incarnation: Option<ConnectionIncarnation>,
        settings: &Arc<AcceptorSettings>,
        accepted_at: Instant,
    ) -> Self {
        let built_at = Instant::now();
        let mut socket = match holder.lock() {
            Ok(mut held) => held.take(),
            Err(poisoned) => {
                tracing::error!(
                    peer_addr = ?peer_addr,
                    error = %poisoned,
                    "websocket connection handoff failed: socket holder mutex was poisoned; \
                     the connection process will start without a socket and stop immediately"
                );
                None
            }
        };
        if let Some(socket) = socket.as_mut() {
            super::install_inbound_bound(
                socket,
                PRE_ADMISSION_FRAME_BOUND.min(settings.message_bound),
            );
        }
        let limits = runtime.limits();
        let outbound_capacity = limits.max_connection_outbound_bytes;
        let pending_replies = super::super::pending_reply::PendingReplyTable::new(
            limits.max_pending_replies_per_conversation,
            limits.max_pending_conversation_replies_per_connection,
            std::time::Duration::from_millis(limits.reply_timeout_ms),
        );
        let participant_publication = connection_incarnation.and_then(|_| {
            runtime
                .participant_service()
                .map(crate::server::participant::InstalledParticipantService::new_publication_inbox)
        });
        let state = ConnectionProcessState {
            connection_incarnation,
            participant_publication,
            pending_replies,
            // Design §10: the door that admitted this connection was the
            // WebSocket acceptor, and it says so here — from its own knowledge,
            // before the first inbound message is read.
            mount: crate::server::mount::MountKind::WebSocket,
            ..ConnectionProcessState::default()
        };
        Self {
            runtime,
            peer_addr,
            socket,
            state,
            outbound: WebSocketOutbound::with_capacity(outbound_capacity),
            readiness_token: None,
            keepalive_interval: settings.ping_interval,
            keepalive: None,
            inbound_budget_exhausted: false,
            admitted_bound: settings.message_bound,
            admitted_bound_installed: false,
            start: Some((accepted_at, built_at)),
        }
    }

    /// Runs one connection scheduler slice with the TCP process's exact
    /// ordering: inbound socket/control work, pending replies, the shared
    /// delivery pump, the keepalive schedule, the outbound drain, timers,
    /// readiness arming, and the final probe.
    fn handle_slice(&mut self, pid: u64, ctx: &mut NativeContext<'_>) -> NativeOutcome {
        #[cfg(test)]
        self.runtime.record_slice(pid);
        if !self.runtime.is_registered(pid) {
            return NativeOutcome::Continue;
        }
        if let Err(error) = self.ensure_participant_publication_registered(pid) {
            return self.fail_slice(pid, &error);
        }
        // ACCESS-005 R6: acted on at the connection's own wake too.
        if self.state.revocation.is_some()
            && let Some(outcome) = self.revocation_sweep(pid)
        {
            return outcome;
        }
        if let Some(outcome) = self.force_close_if_pending(pid) {
            return outcome;
        }
        match self.service_socket(pid) {
            SliceStep::Stop(reason) => return NativeOutcome::Stop(reason),
            SliceStep::Continue => {}
        }
        if let Err(error) = self.service_pending_replies() {
            tracing::warn!(
                connection_pid = pid,
                %error,
                "outbound overflow while writing conversation replies; tearing down"
            );
            return self.stop_crashed(pid);
        }
        if let Some(service) = self.runtime.participant_service() {
            let result = service_participant_publications(
                &mut self.state,
                service,
                &mut self.outbound,
                UNIT2_PUSH_SLICE_BUDGET,
            );
            if let Err(error) = result
                && !error.is_capacity_refusal()
            {
                tracing::error!(
                    connection_pid = pid,
                    %error,
                    "participant publication failed; tearing down the connection"
                );
                return self.stop_crashed(pid);
            }
        }
        let slice_budget = self.runtime.limits().delivery_slice_budget;
        match service_subscriptions(&mut self.state, &mut self.outbound, slice_budget) {
            Ok(shed) => self.shed_subscriptions(shed),
            Err(error) => {
                tracing::warn!(
                    connection_pid = pid,
                    %error,
                    "outbound overflow while delivering; tearing down the connection"
                );
                return self.stop_crashed(pid);
            }
        }
        if let Err(error) = self.service_keepalive(pid) {
            tracing::warn!(
                connection_pid = pid,
                %error,
                "keepalive ping write failed; tearing down the connection"
            );
            return self.stop_crashed(pid);
        }
        let drain = match self.drain_outbound(Some(DRAIN_SLICE_BYTES)) {
            Ok(drain) => drain,
            Err(error) => {
                tracing::warn!(
                    connection_pid = pid,
                    %error,
                    "outbound drain failed; tearing down the connection"
                );
                return self.stop_crashed(pid);
            }
        };
        if let Err(error) = self.sync_deadline_timers(pid, ctx) {
            return self.fail_slice(pid, &error);
        }
        if drain == DrainOutcome::Progress {
            return NativeOutcome::Continue;
        }
        if drain == DrainOutcome::Drained && has_held_participant_head(&self.state) {
            return NativeOutcome::Continue;
        }
        let interest = if drain == DrainOutcome::WouldBlockWithResidue {
            Interest::both()
        } else {
            Interest::READABLE
        };
        if let Some(outcome) = self.force_close_if_pending(pid) {
            return outcome;
        }
        if let Err(error) = self.arm_readiness(pid, ctx, interest) {
            return self.fail_slice(pid, &error);
        }
        self.finish_slice_after_arm(pid, ctx)
    }

    fn finish_slice_after_arm(&mut self, pid: u64, ctx: &NativeContext<'_>) -> NativeOutcome {
        #[cfg(test)]
        let barrier_staged = self.runtime.run_pre_wait_barrier();
        match self.final_probe(pid, ctx) {
            Ok(true) => {
                #[cfg(test)]
                if barrier_staged {
                    self.runtime.record_pre_wait_probe_hit();
                }
                NativeOutcome::Continue
            }
            Ok(false) => {
                #[cfg(test)]
                self.runtime.record_park(pid);
                NativeOutcome::Wait
            }
            Err(error) => self.fail_slice(pid, &error),
        }
    }

    /// Services the inbound half of a slice: applies complete binary messages
    /// up to the slice budget, treating transport control per RFC 6455 and
    /// every contract violation as a typed terminal failure.
    fn service_socket(&mut self, pid: u64) -> SliceStep {
        self.inbound_budget_exhausted = false;
        if self.socket.is_none() {
            self.release_conversations();
            self.runtime
                .mark_crashed(pid, ExitReason::Error, self.peer_addr);
            return SliceStep::Stop(ExitReason::Error);
        }
        let mut processed: usize = 0;
        loop {
            if let Some(outcome) = self.force_close_if_pending(pid) {
                return match outcome {
                    NativeOutcome::Stop(reason) => SliceStep::Stop(reason),
                    _ => SliceStep::Continue,
                };
            }
            if processed >= INBOUND_SLICE_MESSAGE_BUDGET {
                // Messages may remain buffered inside tungstenite: re-queue,
                // never park (the final probe consults this flag).
                self.inbound_budget_exhausted = true;
                return SliceStep::Continue;
            }
            let Some(socket) = self.socket.as_mut() else {
                return SliceStep::Stop(ExitReason::Error);
            };
            let read = socket.read();
            if let Some(outcome) = self.force_close_if_pending(pid) {
                return match outcome {
                    NativeOutcome::Stop(reason) => SliceStep::Stop(reason),
                    _ => SliceStep::Continue,
                };
            }
            let message = match read {
                Ok(message) => message,
                Err(tungstenite::Error::Io(error))
                    if error.kind() == std::io::ErrorKind::WouldBlock
                        || error.kind() == std::io::ErrorKind::Interrupted =>
                {
                    return SliceStep::Continue;
                }
                Err(tungstenite::Error::Capacity(
                    tungstenite::error::CapacityError::MessageTooLong { size, max_size },
                )) if !self.admitted_bound_installed => {
                    return self.refuse_pre_admission(
                        pid,
                        PreAdmissionOversize {
                            attempted: size,
                            bound: max_size,
                        },
                    );
                }
                Err(tungstenite::Error::ConnectionClosed) => {
                    // The close handshake completed: the WS-clean equivalent of
                    // a TCP FIN. Mirror the TCP EOF path exactly.
                    if let Err(error) = self.drain_outbound(None) {
                        tracing::debug!(connection_pid = pid, %error, "the close path's best-effort drain did not finish; the peer's half-close ends it");
                    }
                    if let Err(error) =
                        self.complete_connection_fate(ConnectionFateClass::ConnectionLost)
                    {
                        return self.fail_fate(pid, &error);
                    }
                    self.release_conversations();
                    self.runtime.finish(pid);
                    return SliceStep::Stop(ExitReason::Normal);
                }
                Err(error) => {
                    return self.finish_read_error(pid, &error);
                }
            };
            processed = processed.saturating_add(1);
            if let Some(step) = self.handle_inbound_message(pid, message) {
                return step;
            }
        }
    }

    /// The pending-reply servicing both doors run, against the
    /// message-framed outbound queue.
    fn service_pending_replies(&mut self) -> Result<(), OutboundError> {
        let frames = super::super::reply_service::settled_frames(&mut self.state, &self.runtime)
            .map_err(OutboundError::LinkCalls)?;
        for frame in frames {
            self.outbound.enqueue_frame(&frame)?;
        }
        Ok(())
    }

    /// The TCP process's shed handling, byte-for-byte.
    fn shed_subscriptions(&mut self, shed: Vec<u64>) {
        for subscription_id in shed {
            self.state.delivery_seqs.remove(&subscription_id);
            self.state.held_deliveries.remove(&subscription_id);
            if let Some(subscription) = self.state.subscriptions.remove(&subscription_id) {
                if let Err(error) = self.runtime.services().unsubscribe(subscription) {
                    tracing::warn!(
                        subscription_id,
                        %error,
                        "releasing a shed (inbox-overflowed) subscription failed"
                    );
                }
            }
        }
    }

    fn ensure_participant_publication_registered(&mut self, pid: u64) -> Result<(), ServerError> {
        if self.state.participant_publication_registered {
            return Ok(());
        }
        let (Some(incarnation), Some(inbox), Some(service), Some(waker)) = (
            self.state.connection_incarnation,
            self.state.participant_publication.as_ref(),
            self.runtime.participant_service(),
            self.runtime.ready_waker(pid),
        ) else {
            return Ok(());
        };
        service
            .publication_registry()
            .register(incarnation, inbox, waker)
            .map_err(|error| ServerError::ListenerAccept {
                message: format!("participant publication registry failed: {error}"),
            })?;
        self.state.participant_publication_registered = true;
        Ok(())
    }

    fn release_participant_publication(&mut self) {
        if self.state.participant_publication_registered {
            if let (Some(incarnation), Some(service)) = (
                self.state.connection_incarnation,
                self.runtime.participant_service(),
            ) {
                service.publication_registry().deregister(incarnation);
            }
            self.state.participant_publication_registered = false;
        }
    }

    fn drain_outbound(&mut self, budget: Option<usize>) -> Result<DrainOutcome, OutboundError> {
        let Some(socket) = self.socket.as_mut() else {
            return Ok(DrainOutcome::Drained);
        };
        let outcome = self.outbound.drain(socket, budget);
        if self.outbound.frames_flushed() > 0 {
            if let Some((accepted_at, built_at)) = self.start.take() {
                tracing::debug!(
                    peer_addr = ?self.peer_addr,
                    built_ms = u64::try_from(built_at.duration_since(accepted_at).as_millis())
                        .unwrap_or(u64::MAX),
                    first_frame_ms = u64::try_from(accepted_at.elapsed().as_millis()).unwrap_or(u64::MAX),
                    "websocket first frame sent"
                );
            }
        }
        outcome
    }
}

impl NativeHandler for WebSocketConnectionProcess {
    fn handle(&mut self, ctx: &mut NativeContext<'_>) -> NativeOutcome {
        let pid = ctx.self_pid();
        if let Some(outcome) = self.force_close_if_pending(pid) {
            return outcome;
        }
        while let Some(message) = ctx.recv() {
            if let Some(outcome) = self.handle_message(pid, message) {
                return outcome;
            }
        }
        self.runtime.acknowledge_ready(pid);
        self.handle_slice(pid, ctx)
    }
}

impl Drop for WebSocketConnectionProcess {
    fn drop(&mut self) {
        // Backstop for termination paths that never run another handler slice
        // (external termination, scheduler shutdown) — identical to the TCP
        // process's Drop discipline, plus the keepalive timer.
        self.release_conversations();
        let mut timers = self.state.pending_replies.take_retired_timers();
        if let Some(keepalive) = self.keepalive.as_mut() {
            if let Some((timer, _)) = keepalive.armed.take() {
                timers.push(timer);
            }
        }
        self.runtime.cancel_deadline_timers(timers);
    }
}