scheduled_task_macro/
lib.rs

1use cron::Schedule;
2use darling::ast::NestedMeta;
3use darling::{Error, FromMeta};
4use proc_macro::TokenStream;
5use quote::quote;
6use syn::{parse_macro_input, ItemFn, ReturnType};
7
8
9#[proc_macro_attribute]
10pub fn scheduled(attr: TokenStream, item: TokenStream) -> TokenStream {
11    let ast = parse_macro_input!(item as ItemFn);
12    let args = match NestedMeta::parse_meta_list(attr.into()) {
13        Ok(v) => v,
14        Err(e) => {
15            let e: Error= e.into();
16            return e.write_errors().into()
17        },
18    };
19    let args = match ScheduledArgs::from_list(&args) {
20        Ok(v) => v,
21        Err(e) => return e.write_errors().into(),
22    };
23    println!("{:#?}", args);
24    let rate = &args.rate;
25    let cron = &args.cron;
26
27    if rate.is_none() && cron.is_none() {
28        return Error::custom("cron和rate不能都为空").write_errors().into()
29    }
30
31    if rate.is_some() && cron.is_some() {
32        return Error::custom("cron和rate只能有一个属性有值").write_errors().into()
33    }
34    let func_block = &ast.block;
35    let job_block = if cron.is_some() {
36        let str_cron = cron.as_ref().unwrap().as_str();
37        if Schedule::try_from(str_cron).is_err() {
38            return Error::custom("cron表达式错误").write_errors().into()
39        }
40        quote! {
41            let job = Job::new_async(#str_cron, move|_uuid,_lock|{
42                Box::pin(async move{
43                    #func_block
44                })
45            }).unwrap();
46            sched.add(job).await.unwrap();
47        }
48    } else {
49        let rate = rate.unwrap();
50        if rate == 0 {
51            return Error::custom("rate的值必须大于0").write_errors().into()
52        }
53        quote! {
54            let job = Job::new_repeated_async(std::time::Duration::from_secs(#rate), move|_uuid,_lock|{
55                Box::pin(async move{
56                    #func_block
57                })
58            }).unwrap();
59            sched.add(job).await.unwrap();
60        }
61    };
62    let func_vis = &ast.vis;
63    let func_decl = &ast.sig;
64    let asyncness = func_decl.asyncness;
65    if asyncness.is_none() {
66        return Error::custom("必须是异步函数").write_errors().into()
67    }
68    let func_inputs = &func_decl.inputs;
69    if !func_inputs.is_empty() {
70        return Error::custom("函数不能有参数").write_errors().into()
71    }
72    let return_type = &func_decl.output;
73    if let ReturnType::Type(_, _) = return_type {
74        return Error::custom("函数不能有返回值").write_errors().into()
75    }
76
77
78    let task_start_block_header = quote! {
79        let sched = JobScheduler::new().await.unwrap();
80    };
81    let task_start_block_tail = quote! {
82        sched.start().await.unwrap();
83    };
84    quote! {
85        #func_vis #func_decl{
86            use scheduled_task::{Job, JobScheduler};
87            #task_start_block_header
88            #job_block
89            #task_start_block_tail
90        }
91    }.into()
92}
93
94
95#[derive(Debug, FromMeta)]
96struct ScheduledArgs {
97    /// cron表达式
98    cron: Option<String>,
99    /// 间隔时间，单位：秒
100    rate: Option<u64>,
101}
scheduled_task_macro/lib.rs

scheduled_task_macro/
lib.rs
