framework_proc_timer/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{Error, ItemFn, Meta, ReturnType, Token, parse::Parser};
4
5/// 将异步函数注册为进程内定时任务。
6///
7/// 任务名称固定为函数名，因此同一进程内不能出现同名任务。
8#[proc_macro_attribute]
9pub fn timer(attribute: TokenStream, item: TokenStream) -> TokenStream {
10    match expand_timer(attribute, item) {
11        Ok(tokens) => tokens.into(),
12        Err(error) => error.into_compile_error().into(),
13    }
14}
15
16fn expand_timer(
17    attribute: TokenStream,
18    item: TokenStream,
19) -> Result<proc_macro2::TokenStream, Error> {
20    let function = syn::parse::<ItemFn>(item)?;
21    let (interval, count, retry_immediately_on_panic) = parse_arguments(attribute)?;
22    validate_function(&function)?;
23
24    let function_name = &function.sig.ident;
25    let wrapper_name = syn::Ident::new(
26        &format!("__timer_target_{function_name}"),
27        function_name.span(),
28    );
29    let task_name = function_name.to_string();
30    let retry_immediately_on_panic =
31        retry_immediately_on_panic.map_or_else(|| quote!(None), |value| quote!(Some(#value)));
32
33    Ok(quote! {
34        #function
35
36        fn #wrapper_name() -> ::framework_timer::TimerFuture {
37            Box::pin(#function_name())
38        }
39
40        ::framework_timer::inventory::submit! {
41            ::framework_timer::TimerRegistration {
42                name: #task_name,
43                count: #count,
44                interval: #interval,
45                retry_immediately_on_panic: #retry_immediately_on_panic,
46                target: #wrapper_name,
47            }
48        }
49    })
50}
51
52fn parse_arguments(attribute: TokenStream) -> Result<(i64, usize, Option<bool>), Error> {
53    let arguments =
54        syn::punctuated::Punctuated::<Meta, Token![,]>::parse_terminated.parse(attribute)?;
55    let mut interval = None;
56    let mut count = None;
57    let mut retry_immediately_on_panic = None;
58
59    for argument in arguments {
60        let Meta::NameValue(argument) = argument else {
61            return Err(Error::new_spanned(
62                argument,
63                "timer 参数必须使用 key = value 形式",
64            ));
65        };
66        if argument.path.is_ident("interval") {
67            if interval.is_some() {
68                return Err(Error::new_spanned(argument, "interval 只能指定一次"));
69            }
70            let value = parse_integer(&argument.value, "interval 必须是毫秒整数")?;
71            interval = Some(value);
72        } else if argument.path.is_ident("count") {
73            if count.is_some() {
74                return Err(Error::new_spanned(argument, "count 只能指定一次"));
75            }
76            let value = parse_usize(&argument.value, "count 必须是大于等于 1 的整数")?;
77            if value == 0 {
78                return Err(Error::new_spanned(
79                    argument.value,
80                    "count 必须是大于等于 1 的整数",
81                ));
82            }
83            count = Some(value);
84        } else if argument.path.is_ident("retry_immediately_on_panic") {
85            if retry_immediately_on_panic.is_some() {
86                return Err(Error::new_spanned(
87                    argument,
88                    "retry_immediately_on_panic 只能指定一次",
89                ));
90            }
91            retry_immediately_on_panic = Some(parse_bool(
92                &argument.value,
93                "retry_immediately_on_panic 必须是布尔值",
94            )?);
95        } else {
96            return Err(Error::new_spanned(argument, "不支持的 timer 参数"));
97        }
98    }
99
100    let interval = interval
101        .ok_or_else(|| Error::new(proc_macro2::Span::call_site(), "缺少必填参数 interval"))?;
102    Ok((interval, count.unwrap_or(1), retry_immediately_on_panic))
103}
104
105fn parse_integer(expression: &syn::Expr, message: &str) -> Result<i64, Error> {
106    let syn::Expr::Lit(expression) = expression else {
107        return Err(Error::new_spanned(expression, message));
108    };
109    let syn::Lit::Int(value) = &expression.lit else {
110        return Err(Error::new_spanned(expression, message));
111    };
112    value
113        .base10_parse::<i64>()
114        .map_err(|_| Error::new_spanned(value, message))
115}
116
117fn parse_usize(expression: &syn::Expr, message: &str) -> Result<usize, Error> {
118    let syn::Expr::Lit(expression) = expression else {
119        return Err(Error::new_spanned(expression, message));
120    };
121    let syn::Lit::Int(value) = &expression.lit else {
122        return Err(Error::new_spanned(expression, message));
123    };
124    value
125        .base10_parse::<usize>()
126        .map_err(|_| Error::new_spanned(value, message))
127}
128
129fn parse_bool(expression: &syn::Expr, message: &str) -> Result<bool, Error> {
130    let syn::Expr::Lit(expression) = expression else {
131        return Err(Error::new_spanned(expression, message));
132    };
133    let syn::Lit::Bool(value) = &expression.lit else {
134        return Err(Error::new_spanned(expression, message));
135    };
136    Ok(value.value)
137}
138
139fn validate_function(function: &ItemFn) -> Result<(), Error> {
140    if function.sig.asyncness.is_none() {
141        return Err(Error::new_spanned(
142            &function.sig,
143            "#[timer] 只能标注 async 函数",
144        ));
145    }
146    if !function.sig.inputs.is_empty() {
147        return Err(Error::new_spanned(
148            &function.sig.inputs,
149            "#[timer] 标注的函数不能包含参数",
150        ));
151    }
152    match &function.sig.output {
153        ReturnType::Type(_, output) if matches!(output.as_ref(), syn::Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "TimerResult")) =>
154            {}
155        _ => {
156            return Err(Error::new_spanned(
157                &function.sig.output,
158                "#[timer] 函数必须返回 framework_timer::TimerResult",
159            ));
160        }
161    }
162    Ok(())
163}
framework_proc_timer/lib.rs

framework_proc_timer/
lib.rs
