sausage/sdk/src/lib.rs

//! Small runtime helpers for Sausage actors.

pub use serde_json::{Value, json};

use serde::{Deserialize, Serialize};
use std::io::BufRead;

/// A message that an actor receives from the broker.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct Message {
    /// The message ID.
    pub id: String,
    /// The ID of the sending actor.
    pub from: String,
    /// The ID of the message that this message answers.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reply_to: Option<String>,
    /// The message body.
    pub body: Value,
}

impl Message {
    /// Return the body's `type` string, or an empty string when it is missing.
    pub fn kind(&self) -> &str {
        self.body["type"].as_str().unwrap_or("")
    }
}

fn parse(line: &[u8]) -> Option<Message> {
    serde_json::from_slice::<Message>(line).ok()
}

fn is_shutdown(message: &Message) -> bool {
    message.from == "broker" && message.kind() == "shutdown"
}

/// Return this actor's ID from `SAUSAGE_ACTOR`, or an empty string if unset.
pub fn me() -> String {
    std::env::var("SAUSAGE_ACTOR").unwrap_or_default()
}

/// Read messages from stdin until EOF or a shutdown message from the broker.
///
/// Invalid JSON lines are skipped.
pub fn messages() -> impl Iterator<Item = Message> {
    let stdin = std::io::stdin();
    let mut lines = stdin.lock().lines();
    std::iter::from_fn(move || {
        loop {
            let Ok(line) = lines.next()? else {
                continue;
            };
            let Some(message) = parse(line.as_bytes()) else {
                continue;
            };
            return (!is_shutdown(&message)).then_some(message);
        }
    })
}

/// Send a message to `to` and return its full ID (`<actor>/<own-id>`).
pub fn send(to: &str, body: Value) -> String {
    let id = next_id();
    write_message(json!({"to": to, "id": id, "body": body}));
    format!("{}/{id}", me())
}

/// Reply to `message` and return the new message's full ID.
pub fn reply(message: &Message, body: Value) -> String {
    let id = next_id();
    write_message(json!({
        "to": message.from,
        "id": id,
        "reply_to": message.id,
        "body": body,
    }));
    format!("{}/{id}", me())
}

/// Write a file so that readers see either the old or the new content,
/// never half of it: write a temporary file next to `path`, then rename it.
pub fn write_atomic(path: impl AsRef<std::path::Path>, contents: impl AsRef<[u8]>) -> std::io::Result<()> {
    let path = path.as_ref();
    let tmp = path.with_file_name(format!(".{}.tmp-{}", path.file_name().unwrap_or_default().to_string_lossy(), next_id()));
    std::fs::write(&tmp, contents)?;
    std::fs::rename(&tmp, path)
}

/// Publish a message to this actor's subscribers and return its full ID.
pub fn publish(body: Value) -> String {
    let id = next_id();
    write_message(json!({"id": id, "body": body}));
    format!("{}/{id}", me())
}

fn write_message(message: Value) {
    use std::io::Write;

    let mut stdout = std::io::stdout().lock();
    writeln!(stdout, "{message}").unwrap();
    stdout.flush().unwrap();
}

fn next_id() -> String {
    use std::sync::OnceLock;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::{SystemTime, UNIX_EPOCH};

    static START: OnceLock<u128> = OnceLock::new();
    static SEQUENCE: AtomicUsize = AtomicUsize::new(0);
    let start = START.get_or_init(|| {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos()
    });
    format!("{start}-{}", SEQUENCE.fetch_add(1, Ordering::Relaxed))
}

/// An asynchronous inbox for actors that use HTTP or timers.
#[cfg(feature = "async")]
pub struct Inbox {
    rx: async_channel::Receiver<Message>,
    _reader: wstd::runtime::Task<()>,
}

#[cfg(feature = "async")]
impl Inbox {
    /// Start the stdin reader inside the wstd runtime and return its inbox.
    /// The inbox owns the reader task and must be dropped to cancel it.
    pub fn start() -> Self {
        let (tx, rx) = async_channel::unbounded();
        let reader = wstd::runtime::spawn(read_messages(tx));
        Self { rx, _reader: reader }
    }

    /// Wait for the next message, or return `None` at EOF or shutdown.
    pub async fn recv(&self) -> Option<Message> {
        self.rx.recv().await.ok()
    }

    /// Return the next available message without waiting.
    pub fn try_recv(&self) -> Option<Message> {
        self.rx.try_recv().ok()
    }

    /// Return whether the reader closed the inbox, even when messages remain buffered.
    pub fn is_closed(&self) -> bool {
        self.rx.is_closed()
    }
}

#[cfg(feature = "async")]
async fn read_messages(tx: async_channel::Sender<Message>) {
    use wstd::io::AsyncRead;

    let mut stdin = wstd::io::stdin();
    let mut buf = Vec::new();
    let mut chunk = [0; 8192];
    loop {
        let Ok(n) = stdin.read(&mut chunk).await else {
            break;
        };
        if n == 0 {
            break;
        }
        buf.extend_from_slice(&chunk[..n]);
        while let Some(pos) = buf.iter().position(|b| *b == b'\n') {
            let message = parse(&buf[..pos]);
            buf.drain(..=pos);
            let Some(message) = message else {
                continue;
            };
            if is_shutdown(&message) {
                tx.close();
                return;
            }
            if tx.send(message).await.is_err() {
                return;
            }
        }
    }
    tx.close();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn kind_reads_body_type() {
        let message: Message = serde_json::from_value(json!({
            "id": "broker/1",
            "from": "broker",
            "body": {"type": "ping"}
        }))
        .unwrap();
        assert_eq!(message.kind(), "ping");
        assert_eq!(message.reply_to, None);
    }

    #[test]
    fn kind_is_empty_when_type_is_missing() {
        let message: Message = serde_json::from_value(json!({
            "id": "broker/1",
            "from": "broker",
            "body": {}
        }))
        .unwrap();
        assert_eq!(message.kind(), "");
    }

    #[test]
    fn write_atomic_replaces_the_file() {
        let dir = std::env::temp_dir().join(format!("sausage-actor-{}", next_id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("state.json");
        write_atomic(&path, "old").unwrap();
        write_atomic(&path, "new").unwrap();
        assert_eq!(std::fs::read_to_string(&path).unwrap(), "new");
        assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1);
        std::fs::remove_dir_all(dir).unwrap();
    }

    #[test]
    fn generated_ids_differ() {
        assert_ne!(next_id(), next_id());
    }
}