apalis-file-storage 0.1.0-rc.9

#![cfg_attr(docsrs, feature(doc_cfg))]
#![doc = include_str!("../README.md")]
use serde_json::Value;
use std::{
    collections::BTreeMap,
    fs::{File, OpenOptions},
    io::{BufRead, Write},
    path::PathBuf,
    sync::{Arc, RwLock},
};

use self::util::{TaskKey, TaskWithMeta};
use apalis_core::{
    features_table,
    task::{
        Task,
        status::Status,
        task_id::{RandomId, TaskId},
    },
};
use std::io::{BufReader, BufWriter};

mod backend;
mod meta;
mod shared;
mod sink;
mod util;

pub use self::shared::SharedJsonStore;
pub use meta::JsonMapMetadata;
/// A backend that persists to a file using json encoding
///
/// *Warning*: This backend is not optimized for high-throughput scenarios and is best suited for development, testing, or low-volume workloads.
///
/// # Example
///
/// Creates a temporary JSON storage backend
/// ```rust
/// # use apalis_file_storage::JsonStorage;;
/// # pub fn setup_json_storage() -> JsonStorage<u32> {
/// let mut backend = JsonStorage::new_temp().unwrap();
/// # backend
/// # }
/// ```
#[doc = features_table! {
    setup = r#"
        # {
        #   use apalis_file_storage::JsonStorage;;
        #   let mut backend = JsonStorage::new_temp().unwrap();
        #   backend
        # };
    "#,
    Backend => supported("Basic Backend functionality", true),
    TaskSink => supported("Ability to push new tasks", true),
    Serialization => limited("Serialization support for arguments. Only accepts `json`", false),
    WebUI => not_implemented("Expose a web interface for monitoring tasks"),
    FetchById => not_implemented("Allow fetching a task by its ID"),
    RegisterWorker => not_supported("Allow registering a worker with the backend"),
    "[`PipeExt`]" => supported("Allow other backends to pipe to this backend", false),
    MakeShared => supported("Share the same JSON storage across multiple workers via [`SharedJsonStore`]", false),
    Workflow => supported("Flexible enough to support workflows", true),
    WaitForCompletion => supported("Wait for tasks to complete without blocking", true),
    ResumeById => not_implemented("Resume a task by its ID"),
    ResumeAbandoned => not_implemented("Resume abandoned tasks"),
    ListWorkers => not_supported("List all workers registered with the backend"),
    ListTasks => not_implemented("List all tasks in the backend"),
}]
///
/// [`PipeExt`]: crate::backend::pipe::PipeExt
#[derive(Debug)]
pub struct JsonStorage<Args> {
    tasks: Arc<RwLock<BTreeMap<TaskKey, TaskWithMeta>>>,
    buffer: Vec<Task<Value, JsonMapMetadata, RandomId>>,
    path: PathBuf,
    _marker: std::marker::PhantomData<Args>,
}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
struct StorageEntry {
    task_id: TaskId<RandomId>,
    status: Status,
    task: TaskWithMeta,
}

impl<Args> JsonStorage<Args> {
    /// Creates a new `JsonStorage` instance using the specified file path.
    pub fn new(path: impl Into<PathBuf>) -> std::io::Result<Self> {
        let path = path.into();
        let mut data = BTreeMap::new();

        if path.exists() {
            let file = File::open(&path)?;
            let reader = BufReader::new(file);

            for line in reader.lines() {
                let line = line?;
                if line.trim().is_empty() {
                    continue;
                }

                if let Ok(entry) = serde_json::from_str::<StorageEntry>(&line) {
                    let key = TaskKey {
                        status: entry.status,
                        task_id: entry.task_id,
                        queue: std::any::type_name::<Args>().to_owned(),
                    };
                    data.insert(key, entry.task);
                }
            }
        }

        Ok(Self {
            path,
            tasks: Arc::new(RwLock::new(data)),
            buffer: Vec::new(),
            _marker: std::marker::PhantomData,
        })
    }

    /// Creates a new temporary `JsonStorage` instance.
    pub fn new_temp() -> Result<Self, std::io::Error> {
        let p = std::env::temp_dir().join(format!("apalis-json-store-{}", RandomId::default()));
        Self::new(p)
    }

    fn insert(&self, k: &TaskKey, v: TaskWithMeta) -> Result<(), std::io::Error> {
        self.tasks.try_write().unwrap().insert(k.clone(), v);
        Ok(())
    }

