46db531df9* fix(sse): saturate exponential reconnect backoff ExponentialBackoff::retry computed the reconnect multiplier with 2u32.pow(current_times). With max_times unset, current_times can reach the bit width, panicking in debug builds and wrapping to a zero delay in release builds for long-lived SSE clients. Use saturating_pow and Duration::saturating_mul so the delay stays monotonic and panic-free. * fix(sse): cap exponential reconnect backoff at a bounded max delay Saturating the multiplier alone can still yield decades-long sleeps once current_times reaches the bit width, pinning the stream in tokio::time::sleep without reconnecting or terminating. Add an optional max_delay (default 30s) that clamps the computed delay, keeping the backoff monotonic and panic-free while guaranteeing the client retries. * fix(sse): default max_delay to None to preserve unbounded backoff Per maintainer feedback, leave ExponentialBackoff::default() unbounded so the fix stays a pure overflow bug-fix. The saturating multiplier removes the debug panic / release wrap from #1198, while max_delay stays opt-in for callers that want a bounded reconnect delay.
864 lines
31 KiB
Rust
864 lines
31 KiB
Rust
use std::{
|
|
pin::Pin,
|
|
sync::Arc,
|
|
task::{Poll, ready},
|
|
time::Duration,
|
|
};
|
|
|
|
use bytes::Bytes;
|
|
use futures::{Stream, StreamExt, stream::BoxStream};
|
|
use sse_stream::{Error as SseError, Sse, SseStream};
|
|
use thiserror::Error;
|
|
|
|
use crate::model::ServerJsonRpcMessage;
|
|
|
|
pub type BoxedSseResponse = BoxStream<'static, Result<Sse, SseError>>;
|
|
|
|
/// Maximum raw size of one SSE event accepted from a remote server.
|
|
pub(crate) const DEFAULT_MAX_SSE_EVENT_SIZE: usize = 16 * 1024 * 1024;
|
|
|
|
#[derive(Debug, Error)]
|
|
enum BoundedSseStreamError {
|
|
#[error(transparent)]
|
|
Source(Box<dyn std::error::Error + Send + Sync>),
|
|
#[error("SSE event exceeded the maximum size of {max_size} bytes")]
|
|
EventTooLarge { max_size: usize },
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct SseEventSizeLimiter {
|
|
max_size: usize,
|
|
retained_size: usize,
|
|
line_size: usize,
|
|
line_is_comment: bool,
|
|
previous_was_cr: bool,
|
|
}
|
|
|
|
impl SseEventSizeLimiter {
|
|
fn new(max_size: usize) -> Self {
|
|
Self {
|
|
max_size,
|
|
retained_size: 0,
|
|
line_size: 0,
|
|
line_is_comment: false,
|
|
previous_was_cr: false,
|
|
}
|
|
}
|
|
|
|
fn observe(&mut self, chunk: &[u8]) -> Result<(), ()> {
|
|
for &byte in chunk {
|
|
if self.previous_was_cr {
|
|
self.previous_was_cr = false;
|
|
if byte == b'\n' {
|
|
continue;
|
|
}
|
|
}
|
|
|
|
match byte {
|
|
b'\r' => {
|
|
self.finish_line()?;
|
|
self.previous_was_cr = true;
|
|
}
|
|
b'\n' => self.finish_line()?,
|
|
_ => {
|
|
if self.line_size == 0 {
|
|
self.line_is_comment = byte == b':';
|
|
}
|
|
self.line_size = self.line_size.saturating_add(1);
|
|
self.check_limit()?;
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn finish_line(&mut self) -> Result<(), ()> {
|
|
if self.line_size == 0 {
|
|
self.retained_size = 0;
|
|
} else if !self.line_is_comment {
|
|
// The SSE parser inserts a newline when joining multiple data fields.
|
|
self.retained_size = self
|
|
.retained_size
|
|
.saturating_add(self.line_size)
|
|
.saturating_add(1);
|
|
}
|
|
self.line_size = 0;
|
|
self.line_is_comment = false;
|
|
self.check_limit()
|
|
}
|
|
|
|
fn check_limit(&self) -> Result<(), ()> {
|
|
if self.retained_size.saturating_add(self.line_size) > self.max_size {
|
|
Err(())
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pin_project_lite::pin_project! {
|
|
struct BoundedSseByteStream<S> {
|
|
#[pin]
|
|
inner: S,
|
|
limiter: SseEventSizeLimiter,
|
|
failed: bool,
|
|
}
|
|
}
|
|
|
|
impl<S, E> Stream for BoundedSseByteStream<S>
|
|
where
|
|
S: Stream<Item = Result<Bytes, E>>,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
type Item = Result<Bytes, BoundedSseStreamError>;
|
|
|
|
fn poll_next(
|
|
self: Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> Poll<Option<Self::Item>> {
|
|
let mut this = self.project();
|
|
if *this.failed {
|
|
return Poll::Ready(None);
|
|
}
|
|
|
|
match ready!(this.inner.as_mut().poll_next(cx)) {
|
|
Some(Ok(chunk)) => {
|
|
if this.limiter.observe(&chunk).is_err() {
|
|
*this.failed = true;
|
|
Poll::Ready(Some(Err(BoundedSseStreamError::EventTooLarge {
|
|
max_size: this.limiter.max_size,
|
|
})))
|
|
} else {
|
|
Poll::Ready(Some(Ok(chunk)))
|
|
}
|
|
}
|
|
Some(Err(error)) => {
|
|
*this.failed = true;
|
|
Poll::Ready(Some(Err(BoundedSseStreamError::Source(Box::new(error)))))
|
|
}
|
|
None => Poll::Ready(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn bounded_sse_stream<S, E>(stream: S, max_event_size: usize) -> BoxedSseResponse
|
|
where
|
|
S: Stream<Item = Result<Bytes, E>> + Send + 'static,
|
|
E: std::error::Error + Send + Sync + 'static,
|
|
{
|
|
let stream = BoundedSseByteStream {
|
|
inner: stream,
|
|
limiter: SseEventSizeLimiter::new(max_event_size),
|
|
failed: false,
|
|
};
|
|
SseStream::from_bytes_stream(stream).boxed()
|
|
}
|
|
|
|
fn is_event_too_large_error(error: &SseError) -> bool {
|
|
matches!(
|
|
error,
|
|
SseError::Body(error)
|
|
if matches!(
|
|
error.downcast_ref::<BoundedSseStreamError>(),
|
|
Some(BoundedSseStreamError::EventTooLarge { .. })
|
|
)
|
|
)
|
|
}
|
|
|
|
pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync {
|
|
fn retry(&self, current_times: usize) -> Option<Duration>;
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
#[non_exhaustive]
|
|
pub struct FixedInterval {
|
|
pub max_times: Option<usize>,
|
|
pub duration: Duration,
|
|
}
|
|
|
|
impl SseRetryPolicy for FixedInterval {
|
|
fn retry(&self, current_times: usize) -> Option<Duration> {
|
|
if let Some(max_times) = self.max_times
|
|
&& current_times >= max_times
|
|
{
|
|
return None;
|
|
}
|
|
Some(self.duration)
|
|
}
|
|
}
|
|
|
|
impl FixedInterval {
|
|
pub const DEFAULT_MIN_DURATION: Duration = Duration::from_millis(1000);
|
|
}
|
|
|
|
impl Default for FixedInterval {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_times: None,
|
|
duration: Self::DEFAULT_MIN_DURATION,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
#[non_exhaustive]
|
|
pub struct ExponentialBackoff {
|
|
pub max_times: Option<usize>,
|
|
pub base_duration: Duration,
|
|
/// Optional upper bound on a single reconnect delay. `None` (the default) preserves the
|
|
/// pre-existing unbounded doubling behavior. Once the multiplier saturates near the bit
|
|
/// width that can still produce very long sleeps, so callers that need the client to
|
|
/// actually reconnect can set `Some(...)` to clamp the delay.
|
|
pub max_delay: Option<Duration>,
|
|
}
|
|
|
|
impl ExponentialBackoff {
|
|
pub const DEFAULT_DURATION: Duration = Duration::from_millis(1000);
|
|
}
|
|
|
|
impl Default for ExponentialBackoff {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_times: None,
|
|
base_duration: Self::DEFAULT_DURATION,
|
|
max_delay: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SseRetryPolicy for ExponentialBackoff {
|
|
fn retry(&self, current_times: usize) -> Option<Duration> {
|
|
if let Some(max_times) = self.max_times
|
|
&& current_times >= max_times
|
|
{
|
|
return None;
|
|
}
|
|
// `current_times` is unbounded when `max_times` is unset, so the exponent can reach
|
|
// the bit width. Saturate the multiplier at `u32::MAX` and use saturating multiplication
|
|
// for the base duration so the delay stays monotonic and panic-free instead of an
|
|
// overflow panic (debug) or a wrapped-to-zero backoff (release).
|
|
let multiplier = 2u32.saturating_pow(current_times as u32);
|
|
let delay = self.base_duration.saturating_mul(multiplier);
|
|
Some(match self.max_delay {
|
|
Some(max_delay) => delay.min(max_delay),
|
|
None => delay,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, Default)]
|
|
#[non_exhaustive]
|
|
pub struct NeverRetry;
|
|
|
|
impl SseRetryPolicy for NeverRetry {
|
|
fn retry(&self, _current_times: usize) -> Option<Duration> {
|
|
None
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct NeverReconnect<E> {
|
|
error: Option<E>,
|
|
}
|
|
|
|
impl<E: std::error::Error + Send> SseStreamReconnect for NeverReconnect<E> {
|
|
type Error = E;
|
|
type Future = futures::future::Ready<Result<BoxedSseResponse, Self::Error>>;
|
|
fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future {
|
|
futures::future::ready(Err(self.error.take().expect("should not be called again")))
|
|
}
|
|
}
|
|
|
|
/// Abstraction for SSE reconnection logic. Implementors can hook into
|
|
/// [`handle_control_event`](Self::handle_control_event) to consume control
|
|
/// frames (e.g. `event: endpoint`) that arrive when a server restarts an SSE
|
|
/// stream. The default implementation is a no-op, keeping existing behaviour
|
|
/// intact.
|
|
pub(crate) trait SseStreamReconnect {
|
|
type Error: std::error::Error;
|
|
type Future: Future<Output = Result<BoxedSseResponse, Self::Error>> + Send;
|
|
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future;
|
|
fn handle_control_event(&mut self, _event: &Sse) -> Result<(), Self::Error> {
|
|
Ok(())
|
|
}
|
|
fn handle_stream_error(
|
|
&mut self,
|
|
error: &(dyn std::error::Error + 'static),
|
|
last_event_id: Option<&str>,
|
|
) {
|
|
if let Some(id) = last_event_id {
|
|
tracing::warn!(%id, "sse stream error: {error}");
|
|
} else {
|
|
tracing::warn!("sse stream error: {error}");
|
|
}
|
|
}
|
|
fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
|
|
tracing::warn!("fatal sse stream error: {error}");
|
|
None
|
|
}
|
|
}
|
|
|
|
pin_project_lite::pin_project! {
|
|
pub(crate) struct SseAutoReconnectStream<R>
|
|
where R: SseStreamReconnect
|
|
{
|
|
retry_policy: Arc<dyn SseRetryPolicy>,
|
|
reconnect_only_after_event_id: bool,
|
|
last_event_id: Option<String>,
|
|
server_retry_interval: Option<Duration>,
|
|
connector: R,
|
|
#[pin]
|
|
state: SseAutoReconnectStreamState<R::Future>,
|
|
}
|
|
}
|
|
|
|
impl<R: SseStreamReconnect> SseAutoReconnectStream<R> {
|
|
pub fn new(
|
|
stream: BoxedSseResponse,
|
|
connector: R,
|
|
retry_policy: Arc<dyn SseRetryPolicy>,
|
|
) -> Self {
|
|
Self {
|
|
retry_policy,
|
|
reconnect_only_after_event_id: false,
|
|
last_event_id: None,
|
|
server_retry_interval: None,
|
|
connector,
|
|
state: SseAutoReconnectStreamState::Connected { stream },
|
|
}
|
|
}
|
|
|
|
pub fn new_after_event_id(
|
|
stream: BoxedSseResponse,
|
|
connector: R,
|
|
retry_policy: Arc<dyn SseRetryPolicy>,
|
|
) -> Self {
|
|
Self {
|
|
retry_policy,
|
|
reconnect_only_after_event_id: true,
|
|
last_event_id: None,
|
|
server_retry_interval: None,
|
|
connector,
|
|
state: SseAutoReconnectStreamState::Connected { stream },
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<E: std::error::Error + Send> SseAutoReconnectStream<NeverReconnect<E>> {
|
|
#[allow(dead_code)]
|
|
pub(crate) fn never_reconnect(stream: BoxedSseResponse, error_when_reconnect: E) -> Self {
|
|
Self {
|
|
retry_policy: Arc::new(NeverRetry),
|
|
reconnect_only_after_event_id: false,
|
|
last_event_id: None,
|
|
server_retry_interval: None,
|
|
connector: NeverReconnect {
|
|
error: Some(error_when_reconnect),
|
|
},
|
|
state: SseAutoReconnectStreamState::Connected { stream },
|
|
}
|
|
}
|
|
}
|
|
|
|
pin_project_lite::pin_project! {
|
|
#[project = SseAutoReconnectStreamStateProj]
|
|
#[non_exhaustive]
|
|
pub enum SseAutoReconnectStreamState<F> {
|
|
Connected {
|
|
#[pin]
|
|
stream: BoxedSseResponse,
|
|
},
|
|
Retrying {
|
|
retry_times: usize,
|
|
#[pin]
|
|
retrying: F,
|
|
},
|
|
WaitingNextRetry {
|
|
#[pin]
|
|
sleep: tokio::time::Sleep,
|
|
retry_times: usize,
|
|
},
|
|
Terminated,
|
|
}
|
|
}
|
|
|
|
impl<R> Stream for SseAutoReconnectStream<R>
|
|
where
|
|
R: SseStreamReconnect,
|
|
{
|
|
type Item = Result<ServerJsonRpcMessage, R::Error>;
|
|
fn poll_next(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> Poll<Option<Self::Item>> {
|
|
loop {
|
|
let mut this = self.as_mut().project();
|
|
// let this_state = this.state.as_mut().project()
|
|
let state = this.state.as_mut().project();
|
|
let next_state = match state {
|
|
SseAutoReconnectStreamStateProj::Connected { stream } => {
|
|
match ready!(stream.poll_next(cx)) {
|
|
Some(Ok(sse)) => {
|
|
if let Some(new_server_retry) = sse.retry {
|
|
*this.server_retry_interval =
|
|
Some(Duration::from_millis(new_server_retry));
|
|
}
|
|
if let Some(ref event_id) = sse.id {
|
|
*this.last_event_id = Some(event_id.clone());
|
|
}
|
|
// Only treat blank/`message` events as JSON-RPC payloads.
|
|
// Other control frames (endpoint, ping, etc.) are passed to
|
|
// the reconnection handler.
|
|
let is_message_event =
|
|
matches!(sse.event.as_deref(), None | Some("") | Some("message"));
|
|
if !is_message_event {
|
|
match this.connector.handle_control_event(&sse) {
|
|
Ok(()) => continue,
|
|
Err(e) => {
|
|
this.state.set(SseAutoReconnectStreamState::Terminated);
|
|
return Poll::Ready(Some(Err(e)));
|
|
}
|
|
}
|
|
}
|
|
if let Some(data) = sse.data {
|
|
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
|
|
Err(e) => {
|
|
// Downgrade to debug to avoid noisy logs when servers emit
|
|
// non-JSON payloads as message frames. Include last_event_id
|
|
// to aid troubleshooting while keeping default behaviour.
|
|
let last_id = this.last_event_id.as_deref().unwrap_or("");
|
|
tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}");
|
|
continue;
|
|
}
|
|
Ok(message) => {
|
|
return Poll::Ready(Some(Ok(message)));
|
|
}
|
|
};
|
|
} else {
|
|
continue;
|
|
}
|
|
}
|
|
Some(Err(e)) => {
|
|
if is_event_too_large_error(&e) {
|
|
this.state.set(SseAutoReconnectStreamState::Terminated);
|
|
return Poll::Ready(
|
|
this.connector.map_fatal_stream_error(e).map(Err),
|
|
);
|
|
}
|
|
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
|
|
this.state.set(SseAutoReconnectStreamState::Terminated);
|
|
return Poll::Ready(
|
|
this.connector.map_fatal_stream_error(e).map(Err),
|
|
);
|
|
}
|
|
this.connector
|
|
.handle_stream_error(&e, this.last_event_id.as_deref());
|
|
let retrying = this
|
|
.connector
|
|
.retry_connection(this.last_event_id.as_deref());
|
|
SseAutoReconnectStreamState::Retrying {
|
|
retry_times: 0,
|
|
retrying,
|
|
}
|
|
}
|
|
None => {
|
|
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
|
|
tracing::debug!(
|
|
"sse response ended before an event ID was received; cannot resume"
|
|
);
|
|
this.state.set(SseAutoReconnectStreamState::Terminated);
|
|
return Poll::Ready(None);
|
|
}
|
|
// Per SEP-1699, a graceful stream close is
|
|
// reconnectable. If the server sent a `retry` field
|
|
// we MUST wait that long before reconnecting.
|
|
let interval = this
|
|
.server_retry_interval
|
|
.take()
|
|
.or_else(|| this.retry_policy.retry(0));
|
|
if let Some(interval) = interval {
|
|
tracing::debug!(
|
|
?interval,
|
|
"sse stream ended gracefully, reconnecting"
|
|
);
|
|
SseAutoReconnectStreamState::WaitingNextRetry {
|
|
sleep: tokio::time::sleep(interval),
|
|
retry_times: 0,
|
|
}
|
|
} else {
|
|
tracing::debug!("sse stream terminated, no reconnect policy");
|
|
return Poll::Ready(None);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
SseAutoReconnectStreamStateProj::Retrying {
|
|
retry_times,
|
|
retrying,
|
|
} => {
|
|
let retry_result = ready!(retrying.poll(cx));
|
|
match retry_result {
|
|
Ok(new_stream) => {
|
|
SseAutoReconnectStreamState::Connected { stream: new_stream }
|
|
}
|
|
Err(e) => {
|
|
tracing::debug!("retry sse stream error: {e}");
|
|
*retry_times += 1;
|
|
if let Some(interval) = this.retry_policy.retry(*retry_times) {
|
|
let interval = this
|
|
.server_retry_interval
|
|
.map(|server_retry_interval| {
|
|
server_retry_interval.max(interval)
|
|
})
|
|
.unwrap_or(interval);
|
|
let sleep = tokio::time::sleep(interval);
|
|
SseAutoReconnectStreamState::WaitingNextRetry {
|
|
sleep,
|
|
retry_times: *retry_times,
|
|
}
|
|
} else {
|
|
tracing::error!("sse stream error: {e}, max retry times reached");
|
|
this.state.set(SseAutoReconnectStreamState::Terminated);
|
|
return Poll::Ready(Some(Err(e)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => {
|
|
ready!(sleep.poll(cx));
|
|
let retrying = this
|
|
.connector
|
|
.retry_connection(this.last_event_id.as_deref());
|
|
let retry_times = *retry_times;
|
|
SseAutoReconnectStreamState::Retrying {
|
|
retry_times,
|
|
retrying,
|
|
}
|
|
}
|
|
SseAutoReconnectStreamStateProj::Terminated => {
|
|
return Poll::Ready(None);
|
|
}
|
|
};
|
|
// update the state
|
|
this.state.set(next_state);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
use super::*;
|
|
|
|
#[derive(Debug, Error)]
|
|
enum TestReconnectError {
|
|
#[error("SSE stream error: {0}")]
|
|
Sse(SseError),
|
|
#[error("unexpected reconnect")]
|
|
Reconnect,
|
|
}
|
|
|
|
struct CountingReconnect {
|
|
attempts: Arc<AtomicUsize>,
|
|
}
|
|
|
|
impl SseStreamReconnect for CountingReconnect {
|
|
type Error = TestReconnectError;
|
|
type Future = futures::future::Ready<Result<BoxedSseResponse, Self::Error>>;
|
|
|
|
fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future {
|
|
self.attempts.fetch_add(1, Ordering::Relaxed);
|
|
futures::future::ready(Err(TestReconnectError::Reconnect))
|
|
}
|
|
|
|
fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
|
|
Some(TestReconnectError::Sse(error))
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_rejects_unterminated_event_over_limit() {
|
|
let source = futures::stream::iter([
|
|
Ok::<_, std::io::Error>(Bytes::from_static(b"data: aaaaa")),
|
|
Ok(Bytes::from_static(b"aaaaaa")),
|
|
]);
|
|
let mut stream = bounded_sse_stream(source, 16);
|
|
|
|
let error = stream.next().await.unwrap().unwrap_err();
|
|
|
|
assert!(
|
|
error.to_string().contains("maximum size of 16 bytes"),
|
|
"unexpected error: {error}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_resets_limit_after_event_terminator() {
|
|
let source = futures::stream::iter([
|
|
Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\r")),
|
|
Ok(Bytes::from_static(b"\n\r\ndata: b\n\n")),
|
|
]);
|
|
let mut stream = bounded_sse_stream(source, 8);
|
|
|
|
let first = stream.next().await.unwrap().unwrap();
|
|
let second = stream.next().await.unwrap().unwrap();
|
|
|
|
assert_eq!(
|
|
(first.data.as_deref(), second.data.as_deref()),
|
|
(Some("a"), Some("b"))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_passes_event_at_exact_limit() {
|
|
let source =
|
|
futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"data: ab\n\n"))]);
|
|
let mut stream = bounded_sse_stream(source, 9);
|
|
|
|
let event = stream.next().await.unwrap().unwrap();
|
|
assert_eq!(event.data.as_deref(), Some("ab"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_rejects_oversize_split_across_many_chunks() {
|
|
let source = futures::stream::iter(
|
|
std::iter::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"data: x")))
|
|
.take(20),
|
|
);
|
|
let mut stream = bounded_sse_stream(source, 32);
|
|
|
|
let mut found_error = false;
|
|
while let Some(item) = stream.next().await {
|
|
if item.is_err() {
|
|
found_error = true;
|
|
break;
|
|
}
|
|
}
|
|
assert!(found_error, "expected oversize error");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_handles_crlf_split_across_chunks() {
|
|
let source = futures::stream::iter([
|
|
Ok::<_, std::io::Error>(Bytes::from_static(b"data: hello\r")),
|
|
Ok(Bytes::from_static(b"\n\ndata: world\n\n")),
|
|
]);
|
|
let mut stream = bounded_sse_stream(source, 64);
|
|
|
|
let first = stream.next().await.unwrap().unwrap();
|
|
let second = stream.next().await.unwrap().unwrap();
|
|
|
|
assert_eq!(first.data.as_deref(), Some("hello"));
|
|
assert_eq!(second.data.as_deref(), Some("world"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_discards_completed_comment_lines_from_limit() {
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
|
b": ping\n: ping\n: ping\n",
|
|
))]);
|
|
let mut stream = bounded_sse_stream(source, 6);
|
|
|
|
assert!(stream.next().await.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_comments_do_not_reset_accumulated_data() {
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
|
b"data: a\n: ping\ndata: b\n",
|
|
))]);
|
|
let mut stream = bounded_sse_stream(source, 14);
|
|
|
|
assert!(stream.next().await.unwrap().is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_counts_multiline_data_join_newlines() {
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
|
b"data: aaa\ndata: bbb\n\n",
|
|
))]);
|
|
let mut stream = bounded_sse_stream(source, 18);
|
|
|
|
let error = stream.next().await.unwrap().unwrap_err();
|
|
assert!(
|
|
error.to_string().contains("maximum size"),
|
|
"unexpected error: {error}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn bounded_sse_stream_propagates_source_error() {
|
|
let source = futures::stream::iter([Err::<Bytes, _>(std::io::Error::new(
|
|
std::io::ErrorKind::ConnectionReset,
|
|
"connection reset",
|
|
))]);
|
|
let mut stream = bounded_sse_stream(source, 1024);
|
|
|
|
let error = stream.next().await.unwrap().unwrap_err();
|
|
assert!(
|
|
error.to_string().contains("connection reset"),
|
|
"unexpected error: {error}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn is_event_too_large_error_detects_oversize() {
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]);
|
|
let mut stream = bounded_sse_stream(source, 8);
|
|
|
|
let error = stream.next().await.unwrap().unwrap_err();
|
|
assert!(is_event_too_large_error(&error));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn is_event_too_large_error_rejects_other_errors() {
|
|
let source =
|
|
futures::stream::iter([Err::<Bytes, _>(std::io::Error::other("something else"))]);
|
|
let mut stream = bounded_sse_stream(source, 1024);
|
|
|
|
let error = stream.next().await.unwrap().unwrap_err();
|
|
assert!(!is_event_too_large_error(&error));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn oversized_event_returns_error_without_reconnecting() {
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]);
|
|
let attempts = Arc::new(AtomicUsize::new(0));
|
|
let connector = CountingReconnect {
|
|
attempts: attempts.clone(),
|
|
};
|
|
let stream = SseAutoReconnectStream::new(
|
|
bounded_sse_stream(source, 8),
|
|
connector,
|
|
Arc::new(NeverRetry),
|
|
);
|
|
let mut stream = std::pin::pin!(stream);
|
|
|
|
let result = stream.next().await;
|
|
|
|
assert!(
|
|
matches!(result, Some(Err(TestReconnectError::Sse(_))))
|
|
&& attempts.load(Ordering::Relaxed) == 0
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn skipped_events_do_not_grow_the_stack() {
|
|
// Every frame here fails to deserialize, so each one takes the "skip this
|
|
// event" path. The inner stream is always ready, so nothing yields in
|
|
// between: when that path recursed into poll_next instead of looping, this
|
|
// overflowed the stack and aborted the process.
|
|
const SKIPPED_EVENTS: usize = 50_000;
|
|
let mut payload = Vec::with_capacity(SKIPPED_EVENTS * 9);
|
|
for _ in 0..SKIPPED_EVENTS {
|
|
payload.extend_from_slice(b"data: x\n\n");
|
|
}
|
|
|
|
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(payload))]);
|
|
let attempts = Arc::new(AtomicUsize::new(0));
|
|
let connector = CountingReconnect {
|
|
attempts: attempts.clone(),
|
|
};
|
|
let stream = SseAutoReconnectStream::new(
|
|
bounded_sse_stream(source, 1024),
|
|
connector,
|
|
Arc::new(NeverRetry),
|
|
);
|
|
let mut stream = std::pin::pin!(stream);
|
|
|
|
assert!(stream.next().await.is_none());
|
|
assert_eq!(attempts.load(Ordering::Relaxed), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn response_without_event_id_does_not_reconnect() {
|
|
let attempts = Arc::new(AtomicUsize::new(0));
|
|
let connector = CountingReconnect {
|
|
attempts: attempts.clone(),
|
|
};
|
|
let stream = SseAutoReconnectStream::new_after_event_id(
|
|
futures::stream::empty().boxed(),
|
|
connector,
|
|
Arc::new(FixedInterval {
|
|
max_times: Some(1),
|
|
duration: Duration::ZERO,
|
|
}),
|
|
);
|
|
let mut stream = std::pin::pin!(stream);
|
|
|
|
assert!(stream.next().await.is_none());
|
|
assert_eq!(attempts.load(Ordering::Relaxed), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn exponential_backoff_saturates_at_high_retry_counts() {
|
|
// With `max_times` unset, `current_times` can reach the bit width. The old
|
|
// `2u32.pow(current_times)` panicked in debug builds and wrapped in release;
|
|
// the saturating implementation must return a monotonic, non-zero delay instead.
|
|
let policy = ExponentialBackoff {
|
|
max_times: None,
|
|
base_duration: Duration::from_millis(1),
|
|
max_delay: None,
|
|
};
|
|
let mut previous = Duration::ZERO;
|
|
for current_times in [31usize, 32, 63, 64, 100] {
|
|
let delay = policy
|
|
.retry(current_times)
|
|
.expect("unbounded policy never gives up");
|
|
assert!(
|
|
!delay.is_zero(),
|
|
"delay must stay non-zero at {current_times}"
|
|
);
|
|
assert!(
|
|
delay >= previous,
|
|
"delay must stay monotonic at {current_times}"
|
|
);
|
|
previous = delay;
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn exponential_backoff_caps_delay_at_max_delay() {
|
|
// An explicit cap keeps the unbounded doubling policy from producing decades-long
|
|
// sleeps once the multiplier saturates. The delay must grow monotonically, stop at
|
|
// the configured ceiling, and never exceed it.
|
|
let policy = ExponentialBackoff {
|
|
max_times: None,
|
|
base_duration: Duration::from_secs(1),
|
|
max_delay: Some(Duration::from_secs(30)),
|
|
};
|
|
let mut previous = Duration::ZERO;
|
|
for current_times in [0usize, 1, 2, 3, 4, 5, 10, 32, 64, 100] {
|
|
let delay = policy
|
|
.retry(current_times)
|
|
.expect("unbounded policy never gives up");
|
|
assert!(
|
|
delay >= previous,
|
|
"delay must stay monotonic at {current_times}"
|
|
);
|
|
assert!(
|
|
delay <= Duration::from_secs(30),
|
|
"delay must respect max_delay at {current_times}"
|
|
);
|
|
previous = delay;
|
|
}
|
|
// Beyond the ceiling the delay stays pinned at max_delay.
|
|
assert_eq!(
|
|
policy.retry(100).expect("never gives up"),
|
|
Duration::from_secs(30)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn exponential_backoff_respects_max_times() {
|
|
let policy = ExponentialBackoff {
|
|
max_times: Some(3),
|
|
base_duration: Duration::from_millis(1),
|
|
max_delay: None,
|
|
};
|
|
assert!(policy.retry(0).is_some());
|
|
assert!(policy.retry(2).is_some());
|
|
assert!(policy.retry(3).is_none());
|
|
}
|
|
}
|