karyon_jsonrpc/server/
quic.rs

1//! QUIC transport: one request per stream, with separate
2//! notification streaming for pubsub.
3
4use std::sync::Arc;
5
6use log::{debug, error};
7
8use karyon_core::async_util::{select, Either, TaskResult};
9
10use karyon_net::{
11    framed,
12    quic::{QuicConn, QuicIncoming},
13    FramedConn, StreamMux,
14};
15
16use crate::{
17    codec::JsonCodec,
18    error::Result,
19    message,
20    server::{
21        channel::Channel,
22        dispatch::{sanity_check, Handler, NewRequest, SanityCheckResult},
23        Server, CHANNEL_SUBSCRIPTION_BUFFER_SIZE,
24    },
25};
26
27impl Server {
28    /// Run the handshake off the accept loop, then serve the connection.
29    pub(super) fn handle_quic_incoming(self: &Arc<Self>, incoming: QuicIncoming) {
30        let peer = incoming.peer_endpoint();
31        let server = self.clone();
32        self.task_group.spawn(async move {
33            // Bounded by the QUIC idle timeout (QuicConfig::idle_timeout).
34            match incoming.handshake().await {
35                Ok(quic_conn) => server.handle_quic_conn(quic_conn),
36                Err(err) => debug!("QUIC handshake with {peer} failed: {err}"),
37            }
38        });
39    }
40
41    /// Serve a QUIC connection; each incoming stream is handled
42    /// as an independent request.
43    fn handle_quic_conn(self: &Arc<Self>, quic_conn: QuicConn) {
44        let peer = quic_conn.peer_endpoint().ok();
45        debug!("Handle QUIC connection {peer:?}");
46
47        self.task_group.spawn_then(
48            quic_accept_streams_task(self.clone(), Arc::new(quic_conn)),
49            move |result: TaskResult<Result<()>>| async move {
50                if let TaskResult::Completed(Err(err)) = result {
51                    debug!("QUIC conn {peer:?} dropped: {err}");
52                } else {
53                    debug!("QUIC conn {peer:?} dropped");
54                }
55            },
56        );
57    }
58}
59
60async fn quic_accept_streams_task(server: Arc<Server>, quic_conn: Arc<QuicConn>) -> Result<()> {
61    let codec = JsonCodec::default();
62    loop {
63        let stream = quic_conn.accept_stream().await?;
64        let conn = framed(stream, codec.clone());
65        server
66            .task_group
67            .spawn(quic_handle_stream_task(server.clone(), conn));
68    }
69}
70
71async fn quic_handle_stream_task(server: Arc<Server>, conn: FramedConn<JsonCodec>) -> Result<()> {
72    if let Err(err) = handle_quic_stream(server, conn).await {
73        error!("Handle QUIC stream: {err}");
74    }
75    Ok(())
76}
77
78/// Handle a single QUIC stream: one request, one response.
79/// Upgrades to pubsub notification streaming if the method matches.
80async fn handle_quic_stream(server: Arc<Server>, mut conn: FramedConn<JsonCodec>) -> Result<()> {
81    let msg = conn.recv_msg().await?;
82
83    let req = match sanity_check(msg) {
84        SanityCheckResult::NewReq(req) => req,
85        SanityCheckResult::ErrRes(res) => {
86            conn.send_msg(serde_json::json!(res)).await?;
87            return Ok(());
88        }
89    };
90
91    // Pubsub handlers need a dedicated stream for streaming notifications.
92    // Everything else (regular RPC, method not found) can reuse the
93    // shared dispatch path.
94    let is_pubsub = matches!(
95        server.resolve_handler(&req.srvc_name, &req.method_name, true),
96        Handler::Pubsub(_)
97    );
98
99    if is_pubsub {
100        return handle_quic_subscription(server, conn, req).await;
101    }
102
103    let msg = serde_json::to_value(&req.msg).expect("serializable request");
104    let response = server.handle_request(None, msg).await;
105    debug!("--> {response}");
106    conn.send_msg(serde_json::json!(response)).await?;
107    Ok(())
108}
109
110/// Pubsub over QUIC: split the stream so notifications stream
111/// from the writer while the reader waits for an unsubscribe.
112async fn handle_quic_subscription(
113    server: Arc<Server>,
114    conn: FramedConn<JsonCodec>,
115    req: NewRequest,
116) -> Result<()> {
117    let (ch_tx, ch_rx) = async_channel::bounded(CHANNEL_SUBSCRIPTION_BUFFER_SIZE);
118    let channel = Channel::new(ch_tx);
119
120    let method = match server.resolve_handler(&req.srvc_name, &req.method_name, true) {
121        Handler::Pubsub(m) => m,
122        _ => unreachable!("pubsub method presence checked by caller"),
123    };
124
125    let params = req.msg.params.unwrap_or(serde_json::json!(()));
126    let result = method(channel.clone(), req.msg.method, params).await;
127
128    let response = match result {
129        Ok(res) => message::Response {
130            result: Some(res),
131            id: Some(req.msg.id),
132            ..Default::default()
133        },
134        Err(err) => {
135            let mut conn = conn;
136            let response = err.to_response(Some(req.msg.id), None);
137            conn.send_msg(serde_json::json!(response)).await?;
138            return Ok(());
139        }
140    };
141
142    debug!("--> {response}");
143
144    let (mut reader, mut writer) = conn.split();
145    writer.send_msg(serde_json::json!(response)).await?;
146
147    let notification_encoder = server.config.notification_encoder;
148
149    loop {
150        match select(ch_rx.recv(), reader.recv_msg()).await {
151            Either::Left(nt) => {
152                let nt = nt?;
153                let notification = notification_encoder(nt);
154                debug!("--> {notification}");
155                writer.send_msg(serde_json::json!(notification)).await?;
156            }
157            Either::Right(msg) => match msg {
158                Ok(msg) => {
159                    let response = server.handle_request(Some(channel.clone()), msg).await;
160                    debug!("--> {response}");
161                    writer.send_msg(serde_json::json!(response)).await?;
162                    channel.close();
163                    return Ok(());
164                }
165                Err(_) => {
166                    channel.close();
167                    return Ok(());
168                }
169            },
170        }
171    }
172}
karyon_jsonrpc/server/quic.rs

karyon_jsonrpc/server/
quic.rs
