tibba_hook/
lib.rs

1// Copyright 2026 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use dashmap::DashMap;
16use std::future::Future;
17use std::pin::Pin;
18use std::sync::Arc;
19use std::sync::LazyLock;
20use tibba_error::Error;
21use tracing::info;
22
23/// 该 crate 所有日志事件的 tracing target。
24/// 可通过 `RUST_LOG=tibba:hook=info`（或 `debug`）进行过滤。
25const LOG_TARGET: &str = "tibba:hook";
26
27type Result<T> = std::result::Result<T, Error>;
28
29/// 装箱的异步 Future，用于 trait object 场景下的异步方法返回类型。
30pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
31
32/// 生命周期钩子 trait，用于在应用启动/关闭时执行自定义逻辑。
33///
34/// - `before`：应用启动前执行（如初始化资源），按优先级从低到高顺序调用。
35/// - `after`：应用关闭后执行（如释放资源），按优先级从高到低顺序调用。
36/// - 返回 `true` 表示该钩子实际执行了操作，会记录耗时日志；返回 `false` 则静默跳过。
37pub trait Task {
38    /// 应用启动前的钩子，默认不执行任何操作。
39    fn before(&self) -> BoxFuture<'_, Result<bool>> {
40        Box::pin(async { Ok(false) })
41    }
42    /// 应用关闭后的钩子，默认不执行任何操作。
43    fn after(&self) -> BoxFuture<'_, Result<bool>> {
44        Box::pin(async { Ok(false) })
45    }
46    /// 执行优先级，数值越小优先级越高（before 阶段），after 阶段反之。默认为 0。
47    fn priority(&self) -> u8 {
48        0
49    }
50}
51
52/// 全局任务注册表，键为任务名称，值为线程安全的任务实例。
53static TASKS: LazyLock<DashMap<String, Arc<dyn Task + Send + Sync>>> = LazyLock::new(DashMap::new);
54
55/// 任务执行阶段：启动前（Before）或关闭后（After）。
56#[derive(Clone, Copy)]
57enum TaskType {
58    Before,
59    After,
60}
61
62impl std::fmt::Display for TaskType {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        match self {
65            TaskType::Before => write!(f, "before"),
66            TaskType::After => write!(f, "after"),
67        }
68    }
69}
70
71/// 按优先级顺序执行所有已注册的钩子任务。
72/// Before 阶段按优先级升序（数值小的先执行），After 阶段按降序。
73/// 任务返回 `true` 时记录耗时日志。
74async fn run_tasks(task_type: TaskType) -> Result<()> {
75    let mut executable_tasks: Vec<_> = TASKS
76        .iter()
77        .map(|item| {
78            (
79                item.key().clone(),      // 任务名称
80                item.value().priority(), // 排序优先级
81                item.value().clone(),    // 任务实例的 Arc 克隆
82            )
83        })
84        .collect();
85
86    match task_type {
87        TaskType::Before => {
88            // 启动前：优先级数值小的先执行
89            executable_tasks.sort_by_key(|k| k.1);
90        }
91        TaskType::After => {
92            // 关闭后：优先级数值大的先执行（与 before 相反）
93            executable_tasks.sort_by_key(|k| std::cmp::Reverse(k.1));
94        }
95    }
96
97    for (name, _, task) in executable_tasks {
98        let start = std::time::Instant::now();
99        let executed = match task_type {
100            TaskType::Before => task.before().await?,
101            TaskType::After => task.after().await?,
102        };
103
104        if executed {
105            info!(
106                target: LOG_TARGET,
107                task_type = %task_type,
108                name,
109                elapsed = start.elapsed().as_millis(),
110            );
111        }
112    }
113
114    Ok(())
115}
116
117/// 注册一个具名钩子任务。同名任务重复注册时，新任务会覆盖旧任务。
118pub fn register_task(name: &str, task: Arc<dyn Task + Send + Sync>) {
119    TASKS.insert(name.to_string(), task);
120}
121
122/// 按优先级升序执行所有已注册的 `before` 钩子（应用启动前调用）。
123pub async fn run_before_tasks() -> Result<()> {
124    run_tasks(TaskType::Before).await
125}
126
127/// 按优先级降序执行所有已注册的 `after` 钩子（应用关闭后调用）。
128pub async fn run_after_tasks() -> Result<()> {
129    run_tasks(TaskType::After).await
130}
tibba_hook/lib.rs

tibba_hook/
lib.rs
