81 lines
2.0 KiB
Rust
81 lines
2.0 KiB
Rust
use chrono::{DateTime, Utc};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::{collections::HashMap, sync::Arc};
|
|
use tokio::sync::{RwLock, broadcast};
|
|
use tokio_util::sync::CancellationToken;
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RunEvent {
|
|
pub run_id: Uuid,
|
|
pub event_type: String,
|
|
pub stage: String,
|
|
pub message: String,
|
|
pub progress_current: i64,
|
|
pub progress_total: Option<i64>,
|
|
pub at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct RunControl {
|
|
tokens: Arc<RwLock<HashMap<Uuid, CancellationToken>>>,
|
|
events: broadcast::Sender<RunEvent>,
|
|
}
|
|
|
|
impl Default for RunControl {
|
|
fn default() -> Self {
|
|
Self::new(1_024)
|
|
}
|
|
}
|
|
|
|
impl RunControl {
|
|
pub fn new(event_capacity: usize) -> Self {
|
|
let (events, _) = broadcast::channel(event_capacity.max(16));
|
|
Self {
|
|
tokens: Arc::new(RwLock::new(HashMap::new())),
|
|
events,
|
|
}
|
|
}
|
|
|
|
pub async fn token(&self, run_id: Uuid) -> CancellationToken {
|
|
let mut tokens = self.tokens.write().await;
|
|
tokens.entry(run_id).or_default().clone()
|
|
}
|
|
|
|
pub async fn cancel(&self, run_id: Uuid) -> bool {
|
|
let token = self.tokens.read().await.get(&run_id).cloned();
|
|
if let Some(token) = token {
|
|
token.cancel();
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
pub async fn forget(&self, run_id: Uuid) {
|
|
self.tokens.write().await.remove(&run_id);
|
|
}
|
|
|
|
pub fn publish(&self, event: RunEvent) {
|
|
let _ = self.events.send(event);
|
|
}
|
|
|
|
pub fn subscribe(&self) -> broadcast::Receiver<RunEvent> {
|
|
self.events.subscribe()
|
|
}
|
|
}
|
|
|
|
impl RunEvent {
|
|
pub fn status(run_id: Uuid, stage: impl Into<String>, message: impl Into<String>) -> Self {
|
|
Self {
|
|
run_id,
|
|
event_type: "status".into(),
|
|
stage: stage.into(),
|
|
message: message.into(),
|
|
progress_current: 0,
|
|
progress_total: None,
|
|
at: Utc::now(),
|
|
}
|
|
}
|
|
}
|