runledger_runtime/
registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4pub use runledger_core::jobs::JobHandler;
5use runledger_core::jobs::{JobHandlerRegistry, JobType};
6use thiserror::Error;
7
8#[derive(Debug, Clone, Error, Eq, PartialEq)]
9#[non_exhaustive]
10pub enum JobRegistryError {
11    #[error("job handler already registered for {job_type}")]
12    DuplicateJobType { job_type: JobType<'static> },
13}
14
15#[derive(Clone, Default)]
16pub struct JobRegistry {
17    handlers: HashMap<JobType<'static>, Arc<dyn JobHandler>>,
18    retry_delay_overrides: HashMap<JobType<'static>, HashMap<&'static str, i32>>,
19}
20
21impl JobRegistry {
22    #[must_use]
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    pub fn register<H>(&mut self, handler: H)
28    where
29        H: JobHandler + 'static,
30    {
31        let handler: Arc<dyn JobHandler> = Arc::new(handler);
32        self.handlers.insert(handler.job_type(), handler);
33    }
34
35    pub fn try_register<H>(&mut self, handler: H) -> Result<(), JobRegistryError>
36    where
37        H: JobHandler + 'static,
38    {
39        self.try_register_boxed(Arc::new(handler))
40    }
41
42    pub fn try_register_boxed(
43        &mut self,
44        handler: Arc<dyn JobHandler>,
45    ) -> Result<(), JobRegistryError> {
46        let job_type = handler.job_type();
47        if self.handlers.contains_key(job_type.as_str()) {
48            return Err(JobRegistryError::DuplicateJobType { job_type });
49        }
50
51        self.handlers.insert(job_type, handler);
52        Ok(())
53    }
54
55    /// Registers the policy retry delay for one job type and failure code.
56    ///
57    /// A lower bound attached directly to [`runledger_core::jobs::JobFailure`]
58    /// may extend this delay but can never shorten it. The worker uses its
59    /// exponential backoff when no matching override exists.
60    ///
61    /// # Panics
62    ///
63    /// Panics when `retry_delay_ms` is not positive.
64    pub fn register_retry_delay_override(
65        &mut self,
66        job_type: JobType<'static>,
67        failure_code: &'static str,
68        retry_delay_ms: i32,
69    ) {
70        assert!(retry_delay_ms > 0, "retry delay override must be positive");
71
72        self.retry_delay_overrides
73            .entry(job_type)
74            .or_default()
75            .insert(failure_code, retry_delay_ms);
76    }
77
78    #[must_use]
79    pub fn get(&self, job_type: JobType<'_>) -> Option<Arc<dyn JobHandler>> {
80        self.handlers.get(job_type.as_str()).cloned()
81    }
82
83    /// Returns the configured policy retry delay for an exact job type and
84    /// failure code.
85    ///
86    /// Handler retry timing is a lower bound that may extend this policy delay
87    /// but cannot shorten it.
88    #[must_use]
89    pub fn retry_delay_override(&self, job_type: JobType<'_>, failure_code: &str) -> Option<i32> {
90        self.retry_delay_overrides
91            .get(job_type.as_str())
92            .and_then(|overrides| overrides.get(failure_code).copied())
93    }
94
95    #[must_use]
96    pub fn registered_types(&self) -> Vec<JobType<'_>> {
97        let mut keys: Vec<JobType<'_>> = self.handlers.keys().copied().collect();
98        keys.sort_unstable();
99        keys
100    }
101}
102
103impl JobHandlerRegistry for JobRegistry {
104    fn register_boxed(&mut self, handler: Arc<dyn JobHandler>) {
105        self.handlers.insert(handler.job_type(), handler);
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use std::sync::Arc;
112
113    use async_trait::async_trait;
114    use runledger_core::jobs::{
115        JobCompletion, JobContext, JobFailure, JobHandlerRegistry, JobType,
116    };
117    use serde_json::{Value, json};
118    use uuid::Uuid;
119
120    use super::{JobHandler, JobRegistry, JobRegistryError};
121
122    struct ExampleHandler;
123
124    #[async_trait]
125    impl JobHandler for ExampleHandler {
126        fn job_type(&self) -> JobType<'static> {
127            JobType::new("jobs.example")
128        }
129
130        async fn execute(
131            &self,
132            _context: JobContext,
133            _payload: Value,
134        ) -> Result<JobCompletion, JobFailure> {
135            Ok(JobCompletion::success())
136        }
137    }
138
139    struct OutputHandler(&'static str);
140
141    #[async_trait]
142    impl JobHandler for OutputHandler {
143        fn job_type(&self) -> JobType<'static> {
144            JobType::new("jobs.example")
145        }
146
147        async fn execute(
148            &self,
149            _context: JobContext,
150            _payload: Value,
151        ) -> Result<JobCompletion, JobFailure> {
152            Ok(JobCompletion::with_output(json!(self.0)))
153        }
154    }
155
156    fn test_context() -> JobContext {
157        JobContext {
158            job_id: Uuid::now_v7(),
159            run_number: 1,
160            attempt: 1,
161            organization_id: None,
162            worker_id: "registry-test-worker".to_string(),
163            checkpoint: None,
164        }
165    }
166
167    async fn registered_output(registry: &JobRegistry) -> Value {
168        registry
169            .get(JobType::new("jobs.example"))
170            .expect("handler exists")
171            .execute(test_context(), json!({}))
172            .await
173            .expect("registered handler should execute")
174            .output()
175            .cloned()
176            .expect("handler should return output")
177    }
178
179    #[tokio::test]
180    async fn registered_handler_executes_successfully_via_trait_object() {
181        let mut registry = JobRegistry::new();
182        registry.register(ExampleHandler);
183        let handler = registry
184            .get(JobType::new("jobs.example"))
185            .expect("handler exists");
186
187        handler
188            .execute(test_context(), json!({}))
189            .await
190            .expect("registered handler should execute");
191    }
192
193    #[test]
194    fn registered_types_returns_sorted_job_types() {
195        let mut registry = JobRegistry::new();
196        registry.register(ExampleHandler);
197
198        assert_eq!(
199            registry.registered_types(),
200            vec![JobType::new("jobs.example")]
201        );
202    }
203
204    #[tokio::test]
205    async fn try_register_rejects_duplicate_job_type_and_keeps_first_handler() {
206        let mut registry = JobRegistry::new();
207        registry
208            .try_register(OutputHandler("first"))
209            .expect("first handler registration should succeed");
210
211        assert_eq!(
212            registry.try_register(OutputHandler("second")),
213            Err(JobRegistryError::DuplicateJobType {
214                job_type: JobType::new("jobs.example"),
215            })
216        );
217        assert_eq!(registered_output(&registry).await, json!("first"));
218    }
219
220    #[tokio::test]
221    async fn try_register_boxed_rejects_duplicate_job_type_and_keeps_first_handler() {
222        let mut registry = JobRegistry::new();
223        registry
224            .try_register_boxed(Arc::new(OutputHandler("first")))
225            .expect("first handler registration should succeed");
226
227        assert_eq!(
228            registry.try_register_boxed(Arc::new(OutputHandler("second"))),
229            Err(JobRegistryError::DuplicateJobType {
230                job_type: JobType::new("jobs.example"),
231            })
232        );
233        assert_eq!(registered_output(&registry).await, json!("first"));
234    }
235
236    #[tokio::test]
237    async fn register_overwrites_duplicate_job_type() {
238        let mut registry = JobRegistry::new();
239        registry.register(OutputHandler("first"));
240        registry.register(OutputHandler("second"));
241
242        assert_eq!(registered_output(&registry).await, json!("second"));
243    }
244
245    #[tokio::test]
246    async fn trait_register_boxed_overwrites_duplicate_job_type() {
247        let mut registry = JobRegistry::new();
248        JobHandlerRegistry::register_boxed(&mut registry, Arc::new(OutputHandler("first")));
249        JobHandlerRegistry::register_boxed(&mut registry, Arc::new(OutputHandler("second")));
250
251        assert_eq!(registered_output(&registry).await, json!("second"));
252    }
253
254    #[test]
255    fn retry_delay_override_matches_job_type_and_failure_code() {
256        let mut registry = JobRegistry::new();
257        registry.register_retry_delay_override(
258            JobType::new("jobs.example"),
259            "job.example.wait",
260            42,
261        );
262
263        assert_eq!(
264            registry.retry_delay_override(JobType::new("jobs.example"), "job.example.wait"),
265            Some(42)
266        );
267        assert_eq!(
268            registry.retry_delay_override(JobType::new("jobs.other"), "job.example.wait"),
269            None
270        );
271        assert_eq!(
272            registry.retry_delay_override(JobType::new("jobs.example"), "job.example.other"),
273            None
274        );
275    }
276
277    #[test]
278    #[should_panic(expected = "retry delay override must be positive")]
279    fn retry_delay_override_rejects_zero_delay() {
280        let mut registry = JobRegistry::new();
281        registry.register_retry_delay_override(JobType::new("jobs.example"), "job.example.wait", 0);
282    }
283}
runledger_runtime/registry.rs

runledger_runtime/
registry.rs
