use core::num::NonZeroUsize;
use core::pin::Pin;
use core::sync::atomic::{AtomicU64, Ordering};
use core::task::{Context, Poll};
use core::time::Duration;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::task::Waker;
use axum::body::{Body, Bytes};
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use futures_core::Stream;
use thiserror::Error;
use super::frame_payload::FramePayload;
use super::resume::{Resume, decide, plan};
use super::sse::SseEvent;
use crate::domain::{EventId, EventRing};
const EVENT_STREAM: &str = "text/event-stream";
const ACCEL_BUFFERING: &str = "x-accel-buffering";
pub const KEEP_ALIVE_FRAME: &str = ":ping\n\n";
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum StreamError {
#[error("the event-stream response could not be assembled: {0}")]
Response(String),
}
impl From<axum::http::Error> for StreamError {
fn from(error: axum::http::Error) -> Self {
Self::Response(error.to_string())
}
}
#[derive(Debug, Default)]
struct Waiters {
entries: Vec<(u64, Waker)>,
}
#[derive(Debug)]
struct Shared<T> {
ring: Mutex<EventRing<T>>,
waiters: Mutex<Waiters>,
next_registration: AtomicU64,
}
impl<T> Shared<T> {
fn ring(&self) -> MutexGuard<'_, EventRing<T>> {
self.ring.lock().unwrap_or_else(PoisonError::into_inner)
}
fn waiters(&self) -> MutexGuard<'_, Waiters> {
self.waiters.lock().unwrap_or_else(PoisonError::into_inner)
}
fn next_registration(&self) -> u64 {
self.next_registration.fetch_add(1, Ordering::Relaxed)
}
}
#[derive(Debug)]
pub struct EventTail<T> {
shared: Arc<Shared<T>>,
}
impl<T> EventTail<T> {
#[must_use]
pub fn new(capacity: NonZeroUsize) -> Self {
Self {
shared: Arc::new(Shared {
ring: Mutex::new(EventRing::new(capacity)),
waiters: Mutex::new(Waiters::default()),
next_registration: AtomicU64::new(0),
}),
}
}
#[must_use]
pub fn len(&self) -> usize {
self.shared.ring().len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.shared.ring().is_empty()
}
#[must_use]
pub fn last(&self) -> EventId {
self.shared.ring().last()
}
}
impl<T> Clone for EventTail<T> {
fn clone(&self) -> Self {
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl<T> Default for EventTail<T> {
fn default() -> Self {
Self::new(crate::domain::RECOMMENDED_CAPACITY.into())
}
}
impl<T> EventTail<T>
where
T: FramePayload,
{
pub fn push(&self, payload: T) -> EventId {
let id = self.shared.ring().push(payload);
let wakers: Vec<Waker> = {
let _ring = self.shared.ring();
self.shared
.waiters()
.entries
.drain(..)
.map(|(_, waker)| waker)
.collect()
};
for waker in &wakers {
waker.wake_by_ref();
}
id
}
pub fn plan(&self, header: Option<&str>) -> Planned<T> {
let ring = self.shared.ring();
let resume = plan(&ring, header);
let cursor = match &resume {
Resume::Gap {
evicted_through, ..
} => *evicted_through,
other => other.last().unwrap_or(EventId::WATERMARK),
};
drop(ring);
Planned {
shared: Arc::clone(&self.shared),
resume,
cursor,
}
}
}
impl<T> EventTail<T>
where
T: FramePayload + Send + 'static,
{
pub fn respond(
&self,
header: Option<&str>,
keep_alive: Duration,
) -> Result<Response, StreamError> {
self.plan(header).into_response(keep_alive)
}
}
#[derive(Debug)]
pub struct Planned<T> {
shared: Arc<Shared<T>>,
resume: Resume,
cursor: EventId,
}
impl<T> Planned<T> {
#[must_use]
pub const fn resume(&self) -> &Resume {
&self.resume
}
#[must_use]
pub const fn cursor(&self) -> EventId {
self.cursor
}
#[must_use]
pub const fn refusal(&self) -> Option<&super::Refusal> {
match &self.resume {
Resume::Refused(refusal) => Some(refusal),
Resume::Replay { .. } | Resume::Live { .. } | Resume::Gap { .. } => None,
}
}
}
impl<T> Planned<T> {
#[must_use]
pub fn refused(tail: &EventTail<T>, refusal: super::Refusal) -> Self {
Self {
shared: Arc::clone(&tail.shared),
resume: Resume::Refused(refusal),
cursor: EventId::WATERMARK,
}
}
}
#[must_use]
pub fn planned_refusal<T>(tail: &EventTail<T>, refusal: super::Refusal) -> Planned<T> {
Planned::refused(tail, refusal)
}
impl<T> Planned<T>
where
T: FramePayload,
{
pub fn try_stream(&self, keep_alive: Duration) -> Result<EventStream<T>, Box<Response>> {
if let Resume::Refused(refusal) = &self.resume {
return Err(Box::new(refusal.clone().into_response()));
}
Ok(EventStream {
shared: Arc::clone(&self.shared),
queued: self.resume.pending().iter().cloned().collect(),
cursor: self.cursor,
keep_alive,
deadline: None,
registration: None,
closed: false,
})
}
pub fn into_draining_response(&self, keep_alive: Duration) -> Result<Response, StreamError> {
let stream = self
.try_stream(keep_alive)
.map_err(|_| StreamError::Response("the stream could not open".to_owned()))?;
stream_response(stream.finishing_when_drained())
}
pub fn into_finishable_response(
&self,
finished: Finishing,
keep_alive: Duration,
) -> Result<Response, StreamError> {
let stream = self
.try_stream(keep_alive)
.map_err(|_| StreamError::Response("the stream could not open".to_owned()))?;
stream_response(stream.finishing_on(finished))
}
}
#[derive(Debug, Clone)]
enum Until {
Drained,
Signalled(Finishing),
}
#[derive(Debug, Clone)]
pub struct Finishing {
inner: Arc<FinishingInner>,
}
impl Finishing {
#[must_use]
pub fn new() -> Self {
Self {
inner: Arc::new(FinishingInner::new()),
}
}
#[must_use]
pub fn is_fired(&self) -> bool {
self.inner.fired.load(Ordering::SeqCst)
}
pub fn fire(&self) {
self.inner.fire();
}
}
impl Default for Finishing {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
struct FinishingInner {
fired: std::sync::atomic::AtomicBool,
wakers: std::sync::Mutex<Vec<Waker>>,
}
impl FinishingInner {
const fn new() -> Self {
Self {
fired: std::sync::atomic::AtomicBool::new(false),
wakers: std::sync::Mutex::new(Vec::new()),
}
}
fn fire(&self) {
self.fired.store(true, Ordering::SeqCst);
let wakers =
std::mem::take(&mut *self.wakers.lock().unwrap_or_else(PoisonError::into_inner));
for waker in wakers {
waker.wake();
}
}
fn subscribe(&self, waker: &Waker) {
if self.fired.load(Ordering::SeqCst) {
waker.wake_by_ref();
return;
}
self.wakers
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(waker.clone());
}
}
struct FinishingStream<T> {
inner: EventStream<T>,
until: Until,
}
impl<T: FramePayload> Stream for FinishingStream<T> {
type Item = Result<Bytes, StreamError>;
fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
match Pin::new(&mut this.inner).poll_next(context) {
Poll::Ready(Some(frame)) => return Poll::Ready(Some(frame)),
Poll::Ready(None) => {
this.inner.finish();
return Poll::Ready(None);
}
Poll::Pending => {}
}
let ended = match &this.until {
Until::Drained => !this.inner.has_pending(),
Until::Signalled(finishing) => {
let ready = finishing.is_fired() && !this.inner.has_pending();
if !ready {
finishing.inner.subscribe(context.waker());
}
ready
}
};
if ended {
this.inner.finish();
return Poll::Ready(None);
}
Poll::Pending
}
}
impl<T> Planned<T>
where
T: FramePayload + Send + 'static,
{
pub fn into_response(self, keep_alive: Duration) -> Result<Response, StreamError> {
if let Resume::Refused(refusal) = &self.resume {
return Ok(refusal.clone().into_response());
}
let stream = EventStream {
shared: Arc::clone(&self.shared),
queued: self.resume.pending().iter().cloned().collect(),
cursor: self.cursor,
keep_alive,
deadline: None,
registration: None,
closed: false,
};
stream_response(stream)
}
}
#[derive(Debug)]
pub struct EventStream<T> {
shared: Arc<Shared<T>>,
queued: VecDeque<SseEvent>,
cursor: EventId,
keep_alive: Duration,
deadline: Option<Pin<Box<tokio::time::Sleep>>>,
registration: Option<u64>,
closed: bool,
}
impl<T> EventStream<T> {
#[must_use]
pub const fn cursor(&self) -> EventId {
self.cursor
}
#[must_use]
pub const fn is_subscribed(&self) -> bool {
self.registration.is_some()
}
#[must_use]
pub fn has_pending(&self) -> bool {
!self.queued.is_empty()
}
fn finish(&mut self) {
self.unregister();
self.closed = true;
}
const fn finishing_when_drained(self) -> FinishingStream<T> {
FinishingStream {
inner: self,
until: Until::Drained,
}
}
const fn finishing_on(self, finished: Finishing) -> FinishingStream<T> {
FinishingStream {
inner: self,
until: Until::Signalled(finished),
}
}
fn unregister(&mut self) {
let Some(id) = self.registration.take() else {
return;
};
self.shared
.waiters()
.entries
.retain(|(entry, _)| *entry != id);
}
}
impl<T> EventStream<T>
where
T: FramePayload,
{
fn take(&mut self, context: &Context<'_>) -> Batch {
let ring = self.shared.ring();
let decision = decide(&ring.replay_from(Some(self.cursor)));
let batch = match decision {
Resume::Replay { frames, last } => Batch {
frames,
cursor: last,
parked: false,
closed: false,
},
Resume::Live { .. } => Batch::parked(),
Resume::Gap {
notice,
resume_from,
..
} => Batch {
frames: vec![notice],
cursor: resume_from,
parked: false,
closed: false,
},
Resume::Refused(_) => Batch {
frames: Vec::new(),
cursor: self.cursor,
parked: false,
closed: true,
},
};
if batch.parked {
let id = self.registration(context);
self.registration = Some(id);
}
drop(ring);
batch
}
fn registration(&self, context: &Context<'_>) -> u64 {
let waker = context.waker();
let mut waiters = self.shared.waiters();
if let Some(id) = self.registration
&& let Some((_, registered)) =
waiters.entries.iter_mut().find(|(entry, _)| *entry == id)
{
if !registered.will_wake(waker) {
registered.clone_from(waker);
}
return id;
}
let fresh = self.shared.next_registration();
waiters.entries.push((fresh, waker.clone()));
fresh
}
fn arm_keep_alive(&mut self) {
if self.deadline.is_some() {
return;
}
if tokio::runtime::Handle::try_current().is_err() {
return;
}
self.deadline = Some(Box::pin(tokio::time::sleep(self.keep_alive)));
}
}
impl<T> Drop for EventStream<T> {
fn drop(&mut self) {
self.unregister();
}
}
struct Batch {
frames: Vec<SseEvent>,
cursor: EventId,
parked: bool,
closed: bool,
}
impl Batch {
const fn parked() -> Self {
Self {
frames: Vec::new(),
cursor: EventId::WATERMARK,
parked: true,
closed: false,
}
}
}
impl<T> Stream for EventStream<T>
where
T: FramePayload,
{
type Item = Result<Bytes, StreamError>;
fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if this.closed {
return Poll::Ready(None);
}
if let Some(frame) = this.queued.pop_front() {
this.unregister();
return Poll::Ready(Some(Ok(Bytes::from(frame.encode()))));
}
let batch = this.take(context);
if batch.closed {
this.unregister();
this.closed = true;
return Poll::Ready(None);
}
if !batch.frames.is_empty() {
this.cursor = batch.cursor;
this.queued.extend(batch.frames);
continue;
}
this.arm_keep_alive();
if let Some(deadline) = this.deadline.as_mut()
&& deadline.as_mut().poll(context).is_ready()
{
this.deadline = None;
this.unregister();
return Poll::Ready(Some(Ok(Bytes::from_static(KEEP_ALIVE_FRAME.as_bytes()))));
}
return Poll::Pending;
}
}
}
pub(super) fn stream_response<S>(stream: S) -> Result<Response, StreamError>
where
S: Stream<Item = Result<Bytes, StreamError>> + Send + 'static,
{
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, EVENT_STREAM)
.header(header::CACHE_CONTROL, "no-cache")
.header(ACCEL_BUFFERING, "no")
.body(Body::from_stream(stream))
.map_err(StreamError::from)
}
#[must_use]
pub fn parked_streams<T>(tail: &EventTail<T>) -> usize {
tail.shared.waiters().entries.len()
}
#[cfg(test)]
mod tests {
use core::num::NonZeroUsize;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};
use core::time::Duration;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::Wake;
use axum::body::Bytes;
use axum::http::StatusCode;
use futures_core::Stream;
use serde_json::{Value, json};
use super::{
EventStream, EventTail, FramePayload, KEEP_ALIVE_FRAME, StreamError, parked_streams,
};
use crate::adapters::primary::transport::http::{
Refusal, Resume, SseError, SseEvent, planned_refusal,
};
use crate::domain::EventId;
#[derive(Debug)]
enum Failure {
ZeroCapacity,
UnexpectedStream,
Sse(SseError),
Wrong(String),
}
impl fmt::Display for Failure {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ZeroCapacity => {
write!(formatter, "a ring capacity of zero is not a capacity")
}
Self::UnexpectedStream => write!(formatter, "expected a refusal, got a stream"),
Self::Sse(error) => write!(formatter, "sse: {error}"),
Self::Wrong(detail) => write!(formatter, "{detail}"),
}
}
}
impl std::error::Error for Failure {}
impl From<SseError> for Failure {
fn from(error: SseError) -> Self {
Self::Sse(error)
}
}
impl From<StreamError> for Failure {
fn from(error: StreamError) -> Self {
Self::Wrong(error.to_string())
}
}
type TestError = Failure;
#[derive(Default)]
struct Counting {
wakes: AtomicUsize,
}
impl Counting {
fn wakes(&self) -> usize {
self.wakes.load(Ordering::SeqCst)
}
fn waker(self: &Arc<Self>) -> Waker {
Waker::from(Arc::clone(self))
}
}
impl Wake for Counting {
fn wake(self: Arc<Self>) {
self.wakes.fetch_add(1, Ordering::SeqCst);
}
fn wake_by_ref(self: &Arc<Self>) {
self.wakes.fetch_add(1, Ordering::SeqCst);
}
}
fn tail(capacity: usize) -> Result<EventTail<Value>, TestError> {
let capacity = NonZeroUsize::new(capacity).ok_or(Failure::ZeroCapacity)?;
Ok(EventTail::new(capacity))
}
fn poll_once<T: FramePayload>(
stream: &mut EventStream<T>,
waker: &Waker,
) -> Poll<Option<Result<Bytes, StreamError>>> {
let mut context = Context::from_waker(waker);
core::pin::Pin::new(stream).poll_next(&mut context)
}
fn bytes_of(poll: Poll<Option<Result<Bytes, StreamError>>>) -> Result<String, TestError> {
match poll {
Poll::Ready(Some(Ok(bytes))) => Ok(String::from_utf8_lossy(&bytes).into_owned()),
Poll::Ready(Some(Err(error))) => Err(Failure::Wrong(error.to_string())),
Poll::Ready(None) => Err(Failure::Wrong("the stream ended".to_owned())),
Poll::Pending => Err(Failure::Wrong("the stream parked".to_owned())),
}
}
fn event_of(text: &str) -> Result<SseEvent, TestError> {
Ok(SseEvent::decode(text)?)
}
fn id_of(event: &SseEvent) -> EventId {
event.id().unwrap_or(EventId::WATERMARK)
}
fn assert_payload(event: &SseEvent, expected: &Value) -> Result<(), TestError> {
let actual: Value = serde_json::from_str(event.data())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if &actual == expected {
Ok(())
} else {
Err(Failure::Wrong(format!(
"expected payload {expected}, got {actual}"
)))
}
}
fn no_keep_alive() -> Duration {
Duration::from_mins(15)
}
#[test]
fn the_buffered_frames_come_out_in_order() -> Result<(), TestError> {
let tail = tail(8)?;
for index in 1..=3 {
tail.push(json!({ "index": index }));
}
let planned = tail.plan(Some("0"));
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
for index in 1..=3u64 {
let frame = bytes_of(poll_once(&mut stream, &waker))?;
let event = event_of(&frame)?;
if id_of(&event) != EventId::new(index) {
return Err(Failure::Wrong(format!(
"expected frame {index}, got {}",
id_of(&event).get()
)));
}
assert_payload(&event, &json!({ "index": index }))?;
}
Ok(())
}
#[test]
fn the_gap_notice_precedes_the_live_frames_and_carries_no_id() -> Result<(), TestError> {
let tail = tail(2)?;
for index in 1..=3 {
tail.push(json!({ "index": index }));
}
let planned = tail.plan(Some("0"));
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
let notice = event_of(&bytes_of(poll_once(&mut stream, &waker))?)?;
if notice.id().is_some() {
return Err(Failure::Wrong(
"a gap notice must carry no id: it would move the client's \
resume cursor past events it never received"
.to_owned(),
));
}
let payload: Value = serde_json::from_str(notice.data())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if payload.get("type").and_then(Value::as_str) != Some("resume-gap") {
return Err(Failure::Wrong(format!(
"expected a resume-gap notice, got {payload}"
)));
}
for index in 2..=3u64 {
let event = event_of(&bytes_of(poll_once(&mut stream, &waker))?)?;
if id_of(&event) != EventId::new(index) {
return Err(Failure::Wrong(format!(
"expected frame {index} after the notice, got {}",
id_of(&event).get()
)));
}
}
Ok(())
}
#[test]
fn a_client_that_falls_behind_mid_stream_is_told_rather_than_skipped() -> Result<(), TestError>
{
let tail = tail(2)?;
tail.push(json!({ "index": 1 }));
tail.push(json!({ "index": 2 }));
let planned = tail.plan(Some("2"));
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
if !matches!(poll_once(&mut stream, &waker), Poll::Pending) {
return Err(Failure::Wrong(
"a current client should have parked, not produced bytes".to_owned(),
));
}
for index in 3..=5 {
tail.push(json!({ "index": index }));
}
let notice = event_of(&bytes_of(poll_once(&mut stream, &waker))?)?;
if notice.id().is_some() {
return Err(Failure::Wrong(
"a mid-stream gap notice must carry no id either".to_owned(),
));
}
let payload: Value = serde_json::from_str(notice.data())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if payload.get("type").and_then(Value::as_str) != Some("resume-gap") {
return Err(Failure::Wrong(format!(
"a client that fell behind mid-stream was not told: got {payload}"
)));
}
Ok(())
}
#[test]
fn a_live_stream_is_silent_until_something_happens() -> Result<(), TestError> {
let tail = tail(4)?;
tail.push(json!({ "index": 1 }));
tail.push(json!({ "index": 2 }));
let planned = tail.plan(Some("2"));
if !planned.resume().pending().is_empty() {
return Err(Failure::Wrong(
"a current client must have nothing queued to write".to_owned(),
));
}
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
for _ in 0..3 {
match poll_once(&mut stream, &waker) {
Poll::Pending => {}
Poll::Ready(_) => {
return Err(Failure::Wrong(
"a current client received bytes before anything happened".to_owned(),
));
}
}
}
Ok(())
}
#[tokio::test]
async fn a_refusal_is_a_status_and_a_json_rpc_body_not_a_stream() -> Result<(), TestError> {
let tail = tail(4)?;
let response = tail
.respond(Some("not-a-cursor"), no_keep_alive())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if response.status() != StatusCode::BAD_REQUEST {
return Err(Failure::Wrong(format!(
"a malformed cursor is a 400, got {}",
response.status()
)));
}
let content_type = response
.headers()
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
if content_type.starts_with(super::EVENT_STREAM) {
return Err(Failure::Wrong(format!(
"a refusal must not be served as {content_type}"
)));
}
let body = axum::body::to_bytes(response.into_body(), 4096)
.await
.map_err(|error| Failure::Wrong(error.to_string()))?;
let document: Value =
serde_json::from_slice(&body).map_err(|error| Failure::Wrong(error.to_string()))?;
if document.get("jsonrpc").and_then(Value::as_str) != Some("2.0") {
return Err(Failure::Wrong(format!(
"a refusal body must be a JSON-RPC document, got {document}"
)));
}
let error = document
.get("error")
.ok_or_else(|| Failure::Wrong("a refusal body must carry an error".to_owned()))?;
if error.get("code").and_then(Value::as_i64).is_none() {
return Err(Failure::Wrong(
"a refusal body must carry a numeric error code".to_owned(),
));
}
if error.get("message").and_then(Value::as_str).is_none() {
return Err(Failure::Wrong(
"a refusal body must carry an error message".to_owned(),
));
}
Ok(())
}
#[test]
fn every_refusal_becomes_a_non_200() -> Result<(), TestError> {
let refusals = [
Refusal::MalformedCursor,
Refusal::MalformedSessionId,
Refusal::MissingSession,
Refusal::Unauthenticated,
Refusal::Forbidden,
Refusal::UnknownCursor {
requested: EventId::new(4),
last: EventId::new(9),
},
];
for refusal in refusals {
let tail = tail(4)?;
let planned = planned_refusal(&tail, refusal.clone());
if !planned.resume().pending().is_empty() {
return Err(Failure::Wrong(format!(
"{refusal} queued frames for a stream it will never open"
)));
}
let response = planned
.into_response(no_keep_alive())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if response.status().is_success() {
return Err(Failure::Wrong(format!("{refusal} was served as a success")));
}
let content_type = response
.headers()
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
if content_type.starts_with(super::EVENT_STREAM) {
return Err(Failure::Wrong(format!(
"{refusal} was served as {content_type}"
)));
}
}
Ok(())
}
#[test]
fn a_frame_pushed_after_the_replay_is_not_dropped() -> Result<(), TestError> {
let tail = tail(8)?;
for index in 1..=3 {
tail.push(json!({ "index": index }));
}
let planned = tail.plan(Some("0"));
if planned.cursor() != EventId::new(3) {
return Err(Failure::Wrong(format!(
"the cursor should be the watermark 3, got {}",
planned.cursor().get()
)));
}
tail.push(json!({ "index": 4 }));
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
for index in 1..=3u64 {
let event = event_of(&bytes_of(poll_once(&mut stream, &waker))?)?;
if id_of(&event) != EventId::new(index) {
return Err(Failure::Wrong(format!(
"expected replayed frame {index}, got {}",
id_of(&event).get()
)));
}
}
let raced = event_of(&bytes_of(poll_once(&mut stream, &waker))?)?;
if id_of(&raced) != EventId::new(4) {
return Err(Failure::Wrong(format!(
"the frame published in the replay-to-tail window was lost: \
expected 4, got {}",
id_of(&raced).get()
)));
}
assert_payload(&raced, &json!({ "index": 4 }))
}
#[tokio::test]
async fn a_parked_stream_wakes_when_a_frame_is_pushed() -> Result<(), TestError> {
let tail = tail(8)?;
let planned = tail.plan(None);
let mut stream = planned
.try_stream(Duration::from_mins(15))
.map_err(|_| Failure::UnexpectedStream)?;
let publisher = tail.clone();
let pusher = tokio::spawn(async move {
tokio::task::yield_now().await;
publisher.push(json!({ "woken": true }));
});
let frame =
core::future::poll_fn(|context| core::pin::Pin::new(&mut stream).poll_next(context))
.await
.ok_or_else(|| Failure::Wrong("the stream ended without a frame".to_owned()))?
.map_err(|error| Failure::Wrong(error.to_string()))?;
pusher
.await
.map_err(|error| Failure::Wrong(error.to_string()))?;
let text = String::from_utf8_lossy(&frame).into_owned();
let event = event_of(&text)?;
assert_payload(&event, &json!({ "woken": true }))
}
#[test]
fn dropping_the_body_releases_the_subscription() -> Result<(), TestError> {
let tail = tail(4)?;
let planned = tail.plan(None);
let mut stream = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
if !matches!(poll_once(&mut stream, &waker), Poll::Pending) {
return Err(Failure::Wrong("the stream did not park".to_owned()));
}
if !stream.is_subscribed() {
return Err(Failure::Wrong(
"a parked stream must hold a subscription".to_owned(),
));
}
if parked_streams(&tail) != 1 {
return Err(Failure::Wrong(format!(
"expected one parked stream, got {}",
parked_streams(&tail)
)));
}
drop(stream);
if parked_streams(&tail) != 0 {
return Err(Failure::Wrong(format!(
"a dropped stream left {} subscription(s) behind",
parked_streams(&tail)
)));
}
Ok(())
}
#[test]
fn one_streams_drop_does_not_silence_another() -> Result<(), TestError> {
let tail = tail(4)?;
let planned = tail.plan(None);
let mut first = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let mut second = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
for stream in [&mut first, &mut second] {
if !matches!(poll_once(stream, &waker), Poll::Pending) {
return Err(Failure::Wrong("a stream did not park".to_owned()));
}
}
if parked_streams(&tail) != 2 {
return Err(Failure::Wrong(format!(
"expected two parked streams on one waker, got {}",
parked_streams(&tail)
)));
}
drop(first);
if parked_streams(&tail) != 1 {
return Err(Failure::Wrong(format!(
"dropping one stream disturbed the other: {} left",
parked_streams(&tail)
)));
}
if !second.is_subscribed() {
return Err(Failure::Wrong(
"the surviving stream lost its subscription".to_owned(),
));
}
drop(second);
Ok(())
}
#[test]
fn a_push_wakes_the_parked_streams_exactly_once_each() -> Result<(), TestError> {
let tail = tail(4)?;
let planned = tail.plan(None);
let mut first = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let mut second = planned
.try_stream(no_keep_alive())
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
for stream in [&mut first, &mut second] {
if !matches!(poll_once(stream, &waker), Poll::Pending) {
return Err(Failure::Wrong("a stream did not park".to_owned()));
}
}
if counting.wakes() != 0 {
return Err(Failure::Wrong(format!(
"parking should not have woken anything, got {}",
counting.wakes()
)));
}
tail.push(json!({ "index": 1 }));
if counting.wakes() != 2 {
return Err(Failure::Wrong(format!(
"one push to two parked streams should wake twice, got {}",
counting.wakes()
)));
}
let event = event_of(&bytes_of(poll_once(&mut first, &waker))?)?;
assert_payload(&event, &json!({ "index": 1 }))
}
#[test]
fn the_keep_alive_is_a_comment_the_client_dispatches_nothing_for() -> Result<(), TestError> {
if KEEP_ALIVE_FRAME.lines().next() != Some(":ping") {
return Err(Failure::Wrong(format!(
"a keep-alive must be a comment line, got {KEEP_ALIVE_FRAME:?}"
)));
}
if let Some(event) = SseEvent::dispatch(KEEP_ALIVE_FRAME)? {
return Err(Failure::Wrong(format!(
"a keep-alive dispatched an event: {}",
event.data()
)));
}
Ok(())
}
#[tokio::test(start_paused = true)]
async fn the_keep_alive_fires_after_its_interval_and_nothing_else() -> Result<(), TestError> {
let tail = tail(4)?;
let planned = tail.plan(None);
let mut stream = planned
.try_stream(Duration::from_secs(15))
.map_err(|_| Failure::UnexpectedStream)?;
let counting = Arc::new(Counting::default());
let waker = counting.waker();
if !matches!(poll_once(&mut stream, &waker), Poll::Pending) {
return Err(Failure::Wrong(
"an idle stream wrote before its keep-alive interval elapsed".to_owned(),
));
}
tokio::time::advance(Duration::from_secs(16)).await;
let frame = bytes_of(poll_once(&mut stream, &waker))?;
if frame != KEEP_ALIVE_FRAME {
return Err(Failure::Wrong(format!(
"expected the keep-alive comment, got {frame:?}"
)));
}
if SseEvent::dispatch(&frame)?.is_some() {
return Err(Failure::Wrong(
"the emitted keep-alive dispatched an event".to_owned(),
));
}
Ok(())
}
#[tokio::test]
async fn the_stream_response_carries_the_headers_a_proxy_needs() -> Result<(), TestError> {
let tail = tail(4)?;
let response = tail
.respond(None, no_keep_alive())
.map_err(|error| Failure::Wrong(error.to_string()))?;
if response.status() != StatusCode::OK {
return Err(Failure::Wrong(format!(
"a resumable stream is a 200, got {}",
response.status()
)));
}
let headers = response.headers();
let content_type = headers
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
if content_type != super::EVENT_STREAM {
return Err(Failure::Wrong(format!(
"expected {}, got {content_type:?}",
super::EVENT_STREAM
)));
}
let cache_control = headers
.get(axum::http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
if cache_control != "no-cache" {
return Err(Failure::Wrong(format!(
"expected no-cache, got {cache_control:?}"
)));
}
let buffering = headers
.get(super::ACCEL_BUFFERING)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
if buffering != "no" {
return Err(Failure::Wrong(format!(
"expected the proxy buffering opt-out, got {buffering:?}"
)));
}
Ok(())
}
#[tokio::test]
async fn the_response_body_carries_the_frames_the_replay_queued() -> Result<(), TestError> {
let tail = tail(8)?;
for index in 1..=2 {
tail.push(json!({ "index": index }));
}
let response = tail
.respond(Some("0"), no_keep_alive())
.map_err(|error| Failure::Wrong(error.to_string()))?;
let mut chunks = response.into_body().into_data_stream();
let first = tokio::time::timeout(
Duration::from_secs(5),
core::future::poll_fn(|context| Pin::new(&mut chunks).poll_next(context)),
)
.await
.map_err(|_| Failure::Wrong("the body never produced a chunk".to_owned()))?
.ok_or_else(|| Failure::Wrong("the body ended without a chunk".to_owned()))?
.map_err(|error| Failure::Wrong(error.to_string()))?;
let text = String::from_utf8_lossy(&first).into_owned();
if !text.starts_with("id: 1\n") {
return Err(Failure::Wrong(format!(
"expected the body to start with frame 1, got {text:?}"
)));
}
Ok(())
}
#[test]
fn a_refusal_has_a_response_and_not_a_stream() -> Result<(), TestError> {
let tail = tail(4)?;
let planned = tail.plan(Some("garbage"));
if planned.refusal().is_none() {
return Err(Failure::Wrong(
"an unreadable cursor should have been refused".to_owned(),
));
}
match planned.try_stream(no_keep_alive()) {
Ok(_) => Err(Failure::Wrong("a refusal produced a stream".to_owned())),
Err(response) => {
if response.status().is_success() {
return Err(Failure::Wrong(
"a refusal produced a success status".to_owned(),
));
}
Ok(())
}
}
}
#[test]
fn the_queued_frames_are_exactly_the_decisions_pending_frames() -> Result<(), TestError> {
let source = tail(4)?;
source.push(json!({ "index": 1 }));
let planned = source.plan(Some("0"));
let expected = planned.resume().pending().len();
if expected != 1 {
return Err(Failure::Wrong(format!(
"expected one pending frame for a fresh client, got {expected}"
)));
}
let evicting = tail(1)?;
evicting.push(json!({ "index": 1 }));
evicting.push(json!({ "index": 2 }));
let behind = evicting.plan(Some("0"));
if !matches!(behind.resume(), Resume::Gap { .. }) {
return Err(Failure::Wrong(
"a client behind the eviction point should be a gap".to_owned(),
));
}
Ok(())
}
}