orchai/src-tauri/src/services/process_registry.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::process::Child;
use tokio::sync::Mutex as AsyncMutex;

#[derive(Clone)]
struct ActiveProcess {
    child: Arc<AsyncMutex<Child>>,
    cancellation_requested: Arc<AtomicBool>,
}

impl ActiveProcess {
    fn new(child: Arc<AsyncMutex<Child>>, cancellation_requested: Arc<AtomicBool>) -> Self {
        Self {
            child,
            cancellation_requested,
        }
    }
}

#[derive(Clone, Default)]
pub struct ProcessRegistry {
    ticket_processes: Arc<Mutex<HashMap<String, ActiveProcess>>>,
    task_processes: Arc<Mutex<HashMap<String, ActiveProcess>>>,
}

impl ProcessRegistry {
    pub fn register_ticket(
        &self,
        ticket_id: &str,
        child: Arc<AsyncMutex<Child>>,
        cancellation_requested: Arc<AtomicBool>,
    ) {
        let mut processes = self
            .ticket_processes
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        processes.insert(
            ticket_id.to_string(),
            ActiveProcess::new(child, cancellation_requested),
        );
    }

    pub fn unregister_ticket(&self, ticket_id: &str) {
        let mut processes = self
            .ticket_processes
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        processes.remove(ticket_id);
    }

    pub async fn cancel_ticket(&self, ticket_id: &str) -> bool {
        self.cancel_process(&self.ticket_processes, ticket_id).await
    }

    pub fn register_task(
        &self,
        task_id: &str,
        child: Arc<AsyncMutex<Child>>,
        cancellation_requested: Arc<AtomicBool>,
    ) {
        let mut processes = self
            .task_processes
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        processes.insert(
            task_id.to_string(),
            ActiveProcess::new(child, cancellation_requested),
        );
    }

    pub fn unregister_task(&self, task_id: &str) {
        let mut processes = self
            .task_processes
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        processes.remove(task_id);
    }

    pub async fn cancel_task(&self, task_id: &str) -> bool {
        self.cancel_process(&self.task_processes, task_id).await
    }

    async fn cancel_process(
        &self,
        map: &Arc<Mutex<HashMap<String, ActiveProcess>>>,
        entity_id: &str,
    ) -> bool {
        let process = {
            let processes = map.lock().unwrap_or_else(|poison| poison.into_inner());
            processes.get(entity_id).cloned()
        };

        let Some(process) = process else {
            return false;
        };

        process.cancellation_requested.store(true, Ordering::SeqCst);
        let mut child = process.child.lock().await;
        let _ = child.start_kill();
        true
    }
}