aion_integrations/jsonrpc/
transport.rs

1//! Newline-delimited JSON-RPC 2.0 framing over an async duplex, with a single serializing
2//! writer and request-id correlation.
3//!
4//! The one net-new piece over the in-tree prior art (§9.4): because this channel is
5//! **bidirectional**, responses and outbound notifications could interleave and corrupt a frame.
6//! [`JsonRpcConnection`] serialises all writes through a single [`tokio::sync::Mutex`]-guarded
7//! writer so every frame is emitted atomically, and allocates monotonic request ids from a second
8//! guarded counter so an adapter can correlate a response to the request it sent.
9
10use std::sync::atomic::{AtomicU64, Ordering};
11
12use serde::Serialize;
13use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, Lines};
14use tokio::sync::Mutex;
15
16use super::envelope::{
17    IncomingMessage, JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse,
18};
19use crate::error::HarnessError;
20
21/// A framed JSON-RPC 2.0 connection over one async read half and one async write half.
22///
23/// Generic over any [`AsyncRead`] + [`AsyncWrite`] pair — a child's stdout/stdin, an in-memory
24/// duplex in tests, or a socket. The read and write halves are independent, so a caller may read
25/// inbound frames on one task while another task writes outbound frames concurrently: writes are
26/// serialised internally, and the id allocator is shared, so both are safe to use behind a shared
27/// reference (e.g. an [`std::sync::Arc`]).
28pub struct JsonRpcConnection<R, W> {
29    reader: Mutex<Lines<BufReader<R>>>,
30    writer: Mutex<W>,
31    next_id: AtomicU64,
32}
33
34impl<R, W> JsonRpcConnection<R, W>
35where
36    R: AsyncRead + Unpin + Send,
37    W: AsyncWrite + Unpin + Send,
38{
39    /// Wraps a read half and a write half into a framed connection.
40    #[must_use]
41    pub fn new(read_half: R, write_half: W) -> Self {
42        Self {
43            reader: Mutex::new(BufReader::new(read_half).lines()),
44            writer: Mutex::new(write_half),
45            next_id: AtomicU64::new(1),
46        }
47    }
48
49    /// Allocates the next monotonic request id.
50    ///
51    /// Ids are process-local to this connection and used only to correlate a response to the
52    /// request that produced it, so a plain monotonic counter is correct here.
53    #[must_use]
54    pub fn next_request_id(&self) -> JsonRpcId {
55        JsonRpcId::number(self.next_id.fetch_add(1, Ordering::Relaxed))
56    }
57
58    /// Writes one JSON-RPC request as a newline-delimited frame.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`HarnessError::Protocol`] if the request cannot be serialised, or
63    /// [`HarnessError::Transport`] if the underlying write fails.
64    pub async fn send_request(&self, request: &JsonRpcRequest) -> Result<(), HarnessError> {
65        self.write_frame(request).await
66    }
67
68    /// Writes one JSON-RPC notification as a newline-delimited frame.
69    ///
70    /// # Errors
71    ///
72    /// Returns [`HarnessError::Protocol`] if the notification cannot be serialised, or
73    /// [`HarnessError::Transport`] if the underlying write fails.
74    pub async fn send_notification(
75        &self,
76        notification: &JsonRpcNotification,
77    ) -> Result<(), HarnessError> {
78        self.write_frame(notification).await
79    }
80
81    /// Writes one JSON-RPC response as a newline-delimited frame.
82    ///
83    /// # Errors
84    ///
85    /// Returns [`HarnessError::Protocol`] if the response cannot be serialised, or
86    /// [`HarnessError::Transport`] if the underlying write fails.
87    pub async fn send_response(&self, response: &JsonRpcResponse) -> Result<(), HarnessError> {
88        self.write_frame(response).await
89    }
90
91    /// Reads the next inbound frame, classifying it into an [`IncomingMessage`].
92    ///
93    /// Returns `Ok(None)` at end of stream (the peer closed its write half). Blank lines are
94    /// skipped, matching the newline-delimited framing convention.
95    ///
96    /// # Errors
97    ///
98    /// Returns [`HarnessError::Transport`] on an I/O failure and [`HarnessError::Protocol`] on a
99    /// frame that is not a valid JSON-RPC message.
100    pub async fn recv(&self) -> Result<Option<IncomingMessage>, HarnessError> {
101        let mut reader = self.reader.lock().await;
102        loop {
103            let line = reader
104                .next_line()
105                .await
106                .map_err(|source| HarnessError::transport(format!("read failed: {source}")))?;
107            let Some(line) = line else {
108                return Ok(None);
109            };
110            let trimmed = line.trim();
111            if trimmed.is_empty() {
112                continue;
113            }
114            let value: serde_json::Value = serde_json::from_str(trimmed).map_err(|source| {
115                HarnessError::protocol(format!("invalid JSON frame: {source}"))
116            })?;
117            let message = IncomingMessage::from_value(value).map_err(|source| {
118                HarnessError::protocol(format!("frame is not a JSON-RPC message: {source}"))
119            })?;
120            return Ok(Some(message));
121        }
122    }
123
124    /// Serialises `frame` to a single line and writes it atomically through the guarded writer.
125    async fn write_frame<T: Serialize>(&self, frame: &T) -> Result<(), HarnessError> {
126        let mut encoded = serde_json::to_vec(frame).map_err(|source| {
127            HarnessError::protocol(format!("frame cannot be encoded: {source}"))
128        })?;
129        encoded.push(b'\n');
130        let mut writer = self.writer.lock().await;
131        writer
132            .write_all(&encoded)
133            .await
134            .map_err(|source| HarnessError::transport(format!("write failed: {source}")))?;
135        writer
136            .flush()
137            .await
138            .map_err(|source| HarnessError::transport(format!("flush failed: {source}")))?;
139        Ok(())
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use tokio::io::duplex;
146
147    use super::super::envelope::{IncomingMessage, JsonRpcNotification, JsonRpcRequest};
148    use super::JsonRpcConnection;
149    use crate::error::HarnessError;
150
151    #[tokio::test]
152    async fn request_and_notification_round_trip_over_a_duplex() -> Result<(), HarnessError> {
153        // A loopback pair: host writes into `host`, child reads from `child`, and vice versa.
154        let (host_io, child_io) = duplex(4096);
155        let (host_read, host_write) = tokio::io::split(host_io);
156        let (child_read, child_write) = tokio::io::split(child_io);
157        let host = JsonRpcConnection::new(host_read, host_write);
158        let child = JsonRpcConnection::new(child_read, child_write);
159
160        let id = host.next_request_id();
161        let request = JsonRpcRequest::new(id.clone(), "run/execute", None);
162        host.send_request(&request).await?;
163
164        let received = child
165            .recv()
166            .await?
167            .ok_or_else(|| HarnessError::protocol("expected a frame, got end-of-stream"))?;
168        assert!(
169            matches!(&received, IncomingMessage::Request(got) if got.id == id && got.method == "run/execute"),
170            "expected the id-matched run/execute request, got {received:?}"
171        );
172
173        // The child replies with a notification (no id): it must classify as a notification.
174        let notification = JsonRpcNotification::new("event/stop", None);
175        child.send_notification(&notification).await?;
176        let echoed = host
177            .recv()
178            .await?
179            .ok_or_else(|| HarnessError::protocol("expected a frame, got end-of-stream"))?;
180        assert!(
181            matches!(echoed, IncomingMessage::Notification(_)),
182            "a frame with no id must classify as a notification"
183        );
184        Ok(())
185    }
186
187    #[tokio::test]
188    async fn ids_are_monotonic() {
189        let (a, _b) = duplex(64);
190        let (read, write) = tokio::io::split(a);
191        let connection = JsonRpcConnection::new(read, write);
192        let first = connection.next_request_id();
193        let second = connection.next_request_id();
194        assert_ne!(first, second);
195    }
196
197    #[tokio::test]
198    async fn recv_returns_none_at_end_of_stream() -> Result<(), HarnessError> {
199        let (host_io, child_io) = duplex(64);
200        let (host_read, host_write) = tokio::io::split(host_io);
201        let host = JsonRpcConnection::new(host_read, host_write);
202        // Dropping the child's io closes the write half the host reads from.
203        drop(child_io);
204        let end = host.recv().await?;
205        assert!(end.is_none(), "closed stream yields None");
206        Ok(())
207    }
208}
aion_integrations/jsonrpc/transport.rs

aion_integrations/jsonrpc/
transport.rs
