snerd_rust/

queue.rs

use chrono::Utc;
use std::collections::{HashMap, HashSet, BinaryHeap};
use tokio::sync::Semaphore;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::RwLock;

use crate::file_store::FileStore;
use crate::rate_limiter::RateLimiter;
use crate::task::{RetryableTask, PriorityTask, ProgressMessage};
use tokio::sync::broadcast;

pub type TaskHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;
pub type MaxRetryHandler = Arc<dyn Fn(String) -> Result<(), String> + Send + Sync>;

struct ExecutingGuard {
    executing_tasks: Arc<Mutex<HashSet<String>>>,
    task_id: String,
}

impl Drop for ExecutingGuard {
    fn drop(&mut self) {
        if let Ok(mut executing) = self.executing_tasks.lock() {
            executing.remove(&self.task_id);
        }
    }
}

#[derive(Clone)]
pub struct SnerdQueue {
    pub name: String,
    pub file_store: FileStore,
    pub rate_limiter: RateLimiter,
    task_handlers: Arc<RwLock<HashMap<String, TaskHandler>>>,
    max_retry_handlers: Arc<RwLock<HashMap<String, MaxRetryHandler>>>,
    active_hashes: Arc<Mutex<HashSet<String>>>,
    executing_tasks: Arc<Mutex<HashSet<String>>>,
    worker_semaphore: Arc<Semaphore>,
    pub progress_tx: broadcast::Sender<ProgressMessage>,
}

impl SnerdQueue {
    pub fn new(name: &str, file_store: FileStore, rate_limiter: RateLimiter) -> Self {
        let mut initial_hashes = HashSet::new();
        if let Ok(tasks) = file_store.read_tasks() {
            for task in tasks {
                if task.deleted_at.is_none() {
                    if let Some(hash) = task.payload_hash {
                        initial_hashes.insert(hash);
                    }
                }
            }
        }

        let (progress_tx, _) = broadcast::channel(1024);
        Self {
            name: name.to_string(),
            file_store,
            rate_limiter,
            task_handlers: Arc::new(RwLock::new(HashMap::new())),
            max_retry_handlers: Arc::new(RwLock::new(HashMap::new())),
            active_hashes: Arc::new(Mutex::new(initial_hashes)),
            executing_tasks: Arc::new(Mutex::new(HashSet::new())),
            worker_semaphore: Arc::new(Semaphore::new(100)), // Limit to 100 concurrent tasks
            progress_tx,
        }
    }

    pub fn subscribe_progress(&self) -> broadcast::Receiver<ProgressMessage> {
        self.progress_tx.subscribe()
    }

    pub fn yield_progress(&self, task_id: &str, data: &str) {
        let _ = self.progress_tx.send(ProgressMessage {
            task_id: task_id.to_string(),
            data: data.to_string(),
        });
    }

    pub async fn register_task_handler<F>(&self, task_type: &str, handler: F)
    where
        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
    {
        self.task_handlers
            .write()
            .await
            .insert(task_type.to_string(), Arc::new(handler));
    }

    pub async fn register_max_retry_handler<F>(&self, task_type: &str, handler: F)
    where
        F: Fn(String) -> Result<(), String> + Send + Sync + 'static,
    {
        self.max_retry_handlers
            .write()
            .await
            .insert(task_type.to_string(), Arc::new(handler));
    }

    pub fn enqueue(&self, mut task: RetryableTask) -> std::io::Result<()> {
        if let Some(ref hash) = task.payload_hash {
            if let Ok(mut hashes) = self.active_hashes.lock() {
                if hashes.contains(hash) {
                    return Ok(());
                }
                hashes.insert(hash.clone());
            }
        }
        task.deleted_at = None;
        self.file_store.save_task(&task)?;

        if task.retry_after_time <= Utc::now() {
            if let Some(ref group) = task.rate_limit_group {
                if let Some(limit) = task.max_per_minute {
                    match self.rate_limiter.check_and_increment(group, limit) {
                        Ok(true) => {}
                        Ok(false) | Err(_) => {
                            task.retry_after_time = Utc::now() + chrono::Duration::seconds(60);
                            let _ = self.file_store.save_task(&task);
                            return Ok(());
                        }
                    }
                }
            }

            // Lock check before executing
            if let Ok(mut executing) = self.executing_tasks.lock() {
                if executing.contains(&task.task_id) {
                    return Ok(());
                }
                executing.insert(task.task_id.clone());
            }

            let q = self.clone();
            tokio::spawn(async move {
                q.execute_task(task).await;
            });
        }
        Ok(())
    }

