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        quote! {
51            let job = Job::new_repeated_async(std::time::Duration::from_secs(#rate), move|_uuid,_lock|{
52                Box::pin(async move{
53                    #func_block
54                })
55            }).unwrap();
56            sched.add(job).await.unwrap();
57        }
58    };
59    let func_vis = &ast.vis;
60    let func_decl = &ast.sig;
61    let asyncness = func_decl.asyncness;
62    if asyncness.is_none() {
63        return Error::custom("必须是异步函数").write_errors().into()
64    }
65    let func_inputs = &func_decl.inputs;
66    if !func_inputs.is_empty() {
67        return Error::custom("函数不能有参数").write_errors().into()
68    }
69    let return_type = &func_decl.output;
70    if let ReturnType::Type(_, _) = return_type {
71        return Error::custom("函数不能有返回值").write_errors().into()
72    }
73
74
75    let task_start_block_header = quote! {
76        let sched = JobScheduler::new().await.unwrap();
77    };
78    let task_start_block_tail = quote! {
79        sched.start().await.unwrap();
80    };
81    quote! {
82        #func_vis #func_decl{
83            use scheduled_task::{Job, JobScheduler};
84            #task_start_block_header
85            #job_block
86            #task_start_block_tail
87        }
88    }.into()
89}
90
91
92#[derive(Debug, FromMeta)]
93struct ScheduledArgs {
94    /// cron表达式
95    cron: Option<String>,
96    /// 间隔时间，单位：秒
97    rate: Option<u64>,
98}
scheduled_task_macro/lib.rs

scheduled_task_macro/
lib.rs
