sausage/examples/fetch/src/main.rs

//! Fetches URLs over HTTP(S) and writes the body to /public/fetch/<n>.
//! Reads stdin and runs requests at the same time: while a fetch runs,
//! `ping` still answers at once.
use serde_json::{Value, json};
use std::cell::RefCell;
use std::rc::Rc;
use wstd::http::{Body, Client, Request};
use wstd::io::{AsyncRead, AsyncWrite};

type Out = Rc<RefCell<Vec<String>>>;

async fn flush(out: &Out) {
    let lines: Vec<String> = out.borrow_mut().drain(..).collect();
    let mut stdout = wstd::io::stdout();
    for l in lines {
        let _ = stdout.write_all(format!("{l}\n").as_bytes()).await;
    }
    let _ = stdout.flush().await;
}

async fn fetch(n: u64, msg: Value, out: Out) {
    let url = msg["body"]["url"].as_str().unwrap_or_default().to_string();
    let started = wstd::time::Instant::now();
    let result = async {
        let req = Request::get(&url).body(Body::empty()).map_err(|e| e.to_string())?;
        let resp = Client::new().send(req).await.map_err(|e| format!("{e:?}"))?;
        let status = resp.status().as_u16();
        let bytes = resp.into_body().contents().await.map_err(|e| format!("{e:?}"))?.to_vec();
        std::fs::create_dir_all("/public/fetch").map_err(|e| e.to_string())?;
        let path = format!("/public/fetch/{n}");
        std::fs::write(&path, &bytes).map_err(|e| e.to_string())?;
        Ok::<_, String>((status, bytes.len(), path))
    }
    .await;
    let ms = started.elapsed().as_millis();
    let body = match result {
        Ok((status, size, path)) => {
            let id = std::env::var("SAUSAGE_ACTOR").unwrap_or("fetch".into());
            json!({"type": "fetched", "status": status, "bytes": size, "path": path.replacen("/public", &format!("/actor/{id}/public"), 1), "ms": ms})
        }
        Err(e) => json!({"type": "error", "reason": e, "ms": ms}),
    };
    out.borrow_mut().push(json!({"to": msg["from"], "reply_to": msg["id"], "body": body}).to_string());
    flush(&out).await;
}

fn main() {
    wstd::runtime::block_on(async {
        let out: Out = Rc::new(RefCell::new(Vec::new()));
        let mut stdin = wstd::io::stdin();
        let mut buf = Vec::new();
        let mut chunk = vec![0u8; 8192];
        let mut n = 0u64;
        loop {
            let Ok(len) = stdin.read(&mut chunk).await else { break };
            if len == 0 {
                break;
            }
            buf.extend_from_slice(&chunk[..len]);
            while let Some(pos) = buf.iter().position(|b| *b == b'\n') {
                let line: Vec<u8> = buf.drain(..=pos).collect();
                let Ok(msg) = serde_json::from_slice::<Value>(&line) else { continue };
                match msg["body"]["type"].as_str() {
                    Some("shutdown") => return,
                    Some("fetch") => {
                        n += 1;
                        wstd::runtime::spawn(fetch(n, msg, out.clone())).detach();
                    }
                    Some("ping") => {
                        out.borrow_mut().push(json!({"to": msg["from"], "reply_to": msg["id"], "body": {"type": "pong"}}).to_string());
                        flush(&out).await;
                    }
                    _ => {}
                }
            }
        }
    });
}