escouade_plugin_kit/serve/
process.rs

1//! The loop in process: on the standard input and output, on tokio's threads.
2
3use super::{arrived, called, Cx, Plugin};
4use std::sync::Arc;
5use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
6use tokio::sync::mpsc;
7
8/// Runs `plugin` on `input` and `output` until `input` ends or `shutdown` is answered.
9pub async fn serve_on<P, R, W>(plugin: Arc<P>, input: R, mut output: W) -> std::io::Result<()>
10where
11    P: Plugin,
12    R: tokio::io::AsyncRead + Unpin,
13    W: AsyncWrite + Unpin + Send + 'static,
14{
15    let (tx, mut rx) = mpsc::unbounded_channel::<String>();
16    let writer = tokio::spawn(async move {
17        while let Some(line) = rx.recv().await {
18            if output.write_all(line.as_bytes()).await.is_err()
19                || output.write_all(b"\n").await.is_err()
20                || output.flush().await.is_err()
21            {
22                break;
23            }
24        }
25    });
26    let out = tx.clone();
27    let cx = Arc::new(Cx::new(move |line| out.send(line).is_ok()));
28    let (stop_tx, mut stop_rx) = mpsc::unbounded_channel::<()>();
29    let mut lines = BufReader::new(input).lines();
30    loop {
31        let line = tokio::select! {
32            line = lines.next_line() => line?,
33            _ = stop_rx.recv() => break,
34        };
35        let Some(line) = line else { break };
36        let Some(msg) = arrived(&cx, &line) else {
37            continue;
38        };
39        let (plugin, cx, tx, stop) = (plugin.clone(), cx.clone(), tx.clone(), stop_tx.clone());
40        tokio::spawn(async move {
41            let (out, shut) = called(&*plugin, &cx, msg).await;
42            if let Some(out) = out {
43                let _ = tx.send(out.to_string());
44            }
45            if shut {
46                let _ = stop.send(());
47            }
48        });
49    }
50    drop(tx);
51    drop(cx);
52    // What was answered is written before the plugin ends.
53    let _ = tokio::time::timeout(std::time::Duration::from_secs(2), writer).await;
54    Ok(())
55}
56
57/// The whole life of a plugin of `escouade-plugins`: its language, then its loop on the standard input and output, on
58/// a runtime of its own. Its exit code.
59pub fn run<P: Plugin>(name: &str, plugin: P) -> i32 {
60    crate::i18n::from_env();
61    let runtime = match tokio::runtime::Builder::new_multi_thread()
62        .worker_threads(2)
63        .enable_all()
64        .build()
65    {
66        Ok(r) => r,
67        Err(e) => {
68            eprintln!("escouade-plugins {name}: {e}");
69            return 1;
70        }
71    };
72    let result = runtime.block_on(serve_on(
73        Arc::new(plugin),
74        tokio::io::stdin(),
75        tokio::io::stdout(),
76    ));
77    // A call still running (a slow host) does not keep the process.
78    runtime.shutdown_timeout(std::time::Duration::from_millis(200));
79    match result {
80        Ok(()) => 0,
81        Err(e) => {
82            eprintln!("escouade-plugins {name}: {e}");
83            1
84        }
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::forge::BoxFuture;
92    use escouade_plugin_protocol::{code, Refusal};
93    use serde_json::{json, Value};
94
95    struct Echo;
96
97    impl Plugin for Echo {
98        fn call<'a>(
99            &'a self,
100            cx: &'a Cx,
101            method: &'a str,
102            params: Value,
103        ) -> BoxFuture<'a, Result<Value, Refusal>> {
104            Box::pin(async move {
105                match method {
106                    "echo" => Ok(params),
107                    "slow" => {
108                        crate::time::sleep(std::time::Duration::from_millis(300)).await;
109                        Ok(json!("slow"))
110                    }
111                    "secret" => Ok(json!(cx.secret("account:x").await?)),
112                    other => Err(Refusal::not_found(other)),
113                }
114            })
115        }
116    }
117
118    /// The plugin's lines, as Escouade reads them, and what Escouade writes to it.
119    async fn session() -> (
120        tokio::io::DuplexStream,
121        tokio::io::Lines<BufReader<tokio::io::DuplexStream>>,
122    ) {
123        let (to_plugin, plugin_in) = tokio::io::duplex(1 << 16);
124        let (plugin_out, from_plugin) = tokio::io::duplex(1 << 16);
125        tokio::spawn(serve_on(Arc::new(Echo), plugin_in, plugin_out));
126        (to_plugin, BufReader::new(from_plugin).lines())
127    }
128
129    async fn next(lines: &mut tokio::io::Lines<BufReader<tokio::io::DuplexStream>>) -> Value {
130        serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap()
131    }
132
133    #[tokio::test]
134    async fn a_slow_call_does_not_hold_back_the_next() {
135        let (mut w, mut r) = session().await;
136        w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"slow\"}\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"echo\",\"params\":{\"a\":1}}\n")
137            .await
138            .unwrap();
139        let first = next(&mut r).await;
140        assert_eq!(first["id"], 2, "{first}");
141        assert_eq!(next(&mut r).await["id"], 1);
142    }
143
144    #[tokio::test]
145    async fn the_plugin_asks_escouade_while_it_answers() {
146        let (mut w, mut r) = session().await;
147        w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"secret\"}\n")
148            .await
149            .unwrap();
150        let ask = next(&mut r).await;
151        assert_eq!(ask["method"], "secrets/get");
152        assert_eq!(ask["params"]["key"], "account:x");
153        let reply = json!({ "jsonrpc": "2.0", "id": ask["id"], "result": "ghp_x" });
154        w.write_all(format!("{reply}\n").as_bytes()).await.unwrap();
155        let done = next(&mut r).await;
156        assert_eq!(
157            (done["id"].clone(), done["result"].clone()),
158            (json!(7), json!("ghp_x"))
159        );
160    }
161
162    #[tokio::test]
163    async fn the_handshake_garbage_and_shutdown() {
164        let (mut w, mut r) = session().await;
165        w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}\nnot json\n[{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"echo\",\"params\":1},{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"nope\"}]\n")
166            .await
167            .unwrap();
168        // Answered as each ends: in any order.
169        let mut got = [next(&mut r).await, next(&mut r).await, next(&mut r).await];
170        got.sort_by_key(|v| (v.is_array(), v["id"].as_u64()));
171        assert_eq!(got[0]["error"]["code"], code::PARSE);
172        assert_eq!(got[1]["result"]["protocolVersion"], 1);
173        let batch = got[2].clone();
174        assert_eq!(batch[0]["result"], 1);
175        assert_eq!(batch[1]["error"]["code"], code::METHOD_NOT_FOUND);
176        w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":9,\"method\":\"shutdown\"}\n")
177            .await
178            .unwrap();
179        assert_eq!(next(&mut r).await["id"], 9);
180        assert!(
181            r.next_line().await.unwrap().is_none(),
182            "it ends after shutdown"
183        );
184    }
185}
escouade_plugin_kit/serve/process.rs

escouade_plugin_kit/serve/
process.rs
