//! 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, /// 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 { serde_json::from_slice::(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 { 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 (`/`). 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, 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 = 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, _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 { self.rx.recv().await.ok() } /// Return the next available message without waiting. pub fn try_recv(&self) -> Option { 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) { 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()); } }