Skip to main content

chronon_backend_mem/store/
mod.rs

1//! In-memory [`SchedulerStore`] implementation split by concern.
2
3mod claims;
4mod coordinator;
5mod jobs;
6mod runs;
7mod trait_impl;
8
9use std::collections::HashMap;
10use std::sync::RwLock;
11
12use chrono::{DateTime, Utc};
13use chronon_core::error::{ChrononError, Result};
14use chronon_core::models::{
15    Job, JobRevision, PartitionAssignment, Run, SchedulerLeader, Script, Worker,
16};
17
18/// Fixed primary key for the singleton leader election row.
19pub const LEADER_ROW_ID: &str = "singleton";
20
21/// Thread-safe in-memory persistence for jobs, runs, and coordinator metadata.
22///
23/// Process-local and **non-durable** — suitable for Mode 1 embedded experiments, unit tests,
24/// examples, and benchmarks. Do **not** share across processes; use SQLite or Postgres for
25/// durable / Mode 2 topologies.
26///
27/// Enable the facade `mem` feature to re-export this type from `chronon`. Wire with
28/// `ChrononBuilder::scheduler_store(Arc::new(InMemorySchedulerStore::new()))` or
29/// [`crate::install_default_mem_store`].
30///
31/// # Examples
32///
33/// ```
34/// use std::sync::Arc;
35/// use chronon_backend_mem::InMemorySchedulerStore;
36/// use chronon_core::{Job, SchedulerStore};
37///
38/// # #[tokio::main]
39/// # async fn main() -> chronon_core::Result<()> {
40/// let store = Arc::new(InMemorySchedulerStore::new());
41/// store.upsert_job(&Job::new("demo", "noop")).await?;
42/// assert_eq!(store.list_jobs().await?.len(), 1);
43/// # Ok(())
44/// # }
45/// ```
46#[derive(Default)]
47pub struct InMemorySchedulerStore {
48    /// Jobs keyed by [`Job::job_id`].
49    pub(super) jobs: RwLock<HashMap<String, Job>>,
50    /// Secondary index from [`Job::job_name`] to job id.
51    pub(super) jobs_by_name: RwLock<HashMap<String, String>>,
52    /// Runs keyed by [`Run::run_id`].
53    pub(super) runs: RwLock<HashMap<String, Run>>,
54    /// Job revision history keyed by job id.
55    pub(super) revisions: RwLock<HashMap<String, Vec<JobRevision>>>,
56    /// Script metadata keyed by script name.
57    pub(super) scripts: RwLock<HashMap<String, Script>>,
58    /// Singleton scheduler leader row.
59    pub(super) leader: RwLock<Option<SchedulerLeader>>,
60    /// Partition assignments keyed by partition id.
61    pub(super) partitions: RwLock<HashMap<String, PartitionAssignment>>,
62    /// Registered workers keyed by worker id.
63    pub(super) workers: RwLock<HashMap<String, Worker>>,
64}
65
66impl InMemorySchedulerStore {
67    /// Creates an empty store.
68    ///
69    /// # Examples
70    ///
71    /// ```
72    /// # #[tokio::main]
73    /// # async fn main() {
74    /// use chronon_backend_mem::InMemorySchedulerStore;
75    /// use chronon_core::SchedulerStore;
76    ///
77    /// let store = InMemorySchedulerStore::new();
78    /// assert!(store.list_jobs().await.unwrap().is_empty());
79    /// # }
80    /// ```
81    pub fn new() -> Self {
82        Self::default()
83    }
84
85    /// Insert or replace a job and update the name index.
86    pub(super) fn write_job(&self, job: Job) -> Result<()> {
87        let mut jobs = self.jobs.write().expect("jobs lock");
88        let mut by_name = self.jobs_by_name.write().expect("jobs_by_name lock");
89        by_name.insert(job.job_name.clone(), job.job_id.clone());
90        jobs.insert(job.job_id.clone(), job);
91        Ok(())
92    }
93
94    /// Mutate an existing job and stamp `updated_at` to now.
95    pub(super) fn mutate_job<F>(&self, job_id: &str, f: F) -> Result<()>
96    where
97        F: FnOnce(&mut Job),
98    {
99        let mut jobs = self.jobs.write().expect("jobs lock");
100        let job = jobs
101            .get_mut(job_id)
102            .ok_or_else(|| ChrononError::JobNotFound(job_id.to_string()))?;
103        f(job);
104        job.updated_at = Utc::now();
105        Ok(())
106    }
107
108    /// Mutate an existing job and stamp `updated_at` to `now`.
109    pub(super) fn mutate_job_at<F>(
110        &self,
111        job_id: &str,
112        now: DateTime<Utc>,
113        f: F,
114    ) -> Result<()>
115    where
116        F: FnOnce(&mut Job),
117    {
118        let mut jobs = self.jobs.write().expect("jobs lock");
119        let job = jobs
120            .get_mut(job_id)
121            .ok_or_else(|| ChrononError::JobNotFound(job_id.to_string()))?;
122        f(job);
123        job.updated_at = now;
124        Ok(())
125    }
126}