use std::error::Error as StdError;
use std::io::ErrorKind;
use std::time::Duration;
use std::time::Instant;
use async_stream::stream;
use futures_util::StreamExt;
use http::header::{
HeaderName,
HeaderValue,
CONTENT_TYPE,
};
use tokio_util::sync::CancellationToken;
use super::{
SseControl,
SseMessageStream,
SseReconnectOptions,
SseRecord,
DEFAULT_SSE_MAX_RECONNECT_DELAY,
};
use crate::{
content_type,
HttpClient,
HttpError,
HttpErrorKind,
HttpRequest,
HttpResponse,
HttpResult,
RetryDelay,
RetryHint,
RetryOptions,
};
const LAST_EVENT_ID_HEADER: &str = "last-event-id";
enum ReconnectDecision {
Allowed,
MaxReconnectsReached,
MaxElapsedExceeded {
elapsed: Duration,
max_elapsed: Duration,
},
}
struct ReconnectRuntime<'a> {
retry_options: &'a RetryOptions,
options: &'a SseReconnectOptions,
max_reconnects: u32,
started_at: Instant,
cancellation_token: Option<&'a CancellationToken>,
request_method: &'a http::Method,
request_url: Option<&'a url::Url>,
}
enum ReconnectAction {
Continue,
Stop,
Fail(Box<HttpError>),
}
struct ReconnectState {
count: u32,
backoff_delay: Duration,
pending_server_retry_delay: Option<Duration>,
}
impl ReconnectState {
fn new(retry_options: &RetryOptions) -> Self {
Self {
count: 0,
backoff_delay: initial_reconnect_delay(retry_options),
pending_server_retry_delay: None,
}
}
fn set_server_retry_delay(&mut self, delay: Duration) {
self.pending_server_retry_delay = Some(delay);
}
async fn after_error(&mut self, error: HttpError, runtime: &ReconnectRuntime<'_>) -> ReconnectAction {
let sleep_delay = self.sleep_delay(runtime);
match reconnect_decision(
self.count,
runtime.max_reconnects,
runtime.started_at,
runtime.retry_options,
sleep_delay,
) {
ReconnectDecision::Allowed => self.sleep_and_advance(sleep_delay, runtime).await,
ReconnectDecision::MaxElapsedExceeded { elapsed, max_elapsed } => {
ReconnectAction::Fail(Box::new(max_elapsed_exceeded_error_with_last_error(
error,
elapsed,
max_elapsed,
runtime.request_method,
runtime.request_url,
)))
}
ReconnectDecision::MaxReconnectsReached => ReconnectAction::Fail(Box::new(error)),
}
}
async fn after_eof(&mut self, runtime: &ReconnectRuntime<'_>) -> ReconnectAction {
let sleep_delay = self.sleep_delay(runtime);
match reconnect_decision(
self.count,
runtime.max_reconnects,
runtime.started_at,
runtime.retry_options,
sleep_delay,
) {
ReconnectDecision::Allowed => self.sleep_and_advance(sleep_delay, runtime).await,
ReconnectDecision::MaxElapsedExceeded { elapsed, max_elapsed } => ReconnectAction::Fail(Box::new(
max_elapsed_exceeded_error(elapsed, max_elapsed, runtime.request_method, runtime.request_url),
)),
ReconnectDecision::MaxReconnectsReached => ReconnectAction::Stop,
}
}
fn sleep_delay(&self, runtime: &ReconnectRuntime<'_>) -> Duration {
reconnect_sleep_delay(
self.backoff_delay,
self.pending_server_retry_delay,
runtime.retry_options,
runtime.options,
)
}
async fn sleep_and_advance(&mut self, sleep_delay: Duration, runtime: &ReconnectRuntime<'_>) -> ReconnectAction {
self.count += 1;
if let Err(error) = sleep_reconnect_delay(
sleep_delay,
runtime.cancellation_token,
runtime.request_method,
runtime.request_url,
)
.await
{
return ReconnectAction::Fail(Box::new(error));
}
self.backoff_delay = next_reconnect_delay(runtime.retry_options, self.backoff_delay);
self.pending_server_retry_delay = None;
ReconnectAction::Continue
}
}
pub(crate) struct SseReconnectRunner {
client: HttpClient,
request_template: HttpRequest,
options: SseReconnectOptions,
}
impl SseReconnectRunner {
pub(crate) fn new(client: HttpClient, request: HttpRequest, options: SseReconnectOptions) -> Self {
Self {
client,
request_template: request,
options,
}
}
pub(crate) fn run(self) -> SseMessageStream {
let client = self.client;
let request_template = self.request_template;
let options = self.options;
let output = stream! {
let retry_options = options.retry.clone();
let max_reconnects = retry_options.max_attempts().saturating_sub(1);
let request_url = request_template.resolved_url().ok();
let request_method = request_template.method().clone();
let cancellation_token = request_template.cancellation_token().cloned();
let started_at = Instant::now();
let runtime = ReconnectRuntime {
retry_options: &retry_options,
options: &options,
max_reconnects,
started_at,
cancellation_token: cancellation_token.as_ref(),
request_method: &request_method,
request_url: request_url.as_ref(),
};
let mut reconnect_state = ReconnectState::new(&retry_options);
let mut last_event_id: Option<String> = None;
loop {
let mut request = request_template.clone();
let retry_override = request.retry_override().clone().force_disable();
request.set_retry_override(retry_override);
if let Some(last_event_id) = last_event_id.as_deref() {
if let Err(error) = apply_last_event_id_header(&mut request, last_event_id) {
yield Err(error);
return;
}
}
let response = match client.execute_once(request).await {
Ok(response) => response,
Err(error) => {
if should_reconnect_sse_error(&error) {
match reconnect_state.after_error(error, &runtime).await {
ReconnectAction::Continue => continue,
ReconnectAction::Stop => return,
ReconnectAction::Fail(error) => {
yield Err(*error);
return;
}
}
}
yield Err(error);
return;
}
};
if let Err(error) = validate_sse_response_content_type(&response) {
yield Err(error);
return;
}
let mut records = response.sse_records();
let mut stream_error: Option<HttpError> = None;
while let Some(item) = records.next().await {
match item {
Ok(SseRecord::Dispatch(message)) => {
if let Some(id) = message.last_event_id.clone() {
last_event_id = Some(id);
}
yield Ok(message);
}
Ok(SseRecord::Control(SseControl::ReconnectDelayMs(retry_ms))) => {
if options.honor_server_retry {
reconnect_state.set_server_retry_delay(server_retry_delay(
retry_ms,
&retry_options,
&options,
));
}
}
Ok(SseRecord::Control(SseControl::LastEventId(id))) => {
last_event_id = Some(id);
}
Err(error) => {
stream_error = Some(error);
break;
}
}
}
if let Some(error) = stream_error {
if should_reconnect_sse_error(&error) {
match reconnect_state.after_error(error, &runtime).await {
ReconnectAction::Continue => continue,
ReconnectAction::Stop => return,
ReconnectAction::Fail(error) => {
yield Err(*error);
return;
}
}
}
yield Err(error);
return;
}
if options.reconnect_on_eof {
match reconnect_state.after_eof(&runtime).await {
ReconnectAction::Continue => continue,
ReconnectAction::Stop => return,
ReconnectAction::Fail(error) => {
yield Err(*error);
return;
}
}
}
return;
}
};
Box::pin(output)
}
}
fn apply_last_event_id_header(request: &mut HttpRequest, last_event_id: &str) -> HttpResult<()> {
let header_value = HeaderValue::from_str(last_event_id).map_err(|error| {
HttpError::other(format!(
"Invalid Last-Event-ID header value ({} bytes): {error}",
last_event_id.len()
))
})?;
request.set_typed_header(HeaderName::from_static(LAST_EVENT_ID_HEADER), header_value);
Ok(())
}
fn should_reconnect_sse_error(error: &HttpError) -> bool {
if error.kind == HttpErrorKind::Cancelled {
return false;
}
matches!(error.retry_hint(), RetryHint::Retryable) || is_unexpected_eof_error(error)
}
fn next_reconnect_delay(retry_options: &RetryOptions, current: Duration) -> Duration {
retry_options
.next_base_delay_from_current(current)
.max(Duration::from_millis(1))
}
fn reconnect_decision(
count: u32,
max_reconnects: u32,
started_at: Instant,
retry_options: &RetryOptions,
sleep_delay: Duration,
) -> ReconnectDecision {
if count >= max_reconnects {
return ReconnectDecision::MaxReconnectsReached;
}
if let Some(max_elapsed) = retry_options.max_total_elapsed() {
let elapsed = started_at.elapsed();
if (elapsed >= max_elapsed) || will_exceed_elapsed(elapsed, sleep_delay, max_elapsed) {
return ReconnectDecision::MaxElapsedExceeded { elapsed, max_elapsed };
}
}
ReconnectDecision::Allowed
}
fn will_exceed_elapsed(elapsed: Duration, sleep_delay: Duration, max_elapsed: Duration) -> bool {
elapsed
.checked_add(sleep_delay)
.is_none_or(|next_elapsed| next_elapsed >= max_elapsed)
}
fn initial_reconnect_delay(retry_options: &RetryOptions) -> Duration {
retry_options.base_delay_for_attempt(1).max(Duration::from_millis(1))
}
fn reconnect_sleep_delay(
backoff_delay: Duration,
pending_server_retry_delay: Option<Duration>,
retry_options: &RetryOptions,
options: &SseReconnectOptions,
) -> Duration {
let delay = if let Some(server_delay) = pending_server_retry_delay {
if options.apply_jitter_to_server_retry {
retry_options.jittered_delay(server_delay)
} else {
server_delay
}
} else {
retry_options.jittered_delay(backoff_delay)
};
delay.max(Duration::from_millis(1))
}
async fn sleep_reconnect_delay(
delay: Duration,
cancellation_token: Option<&CancellationToken>,
request_method: &http::Method,
request_url: Option<&url::Url>,
) -> HttpResult<()> {
if let Some(token) = cancellation_token {
tokio::select! {
_ = token.cancelled() => {
let mut error = HttpError::cancelled(
"SSE reconnect cancelled while waiting before next attempt",
)
.with_method(request_method);
if let Some(url) = request_url {
error = error.with_url(url);
}
Err(error)
}
_ = tokio::time::sleep(delay) => Ok(()),
}
} else {
tokio::time::sleep(delay).await;
Ok(())
}
}
fn server_retry_delay(retry_ms: u64, retry_options: &RetryOptions, options: &SseReconnectOptions) -> Duration {
let raw = Duration::from_millis(retry_ms.max(1));
let cap = server_retry_max_delay(retry_options, options);
raw.min(cap).max(Duration::from_millis(1))
}
fn server_retry_max_delay(retry_options: &RetryOptions, options: &SseReconnectOptions) -> Duration {
options
.server_retry_max_delay
.unwrap_or_else(|| default_server_retry_max_delay(retry_options))
.max(Duration::from_millis(1))
}
fn default_server_retry_max_delay(retry_options: &RetryOptions) -> Duration {
match retry_options.delay() {
RetryDelay::None | RetryDelay::Fixed(_) => DEFAULT_SSE_MAX_RECONNECT_DELAY,
RetryDelay::Random { max, .. } | RetryDelay::Exponential { max, .. } => *max,
}
.max(Duration::from_millis(1))
}
fn max_elapsed_exceeded_error(
elapsed: Duration,
max_elapsed: Duration,
request_method: &http::Method,
request_url: Option<&url::Url>,
) -> HttpError {
let mut error = HttpError::retry_max_elapsed_exceeded(format!(
"SSE reconnect max duration exceeded: {elapsed:?}/{max_elapsed:?}"
))
.with_method(request_method);
if let Some(url) = request_url {
error = error.with_url(url);
}
error
}
fn max_elapsed_exceeded_error_with_last_error(
last_error: HttpError,
elapsed: Duration,
max_elapsed: Duration,
request_method: &http::Method,
request_url: Option<&url::Url>,
) -> HttpError {
let mut error = max_elapsed_exceeded_error(elapsed, max_elapsed, request_method, request_url);
if let Some(method) = last_error.method.as_ref() {
error = error.with_method(method);
}
if let Some(url) = last_error.url.as_ref() {
error = error.with_url(url);
}
if let Some(status) = last_error.status {
error = error.with_status(status);
}
let mut message = format!("{}; last retryable error: {}", error.message, last_error.message);
if let Some(status) = last_error.status {
message = format!("{message} (status: {status})");
}
error.message = message;
error.source = Some(Box::new(last_error));
error
}
fn validate_sse_response_content_type(response: &HttpResponse) -> HttpResult<()> {
let method = response.meta().method().clone();
let url = response.request_url().clone();
let Some(value) = response.headers().get(CONTENT_TYPE) else {
return Err(HttpError::sse_protocol("Missing Content-Type header for SSE response")
.with_status(response.status())
.with_method(&method)
.with_url(&url));
};
let content_type = value.to_str().map_err(|_| {
HttpError::sse_protocol("Invalid non-UTF8 Content-Type header for SSE response")
.with_status(response.status())
.with_method(&method)
.with_url(&url)
})?;
if content_type::is_sse(content_type) {
return Ok(());
}
Err(HttpError::sse_protocol(format!(
"Expected Content-Type 'text/event-stream' for SSE response, got '{content_type}'"
))
.with_status(response.status())
.with_method(&method)
.with_url(&url))
}
fn is_unexpected_eof_error(error: &HttpError) -> bool {
let contains_unexpected_eof = |text: &str| text.to_ascii_lowercase().contains("unexpected eof");
if contains_unexpected_eof(&error.message) {
return true;
}
error.source.as_ref().is_some_and(|source| {
has_unexpected_eof_in_error_chain(source.as_ref())
|| contains_unexpected_eof(&source.to_string())
|| contains_unexpected_eof(&format!("{source:?}"))
})
}
fn has_unexpected_eof_in_error_chain(error: &(dyn StdError + 'static)) -> bool {
let mut current: Option<&(dyn StdError + 'static)> = Some(error);
while let Some(item) = current {
if item
.downcast_ref::<std::io::Error>()
.is_some_and(|io_error| io_error.kind() == ErrorKind::UnexpectedEof)
{
return true;
}
current = item.source();
}
false
}