use crate::prelude::*;
use axum::{
extract::State,
response::sse::{Event, KeepAlive, Sse},
Json,
};
use chrono::{DateTime, Utc};
use futures::stream::Stream;
use std::collections::VecDeque;
use std::convert::Infallible;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;
use super::state::WebState;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OpKind {
GitScan,
GitMutation,
Sync,
FileOperation,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OpStatus {
Running,
Completed,
Failed,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Op {
pub id: Uuid,
pub kind: OpKind,
pub label: String,
pub status: OpStatus,
pub progress: f32,
pub error: Option<String>,
pub started_at: DateTime<Utc>,
pub ended_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum OpsEvent {
Started(Op),
Progress {
id: Uuid,
progress: f32,
label: Option<String>,
},
Completed(Op),
Failed(Op),
Snapshot {
ops: Vec<Op>,
},
}
struct Inner {
ops: VecDeque<Op>,
}
const HISTORY_CAPACITY: usize = 100;
impl Inner {
fn new() -> Self {
Self {
ops: VecDeque::new(),
}
}
fn insert(&mut self, op: Op) {
self.ops.push_back(op);
}
fn update_progress(&mut self, id: Uuid, progress: f32, label: Option<String>) {
if let Some(op) = self.ops.iter_mut().find(|o| o.id == id) {
op.progress = progress.clamp(0.0, 1.0);
if let Some(ref l) = label {
op.label = l.clone();
}
}
}
fn resolve(&mut self, id: Uuid, status: OpStatus, error: Option<String>) {
if let Some(op) = self.ops.iter_mut().find(|o| o.id == id) {
op.status = status;
op.error = error;
op.ended_at = Some(Utc::now());
if status != OpStatus::Running {
op.progress = if status == OpStatus::Completed {
1.0
} else {
op.progress
};
}
}
while self
.ops
.iter()
.filter(|o| o.status != OpStatus::Running)
.count()
> HISTORY_CAPACITY
{
if let Some(idx) = self.ops.iter().position(|o| o.status != OpStatus::Running) {
self.ops.remove(idx);
} else {
break;
}
}
}
fn snapshot(&self) -> Vec<Op> {
self.ops.iter().cloned().collect()
}
}
#[derive(Clone)]
pub struct OpsRegistry {
inner: Arc<RwLock<Inner>>,
tx: broadcast::Sender<OpsEvent>,
}
impl OpsRegistry {
pub fn new() -> Self {
let (tx, _) = broadcast::channel(256);
Self {
inner: Arc::new(RwLock::new(Inner::new())),
tx,
}
}
pub async fn start_op(&self, kind: OpKind, label: impl Into<String>) -> OpHandle {
let op = Op {
id: Uuid::new_v4(),
kind,
label: label.into(),
status: OpStatus::Running,
progress: 0.0,
error: None,
started_at: Utc::now(),
ended_at: None,
};
let id = op.id;
let event = OpsEvent::Started(op.clone());
{
let mut inner = self.inner.write().await;
inner.insert(op.clone());
}
let _ = self.tx.send(event);
OpHandle {
id,
registry: self.clone(),
resolved: false,
}
}
pub async fn report_progress(&self, id: Uuid, progress: f32, label: Option<String>) {
{
let mut inner = self.inner.write().await;
inner.update_progress(id, progress, label.clone());
}
let _ = self.tx.send(OpsEvent::Progress {
id,
progress,
label,
});
}
async fn complete(&self, id: Uuid) {
let op = {
let mut inner = self.inner.write().await;
inner.resolve(id, OpStatus::Completed, None);
inner.ops.iter().find(|o| o.id == id).cloned()
};
if let Some(op) = op {
let _ = self.tx.send(OpsEvent::Completed(op));
}
}
async fn fail(&self, id: Uuid, error: String) {
let op = {
let mut inner = self.inner.write().await;
inner.resolve(id, OpStatus::Failed, Some(error.clone()));
inner.ops.iter().find(|o| o.id == id).cloned()
};
if let Some(op) = op {
let _ = self.tx.send(OpsEvent::Failed(op));
}
}
pub async fn snapshot(&self) -> Vec<Op> {
self.inner.read().await.snapshot()
}
pub fn subscribe(&self) -> broadcast::Receiver<OpsEvent> {
self.tx.subscribe()
}
}
pub struct OpHandle {
id: Uuid,
registry: OpsRegistry,
resolved: bool,
}
impl OpHandle {
pub fn id(&self) -> Uuid {
self.id
}
pub async fn progress(&self, progress: f32) {
self.registry.report_progress(self.id, progress, None).await;
}
pub async fn progress_labeled(&self, progress: f32, label: impl Into<String>) {
self.registry
.report_progress(self.id, progress, Some(label.into()))
.await;
}
pub async fn complete(mut self) {
self.resolved = true;
self.registry.complete(self.id).await;
}
pub async fn fail(mut self, error: impl Into<String>) {
self.resolved = true;
self.registry.fail(self.id, error.into()).await;
}
}
impl Drop for OpHandle {
fn drop(&mut self) {
if self.resolved {
return;
}
let id = self.id;
let registry = self.registry.clone();
tokio::spawn(async move {
registry
.fail(id, "dropped without completion".to_string())
.await;
});
}
}
pub async fn ops_snapshot(State(state): State<WebState>) -> Json<Vec<Op>> {
Json(state.ops.snapshot().await)
}
pub async fn ops_stream(
State(state): State<WebState>,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
let registry = state.ops.clone();
let mut rx = registry.subscribe();
let snapshot = registry.snapshot().await;
let stream = async_stream::stream! {
let init = OpsEvent::Snapshot { ops: snapshot };
if let Ok(json) = serde_json::to_string(&init) {
yield Ok(Event::default().event("ops").data(json));
}
loop {
match rx.recv().await {
Ok(event) => {
match serde_json::to_string(&event) {
Ok(json) => yield Ok(Event::default().event("ops").data(json)),
Err(e) => {
yield Ok(Event::default().event("error").data(format!("ops serialize error: {}", e)));
}
}
}
Err(broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!("ops SSE client lagged by {} events, re-sending snapshot", n);
let ops = registry.snapshot().await;
let recovery = OpsEvent::Snapshot { ops };
if let Ok(json) = serde_json::to_string(&recovery) {
yield Ok(Event::default().event("ops").data(json));
}
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
};
Sse::new(stream).keep_alive(KeepAlive::default())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn op_lifecycle_running_to_completed() {
let registry = OpsRegistry::new();
let mut rx = registry.subscribe();
let handle = registry.start_op(OpKind::Sync, "sync test").await;
let id = handle.id();
let started_event = rx.try_recv().expect("Started event should be immediate");
match started_event {
OpsEvent::Started(ref op) => {
assert_eq!(op.id, id);
assert_eq!(op.status, OpStatus::Running);
assert_eq!(op.label, "sync test");
}
other => panic!("expected Started, got {:?}", other),
}
handle.progress(0.5).await;
let progress_event = rx.try_recv().expect("Progress event should be broadcast");
match progress_event {
OpsEvent::Progress {
id: eid, progress, ..
} => {
assert_eq!(eid, id);
assert!((progress - 0.5).abs() < f32::EPSILON);
}
other => panic!("expected Progress, got {:?}", other),
}
handle.complete().await;
let completed_event = rx.try_recv().expect("Completed event should be broadcast");
match completed_event {
OpsEvent::Completed(ref op) => {
assert_eq!(op.id, id);
assert_eq!(op.status, OpStatus::Completed);
assert!((op.progress - 1.0).abs() < f32::EPSILON);
assert!(op.ended_at.is_some());
}
other => panic!("expected Completed, got {:?}", other),
}
let snap = registry.snapshot().await;
let found = snap
.iter()
.find(|o| o.id == id)
.expect("op must be in snapshot");
assert_eq!(found.status, OpStatus::Completed);
}
#[tokio::test]
async fn op_lifecycle_running_to_failed() {
let registry = OpsRegistry::new();
let mut rx = registry.subscribe();
let handle = registry.start_op(OpKind::GitMutation, "push origin").await;
let id = handle.id();
let _ = rx.try_recv();
handle.fail("authentication failure").await;
let event = rx.try_recv().expect("Failed event should be broadcast");
match event {
OpsEvent::Failed(ref op) => {
assert_eq!(op.id, id);
assert_eq!(op.status, OpStatus::Failed);
assert_eq!(op.error.as_deref(), Some("authentication failure"));
}
other => panic!("expected Failed, got {:?}", other),
}
}
#[tokio::test]
async fn unresolved_drop_fails_the_op() {
let registry = OpsRegistry::new();
let mut rx = registry.subscribe();
let id = {
let handle = registry.start_op(OpKind::GitScan, "scan").await;
let id = handle.id();
let _ = rx.try_recv(); id
};
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
let snap = registry.snapshot().await;
let found = snap
.iter()
.find(|o| o.id == id)
.expect("op must still be in snapshot");
assert_eq!(
found.status,
OpStatus::Failed,
"dropped op must be Failed, not {:?}",
found.status
);
assert_eq!(found.error.as_deref(), Some("dropped without completion"));
}
#[tokio::test]
async fn concurrent_ops_are_independent() {
let registry = OpsRegistry::new();
let mut handles: Vec<OpHandle> = Vec::new();
let mut ids: Vec<Uuid> = Vec::new();
for i in 0..10 {
let h = registry
.start_op(OpKind::GitScan, format!("scan {}", i))
.await;
ids.push(h.id());
handles.push(h);
}
for h in handles {
h.progress(1.0).await;
h.complete().await;
}
let snap = registry.snapshot().await;
for id in &ids {
let op = snap
.iter()
.find(|o| o.id == *id)
.expect("every op must be in snapshot");
assert_eq!(
op.status,
OpStatus::Completed,
"op {} must be Completed",
id
);
}
}
#[tokio::test]
async fn history_ring_buffer_evicts_old_completed() {
let registry = OpsRegistry::new();
for i in 0..(HISTORY_CAPACITY + 20) {
let h = registry.start_op(OpKind::Sync, format!("sync {}", i)).await;
h.complete().await;
}
let snap = registry.snapshot().await;
let completed: Vec<_> = snap
.iter()
.filter(|o| o.status == OpStatus::Completed)
.collect();
assert!(
completed.len() <= HISTORY_CAPACITY,
"completed ops {} must not exceed capacity {}",
completed.len(),
HISTORY_CAPACITY
);
}
#[tokio::test]
async fn snapshot_includes_running_ops() {
let registry = OpsRegistry::new();
let h1 = registry.start_op(OpKind::GitScan, "scan").await;
let h2 = registry.start_op(OpKind::Sync, "sync").await;
let id1 = h1.id();
let id2 = h2.id();
let snap = registry.snapshot().await;
let ids_in_snap: Vec<Uuid> = snap.iter().map(|o| o.id).collect();
assert!(
ids_in_snap.contains(&id1),
"snapshot must include running op1"
);
assert!(
ids_in_snap.contains(&id2),
"snapshot must include running op2"
);
h1.complete().await;
h2.fail("err").await;
}
}