tibba_hook/
lib.rs

1// Copyright 2025 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 async_trait::async_trait;
16use dashmap::DashMap;
17use once_cell::sync::Lazy;
18use std::sync::Arc;
19use tibba_error::Error;
20use tracing::info;
21
22type Result<T> = std::result::Result<T, Error>;
23
24#[async_trait]
25pub trait Task {
26    async fn before(&self) -> Result<bool> {
27        Ok(false)
28    }
29    async fn after(&self) -> Result<bool> {
30        Ok(false)
31    }
32    fn priority(&self) -> u8 {
33        0
34    }
35}
36
37static TASKS: Lazy<DashMap<String, Arc<dyn Task + Send + Sync>>> = Lazy::new(DashMap::new);
38
39#[derive(Clone, Copy)]
40enum TaskType {
41    Before,
42    After,
43}
44
45impl std::fmt::Display for TaskType {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        match self {
48            TaskType::Before => write!(f, "before"),
49            TaskType::After => write!(f, "after"),
50        }
51    }
52}
53
54// Internal function to execute a set of tasks in priority order
55// Parameters:
56// - tasks: reference to a DashMap containing the tasks to execute
57async fn run_tasks(task_type: TaskType) -> Result<()> {
58    let mut executable_tasks: Vec<_> = TASKS
59        .iter()
60        .map(|item| {
61            (
62                item.key().clone(),      // Task name
63                item.value().priority(), // Priority for sorting
64                item.value().clone(),    // Cloned Arc to the task
65            )
66        })
67        .collect();
68
69    match task_type {
70        TaskType::Before => {
71            executable_tasks.sort_by_key(|k| k.1);
72        }
73        TaskType::After => {
74            executable_tasks.sort_by_key(|k| std::cmp::Reverse(k.1));
75        }
76    }
77
78    // Execute tasks in the sorted order.
79    for (name, _, task) in executable_tasks {
80        let start = std::time::Instant::now();
81        let executed = match task_type {
82            TaskType::Before => task.before().await?,
83            TaskType::After => task.after().await?,
84        };
85
86        if executed {
87            info!(
88                category = "task",
89                task_type = task_type.to_string(),
90                name,
91                elapsed = start.elapsed().as_millis(),
92            );
93        }
94    }
95
96    Ok(())
97}
98
99pub fn register_task(name: &str, task: Arc<dyn Task + Send + Sync>) {
100    TASKS.insert(name.to_string(), task);
101}
102
103// Executes all registered "before" tasks in priority order
104pub async fn run_before_tasks() -> Result<()> {
105    run_tasks(TaskType::Before).await
106}
107
108// Executes all registered "after" tasks in priority order
109pub async fn run_after_tasks() -> Result<()> {
110    run_tasks(TaskType::After).await
111}
tibba_hook/lib.rs

tibba_hook/
lib.rs
