rust_task_queue_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, DeriveInput, Error};
4
5/// Automatically register a task type with the task registry.
6///
7/// This derive macro generates code that uses the inventory pattern to automatically
8/// register task types at runtime. The task will be registered using the name returned
9/// by its `name()` method.
10///
11/// This macro uses the inventory pattern to automatically register task types at runtime.
12///
13/// # Example
14///
15/// ```rust
16/// #[derive(Debug, Serialize, Deserialize, AutoRegisterTask)]
17/// struct MyTask {
18///     data: String,
19/// }
20///
21/// #[async_trait]
22/// impl Task for MyTask {
23///     async fn execute(&self) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
24///         // task implementation  
25///         use serde::Serialize;
26///         #[derive(Serialize)]
27///         struct Response { status: String }
28///         let response = Response { status: "completed".to_string() };
29///         Ok(rmp_serde::to_vec(&response)?)
30///     }
31///     
32///     fn name(&self) -> &str {
33///         "my_task"
34///     }
35/// }
36/// ```
37#[proc_macro_derive(AutoRegisterTask)]
38pub fn auto_register_task(input: TokenStream) -> TokenStream {
39    let input = parse_macro_input!(input as DeriveInput);
40
41    impl_auto_register_task(&input).unwrap_or_else(|err| err.to_compile_error().into())
42}
43
44fn impl_auto_register_task(input: &DeriveInput) -> Result<TokenStream, Error> {
45    let type_name = &input.ident;
46
47    // Generate the inventory submission code
48    let expanded = quote! {
49        // Submit this task type to the inventory for automatic registration
50        ::rust_task_queue::inventory::submit! {
51            ::rust_task_queue::TaskRegistration {
52                type_name: stringify!(#type_name),
53                register_fn: |registry: &::rust_task_queue::TaskRegistry| -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
54                    // We need to get the task name from the type
55                    // For this we create a temporary instance with Default::default()
56                    // This is a limitation but keeps the implementation simple
57                    let temp_instance = <#type_name as Default>::default();
58                    let task_name = temp_instance.name();
59
60                    registry.register_with_name::<#type_name>(task_name)
61                },
62            }
63        }
64    };
65
66    Ok(expanded.into())
67}
68
69/// Attribute macro for registering tasks with a custom name.
70///
71/// This is useful when you want to register a task with a specific name
72/// rather than using the name returned by the `name()` method.
73///
74/// # Example
75///
76/// ```rust
77/// #[register_task("custom_task_name")]
78/// #[derive(Debug, Serialize, Deserialize)]
79/// struct MyTask {
80///     data: String,
81/// }
82/// ```
83#[proc_macro_attribute]
84pub fn register_task(args: TokenStream, input: TokenStream) -> TokenStream {
85    let task_name =
86        if args.is_empty() {
87            None
88        } else {
89            // Parse the task name from the attribute arguments
90            match syn::parse::<syn::LitStr>(args.clone()) {
91            Ok(lit) => Some(lit.value()),
92            Err(_) => return Error::new_spanned(
93                proc_macro2::TokenStream::from(args),
94                "Expected string literal for task name, e.g., #[register_task(\"my_task_name\")]"
95            ).to_compile_error().into(),
96        }
97        };
98
99    let input = parse_macro_input!(input as DeriveInput);
100
101    impl_register_task_with_name(&input, task_name)
102        .unwrap_or_else(|err| err.to_compile_error().into())
103}
104
105fn impl_register_task_with_name(
106    input: &DeriveInput,
107    custom_name: Option<String>,
108) -> Result<TokenStream, Error> {
109    let type_name = &input.ident;
110
111    let registration_code = if let Some(name) = custom_name {
112        // Use the custom name provided in the attribute
113        quote! {
114            ::rust_task_queue::inventory::submit! {
115                ::rust_task_queue::TaskRegistration {
116                    type_name: stringify!(#type_name),
117                    register_fn: |registry: &::rust_task_queue::TaskRegistry| -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
118                        registry.register_with_name::<#type_name>(#name)
119                    },
120                }
121            }
122        }
123    } else {
124        // Use the name from the Task::name() method
125        quote! {
126            ::rust_task_queue::inventory::submit! {
127                ::rust_task_queue::TaskRegistration {
128                    type_name: stringify!(#type_name),
129                    register_fn: |registry: &::rust_task_queue::TaskRegistry| -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
130                        let temp_instance = <#type_name as Default>::default();
131                        let task_name = temp_instance.name();
132                        registry.register_with_name::<#type_name>(task_name)
133                    },
134                }
135            }
136        }
137    };
138
139    // Include the original struct/enum definition along with the registration
140    let expanded = quote! {
141        #input
142
143        #registration_code
144    };
145
146    Ok(expanded.into())
147}
rust_task_queue_macro/lib.rs

rust_task_queue_macro/
lib.rs
