1use ahash::HashMap;
64use jsonrpsee::{
65 ConnectionId, MethodResponse, MethodSink,
66 server::{
67 IntoSubscriptionCloseResponse, MethodCallback, Methods, RegisterMethodError,
68 ResponsePayload,
69 },
70 types::{ErrorObjectOwned, Id, Params, error::ErrorCode},
71};
72use parking_lot::Mutex;
73use serde_json::value::{RawValue, to_raw_value};
74use std::sync::Arc;
75use std::sync::atomic::{AtomicU64, Ordering};
76use tokio::sync::broadcast::error::RecvError;
77use tokio::sync::{mpsc, oneshot};
78
79use super::error::ServerError;
80
81pub const NOTIF_METHOD_NAME: &str = "xrpc.ch.val";
82pub const CANCEL_METHOD_NAME: &str = "xrpc.cancel";
83
84pub type ChannelId = u64;
85
86pub type Subscribers =
88 Arc<Mutex<HashMap<(ConnectionId, Id<'static>), (MethodSink, mpsc::Receiver<()>, ChannelId)>>>;
89
90#[derive(Debug)]
97#[must_use = "PendingSubscriptionSink does nothing unless `accept` or `reject` is called"]
98pub struct PendingSubscriptionSink {
99 pub(crate) inner: MethodSink,
101 pub(crate) method: &'static str,
103 pub(crate) subscribers: Subscribers,
105 pub(crate) id: Id<'static>,
108 pub(crate) subscribe: oneshot::Sender<MethodResponse>,
110 pub(crate) channel_id: ChannelId,
112 pub(crate) connection_id: ConnectionId,
114}
115
116impl PendingSubscriptionSink {
117 pub async fn accept(self) -> Result<SubscriptionSink, String> {
123 let channel_id = self.channel_id();
124 let id = self.id.clone();
125 let response = MethodResponse::subscription_response(
126 self.id,
127 ResponsePayload::success_borrowed(&channel_id),
128 self.inner.max_response_size() as usize,
129 );
130 let success = response.is_success();
131
132 self.inner
137 .send(response.to_json())
138 .await
139 .map_err(|e| e.to_string())?;
140 self.subscribe
141 .send(response)
142 .map_err(|e| format!("accept error: {}", e.as_json()))?;
143
144 if success {
145 let (tx, rx) = mpsc::channel(1);
146 self.subscribers.lock().insert(
147 (self.connection_id, id),
148 (self.inner.clone(), rx, self.channel_id),
149 );
150 tracing::debug!(
151 "Accepting subscription (conn_id={}, chann_id={})",
152 self.connection_id.0,
153 self.channel_id
154 );
155 Ok(SubscriptionSink {
156 inner: self.inner,
157 method: self.method,
158 unsubscribe: IsUnsubscribed(tx),
159 channel_id: self.channel_id,
160 })
161 } else {
162 panic!(
163 "The subscription response was too big; adjust the `max_response_size` or change Subscription ID generation"
164 );
165 }
166 }
167
168 pub fn channel_id(&self) -> ChannelId {
170 self.channel_id
171 }
172}
173
174#[derive(Debug, Clone)]
176pub struct IsUnsubscribed(mpsc::Sender<()>);
177
178impl IsUnsubscribed {
179 pub async fn unsubscribed(&self) {
181 self.0.closed().await;
182 }
183}
184
185#[derive(Debug, Clone)]
187pub struct SubscriptionSink {
188 inner: MethodSink,
190 method: &'static str,
192 unsubscribe: IsUnsubscribed,
194 channel_id: ChannelId,
196}
197
198impl SubscriptionSink {
199 pub fn method_name(&self) -> &str {
201 self.method
202 }
203
204 pub fn channel_id(&self) -> ChannelId {
206 self.channel_id
207 }
208
209 pub async fn send(&self, msg: Box<serde_json::value::RawValue>) -> Result<(), String> {
220 if self.is_closed() {
222 return Err(format!("disconnect error: {msg}"));
223 }
224
225 self.inner.send(msg).await.map_err(|e| e.to_string())
226 }
227
228 pub fn is_closed(&self) -> bool {
230 self.inner.is_closed()
231 }
232
233 pub async fn closed(&self) {
235 tokio::select! {
237 _ = self.inner.closed() => (),
238 _ = self.unsubscribe.unsubscribed() => (),
239 }
240 }
241}
242
243fn create_notif_message(
244 sink: &SubscriptionSink,
245 result: &impl serde::Serialize,
246) -> anyhow::Result<Box<RawValue>> {
247 let method = sink.method_name();
248 let channel_id = sink.channel_id();
249 let result = serde_json::to_value(result)?;
250 let msg = serde_json::json!({
251 "jsonrpc": "2.0",
252 "method": method,
253 "params": [channel_id, result]
254 });
255
256 tracing::debug!("Sending notification: {}", msg);
257
258 Ok(to_raw_value(&msg)?)
259}
260
261fn close_payload(channel_id: ChannelId) -> serde_json::Value {
262 serde_json::json!({
263 "jsonrpc":"2.0",
264 "method":"xrpc.ch.close",
265 "params":[channel_id]
266 })
267}
268
269fn close_channel_response(channel_id: ChannelId) -> MethodResponse {
270 MethodResponse::response(
271 Id::Null,
272 ResponsePayload::success(close_payload(channel_id)),
273 1024,
274 )
275}
276
277async fn send_close(sink: &SubscriptionSink) {
280 if let Ok(payload) = to_raw_value(&close_payload(sink.channel_id())) {
281 let _ = sink.send(payload).await;
282 }
283}
284
285#[derive(Debug, Clone)]
286pub struct RpcModule {
287 id_provider: Arc<AtomicU64>,
288 channels: Subscribers,
289 methods: Methods,
290}
291
292impl From<RpcModule> for Methods {
293 fn from(module: RpcModule) -> Methods {
294 module.methods
295 }
296}
297
298impl Default for RpcModule {
299 fn default() -> Self {
300 let mut methods = Methods::default();
301
302 let channels = Subscribers::default();
303 methods
304 .verify_and_insert(
305 CANCEL_METHOD_NAME,
306 MethodCallback::Unsubscription(Arc::new({
307 let channels = channels.clone();
308 move |id,
309 params: Params,
310 connection_id: ConnectionId,
311 _max_response,
312 _extensions| {
313 let cb = || {
314 let [id]: [Id<'_>; 1] = params.parse()?;
315 let sub_id = id.into_owned();
316
317 tracing::debug!("Got cancel request (id={sub_id})");
318
319 let opt = channels.lock().remove(&(connection_id, sub_id));
320 match opt {
321 Some((_, _, channel_id)) => {
322 Ok::<ChannelId, ServerError>(channel_id)
323 }
324 None => Err::<ChannelId, ServerError>(ServerError::from(
325 anyhow::anyhow!("channel not found"),
326 )),
327 }
328 };
329 let result = cb();
330 match result {
331 Ok(channel_id) => {
332 let resp = close_channel_response(channel_id);
333 tracing::debug!("Sending close message: {}", resp.as_json());
334 resp
335 }
336 Err(e) => {
337 let error: ErrorObjectOwned = e.into();
338 MethodResponse::error(id, error)
339 }
340 }
341 }
342 })),
343 )
344 .expect("Inserting a method into an empty methods map is infallible.");
345
346 Self {
347 id_provider: Arc::new(AtomicU64::new(0)),
348 channels,
349 methods,
350 }
351 }
352}
353
354impl RpcModule {
355 pub fn register_channel<R, F>(
356 &mut self,
357 subscribe_method_name: &'static str,
358 callback: F,
359 ) -> Result<&mut MethodCallback, RegisterMethodError>
360 where
361 F: (Fn(Params) -> tokio::sync::broadcast::Receiver<R>) + Send + Sync + 'static,
362 R: serde::Serialize + Clone + Send + 'static,
363 {
364 self.register_channel_raw(subscribe_method_name, {
365 move |params, pending| {
366 let mut receiver = callback(params);
367 tokio::spawn(async move {
368 let sink = if let Ok(sink) = pending.accept().await {
369 sink
370 } else {
371 tracing::error!("Failed to accept subscription");
372 return;
373 };
374 tracing::debug!("Channel created: chann_id={}", sink.channel_id);
375
376 loop {
377 tokio::select! {
378 action = receiver.recv() => {
379 match action {
380 Ok(msg) => {
381 match create_notif_message(&sink, &msg) {
382 Ok(msg) => {
383 if let Err(e) = sink.send(msg).await {
384 tracing::error!("Failed to send message: {:?}", e);
385 break;
386 }
387 }
388 Err(e) => {
389 tracing::error!("Failed to serialize channel message: {:?}", e);
390 break;
391 }
392 }
393 }
394 Err(RecvError::Closed) => {
395 send_close(&sink).await;
396 break;
397 }
398 Err(RecvError::Lagged(n)) => {
399 tracing::warn!(
403 "closing channel {}: subscriber lagged by {n} messages",
404 sink.channel_id()
405 );
406 send_close(&sink).await;
407 break;
408 }
409 }
410 },
411 _ = sink.closed() => {
412 break;
413 }
414 }
415 }
416
417 tracing::debug!("Send notification task ended (chann_id={})", sink.channel_id);
418 });
419 }
420 })
421 }
422
423 fn register_channel_raw<R, F>(
424 &mut self,
425 subscribe_method_name: &'static str,
426 callback: F,
427 ) -> Result<&mut MethodCallback, RegisterMethodError>
428 where
429 F: (Fn(Params, PendingSubscriptionSink) -> R) + Send + Sync + 'static,
430 R: IntoSubscriptionCloseResponse,
431 {
432 self.methods.verify_method_name(subscribe_method_name)?;
433 let subscribers = self.channels.clone();
434
435 self.methods.verify_and_insert(
437 subscribe_method_name,
438 MethodCallback::Subscription(Arc::new({
439 let id_provider = self.id_provider.clone();
440 move |id, params, method_sink, conn, _extensions| {
441 let channel_id = id_provider.fetch_add(1, Ordering::Relaxed);
442
443 let (tx, rx) = oneshot::channel();
445
446 let sink = PendingSubscriptionSink {
447 inner: method_sink,
448 method: NOTIF_METHOD_NAME,
449 subscribers: subscribers.clone(),
450 id: id.clone().into_owned(),
451 subscribe: tx,
452 channel_id,
453 connection_id: conn.conn_id,
454 };
455
456 callback(params, sink);
457
458 let id = id.into_owned();
459
460 Box::pin(async move {
461 match rx.await {
462 Ok(rp) => rp,
463 Err(_) => MethodResponse::error(id, ErrorCode::InternalError),
464 }
465 })
466 }
467 })),
468 )
469 }
470}
471
472#[cfg(test)]
473mod tests {
474 use super::*;
475 use serde_json::{Value, json};
476 use std::time::Duration;
477 use tokio::sync::broadcast;
478
479 const TEST_METHOD: &str = "test.channel";
480 const RECV_TIMEOUT: Duration = Duration::from_secs(1);
483 const SOURCE_CAPACITY: usize = 4;
486 const STREAM_BUF_SIZE: usize = 256;
488
489 fn test_methods(events: &broadcast::Sender<String>) -> Methods {
496 let mut module = RpcModule::default();
497 let prototype = events.subscribe();
498 module
499 .register_channel(TEST_METHOD, move |_params| prototype.resubscribe())
500 .unwrap();
501 module.into()
502 }
503
504 async fn subscribe(
513 methods: &Methods,
514 request_id: u64,
515 ) -> (ChannelId, mpsc::Receiver<Box<RawValue>>) {
516 let request = format!(
517 r#"{{"jsonrpc":"2.0","id":{request_id},"method":"{TEST_METHOD}","params":[]}}"#
518 );
519 let (response, frames) = methods
520 .raw_json_request(&request, STREAM_BUF_SIZE)
521 .await
522 .unwrap();
523 let response: Value = serde_json::from_str(response.get()).unwrap();
524 assert_eq!(response.get("id"), Some(&json!(request_id)));
525 let channel_id = response
526 .get("result")
527 .and_then(Value::as_u64)
528 .unwrap_or_else(|| panic!("channel id must be a bare u64: {response}"));
529 (channel_id, frames)
530 }
531
532 const CANCEL_REQUEST_ID: u64 = 999;
535
536 async fn cancel(methods: &Methods, target_request_id: u64) -> Value {
539 let request = format!(
540 r#"{{"jsonrpc":"2.0","id":{CANCEL_REQUEST_ID},"method":"{CANCEL_METHOD_NAME}","params":[{target_request_id}]}}"#
541 );
542 let (response, _) = methods
543 .raw_json_request(&request, STREAM_BUF_SIZE)
544 .await
545 .unwrap();
546 serde_json::from_str(response.get()).unwrap()
547 }
548
549 async fn next_frame(frames: &mut mpsc::Receiver<Box<RawValue>>) -> Value {
550 let frame = tokio::time::timeout(RECV_TIMEOUT, frames.recv())
551 .await
552 .expect("timed out waiting for a frame")
553 .expect("stream closed while waiting for a frame");
554 serde_json::from_str(frame.get()).unwrap()
555 }
556
557 async fn assert_stream_closed(frames: &mut mpsc::Receiver<Box<RawValue>>) {
561 let frame = tokio::time::timeout(RECV_TIMEOUT, frames.recv())
562 .await
563 .expect("timed out waiting for the stream to close");
564 assert!(
565 frame.is_none(),
566 "expected the stream to close, got frame: {}",
567 frame.unwrap().get()
568 );
569 }
570
571 fn val_frame(channel_id: ChannelId, payload: &str) -> Value {
572 json!({"jsonrpc": "2.0", "method": NOTIF_METHOD_NAME, "params": [channel_id, payload]})
573 }
574
575 fn close_frame(channel_id: ChannelId) -> Value {
576 json!({"jsonrpc": "2.0", "method": "xrpc.ch.close", "params": [channel_id]})
577 }
578
579 fn close_response(channel_id: ChannelId) -> Value {
582 json!({"jsonrpc": "2.0", "id": null, "result": close_frame(channel_id)})
583 }
584
585 #[tokio::test]
586 async fn subscribe_returns_u64_channel_id() {
587 let (events, _) = broadcast::channel::<String>(SOURCE_CAPACITY);
588 let methods = test_methods(&events);
589
590 let (first_channel, _first_frames) = subscribe(&methods, 1).await;
591 let (second_channel, _second_frames) = subscribe(&methods, 2).await;
592
593 assert_eq!(second_channel, first_channel + 1);
594 }
595
596 #[tokio::test]
597 async fn value_framing_positional() {
598 let (events, _) = broadcast::channel(SOURCE_CAPACITY);
599 let methods = test_methods(&events);
600 let (channel_id, mut frames) = subscribe(&methods, 1).await;
601
602 events.send("head-change".into()).unwrap();
603 drop(events);
604
605 assert_eq!(
609 next_frame(&mut frames).await,
610 val_frame(channel_id, "head-change")
611 );
612 assert_eq!(next_frame(&mut frames).await, close_frame(channel_id));
613 }
614
615 #[tokio::test]
616 async fn two_channels_one_conn_independent() {
617 let (events, _) = broadcast::channel(SOURCE_CAPACITY);
618 let methods = test_methods(&events);
619 let (first_channel, mut first_frames) = subscribe(&methods, 1).await;
620 let (second_channel, mut second_frames) = subscribe(&methods, 2).await;
621 assert_ne!(first_channel, second_channel);
622
623 events.send("both".into()).unwrap();
625 assert_eq!(
626 next_frame(&mut first_frames).await,
627 val_frame(first_channel, "both")
628 );
629 assert_eq!(
630 next_frame(&mut second_frames).await,
631 val_frame(second_channel, "both")
632 );
633
634 assert_eq!(cancel(&methods, 1).await, close_response(first_channel));
637 assert_stream_closed(&mut first_frames).await;
638
639 events.send("second-only".into()).unwrap();
641 assert_eq!(
642 next_frame(&mut second_frames).await,
643 val_frame(second_channel, "second-only")
644 );
645 }
646
647 #[tokio::test]
648 async fn hundred_channel_fanout() {
649 let (events, _) = broadcast::channel(SOURCE_CAPACITY);
650 let methods = test_methods(&events);
651
652 let mut channels = Vec::new();
653 for request_id in 1..=100 {
654 channels.push(subscribe(&methods, request_id).await);
655 }
656
657 events.send("fan-out".into()).unwrap();
658
659 let mut seen = ahash::HashSet::default();
660 for (channel_id, frames) in &mut channels {
661 assert_eq!(next_frame(frames).await, val_frame(*channel_id, "fan-out"));
662 assert!(seen.insert(*channel_id), "channel ids must be unique");
663 }
664 }
665
666 #[tokio::test]
667 async fn cancel_unknown_id_errors() {
668 let (events, _) = broadcast::channel(SOURCE_CAPACITY);
669 let methods = test_methods(&events);
670 let (channel_id, mut frames) = subscribe(&methods, 1).await;
671
672 let response = cancel(&methods, 99).await;
673 assert!(
674 response.get("error").is_some(),
675 "cancelling an unknown id must return an error response: {response}"
676 );
677 assert!(response.get("result").is_none());
678 assert_eq!(response.get("id"), Some(&json!(CANCEL_REQUEST_ID)));
680
681 events.send("still-open".into()).unwrap();
683 assert_eq!(
684 next_frame(&mut frames).await,
685 val_frame(channel_id, "still-open")
686 );
687 }
688
689 #[tokio::test]
692 async fn source_closed_sends_bare_close() {
693 let (events, _) = broadcast::channel::<String>(SOURCE_CAPACITY);
694 let methods = test_methods(&events);
695 let (channel_id, mut frames) = subscribe(&methods, 1).await;
696
697 drop(events);
698
699 assert_eq!(next_frame(&mut frames).await, close_frame(channel_id));
700 }
701
702 #[tokio::test]
707 async fn lagged_consumer_channel_closes() {
708 let (events, lagged_rx) = broadcast::channel(SOURCE_CAPACITY);
709 for n in 0..SOURCE_CAPACITY + 2 {
710 events.send(format!("event-{n}")).unwrap();
711 }
712 let lagged_rx = Mutex::new(Some(lagged_rx));
713 let mut module = RpcModule::default();
714 module
715 .register_channel(TEST_METHOD, move |_params| {
716 lagged_rx.lock().take().expect("single subscriber")
717 })
718 .unwrap();
719 let methods: Methods = module.into();
720
721 let (channel_id, mut frames) = subscribe(&methods, 1).await;
722
723 assert_eq!(next_frame(&mut frames).await, close_frame(channel_id));
725
726 assert!(events.send("after-close".into()).is_err());
731 tokio::task::yield_now().await;
732 assert!(matches!(
733 frames.try_recv(),
734 Err(mpsc::error::TryRecvError::Empty)
735 ));
736 }
737}