use crate::event::{Event, EventHub, LongOperationEvent, Origin};
use anyhow::Result;
use std::collections::{HashMap, HashSet};
use std::sync::{
Arc, Condvar, Mutex,
atomic::{AtomicBool, Ordering},
};
use std::thread;
use std::time::Duration;
#[derive(Debug, Clone, PartialEq)]
pub enum OperationStatus {
Running,
Completed,
Cancelled,
Failed(String),
}
#[derive(Debug, Clone)]
pub struct OperationProgress {
pub percentage: f32, pub message: Option<String>,
}
impl OperationProgress {
pub fn new(percentage: f32, message: Option<String>) -> Self {
Self {
percentage: percentage.clamp(0.0, 100.0),
message,
}
}
}
pub trait LongOperation: Send + 'static {
type Output: Send + Sync + 'static + serde::Serialize;
fn execute(
&self,
progress_callback: Box<dyn Fn(OperationProgress) + Send>,
cancel_flag: Arc<AtomicBool>,
) -> Result<Self::Output>;
}
trait OperationHandleTrait: Send {
fn get_status(&self) -> OperationStatus;
fn get_progress(&self) -> OperationProgress;
fn cancel(&self);
fn is_finished(&self) -> bool;
}
fn lock_or_recover<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
pub struct OperationCompletion {
finished: Mutex<HashSet<String>>,
condvar: Condvar,
}
impl OperationCompletion {
fn new() -> Self {
Self {
finished: Mutex::new(HashSet::new()),
condvar: Condvar::new(),
}
}
fn mark_finished(&self, id: &str) {
let mut finished = lock_or_recover(&self.finished);
finished.insert(id.to_string());
drop(finished);
self.condvar.notify_all();
}
pub fn is_finished(&self, id: &str) -> bool {
lock_or_recover(&self.finished).contains(id)
}
pub fn wait_for(&self, id: &str, timeout: Option<Duration>) -> bool {
let mut finished = lock_or_recover(&self.finished);
loop {
if finished.contains(id) {
return true;
}
match timeout {
Some(t) => {
let (guard, outcome) = self
.condvar
.wait_timeout(finished, t)
.unwrap_or_else(|poisoned| poisoned.into_inner());
finished = guard;
if outcome.timed_out() {
return finished.contains(id);
}
}
None => {
finished = self
.condvar
.wait(finished)
.unwrap_or_else(|poisoned| poisoned.into_inner());
}
}
}
}
}
impl std::fmt::Debug for OperationCompletion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OperationCompletion")
.field("finished_len", &lock_or_recover(&self.finished).len())
.finish()
}
}
struct OperationHandle {
status: Arc<Mutex<OperationStatus>>,
progress: Arc<Mutex<OperationProgress>>,
cancel_flag: Arc<AtomicBool>,
_join_handle: thread::JoinHandle<()>,
}
impl OperationHandleTrait for OperationHandle {
fn get_status(&self) -> OperationStatus {
lock_or_recover(&self.status).clone()
}
fn get_progress(&self) -> OperationProgress {
lock_or_recover(&self.progress).clone()
}
fn cancel(&self) {
self.cancel_flag.store(true, Ordering::Relaxed);
let mut status = lock_or_recover(&self.status);
if matches!(*status, OperationStatus::Running) {
*status = OperationStatus::Cancelled;
}
}
fn is_finished(&self) -> bool {
matches!(
self.get_status(),
OperationStatus::Completed | OperationStatus::Cancelled | OperationStatus::Failed(_)
)
}
}
pub struct LongOperationManager {
operations: Arc<Mutex<HashMap<String, Box<dyn OperationHandleTrait>>>>,
next_id: Arc<Mutex<u64>>,
results: Arc<Mutex<HashMap<String, String>>>, event_hub: Option<Arc<EventHub>>,
completion: Arc<OperationCompletion>,
}
impl std::fmt::Debug for LongOperationManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let operations_len = lock_or_recover(&self.operations).len();
let next_id = *lock_or_recover(&self.next_id);
let results_len = lock_or_recover(&self.results).len();
f.debug_struct("LongOperationManager")
.field("operations_len", &operations_len)
.field("next_id", &next_id)
.field("results_len", &results_len)
.field("event_hub_set", &self.event_hub.is_some())
.finish()
}
}
impl LongOperationManager {
pub fn new() -> Self {
Self {
operations: Arc::new(Mutex::new(HashMap::new())),
next_id: Arc::new(Mutex::new(0)),
results: Arc::new(Mutex::new(HashMap::new())),
event_hub: None,
completion: Arc::new(OperationCompletion::new()),
}
}
pub fn set_event_hub(&mut self, event_hub: &Arc<EventHub>) {
self.event_hub = Some(Arc::clone(event_hub));
}
pub fn start_operation<Op: LongOperation>(&self, operation: Op) -> String {
let id = {
let mut next_id = lock_or_recover(&self.next_id);
*next_id += 1;
format!("op_{}", *next_id)
};
if let Some(event_hub) = &self.event_hub {
event_hub.send_event(Event {
origin: Origin::LongOperation(LongOperationEvent::Started),
ids: vec![],
data: Some(id.clone()),
});
}
let status = Arc::new(Mutex::new(OperationStatus::Running));
let progress = Arc::new(Mutex::new(OperationProgress::new(0.0, None)));
let cancel_flag = Arc::new(AtomicBool::new(false));
let status_clone = status.clone();
let progress_clone = progress.clone();
let cancel_flag_clone = cancel_flag.clone();
let results_clone = self.results.clone();
let id_clone = id.clone();
let event_hub_opt = self.event_hub.clone();
let completion_clone = self.completion.clone();
let join_handle = thread::spawn(move || {
let progress_callback = {
let progress = progress_clone.clone();
let event_hub_opt = event_hub_opt.clone();
let id_for_cb = id_clone.clone();
Box::new(move |prog: OperationProgress| {
*lock_or_recover(&progress) = prog.clone();
if let Some(event_hub) = &event_hub_opt {
let payload = serde_json::json!({
"id": id_for_cb,
"percentage": prog.percentage,
"message": prog.message,
})
.to_string();
event_hub.send_event(Event {
origin: Origin::LongOperation(LongOperationEvent::Progress),
ids: vec![],
data: Some(payload),
});
}
}) as Box<dyn Fn(OperationProgress) + Send>
};
let operation_result =
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
operation.execute(progress_callback, cancel_flag_clone.clone())
})) {
Ok(result) => result,
Err(payload) => {
let msg = payload
.downcast_ref::<&str>()
.map(|s| s.to_string())
.or_else(|| payload.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "operation panicked".to_string());
Err(anyhow::anyhow!("operation panicked: {msg}"))
}
};
let final_status = if cancel_flag_clone.load(Ordering::Relaxed) {
OperationStatus::Cancelled
} else {
match &operation_result {
Ok(result) => {
if let Ok(serialized) = serde_json::to_string(result) {
let mut results = lock_or_recover(&results_clone);
results.insert(id_clone.clone(), serialized);
}
OperationStatus::Completed
}
Err(e) => OperationStatus::Failed(e.to_string()),
}
};
if let Some(event_hub) = &event_hub_opt {
let (event, data) = match &final_status {
OperationStatus::Completed => (
LongOperationEvent::Completed,
serde_json::json!({"id": id_clone}).to_string(),
),
OperationStatus::Cancelled => (
LongOperationEvent::Cancelled,
serde_json::json!({"id": id_clone}).to_string(),
),
OperationStatus::Failed(err) => (
LongOperationEvent::Failed,
serde_json::json!({"id": id_clone, "error": err}).to_string(),
),
OperationStatus::Running => (
LongOperationEvent::Progress,
serde_json::json!({"id": id_clone}).to_string(),
),
};
event_hub.send_event(Event {
origin: Origin::LongOperation(event),
ids: vec![],
data: Some(data),
});
}
*lock_or_recover(&status_clone) = final_status;
completion_clone.mark_finished(&id_clone);
});
let handle = OperationHandle {
status,
progress,
cancel_flag,
_join_handle: join_handle,
};
lock_or_recover(&self.operations).insert(id.clone(), Box::new(handle));
id
}
pub fn get_operation_status(&self, id: &str) -> Option<OperationStatus> {
let operations = lock_or_recover(&self.operations);
operations.get(id).map(|handle| handle.get_status())
}
pub fn get_operation_progress(&self, id: &str) -> Option<OperationProgress> {
let operations = lock_or_recover(&self.operations);
operations.get(id).map(|handle| handle.get_progress())
}
pub fn cancel_operation(&self, id: &str) -> bool {
let operations = lock_or_recover(&self.operations);
if let Some(handle) = operations.get(id) {
handle.cancel();
if let Some(event_hub) = &self.event_hub {
let payload = serde_json::json!({"id": id}).to_string();
event_hub.send_event(Event {
origin: Origin::LongOperation(LongOperationEvent::Cancelled),
ids: vec![],
data: Some(payload),
});
}
true
} else {
false
}
}
pub fn is_operation_finished(&self, id: &str) -> Option<bool> {
let operations = lock_or_recover(&self.operations);
operations.get(id).map(|handle| handle.is_finished())
}
pub fn cleanup_finished_operations(&self) {
let mut operations = lock_or_recover(&self.operations);
operations.retain(|_, handle| !handle.is_finished());
}
pub fn list_operations(&self) -> Vec<String> {
let operations = lock_or_recover(&self.operations);
operations.keys().cloned().collect()
}
pub fn get_operations_summary(&self) -> Vec<(String, OperationStatus, OperationProgress)> {
let operations = lock_or_recover(&self.operations);
operations
.iter()
.map(|(id, handle)| (id.clone(), handle.get_status(), handle.get_progress()))
.collect()
}
pub fn store_operation_result<T: serde::Serialize>(&self, id: &str, result: T) -> Result<()> {
let serialized = serde_json::to_string(&result)?;
let mut results = lock_or_recover(&self.results);
results.insert(id.to_string(), serialized);
Ok(())
}
pub fn get_operation_result(&self, id: &str) -> Option<String> {
let results = lock_or_recover(&self.results);
results.get(id).cloned()
}
pub fn completion_signal(&self) -> Arc<OperationCompletion> {
Arc::clone(&self.completion)
}
pub fn wait_for_operation(&self, id: &str, timeout: Option<Duration>) -> bool {
self.completion.wait_for(id, timeout)
}
}
impl Default for LongOperationManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::anyhow;
use std::time::Duration;
pub struct FileProcessingOperation {
pub _file_path: String,
pub total_files: usize,
}
impl LongOperation for FileProcessingOperation {
type Output = ();
fn execute(
&self,
progress_callback: Box<dyn Fn(OperationProgress) + Send>,
cancel_flag: Arc<AtomicBool>,
) -> Result<Self::Output> {
for i in 0..self.total_files {
if cancel_flag.load(Ordering::Relaxed) {
return Err(anyhow!("Operation was cancelled".to_string()));
}
thread::sleep(Duration::from_millis(500));
let percentage = (i as f32 / self.total_files as f32) * 100.0;
progress_callback(OperationProgress::new(
percentage,
Some(format!("Processing file {} of {}", i + 1, self.total_files)),
));
}
progress_callback(OperationProgress::new(100.0, Some("Completed".to_string())));
Ok(())
}
}
pub struct InstantOperation;
impl LongOperation for InstantOperation {
type Output = u32;
fn execute(
&self,
_progress_callback: Box<dyn Fn(OperationProgress) + Send>,
_cancel_flag: Arc<AtomicBool>,
) -> Result<Self::Output> {
Ok(42)
}
}
#[test]
fn wait_for_returns_as_soon_as_the_operation_finishes() {
let manager = LongOperationManager::new();
let completion = manager.completion_signal();
let started = std::time::Instant::now();
let op_id = manager.start_operation(InstantOperation);
assert!(completion.wait_for(&op_id, Some(Duration::from_secs(5))));
let elapsed = started.elapsed();
assert_eq!(
manager.get_operation_result(&op_id).as_deref(),
Some("42"),
"completion must imply the result is already retrievable"
);
assert!(
elapsed < Duration::from_millis(50),
"a trivial operation must not wait out a polling tick (took {elapsed:?})"
);
}
#[test]
fn wait_for_an_already_finished_operation_returns_immediately() {
let manager = LongOperationManager::new();
let completion = manager.completion_signal();
let op_id = manager.start_operation(InstantOperation);
assert!(completion.wait_for(&op_id, Some(Duration::from_secs(5))));
let started = std::time::Instant::now();
assert!(completion.wait_for(&op_id, Some(Duration::from_secs(5))));
assert!(started.elapsed() < Duration::from_millis(50));
assert!(completion.is_finished(&op_id));
}
#[test]
fn wait_for_times_out_while_the_operation_is_still_running() {
let manager = LongOperationManager::new();
let completion = manager.completion_signal();
let op_id = manager.start_operation(FileProcessingOperation {
_file_path: "/tmp/test".to_string(),
total_files: 5,
});
assert!(
!completion.wait_for(&op_id, Some(Duration::from_millis(50))),
"a running operation has published no completion"
);
assert!(!completion.is_finished(&op_id));
manager.cancel_operation(&op_id);
}
#[test]
fn a_cancelled_operation_still_publishes_completion() {
let manager = LongOperationManager::new();
let completion = manager.completion_signal();
let op_id = manager.start_operation(FileProcessingOperation {
_file_path: "/tmp/test".to_string(),
total_files: 5,
});
manager.cancel_operation(&op_id);
assert!(
completion.wait_for(&op_id, Some(Duration::from_secs(5))),
"cancellation must wake a blocked waiter, not strand it"
);
}
#[test]
fn test_operation_manager() {
let manager = LongOperationManager::new();
let operation = FileProcessingOperation {
_file_path: "/tmp/test".to_string(),
total_files: 5,
};
let op_id = manager.start_operation(operation);
assert_eq!(
manager.get_operation_status(&op_id),
Some(OperationStatus::Running)
);
thread::sleep(Duration::from_millis(100));
let progress = manager.get_operation_progress(&op_id);
assert!(progress.is_some());
assert!(manager.cancel_operation(&op_id));
thread::sleep(Duration::from_millis(100));
assert_eq!(
manager.get_operation_status(&op_id),
Some(OperationStatus::Cancelled)
);
}
}