rpi_extensions/
stdio.rs

1//! Exclusive stdio ownership and completion, independent of the served protocol.
2use std::sync::{Arc, Mutex};
3
4use serde_json::{json, Value};
5use tokio::sync::Notify;
6
7#[derive(Default)]
8struct State {
9    owner: Option<(u64, String)>,
10    exit_code: Option<i32>,
11    conflict: Option<String>,
12}
13
14#[derive(Clone, Default)]
15pub struct StdioLifecycle {
16    state: Arc<Mutex<State>>,
17    changed: Arc<Notify>,
18}
19
20impl StdioLifecycle {
21    pub(crate) fn claim(&self, id: u64, name: &str) -> Result<(), String> {
22        let mut state = self.state.lock().unwrap();
23        match &state.owner {
24            Some((owner, _)) if *owner == id => Ok(()),
25            Some((_, owner)) => {
26                let error =
27                    format!("stdio is already owned by `{owner}`; `{name}` cannot claim it");
28                state.conflict = Some(error.clone());
29                Err(error)
30            }
31            None => {
32                state.owner = Some((id, name.to_owned()));
33                Ok(())
34            }
35        }
36    }
37
38    pub(crate) fn release_failed_registration(&self, id: u64) {
39        let mut state = self.state.lock().unwrap();
40        if state.owner.as_ref().is_some_and(|(owner, _)| *owner == id) {
41            state.owner = None;
42            state.exit_code = None;
43        }
44    }
45
46    pub fn validate(&self) -> Result<(), String> {
47        self.state
48            .lock()
49            .unwrap()
50            .conflict
51            .clone()
52            .map_or(Ok(()), Err)
53    }
54
55    pub(crate) fn handle(&self, id: u64, args: Value) -> Result<Value, String> {
56        let mut state = self.state.lock().unwrap();
57        match args.get("op").and_then(Value::as_str) {
58            Some("status") => Ok(
59                json!({"owner": state.owner.as_ref().map(|(_, name)| name), "exitCode": state.exit_code}),
60            ),
61            Some("finish") => {
62                if !state.owner.as_ref().is_some_and(|(owner, _)| *owner == id) {
63                    return Err("only the stdio owner may finish its lifecycle".to_owned());
64                }
65                let code = args
66                    .get("exitCode")
67                    .and_then(Value::as_i64)
68                    .and_then(|code| i32::try_from(code).ok())
69                    .ok_or("finish requires an i32 exitCode")?;
70                if state.exit_code.is_some_and(|previous| previous != code) {
71                    return Err(
72                        "stdio lifecycle already finished with a different exitCode".to_owned()
73                    );
74                }
75                state.exit_code = Some(code);
76                drop(state);
77                self.changed.notify_waiters();
78                Ok(json!({"ok": true}))
79            }
80            _ => Err("stdio lifecycle op must be status or finish".to_owned()),
81        }
82    }
83
84    /// Completion is retained, including when a plugin finishes before the host waits.
85    pub async fn wait(&self) -> i32 {
86        loop {
87            let changed = self.changed.notified();
88            tokio::pin!(changed);
89            changed.as_mut().enable();
90            if let Some(code) = self.state.lock().unwrap().exit_code {
91                return code;
92            }
93            changed.await;
94        }
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[tokio::test]
103    async fn owner_finishes_before_or_after_host_waits() {
104        let lifecycle = StdioLifecycle::default();
105        lifecycle.claim(1, "test-server").unwrap();
106        assert!(lifecycle
107            .handle(2, json!({"op":"finish","exitCode":0}))
108            .is_err());
109        let waiting = tokio::spawn({
110            let lifecycle = lifecycle.clone();
111            async move { lifecycle.wait().await }
112        });
113        lifecycle
114            .handle(1, json!({"op":"finish","exitCode":7}))
115            .unwrap();
116        assert_eq!(waiting.await.unwrap(), 7);
117        assert_eq!(lifecycle.wait().await, 7);
118    }
119
120    #[test]
121    fn duplicate_names_still_have_exclusive_owners() {
122        let lifecycle = StdioLifecycle::default();
123        lifecycle.claim(1, "server").unwrap();
124        lifecycle.claim(1, "server").unwrap();
125        assert!(lifecycle.claim(2, "server").is_err());
126        assert!(lifecycle.validate().is_err());
127    }
128
129    #[test]
130    fn failed_registration_releases_its_claim() {
131        let lifecycle = StdioLifecycle::default();
132        lifecycle.claim(1, "failed").unwrap();
133        lifecycle.release_failed_registration(1);
134        lifecycle.claim(2, "working").unwrap();
135        assert!(lifecycle.validate().is_ok());
136    }
137}
rpi_extensions/stdio.rs

rpi_extensions/
stdio.rs