    /// Removes a task from the storage.
    pub fn remove(&mut self, key: &TaskKey) -> std::io::Result<Option<TaskWithMeta>> {
        let removed = self.tasks.try_write().unwrap().remove(key);

        if removed.is_some() {
            self.persist_to_disk()?;
        }

        Ok(removed)
    }

    /// Persist all current data to disk by rewriting the file
    fn persist_to_disk(&self) -> std::io::Result<()> {
        let tmp_path = &self.path;
        {
            let tmp_file = OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open(tmp_path)?;
            let mut writer = BufWriter::new(tmp_file);

            for (key, value) in self.tasks.try_read().unwrap().iter() {
                let entry = StorageEntry {
                    status: key.status.clone(),
                    task_id: key.task_id.clone(),
                    task: value.clone(),
                };
                let line = serde_json::to_string(&entry)?;
                writeln!(writer, "{line}")?;
            }

            writer.flush()?;
        } // BufWriter is dropped here, ensuring all data is written

        // Atomically replace the old file with the new one
        std::fs::rename(tmp_path, &self.path)?;
        Ok(())
    }
    /// Reload data from disk, useful if the file was modified externally
    pub fn reload(&mut self) -> std::io::Result<()> {
        let mut new_data = BTreeMap::new();

        if self.path.exists() {
            let file = File::open(&self.path)?;
            let reader = BufReader::new(file);

            for line in reader.lines() {
                let line = line?;
                if line.trim().is_empty() {
                    continue;
                }

                if let Ok(entry) = serde_json::from_str::<StorageEntry>(&line) {
                    let key = TaskKey {
                        status: entry.status,
                        task_id: entry.task_id,
                        queue: std::any::type_name::<Args>().to_owned(),
                    };
                    new_data.insert(key, entry.task);
                }
            }
        }

        *self.tasks.try_write().unwrap() = new_data;
        Ok(())
    }
    /// Clear all data from memory and file
    pub fn clear(&mut self) -> std::io::Result<()> {
        self.tasks.try_write().unwrap().clear();

        // Create an empty file
        let file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&self.path)?;
        drop(file);

        Ok(())
    }

    /// Update the status of an existing key
    pub fn update_status(
        &mut self,
        old_key: &TaskKey,
        new_status: Status,
    ) -> std::io::Result<bool> {
        let mut tasks = self.tasks.try_write().unwrap();
        if let Some(value) = tasks.remove(old_key) {
            let new_key = TaskKey {
                status: new_status,
                task_id: old_key.task_id.clone(),
                queue: old_key.queue.clone(),
            };
            tasks.insert(new_key, value);
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Retrieves a task from the storage.
    #[must_use]
    pub fn get(&self, key: &TaskKey) -> Option<TaskWithMeta> {
        let tasks = self.tasks.try_read().unwrap();
        let res = tasks.get(key);
        res.cloned()
    }

    fn update_result(&self, key: &TaskKey, status: Status, val: Value) -> std::io::Result<bool> {
        let mut tasks = self.tasks.try_write().unwrap();
        if let Some(mut task) = tasks.remove(key) {
            let new_key = TaskKey {
                status,
                task_id: key.task_id.clone(),
                queue: key.queue.clone(),
            };
            task.result = Some(val);

            tasks.insert(new_key, task);
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

impl<Args> Clone for JsonStorage<Args> {
    fn clone(&self) -> Self {
        Self {
            tasks: self.tasks.clone(),
            buffer: Vec::new(),
            path: self.path.clone(),
            _marker: std::marker::PhantomData,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    use apalis_core::{
        backend::TaskSink,
        error::BoxDynError,
        worker::{
            builder::WorkerBuilder, context::WorkerContext, ext::event_listener::EventListenerExt,
        },
    };

    const ITEMS: u32 = 100;

    #[tokio::test]
    async fn basic_worker() {
        let mut json_store = JsonStorage::new_temp().unwrap();
        for i in 0..ITEMS {
            json_store.push(i).await.unwrap();
        }

        async fn task(task: u32, ctx: WorkerContext) -> Result<(), BoxDynError> {
            tokio::time::sleep(Duration::from_secs(1)).await;
            if task == ITEMS - 1 {
                ctx.stop().unwrap();
                return Err("Worker stopped!")?;
            }
            Ok(())
        }

        let worker = WorkerBuilder::new("rango-tango")
            .backend(json_store)
            .on_event(|ctx, ev| {
                println!("On Event = {ev:?} from = {}", ctx.name());
            })
            .build(task);
        worker.run().await.unwrap();
    }
}