kshana 0.37.0

// SPDX-License-Identifier: AGPL-3.0-only
//! Sinks for the generated sentences: a file, a TCP server, UDP unicast or broadcast,
//! paced in real time, accelerated, or as fast as possible.
//!
//! Only text leaves this module. This is the one place the crate reads a monotonic clock,
//! to pace the stream; the generated bytes never depend on it.

use super::gen::Epoch;
use std::io::{self, Write};
use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, Instant};

/// Where epochs go.
pub trait Sink {
    /// Deliver one epoch's sentences.
    fn send(&mut self, epoch: &Epoch) -> io::Result<()>;
    /// Flush at the end of the run.
    fn finish(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A file (or any writer): CRLF-terminated sentences.
pub struct WriterSink<W: Write>(pub W);

impl<W: Write> Sink for WriterSink<W> {
    fn send(&mut self, epoch: &Epoch) -> io::Result<()> {
        for l in &epoch.lines {
            self.0.write_all(l.as_bytes())?;
            self.0.write_all(b"\r\n")?;
        }
        Ok(())
    }
    fn finish(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

/// A TCP server: every connected client receives every epoch from its connection on.
pub struct TcpServer {
    listener_addr: SocketAddr,
    clients: Arc<Mutex<Vec<TcpStream>>>,
    stop: Arc<AtomicBool>,
    write_timeout: Arc<Mutex<Duration>>,
}

/// How long a write to one client may take before that client is dropped.
pub const DEFAULT_WRITE_TIMEOUT: Duration = Duration::from_secs(2);

impl TcpServer {
    /// Bind and start accepting in the background.
    pub fn bind(addr: impl ToSocketAddrs) -> io::Result<Self> {
        let listener = TcpListener::bind(addr)?;
        listener.set_nonblocking(true)?;
        let listener_addr = listener.local_addr()?;
        let clients: Arc<Mutex<Vec<TcpStream>>> = Arc::default();
        let stop = Arc::new(AtomicBool::new(false));
        let write_timeout = Arc::new(Mutex::new(DEFAULT_WRITE_TIMEOUT));
        let (c2, s2, w2) = (clients.clone(), stop.clone(), write_timeout.clone());
        std::thread::spawn(move || {
            while !s2.load(Ordering::Relaxed) {
                match listener.accept() {
                    Ok((s, _)) => {
                        let _ = s.set_nodelay(true);
                        let _ = s.set_nonblocking(false);
                        let to = *w2.lock().unwrap_or_else(PoisonError::into_inner);
                        let _ = s.set_write_timeout(Some(to));
                        c2.lock().unwrap_or_else(PoisonError::into_inner).push(s);
                    }
                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                        std::thread::sleep(Duration::from_millis(10))
                    }
                    Err(_) => break,
                }
            }
        });
        Ok(Self {
            listener_addr,
            clients,
            stop,
            write_timeout,
        })
    }

    /// Set how long a write to one client may take before the client is dropped (for
    /// clients that connect after this call; default [`DEFAULT_WRITE_TIMEOUT`]).
    pub fn set_write_timeout(&self, timeout: Duration) {
        *self
            .write_timeout
            .lock()
            .unwrap_or_else(PoisonError::into_inner) = timeout;
    }

    /// The bound address (useful after binding port 0).
    pub fn local_addr(&self) -> SocketAddr {
        self.listener_addr
    }

    /// Block until `n` clients are connected, or `timeout` passes.
    pub fn wait_for_clients(&self, n: usize, timeout: Option<Duration>) -> bool {
        let t0 = Instant::now();
        loop {
            if self
                .clients
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .len()
                >= n
            {
                return true;
            }
            if timeout.is_some_and(|t| t0.elapsed() > t) {
                return false;
            }
            std::thread::sleep(Duration::from_millis(20));
        }
    }
}

impl Drop for TcpServer {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
    }
}

impl Sink for TcpServer {
    fn send(&mut self, epoch: &Epoch) -> io::Result<()> {
        let mut block = String::new();
        for l in &epoch.lines {
            block.push_str(l);
            block.push_str("\r\n");
        }
        // The client list is taken out of the lock while writing, so a slow client cannot
        // hold up the accept thread. A client that went away, or that does not take a
        // write within the timeout, is dropped; the stream carries on for the others.
        let mut current =
            std::mem::take(&mut *self.clients.lock().unwrap_or_else(PoisonError::into_inner));
        current.retain_mut(|c| c.write_all(block.as_bytes()).is_ok());
        let mut shared = self.clients.lock().unwrap_or_else(PoisonError::into_inner);
        current.append(&mut shared);
        *shared = current;
        Ok(())
    }
    fn finish(&mut self) -> io::Result<()> {
        for c in self
            .clients
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .iter_mut()
        {
            let _ = c.flush();
        }
        Ok(())
    }
}

/// UDP, one sentence to a datagram (unicast, or broadcast when asked).
pub struct UdpSink {
    socket: UdpSocket,
    dest: SocketAddr,
}

impl UdpSink {
    /// Send to `dest`; `broadcast` sets SO_BROADCAST.
    pub fn new(dest: impl ToSocketAddrs, broadcast: bool) -> io::Result<Self> {
        let dest = dest
            .to_socket_addrs()?
            .next()
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "no UDP address"))?;
        let socket = UdpSocket::bind(if dest.is_ipv4() {
            "0.0.0.0:0"
        } else {
            "[::]:0"
        })?;
        if broadcast {
            socket.set_broadcast(true)?;
        }
        Ok(Self { socket, dest })
    }
}

impl Sink for UdpSink {
    fn send(&mut self, epoch: &Epoch) -> io::Result<()> {
        for l in &epoch.lines {
            let mut d = l.clone().into_bytes();
            d.extend_from_slice(b"\r\n");
            self.socket.send_to(&d, self.dest)?;
        }
        Ok(())
    }
}

/// How fast simulated time runs against the wall clock.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Pace {
    /// As fast as the sinks take it.
    Max,
    /// `x` simulated seconds per wall-clock second (1 is real time).
    Speed(f64),
}

/// Write every epoch to every sink, paced.
pub fn run<'a>(epochs: &[Epoch], sinks: &mut [Box<dyn Sink + 'a>], pace: Pace) -> io::Result<()> {
    let t0 = Instant::now();
    for e in epochs {
        if let Pace::Speed(x) = pace {
            let due = Duration::from_secs_f64(e.t_s / x);
            if let Some(wait) = due.checked_sub(t0.elapsed()) {
                std::thread::sleep(wait);
            }
        }
        for s in sinks.iter_mut() {
            s.send(e)?;
        }
    }
    for s in sinks.iter_mut() {
        s.finish()?;
    }
    Ok(())
}