use std::convert::Infallible;
use std::time::Duration;
use axum::extract::{Path, Query, State};
use axum::http::HeaderMap;
use axum::response::IntoResponse;
use axum::response::sse::{Event, KeepAlive, Sse};
use salvor_core::{RunId, RunStatus, derive_state};
use serde::Deserialize;
use serde_json::json;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use uuid::Uuid;
use crate::error::ApiError;
use crate::json;
use crate::state::AppState;
#[derive(Debug, Deserialize)]
pub struct StreamParams {
#[serde(default)]
from_seq: Option<u64>,
}
pub async fn stream(
State(state): State<AppState>,
Path(run_id_text): Path<String>,
Query(params): Query<StreamParams>,
headers: HeaderMap,
) -> Result<impl IntoResponse, ApiError> {
let run_id = Uuid::parse_str(&run_id_text)
.map(RunId::from_uuid)
.map_err(|_| {
ApiError::BadRequest(format!(
"`{run_id_text}` is not a valid run id (expected a UUID)"
))
})?;
let log = state
.store()
.read_log(run_id)
.await
.map_err(|error| ApiError::Internal(format!("store: {error}")))?;
if log.is_empty() && !state.is_run_active(run_id) {
return Err(ApiError::UnknownRun(format!(
"no run {} in this store",
run_id.as_uuid()
)));
}
let start = cursor(&headers, params.from_seq);
let poll = state.poll_interval();
let (tx, rx) = mpsc::channel::<Result<Event, Infallible>>(64);
tokio::spawn(produce(state, run_id, start, poll, tx));
Ok(Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default()))
}
fn cursor(headers: &HeaderMap, from_seq: Option<u64>) -> u64 {
if let Some(last) = headers
.get("last-event-id")
.and_then(|value| value.to_str().ok())
.and_then(|text| text.parse::<u64>().ok())
{
return last + 1;
}
from_seq.unwrap_or(0)
}
async fn produce(
state: AppState,
run_id: RunId,
start: u64,
poll: Duration,
tx: mpsc::Sender<Result<Event, Infallible>>,
) {
let store = state.store();
let mut next = start;
loop {
let log = match store.read_log(run_id).await {
Ok(log) => log,
Err(error) => {
let frame = Event::default()
.event("end")
.data(json!({ "error": format!("store: {error}") }).to_string());
let _ = tx.send(Ok(frame)).await;
return;
}
};
let from = next;
for envelope in log.iter().filter(|envelope| envelope.seq.get() >= from) {
let data = serde_json::to_string(envelope).unwrap_or_default();
let frame = Event::default()
.id(envelope.seq.get().to_string())
.data(data);
if tx.send(Ok(frame)).await.is_err() {
return;
}
next = envelope.seq.get() + 1;
}
let status = derive_state(&log).status;
if is_resting(&status) {
let frame = Event::default()
.event("end")
.data(json!({ "status": json::status(&status) }).to_string());
let _ = tx.send(Ok(frame)).await;
return;
}
if !log.is_empty() && !state.is_run_active(run_id) {
let frame = Event::default()
.event("end")
.data(json!({ "status": json::status(&status), "detached": true }).to_string());
let _ = tx.send(Ok(frame)).await;
return;
}
tokio::time::sleep(poll).await;
}
}
fn is_resting(status: &RunStatus) -> bool {
matches!(
status,
RunStatus::Completed { .. }
| RunStatus::Failed { .. }
| RunStatus::Abandoned { .. }
| RunStatus::Suspended { .. }
| RunStatus::BudgetExceeded { .. }
| RunStatus::NeedsReconciliation
)
}