balter-macros 0.3.0

A load/stress testing framework.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use proc_macro::TokenStream;
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::quote;
use syn::{Ident, ItemFn};

/// Proc macro to denote a Transaction
///
/// NOTE: Currently this macro only works on functions with a `Result<T, E>` return value. This is a
/// restriction which will be lifted soon.
///
/// # Example
/// ```ignore
/// use balter::prelude::*;
///
/// #[transaction]
/// fn my_transaction(arg_1: u32, arg_2: &str) -> Result<String, MyError> {
///     ...
/// }
/// ```
#[proc_macro_attribute]
pub fn transaction(attr: TokenStream, item: TokenStream) -> TokenStream {
    transaction_internal(attr, item).into()
}

fn transaction_internal(_attr: TokenStream, item: TokenStream) -> TokenStream2 {
    let input = syn::parse::<ItemFn>(item).unwrap();

    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = input;
    let stmts = &block.stmts;

    let ident = &sig.ident;
    quote! {
        #(#attrs)* #vis #sig {
            ::balter::transaction::transaction_hook(::balter::core::generate_labels!(#ident), async move {
                #(#stmts)*
            }).await
        }
    }
}

/// Proc macro to denote a Scenario
///
/// NOTE: Currently this macro only works on functions which take no arguments and with no return value.
/// (void functions). This is a restriction which will be lifted soon.
///
/// See the `Scenario` struct for more information on the methods this macro provides on functions.
///
/// # Example
/// ```ignore
/// use balter::prelude::*;
///
/// #[scenario]
/// fn my_scenario() {
/// }
/// ```
#[proc_macro_attribute]
pub fn scenario(attr: TokenStream, item: TokenStream) -> TokenStream {
    scenario_internal(attr, item, false).into()
}

/// Proc macro to denote a Scenario
///
/// NOTE: Currently this macro only works on functions which take no arguments and with no return value.
/// (void functions). This is a restriction which will be lifted soon.
///
/// See the `Scenario` struct for more information on the methods this macro provides on functions.
///
/// # Example
/// ```ignore
/// use balter::prelude::*;
///
/// #[scenario]
/// fn my_scenario() {
/// }
/// ```
#[proc_macro_attribute]
pub fn scenario_linkme(attr: TokenStream, item: TokenStream) -> TokenStream {
    scenario_internal(attr, item, true).into()
}

fn scenario_internal(_attr: TokenStream, item: TokenStream, linkme: bool) -> TokenStream2 {
    let input = syn::parse::<ItemFn>(item).expect("Macro only works on fn() items");

    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = input;
    let stmts = &block.stmts;

    let new_name = Ident::new(&format!("__balter_{}", sig.ident), Span::call_site());
    let mut new_sig = sig.clone();
    new_sig.ident = new_name.clone();

    let mut scen_sig = sig.clone();
    let scen_name = sig.ident.clone();
    scen_sig.asyncness = None;
    scen_sig.output = syn::parse(
        quote! {
            -> impl ::balter::scenario::ConfigurableScenario<::balter::prelude::RunStatistics>
        }
        .into(),
    )
    .expect("Scenario signature is invalid");

    let res = quote! {
        #(#attrs)* #vis #scen_sig {
            ::balter::scenario::Scenario::new(stringify!(#scen_name), #new_name)
        }

        #(#attrs)* #vis #new_sig {
            #(#stmts)*
        }
    };

    if linkme {
        let mut linkme_sig = sig.clone();
        let linkme_name = Ident::new(&format!("__balter_distr_{}", sig.ident), Span::call_site());
        linkme_sig.ident = linkme_name.clone();
        linkme_sig.asyncness = None;
        linkme_sig.output = syn::parse(
            quote! {
                -> ::core::pin::Pin<Box<dyn ::balter::prelude::DistributedScenario<Output=::balter::prelude::RunStatistics>>>
            }
            .into(),
        )
        .expect("Scenario signature is invalid");

        let static_name = Ident::new(
            &format!("__BALTER_{}", sig.ident.to_string().to_ascii_uppercase()),
            Span::call_site(),
        );

        let mut linkme = quote! {
            #[::balter::runtime::distributed_slice(::balter::runtime::BALTER_SCENARIOS)]
            static #static_name: (&'static str, fn() -> ::core::pin::Pin<Box<dyn ::balter::prelude::DistributedScenario<Output=::balter::prelude::RunStatistics>>>) = (stringify!(#scen_name), #linkme_name);

            // TODO: This definition can almost certainly merge with the #scen_sig definition
            #(#attrs)* #vis #linkme_sig {
                Box::pin(::balter::scenario::Scenario::new(stringify!(#scen_name), #new_name))
            }
        };

        linkme.extend(res);
        linkme
    } else {
        res
    }
}