    pub async fn start_processor(&self, interval: Duration) {
        let q = self.clone();
        tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);
            loop {
                interval_timer.tick().await;
                q.process_due_tasks().await;
            }
        });
    }

    pub async fn process_due_tasks(&self) {
        let tasks = match self.file_store.read_tasks() {
            Ok(t) => t,
            Err(_) => return,
        };

        let now = Utc::now();
        let mut heap = BinaryHeap::new();
        
        for task in tasks {
            if task.retry_after_time <= now && task.deleted_at.is_none() {
                heap.push(PriorityTask(task));
            }
        }

        let available = self.worker_semaphore.available_permits();
        for _ in 0..available {
            if let Some(PriorityTask(mut task)) = heap.pop() {
                if let Some(ref group) = task.rate_limit_group {
                    if let Some(limit) = task.max_per_minute {
                        match self.rate_limiter.check_and_increment(group, limit) {
                            Ok(true) => {}
                            Ok(false) | Err(_) => {
                                task.retry_after_time = now + chrono::Duration::seconds(60);
                                let _ = self.file_store.save_task(&task);
                                continue;
                            }
                        }
                    }
                }

                // Lock check before executing
                if let Ok(mut executing) = self.executing_tasks.lock() {
                    if executing.contains(&task.task_id) {
                        continue;
                    }
                    executing.insert(task.task_id.clone());
                }

                if let Ok(permit) = self.worker_semaphore.clone().try_acquire_owned() {
                    let q = self.clone();
                    tokio::spawn(async move {
                        let _p = permit;
                        q.execute_task(task).await;
                    });
                }
            } else {
                break;
            }
        }
    }

    async fn execute_task(&self, mut task: RetryableTask) {
        // Drop guard guarantees removal from executing_tasks
        let _guard = ExecutingGuard {
            executing_tasks: Arc::clone(&self.executing_tasks),
            task_id: task.task_id.clone(),
        };

        let handler = {
            let handlers = self.task_handlers.read().await;
            handlers.get(&task.task_type).cloned()
        };

        if let Some(h) = handler {
            let task_data = task.task_data.clone();

            let result = tokio::task::spawn_blocking(move || h(task_data))
                .await
                .unwrap_or_else(|e| Err(format!("Task panic: {:?}", e)));

            match result {
                Ok(_) => {
                    let _ = self.file_store.delete_task(&task.task_id);
                    if let Some(ref hash) = task.payload_hash {
                        if let Ok(mut hashes) = self.active_hashes.lock() {
                            hashes.remove(hash);
                        }
                    }
                }
                Err(e) => {
                    if task.retry_count < task.max_retries {
                        task.update_retry_config(Some(e));
                        let _ = self.file_store.save_task(&task);
                    } else {
                        // Max retries reached
                        let max_handler = {
                            let max_handlers = self.max_retry_handlers.read().await;
                            max_handlers.get(&task.task_type).cloned()
                        };

                        if let Some(mh) = max_handler {
                            let max_data = task.task_data.clone();
                            let _ = tokio::task::spawn_blocking(move || mh(max_data)).await;
                        }

                        let _ = self.file_store.delete_task(&task.task_id);
                        if let Some(ref hash) = task.payload_hash {
                            if let Ok(mut hashes) = self.active_hashes.lock() {
                                hashes.remove(hash);
                            }
                        }
                    }
                }
            }
        }
    }
}