forge_runtime/jobs/
registry.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use forge_core::job::{ForgeJob, JobContext, JobInfo};
7use forge_core::Result;
8use serde_json::Value;
9
10/// Type alias for boxed job handler function.
11pub type BoxedJobHandler = Arc<
12    dyn Fn(&JobContext, Value) -> Pin<Box<dyn Future<Output = Result<Value>> + Send + '_>>
13        + Send
14        + Sync,
15>;
16
17/// Entry in the job registry.
18pub struct JobEntry {
19    /// Job metadata.
20    pub info: JobInfo,
21    /// Job handler function.
22    pub handler: BoxedJobHandler,
23}
24
25/// Registry of all FORGE jobs.
26#[derive(Clone, Default)]
27pub struct JobRegistry {
28    jobs: HashMap<String, Arc<JobEntry>>,
29}
30
31impl JobRegistry {
32    /// Create a new empty registry.
33    pub fn new() -> Self {
34        Self {
35            jobs: HashMap::new(),
36        }
37    }
38
39    /// Register a job type.
40    pub fn register<J: ForgeJob>(&mut self)
41    where
42        J::Args: serde::de::DeserializeOwned + Send + 'static,
43        J::Output: serde::Serialize + Send + 'static,
44    {
45        let info = J::info();
46        let name = info.name.to_string();
47
48        let handler: BoxedJobHandler = Arc::new(move |ctx, args| {
49            Box::pin(async move {
50                let parsed_args: J::Args = serde_json::from_value(args)
51                    .map_err(|e| forge_core::ForgeError::Validation(e.to_string()))?;
52                let result = J::execute(ctx, parsed_args).await?;
53                serde_json::to_value(result)
54                    .map_err(|e| forge_core::ForgeError::Internal(e.to_string()))
55            })
56        });
57
58        self.jobs.insert(name, Arc::new(JobEntry { info, handler }));
59    }
60
61    /// Get a job entry by name.
62    pub fn get(&self, name: &str) -> Option<Arc<JobEntry>> {
63        self.jobs.get(name).cloned()
64    }
65
66    /// Get job info by name.
67    pub fn info(&self, name: &str) -> Option<&JobInfo> {
68        self.jobs.get(name).map(|e| &e.info)
69    }
70
71    /// Check if a job exists.
72    pub fn exists(&self, name: &str) -> bool {
73        self.jobs.contains_key(name)
74    }
75
76    /// Get all job names.
77    pub fn job_names(&self) -> impl Iterator<Item = &str> {
78        self.jobs.keys().map(|s| s.as_str())
79    }
80
81    /// Get all jobs.
82    pub fn jobs(&self) -> impl Iterator<Item = (&str, &Arc<JobEntry>)> {
83        self.jobs.iter().map(|(k, v)| (k.as_str(), v))
84    }
85
86    /// Get the number of registered jobs.
87    pub fn len(&self) -> usize {
88        self.jobs.len()
89    }
90
91    /// Check if registry is empty.
92    pub fn is_empty(&self) -> bool {
93        self.jobs.is_empty()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_empty_registry() {
103        let registry = JobRegistry::new();
104        assert!(registry.is_empty());
105        assert_eq!(registry.len(), 0);
106        assert!(registry.get("nonexistent").is_none());
107    }
108}