use crate::error::Result;
use hyper::body::Bytes;
use serde::Serialize;
use std::time::Duration;
#[derive(Debug, Clone, Default)]
pub struct SseEvent {
data: String,
event: Option<String>,
id: Option<String>,
retry: Option<u64>,
comment: Option<String>,
}
impl SseEvent {
pub fn new<T: Serialize>(value: &T) -> Result<Self> {
Ok(Self {
data: serde_json::to_string(value)?,
..Default::default()
})
}
pub fn data(text: impl Into<String>) -> Self {
Self {
data: text.into(),
..Default::default()
}
}
pub fn comment(text: impl Into<String>) -> Self {
Self {
comment: Some(text.into()),
..Default::default()
}
}
pub fn event(mut self, name: impl Into<String>) -> Self {
self.event = Some(name.into());
self
}
pub fn id(mut self, id: impl Into<String>) -> Self {
self.id = Some(id.into());
self
}
pub fn retry(mut self, dur: Duration) -> Self {
self.retry = Some(dur.as_millis() as u64);
self
}
pub fn encode(&self) -> Bytes {
let mut out = String::new();
if let Some(c) = &self.comment {
for line in c.split('\n') {
out.push_str(": ");
out.push_str(line);
out.push('\n');
}
}
if let Some(e) = &self.event {
out.push_str("event: ");
out.push_str(e);
out.push('\n');
}
if let Some(id) = &self.id {
out.push_str("id: ");
out.push_str(id);
out.push('\n');
}
if let Some(r) = &self.retry {
out.push_str("retry: ");
out.push_str(&r.to_string());
out.push('\n');
}
if !self.data.is_empty() {
for line in self.data.split('\n') {
out.push_str("data: ");
out.push_str(line);
out.push('\n');
}
}
out.push('\n');
Bytes::from(out)
}
}
#[derive(Debug)]
pub struct SseClosed;
impl std::fmt::Display for SseClosed {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("SSE client disconnected")
}
}
impl std::error::Error for SseClosed {}
#[derive(Clone)]
pub struct SseSender {
tx: tokio::sync::mpsc::UnboundedSender<SseEvent>,
}
impl SseSender {
pub fn send(&self, event: SseEvent) -> std::result::Result<(), SseClosed> {
self.tx.send(event).map_err(|_| SseClosed)
}
}
pub fn sse_channel() -> (
SseSender,
impl futures_util::Stream<Item = SseEvent> + Send + 'static,
) {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseEvent>();
let stream =
futures_util::stream::unfold(
rx,
|mut rx| async move { rx.recv().await.map(|ev| (ev, rx)) },
);
(SseSender { tx }, stream)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn s(ev: &SseEvent) -> String {
String::from_utf8(ev.encode().to_vec()).unwrap()
}
#[test]
fn typed_payload_encodes_data_line() {
let ev = SseEvent::new(&json!({"n": 1})).unwrap();
assert_eq!(s(&ev), "data: {\"n\":1}\n\n");
}
#[test]
fn event_id_retry_lines_emitted() {
let ev = SseEvent::data("hi")
.event("message")
.id("42")
.retry(Duration::from_secs(3));
assert_eq!(s(&ev), "event: message\nid: 42\nretry: 3000\ndata: hi\n\n");
}
#[test]
fn multiline_data_splits_into_multiple_lines() {
let ev = SseEvent::data("a\nb");
assert_eq!(s(&ev), "data: a\ndata: b\n\n");
}
#[test]
fn comment_encodes_as_colon_line() {
assert_eq!(s(&SseEvent::comment("ping")), ": ping\n\n");
}
}