use crate::log::error;
use crate::sse::SSEEvent;
use crate::{Result, SilentError, StatusCode};
use async_io::Timer;
use futures_util::{Stream, TryStream};
use pin_project::pin_project;
use std::borrow::Cow;
use std::error::Error as StdError;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use std::time::Instant;
#[derive(Debug)]
pub struct KeepAlive {
comment_text: Cow<'static, str>,
max_interval: Duration,
}
impl Default for KeepAlive {
fn default() -> Self {
Self {
comment_text: Cow::Borrowed(""),
max_interval: Duration::from_secs(15),
}
}
}
impl KeepAlive {
pub fn new() -> Self {
Self::default()
}
pub fn interval(mut self, time: Duration) -> Self {
self.max_interval = time;
self
}
pub fn comment_text(mut self, text: impl Into<Cow<'static, str>>) -> Self {
self.comment_text = text.into();
self
}
pub fn stream<S>(
self,
event_stream: S,
) -> impl TryStream<Ok = SSEEvent, Error = impl StdError + Send + Sync + 'static> + Send + 'static
where
S: TryStream<Ok = SSEEvent> + Send + 'static,
S::Error: StdError + Send + Sync + 'static,
{
let alive_timer = Timer::after(self.max_interval);
SseKeepAlive {
event_stream,
comment_text: self.comment_text,
max_interval: self.max_interval,
alive_timer,
}
}
}
#[pin_project]
struct SseKeepAlive<S> {
#[pin]
event_stream: S,
comment_text: Cow<'static, str>,
max_interval: Duration,
#[pin]
alive_timer: Timer,
}
impl<S> Stream for SseKeepAlive<S>
where
S: TryStream<Ok = SSEEvent> + Send + 'static,
S::Error: StdError + Send + Sync + 'static,
{
type Item = Result<SSEEvent>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut pin = self.project();
match pin.event_stream.try_poll_next(cx) {
Poll::Pending => match Pin::new(&mut pin.alive_timer).poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(_) => {
let next = Instant::now() + *pin.max_interval;
*pin.alive_timer = Timer::at(next);
let comment_str = pin.comment_text.clone();
let event = SSEEvent::default().comment(comment_str);
Poll::Ready(Some(Ok(event)))
}
},
Poll::Ready(Some(Ok(event))) => {
let next = Instant::now() + *pin.max_interval;
*pin.alive_timer = Timer::at(next);
Poll::Ready(Some(Ok(event)))
}
Poll::Ready(None) => Poll::Ready(None),
Poll::Ready(Some(Err(error))) => {
error!("sse::keep error: {}", error);
Poll::Ready(Some(Err(SilentError::BusinessError {
code: StatusCode::INTERNAL_SERVER_ERROR,
msg: "sse::keep error".to_string(),
})))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_keep_alive_default() {
let keep_alive = KeepAlive::default();
assert_eq!(keep_alive.max_interval, Duration::from_secs(15));
assert_eq!(keep_alive.comment_text, Cow::Borrowed(""));
}
#[test]
fn test_keep_alive_new() {
let keep_alive = KeepAlive::new();
assert_eq!(keep_alive.max_interval, Duration::from_secs(15));
assert_eq!(keep_alive.comment_text, Cow::Borrowed(""));
}
#[test]
fn test_keep_alive_interval_custom() {
let keep_alive = KeepAlive::new().interval(Duration::from_secs(30));
assert_eq!(keep_alive.max_interval, Duration::from_secs(30));
}
#[test]
fn test_keep_alive_interval_zero() {
let keep_alive = KeepAlive::new().interval(Duration::ZERO);
assert_eq!(keep_alive.max_interval, Duration::ZERO);
}
#[test]
fn test_keep_alive_interval_millis() {
let keep_alive = KeepAlive::new().interval(Duration::from_millis(500));
assert_eq!(keep_alive.max_interval, Duration::from_millis(500));
}
#[test]
fn test_keep_alive_comment_text_string() {
let keep_alive = KeepAlive::new().comment_text("keep-alive");
assert_eq!(keep_alive.comment_text, Cow::Borrowed("keep-alive"));
}
#[test]
fn test_keep_alive_comment_text_empty() {
let keep_alive = KeepAlive::new().comment_text("");
assert_eq!(keep_alive.comment_text, Cow::Borrowed(""));
}
#[test]
fn test_keep_alive_comment_text_owned() {
let keep_alive = KeepAlive::new().comment_text(String::from("owned"));
assert_eq!(
keep_alive.comment_text,
Cow::Owned::<str>(String::from("owned"))
);
}
#[test]
fn test_keep_alive_chain() {
let keep_alive = KeepAlive::new()
.interval(Duration::from_secs(10))
.comment_text("ping");
assert_eq!(keep_alive.max_interval, Duration::from_secs(10));
assert_eq!(keep_alive.comment_text, Cow::Borrowed("ping"));
}
#[test]
fn test_keep_alive_chain_reverse() {
let keep_alive = KeepAlive::new()
.comment_text("ping")
.interval(Duration::from_secs(20));
assert_eq!(keep_alive.max_interval, Duration::from_secs(20));
assert_eq!(keep_alive.comment_text, Cow::Borrowed("ping"));
}
#[test]
fn test_keep_alive_debug() {
let keep_alive = KeepAlive::new();
let debug_str = format!("{:?}", keep_alive);
assert!(debug_str.contains("KeepAlive"));
}
#[test]
fn test_keep_alive_override_interval() {
let keep_alive = KeepAlive::new()
.interval(Duration::from_secs(5))
.interval(Duration::from_secs(10));
assert_eq!(keep_alive.max_interval, Duration::from_secs(10));
}
#[test]
fn test_keep_alive_override_comment() {
let keep_alive = KeepAlive::new()
.comment_text("first")
.comment_text("second");
assert_eq!(keep_alive.comment_text, Cow::Borrowed("second"));
}
}