Skip to main content

alun_task/
registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use crate::handler::TaskHandler;
5use crate::TaskConfig;
6
7/// 处理器注册中心
8///
9/// 按 `task_type` 注册 `TaskHandler` 实例及其 `TaskConfig`。
10/// 在应用启动时一次性完成注册,运行时只读访问。
11#[derive(Clone)]
12pub struct HandlerRegistry {
13    /// task_type → (handler, config) 映射
14    handlers: HashMap<i16, (Arc<dyn TaskHandler>, TaskConfig)>,
15}
16
17impl HandlerRegistry {
18    /// 创建空的处理器注册中心
19    pub fn new() -> Self {
20        Self {
21            handlers: HashMap::new(),
22        }
23    }
24
25    /// 注册一个处理器及其配置
26    pub fn register(
27        &mut self,
28        handler: impl TaskHandler + 'static,
29        config: TaskConfig,
30    ) -> &mut Self {
31        let task_type = handler.task_type();
32        if task_type != config.task_type {
33            tracing::warn!(
34                "Handler task_type ({}) 与 config.task_type ({}) 不一致,将使用 config 中的值",
35                task_type, config.task_type
36            );
37        }
38        self.handlers
39            .insert(config.task_type, (Arc::new(handler), config));
40        self
41    }
42
43    /// 从编译期自动发现的 handler 注册(配合 `#[task_handler]` 宏)
44    ///
45    /// 读取 `TASK_HANDLERS` 分布式切片中的所有条目并注册。
46    /// 可安全调用多次(后续条目覆盖先前的同名 task_type)。
47    pub fn from_discovered(&mut self) -> &mut Self {
48        for entry in crate::TASK_HANDLERS {
49            let handler = (entry.handler_fn)();
50            let config = (entry.config_fn)();
51            self.handlers
52                .insert(entry.task_type, (Arc::from(handler), config));
53        }
54        self
55    }
56
57    /// 按 task_type 获取处理器和配置
58    pub fn get(&self, task_type: i16) -> Option<(Arc<dyn TaskHandler>, &TaskConfig)> {
59        self.handlers
60            .get(&task_type)
61            .map(|(h, c)| (Arc::clone(h), c))
62    }
63
64    /// 获取所有已注册的 task_type
65    pub fn task_types(&self) -> Vec<i16> {
66        self.handlers.keys().copied().collect()
67    }
68
69    /// 获取配置
70    pub fn get_config(&self, task_type: i16) -> Option<&TaskConfig> {
71        self.handlers.get(&task_type).map(|(_, c)| c)
72    }
73
74    /// 已注册的 handler 数量
75    pub fn len(&self) -> usize {
76        self.handlers.len()
77    }
78
79    /// 是否为空
80    pub fn is_empty(&self) -> bool {
81        self.handlers.is_empty()
82    }
83}
84
85impl Default for HandlerRegistry {
86    fn default() -> Self {
87        Self::new()
88    }
89}