baizekit_app/
application.rs

1use super::command::{Cli, EmptyCommand};
2use super::component::DynComponent;
3use super::component_factory::ComponentFactoryManager;
4use super::signal::shutdown_signal;
5use super::version::GLOBAL_VERSION_PRINTER;
6use anyhow::{Context, Result};
7use clap::{Parser, Subcommand};
8use config::{Config, File};
9use std::any::TypeId;
10use std::collections::HashMap;
11use std::fs;
12use std::future::Future;
13use std::marker::PhantomData;
14use std::path::PathBuf;
15use std::pin::Pin;
16use std::sync::{Arc, Mutex as StdMutex, RwLock as StdRwLock};
17use std::time::Duration;
18use tokio::sync::RwLock;
19use tokio::time::sleep;
20use tracing::info;
21
22/// 组件唯一标识键，由类型 ID 和标签组成
23#[derive(Debug, Clone, PartialEq, Eq, Hash)]
24pub struct ComponentKey {
25    pub type_id: TypeId,
26    pub label: String,
27}
28
29/// 应用内部状态
30#[derive(Default)]
31pub struct ApplicationInner {
32    config: RwLock<Config>,
33    components: RwLock<HashMap<ComponentKey, Arc<dyn DynComponent>>>,
34    wait_signal: StdRwLock<bool>,
35}
36
37impl ApplicationInner {
38    /// 获取组件（返回 Option<Arc<C>>）
39    pub async fn get_component<C: DynComponent + 'static>(&self, label: Option<&str>) -> Option<Arc<C>> {
40        let type_id = TypeId::of::<C>();
41        let label = label.unwrap_or("default").to_string();
42        let key = ComponentKey { type_id, label };
43
44        let components = self.components.read().await;
45        components
46            .get(&key)
47            .cloned()
48            // 将 Arc<dyn DynComponent> 向下转型为 Arc<C>
49            .and_then(|arc| Arc::downcast(arc).ok())
50    }
51
52    /// 强制获取组件（返回 Result<Arc<C>>）
53    pub async fn must_get_component<C: DynComponent + 'static>(&self, label: Option<&str>) -> Result<Arc<C>> {
54        let type_name = std::any::type_name::<C>();
55        let label_str = label.unwrap_or("default");
56
57        self.get_component(label).await.ok_or_else(|| {
58            anyhow::anyhow!("component not found or type mismatch: type={}, label={}", type_name, label_str)
59        })
60    }
61
62    /// 获取配置（返回读锁 Guard）
63    pub async fn config(&self) -> tokio::sync::RwLockReadGuard<'_, Config> {
64        self.config.read().await
65    }
66
67    /// 获取配置可变引用（返回写锁 Guard）
68    pub async fn config_mut(&self) -> tokio::sync::RwLockWriteGuard<'_, Config> {
69        self.config.write().await
70    }
71
72    pub fn set_wait_signal(&self, wait: bool) {
73        let mut wait_signal = self.wait_signal.write().expect("Failed to write wait_signal RwLock");
74        *wait_signal = wait;
75    }
76}
77
78/// 组件初始化策略
79#[derive(Debug, Clone)]
80pub enum InitStrategy {
81    /// 初始化所有组件
82    All,
83    /// 不初始化任何组件
84    None,
85    /// 只初始化指定标签组件
86    Only(Vec<ComponentKey>),
87    /// 初始化除指定标签外的组件
88    Deny(Vec<ComponentKey>),
89}
90
91// 定义通用的处理器Future类型别名，简化重复书写
92type HandlerFuture = Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>;
93
94// 命令处理器类型定义（使用别名简化）
95type CommandHandler<T> = Box<
96    dyn Fn(T, Arc<ApplicationInner>, Arc<ComponentFactoryManager>) -> (InitStrategy, HandlerFuture)
97        + Send
98        + Sync
99        + 'static,
100>;
101
102// 默认处理器类型定义（使用别名简化）
103type DefaultHandler = Box<
104    dyn Fn(Arc<ApplicationInner>, Arc<ComponentFactoryManager>) -> (InitStrategy, HandlerFuture)
105        + Send
106        + Sync
107        + 'static,
108>;
109
110/// 应用核心结构
111pub struct App<T: Subcommand + Clone + 'static = EmptyCommand> {
112    default_handler: StdMutex<Option<DefaultHandler>>,
113    command_handler: StdMutex<Option<CommandHandler<T>>>,
114    component_factories: Arc<ComponentFactoryManager>,
115    inner: Arc<ApplicationInner>,
116    phantom: PhantomData<T>,
117}
118
119impl<T: Subcommand + Clone + 'static> App<T> {
120    /// 创建新的应用实例
121    pub fn new() -> Self {
122        Self {
123            command_handler: StdMutex::new(None),
124            component_factories: Arc::new(ComponentFactoryManager::new()),
125            inner: Arc::new(ApplicationInner { wait_signal: StdRwLock::new(true), ..ApplicationInner::default() }),
126            default_handler: StdMutex::new(None),
127            phantom: PhantomData,
128        }
129    }
130
131    /// 注册命令处理器
132    pub fn register_command_handler<F, Fut>(&self, handler: F) -> &Self
133    where
134        F: Fn(T, Arc<ApplicationInner>, Arc<ComponentFactoryManager>) -> (InitStrategy, Fut) + Send + Sync + 'static,
135        Fut: Future<Output = Result<()>> + Send + 'static,
136    {
137        // 包装处理器，显式转换为动态Future类型
138        let wrapped = move |cmd, inner, factories| {
139            let (strategy, fut) = handler(cmd, inner, factories);
140            (strategy, Box::pin(fut) as HandlerFuture)
141        };
142        let mut cmd_handler = self.command_handler.lock().expect("Failed to lock command_handler mutex");
143        *cmd_handler = Some(Box::new(wrapped) as CommandHandler<T>);
144        self
145    }
146
147    /// 注册组件工厂
148    pub fn register_component_factory<Comp, F, Fut>(&self, label: Option<String>, factory: F) -> &Self
149    where
150        Comp: DynComponent + 'static,
151        F: Fn(Arc<ApplicationInner>, String) -> Fut + Send + Sync + 'static,
152        Fut: Future<Output = Result<Comp>> + Send + 'static,
153    {
154        self.component_factories.register_component_factory(label, factory);
155        self
156    }
157
158    /// 设置默认处理器
159    pub fn set_default_handler<F, Fut>(&self, handler: F) -> &Self
160    where
161        F: Fn(Arc<ApplicationInner>, Arc<ComponentFactoryManager>) -> (InitStrategy, Fut) + Send + Sync + 'static,
162        Fut: Future<Output = Result<()>> + Send + 'static,
163    {
164        // 包装处理器，显式转换为动态Future类型
165        let wrapped = move |inner, factories| {
166            let (strategy, fut) = handler(inner, factories);
167            (strategy, Box::pin(fut) as HandlerFuture)
168        };
169        let mut def_handler = self.default_handler.lock().expect("Failed to lock default_handler mutex");
170        *def_handler = Some(Box::new(wrapped) as DefaultHandler);
171        self
172    }
173
174    /// 设置是否等待关闭信号
175    pub fn set_wait_signal(&self, wait: bool) -> &Self {
176        self.inner.set_wait_signal(wait);
177        self
178    }
179
180    /// 获取所有已初始化组件的键
181    pub async fn get_all_component_keys(&self) -> Vec<ComponentKey> {
182        let components = self.inner.components.read().await;
183        components.keys().cloned().collect()
184    }
185
186    /// 应用主入口点
187    pub async fn run(&self) -> Result<()> {
188        let cli = Cli::<T>::parse();
189        self.load_config(&cli.config).await?;
190        let inner_arc = self.inner.clone();
191        let factories_arc = self.component_factories.clone();
192
193        // 根据命令或默认情况获取处理策略和执行未来
194        let (init_strategy, execute_future) = match &cli.command {
195            Some(command) => {
196                let handler = self
197                    .command_handler
198                    .lock()
199                    .map_err(|e| anyhow::anyhow!("get command handler lock failed: {}", e))?;
200                let handler = handler.as_ref().context("command handler not registered")?;
201                handler(command.clone(), inner_arc.clone(), factories_arc.clone())
202            }
203            None => {
204                let default_handler = self
205                    .default_handler
206                    .lock()
207                    .map_err(|e| anyhow::anyhow!("get default handler lock failed: {}", e))?;
208                if cli.version {
209                    self.set_wait_signal(false);
210                    let fut: HandlerFuture = Box::pin(async {
211                        if let Some(print_version) = GLOBAL_VERSION_PRINTER.get() {
212                            print_version();
213                        }
214                        Ok(())
215                    });
216                    (InitStrategy::None, fut)
217                } else {
218                    default_handler
219                        .as_ref()
220                        .map(|h| h(inner_arc.clone(), factories_arc.clone()))
221                        .unwrap_or_else(|| {
222                            let fut: HandlerFuture = Box::pin(async { Ok(()) });
223                            (InitStrategy::All, fut)
224                        })
225                }
226            }
227        };
228
229        // 初始化组件
230        self.init_components_with_strategy(init_strategy)
231            .await
232            .context("component init failed")?;
233
234        // 执行主逻辑
235        execute_future.await?;
236
237        // 处理等待关闭信号
238        let wait_signal = self
239            .inner
240            .wait_signal
241            .read()
242            .map_err(|e| anyhow::anyhow!("get wait signal lock failed: {}", e))?;
243        if *wait_signal {
244            info!("等待 ctrl+c 信号...");
245            shutdown_signal().await;
246            info!("收到 ctrl+c 信号，正在关闭应用...");
247        }
248
249        // 关闭组件
250        self.shutdown_components().await.context("shutdown component failed.")?;
251        println!("应用已退出");
252        sleep(Duration::from_millis(100)).await;
253        Ok(())
254    }
255
256    /// 根据策略初始化组件
257    async fn init_components_with_strategy(&self, strategy: InitStrategy) -> Result<()> {
258        let inner = self.inner.clone();
259        // 取出所有工厂并清空内部存储
260        let (factories, _) = self.component_factories.take_factories();
261
262        // 过滤需要初始化的工厂
263        let filtered_factories: Vec<_> = match strategy {
264            InitStrategy::All => factories.iter().collect(),
265            InitStrategy::None => Vec::new(),
266            InitStrategy::Only(keys_to_init) => {
267                factories.iter().filter(|(key, _)| keys_to_init.contains(key)).collect()
268            }
269            InitStrategy::Deny(keys_to_deny) => {
270                factories.iter().filter(|(key, _)| !keys_to_deny.contains(key)).collect()
271            }
272        };
273
274        // 合并后的组件创建与初始化阶段（按注册顺序）
275        for (key, factory) in &filtered_factories {
276            // 1. 创建组件实例（此时为可变状态）
277            let mut component = factory.create(inner.clone(), key.label.clone()).await?;
278            let type_name = component.type_name();
279
280            // 2. 立即初始化组件（使用可变引用）
281            let config = inner.config.read().await;
282            component.init(&config, key.label.clone()).await?;
283            info!(com = type_name, label = &key.label, "组件创建并初始化成功");
284
285            // 3. 转换为Arc并存储（初始化完成后转为不可变共享）
286            let arc_component: Arc<dyn DynComponent> = Arc::from(component);
287            //因为上面组件init或者factory都有可能会获取component,所以只能在循环中获取写锁
288            let mut components = inner.components.write().await;
289            components.insert(key.clone(), arc_component);
290        }
291
292        Ok(())
293    }
294
295    /// 关闭所有组件
296    async fn shutdown_components(&self) -> Result<()> {
297        let inner = self.inner.clone();
298        let mut components = inner.components.write().await;
299        let mut component_keys: Vec<ComponentKey> = components.keys().cloned().collect();
300        component_keys.reverse(); // 反向关闭（依赖顺序相反）
301
302        for key in component_keys {
303            if let Some(component) = components.get(&key) {
304                let type_name = component.type_name();
305                component.shutdown().await?;
306                info!(com = type_name, label = &key.label, "组件关闭成功");
307            }
308        }
309
310        // 清除组件
311        components.clear();
312        Ok(())
313    }
314
315    /// 加载配置文件
316    async fn load_config(&self, config_path: &Option<PathBuf>) -> Result<()> {
317        let mut config_builder = Config::builder();
318        if let Some(path) = config_path {
319            let path = fs::canonicalize(path).context(format!("canonicalize config path failed: {:?}", path))?;
320            config_builder = config_builder.add_source(File::from(path));
321        }
322        let config = config_builder.build().context("config load failed")?;
323        *self.inner.config.write().await = config;
324        Ok(())
325    }
326}
327
328impl<T: Subcommand + Clone + 'static> Default for App<T> {
329    fn default() -> Self {
330        Self::new()
331    }
332}
333
334impl App<EmptyCommand> {
335    /// 创建带有空命令的应用实例
336    pub fn with_empty_command() -> Self {
337        Self::new()
338    }
339}
baizekit_app/application.rs

baizekit_app/
application.rs
