1use crate::event::{Event, EventHub, LongOperationEvent, Origin};
114use anyhow::Result;
115use std::collections::HashMap;
116use std::sync::{
117 Arc, Mutex,
118 atomic::{AtomicBool, Ordering},
119};
120use std::thread;
121
122#[derive(Debug, Clone, PartialEq)]
124pub enum OperationStatus {
125 Running,
126 Completed,
127 Cancelled,
128 Failed(String),
129}
130
131#[derive(Debug, Clone)]
133pub struct OperationProgress {
134 pub percentage: f32, pub message: Option<String>,
136}
137
138impl OperationProgress {
139 pub fn new(percentage: f32, message: Option<String>) -> Self {
140 Self {
141 percentage: percentage.clamp(0.0, 100.0),
142 message,
143 }
144 }
145}
146
147pub trait LongOperation: Send + 'static {
149 type Output: Send + Sync + 'static + serde::Serialize;
150
151 fn execute(
152 &self,
153 progress_callback: Box<dyn Fn(OperationProgress) + Send>,
154 cancel_flag: Arc<AtomicBool>,
155 ) -> Result<Self::Output>;
156}
157
158trait OperationHandleTrait: Send {
160 fn get_status(&self) -> OperationStatus;
161 fn get_progress(&self) -> OperationProgress;
162 fn cancel(&self);
163 fn is_finished(&self) -> bool;
164}
165
166struct OperationHandle {
168 status: Arc<Mutex<OperationStatus>>,
169 progress: Arc<Mutex<OperationProgress>>,
170 cancel_flag: Arc<AtomicBool>,
171 _join_handle: thread::JoinHandle<()>,
172}
173
174fn lock_or_recover<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
180 mutex.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
181}
182
183impl OperationHandleTrait for OperationHandle {
184 fn get_status(&self) -> OperationStatus {
185 lock_or_recover(&self.status).clone()
186 }
187
188 fn get_progress(&self) -> OperationProgress {
189 lock_or_recover(&self.progress).clone()
190 }
191
192 fn cancel(&self) {
193 self.cancel_flag.store(true, Ordering::Relaxed);
194 let mut status = lock_or_recover(&self.status);
195 if matches!(*status, OperationStatus::Running) {
196 *status = OperationStatus::Cancelled;
197 }
198 }
199
200 fn is_finished(&self) -> bool {
201 matches!(
202 self.get_status(),
203 OperationStatus::Completed | OperationStatus::Cancelled | OperationStatus::Failed(_)
204 )
205 }
206}
207
208pub struct LongOperationManager {
210 operations: Arc<Mutex<HashMap<String, Box<dyn OperationHandleTrait>>>>,
211 next_id: Arc<Mutex<u64>>,
212 results: Arc<Mutex<HashMap<String, String>>>, event_hub: Option<Arc<EventHub>>,
214}
215
216impl LongOperationManager {
217 pub fn new() -> Self {
218 Self {
219 operations: Arc::new(Mutex::new(HashMap::new())),
220 next_id: Arc::new(Mutex::new(0)),
221 results: Arc::new(Mutex::new(HashMap::new())),
222 event_hub: None,
223 }
224 }
225
226 pub fn set_event_hub(&mut self, event_hub: &Arc<EventHub>) {
228 self.event_hub = Some(Arc::clone(event_hub));
229 }
230
231 pub fn start_operation<Op: LongOperation>(&self, operation: Op) -> String {
233 let id = {
234 let mut next_id = lock_or_recover(&self.next_id);
235 *next_id += 1;
236 format!("op_{}", *next_id)
237 };
238
239 if let Some(event_hub) = &self.event_hub {
241 event_hub.send_event(Event {
242 origin: Origin::LongOperation(LongOperationEvent::Started),
243 ids: vec![],
244 data: Some(id.clone()),
245 });
246 }
247
248 let status = Arc::new(Mutex::new(OperationStatus::Running));
249 let progress = Arc::new(Mutex::new(OperationProgress::new(0.0, None)));
250 let cancel_flag = Arc::new(AtomicBool::new(false));
251
252 let status_clone = status.clone();
253 let progress_clone = progress.clone();
254 let cancel_flag_clone = cancel_flag.clone();
255 let results_clone = self.results.clone();
256 let id_clone = id.clone();
257 let event_hub_opt = self.event_hub.clone();
258
259 let join_handle = thread::spawn(move || {
260 let progress_callback = {
261 let progress = progress_clone.clone();
262 let event_hub_opt = event_hub_opt.clone();
263 let id_for_cb = id_clone.clone();
264 Box::new(move |prog: OperationProgress| {
265 *lock_or_recover(&progress) = prog.clone();
266 if let Some(event_hub) = &event_hub_opt {
267 let payload = serde_json::json!({
268 "id": id_for_cb,
269 "percentage": prog.percentage,
270 "message": prog.message,
271 })
272 .to_string();
273 event_hub.send_event(Event {
274 origin: Origin::LongOperation(LongOperationEvent::Progress),
275 ids: vec![],
276 data: Some(payload),
277 });
278 }
279 }) as Box<dyn Fn(OperationProgress) + Send>
280 };
281
282 let operation_result = operation.execute(progress_callback, cancel_flag_clone.clone());
283
284 let final_status = if cancel_flag_clone.load(Ordering::Relaxed) {
285 OperationStatus::Cancelled
286 } else {
287 match &operation_result {
288 Ok(result) => {
289 if let Ok(serialized) = serde_json::to_string(result) {
291 let mut results = lock_or_recover(&results_clone);
292 results.insert(id_clone.clone(), serialized);
293 }
294 OperationStatus::Completed
295 }
296 Err(e) => OperationStatus::Failed(e.to_string()),
297 }
298 };
299
300 if let Some(event_hub) = &event_hub_opt {
302 let (event, data) = match &final_status {
303 OperationStatus::Completed => (
304 LongOperationEvent::Completed,
305 serde_json::json!({"id": id_clone}).to_string(),
306 ),
307 OperationStatus::Cancelled => (
308 LongOperationEvent::Cancelled,
309 serde_json::json!({"id": id_clone}).to_string(),
310 ),
311 OperationStatus::Failed(err) => (
312 LongOperationEvent::Failed,
313 serde_json::json!({"id": id_clone, "error": err}).to_string(),
314 ),
315 OperationStatus::Running => (
316 LongOperationEvent::Progress,
317 serde_json::json!({"id": id_clone}).to_string(),
318 ),
319 };
320 event_hub.send_event(Event {
321 origin: Origin::LongOperation(event),
322 ids: vec![],
323 data: Some(data),
324 });
325 }
326
327 *lock_or_recover(&status_clone) = final_status;
328 });
329
330 let handle = OperationHandle {
331 status,
332 progress,
333 cancel_flag,
334 _join_handle: join_handle,
335 };
336
337 lock_or_recover(&self.operations)
338 .insert(id.clone(), Box::new(handle));
339
340 id
341 }
342
343 pub fn get_operation_status(&self, id: &str) -> Option<OperationStatus> {
345 let operations = lock_or_recover(&self.operations);
346 operations.get(id).map(|handle| handle.get_status())
347 }
348
349 pub fn get_operation_progress(&self, id: &str) -> Option<OperationProgress> {
351 let operations = lock_or_recover(&self.operations);
352 operations.get(id).map(|handle| handle.get_progress())
353 }
354
355 pub fn cancel_operation(&self, id: &str) -> bool {
357 let operations = lock_or_recover(&self.operations);
358 if let Some(handle) = operations.get(id) {
359 handle.cancel();
360 if let Some(event_hub) = &self.event_hub {
362 let payload = serde_json::json!({"id": id}).to_string();
363 event_hub.send_event(Event {
364 origin: Origin::LongOperation(LongOperationEvent::Cancelled),
365 ids: vec![],
366 data: Some(payload),
367 });
368 }
369 true
370 } else {
371 false
372 }
373 }
374
375 pub fn is_operation_finished(&self, id: &str) -> Option<bool> {
377 let operations = lock_or_recover(&self.operations);
378 operations.get(id).map(|handle| handle.is_finished())
379 }
380
381 pub fn cleanup_finished_operations(&self) {
383 let mut operations = lock_or_recover(&self.operations);
384 operations.retain(|_, handle| !handle.is_finished());
385 }
386
387 pub fn list_operations(&self) -> Vec<String> {
389 let operations = lock_or_recover(&self.operations);
390 operations.keys().cloned().collect()
391 }
392
393 pub fn get_operations_summary(&self) -> Vec<(String, OperationStatus, OperationProgress)> {
395 let operations = lock_or_recover(&self.operations);
396 operations
397 .iter()
398 .map(|(id, handle)| (id.clone(), handle.get_status(), handle.get_progress()))
399 .collect()
400 }
401
402 pub fn store_operation_result<T: serde::Serialize>(&self, id: &str, result: T) -> Result<()> {
404 let serialized = serde_json::to_string(&result)?;
405 let mut results = lock_or_recover(&self.results);
406 results.insert(id.to_string(), serialized);
407 Ok(())
408 }
409
410 pub fn get_operation_result(&self, id: &str) -> Option<String> {
412 let results = lock_or_recover(&self.results);
413 results.get(id).cloned()
414 }
415}
416
417impl Default for LongOperationManager {
418 fn default() -> Self {
419 Self::new()
420 }
421}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426 use anyhow::anyhow;
427 use std::time::Duration;
428
429 pub struct FileProcessingOperation {
431 pub file_path: String,
432 pub total_files: usize,
433 }
434
435 impl LongOperation for FileProcessingOperation {
436 type Output = ();
437
438 fn execute(
439 &self,
440 progress_callback: Box<dyn Fn(OperationProgress) + Send>,
441 cancel_flag: Arc<AtomicBool>,
442 ) -> Result<Self::Output> {
443 for i in 0..self.total_files {
444 if cancel_flag.load(Ordering::Relaxed) {
446 return Err(anyhow!("Operation was cancelled".to_string()));
447 }
448
449 thread::sleep(Duration::from_millis(500));
451
452 let percentage = (i as f32 / self.total_files as f32) * 100.0;
454 progress_callback(OperationProgress::new(
455 percentage,
456 Some(format!("Processing file {} of {}", i + 1, self.total_files)),
457 ));
458 }
459
460 progress_callback(OperationProgress::new(100.0, Some("Completed".to_string())));
462 Ok(())
463 }
464 }
465
466 #[test]
467 fn test_operation_manager() {
468 let manager = LongOperationManager::new();
469
470 let operation = FileProcessingOperation {
471 file_path: "/tmp/test".to_string(),
472 total_files: 5,
473 };
474
475 let op_id = manager.start_operation(operation);
476
477 assert_eq!(
479 manager.get_operation_status(&op_id),
480 Some(OperationStatus::Running)
481 );
482
483 thread::sleep(Duration::from_millis(100));
485 let progress = manager.get_operation_progress(&op_id);
486 assert!(progress.is_some());
487
488 assert!(manager.cancel_operation(&op_id));
490 thread::sleep(Duration::from_millis(100));
491 assert_eq!(
492 manager.get_operation_status(&op_id),
493 Some(OperationStatus::Cancelled)
494 );
495 }
496}