cntryl_stress_macros/
lib.rs

1//! Proc macros for cntryl-stress.
2
3use proc_macro::TokenStream;
4use proc_macro2::{Span, TokenStream as TokenStream2};
5use proc_macro_crate::{crate_name, FoundCrate};
6use quote::quote;
7use std::collections::HashSet;
8use syn::parse::Parser;
9use syn::{
10    parse_macro_input, Expr, ExprLit, FnArg, ItemFn, Lit, Meta, MetaNameValue, Signature, Token,
11    Type,
12};
13
14const MAX_TIER: u32 = 6;
15
16/// Mark a function as a stress benchmark.
17///
18/// Benchmark functions take exactly one `&mut StressContext`, may be synchronous
19/// or asynchronous, and return `()`, `Result<(), E>`, or `StressResult`.
20///
21/// Common usage stays intentionally small:
22///
23/// ```rust,ignore
24/// # use stress_alias as cntryl_stress;
25/// use cntryl_stress::{black_box, stress, StressContext};
26///
27/// #[stress(tier = 1)]
28/// fn parse_hot_path(ctx: &mut StressContext) {
29///     let header = b"content-type:application/json";
30///     ctx.measure("colon lookup", || black_box(header.iter().position(|byte| *byte == b':')));
31/// }
32///
33/// #[stress(tier = 2)]
34/// fn write_batch(ctx: &mut StressContext) {
35///     ctx.measure("write batch", || black_box([1_u8, 2, 3]));
36/// }
37/// ```
38///
39/// Supported attributes:
40///
41/// - `tier = 1` through `6` (defaults to `2`)
42/// - `name = "custom_name"`
43/// - `ignore`
44/// - `role = "gate"`, `"diagnostic"`, or `"experimental"`
45/// - `max_ns_per_op = 1000`
46/// - `max_allocs_per_op = 0`
47/// - `max_bytes_per_op = 0`
48/// - `max_regression_pct = 5`
49/// - `max_rsd_pct = 10`
50/// - `max_peak_rss_mb = 512` (process-wide peak RSS; emits a
51///   `peak_rss_exceeded` warning diagnostic and never fails the budget gate)
52/// - `runtime = "tokio"` or `"tokio-multi"` (async functions only; needs the
53///   `tokio` feature of cntryl-stress) drives the benchmark on a fresh
54///   current-thread or multi-thread tokio runtime instead of the built-in executor
55/// - `metadata(owner = "storage", scenario = "fanout")`
56#[proc_macro_attribute]
57pub fn stress(attr: TokenStream, item: TokenStream) -> TokenStream {
58    let input = parse_macro_input!(item as ItemFn);
59    let fn_name = &input.sig.ident;
60    let fn_name_str = syn::ext::IdentExt::unraw(fn_name).to_string();
61    let cfg_attrs = propagated_cfg_attrs(&input.attrs);
62    let is_async = input.sig.asyncness.is_some();
63
64    let attrs = match StressAttrs::parse(attr.into()) {
65        Ok(attrs) => attrs,
66        Err(error) => return error.to_compile_error().into(),
67    };
68    if let Err(error) = validate_stress_signature(&input.sig) {
69        return error.to_compile_error().into();
70    }
71    if let Some(error) = tier_error(attrs.tier) {
72        return syn::Error::new_spanned(fn_name, error)
73            .to_compile_error()
74            .into();
75    }
76    let stress_crate = match stress_crate_path() {
77        Ok(path) => path,
78        Err(error) => return error.to_compile_error().into(),
79    };
80
81    let benchmark_name = attrs.name.unwrap_or_else(|| fn_name_str.clone());
82    let is_ignored = attrs.ignore;
83    let tier = attrs.tier;
84    let Some(derived_mode) = mode_for_tier(tier) else {
85        return syn::Error::new_spanned(fn_name, "cntryl-stress tiers are 1 through 6")
86            .to_compile_error()
87            .into();
88    };
89    let mode_kind = derived_mode;
90    let mode = mode_kind.tokens(&stress_crate);
91    let budget_setters = [
92        ("with_max_ns_per_op", attrs.budgets.ns_per_op),
93        ("with_max_allocs_per_op", attrs.budgets.allocs_per_op),
94        ("with_max_bytes_per_op", attrs.budgets.bytes_per_op),
95        ("with_max_regression_pct", attrs.budgets.regression_pct),
96        ("with_max_rsd_pct", attrs.budgets.rsd_pct),
97        ("with_max_peak_rss_mb", attrs.budgets.peak_rss_mb),
98    ]
99    .into_iter()
100    .filter_map(|(setter, value)| {
101        value.map(|value| {
102            let setter = syn::Ident::new(setter, Span::call_site());
103            quote! { .#setter(#value) }
104        })
105    });
106    let mut metadata = attrs.metadata;
107    if let Some(role) = attrs.role {
108        metadata.push(("trust_class".to_string(), role));
109    }
110    let metadata_keys = metadata.iter().map(|(key, _)| key);
111    let metadata_values = metadata.iter().map(|(_, value)| value);
112    let submit_ident = syn::Ident::new(&format!("__stress_bench_{fn_name_str}"), fn_name.span());
113    let wrapper_ident = syn::Ident::new(&format!("__stress_wrapper_{fn_name_str}"), fn_name.span());
114    let invocation = match benchmark_invocation(fn_name, is_async, attrs.runtime, &stress_crate) {
115        Ok(tokens) => tokens,
116        Err(error) => return error.to_compile_error().into(),
117    };
118    let wrapper = quote! {
119        #(#cfg_attrs)*
120        #[allow(non_snake_case)]
121        fn #wrapper_ident(ctx: &mut #stress_crate::StressContext) -> #stress_crate::StressResult {
122            #stress_crate::__private::IntoStressResult::into_stress_result(#invocation)
123        }
124    };
125
126    quote! {
127        #input
128        #wrapper
129
130        #(#cfg_attrs)*
131        #[allow(non_upper_case_globals)]
132        #[#stress_crate::__private::linkme::distributed_slice(#stress_crate::__private::STRESS_BENCHMARKS)]
133        #[linkme(crate = #stress_crate::__private::linkme)]
134        static #submit_ident: #stress_crate::__private::BenchmarkEntry = #stress_crate::__private::BenchmarkEntry {
135            name: #benchmark_name,
136            function_name: #fn_name_str,
137            func: #wrapper_ident,
138            ignored: #is_ignored,
139            module_path: ::core::module_path!(),
140            tier: #tier,
141            mode: #mode,
142            budgets: #stress_crate::artifact::BenchmarkBudgets::new()#(#budget_setters)*,
143            metadata: &[#((#metadata_keys, #metadata_values)),*],
144            file: ::core::file!(),
145            line: ::core::line!(),
146        };
147    }
148    .into()
149}
150
151#[derive(Debug)]
152struct StressAttrs {
153    name: Option<String>,
154    tier: u32,
155    ignore: bool,
156    role: Option<String>,
157    budgets: StressBudgets,
158    metadata: Vec<(String, String)>,
159    runtime: Option<RuntimeKind>,
160}
161
162#[derive(Debug, Clone, Copy, PartialEq, Eq)]
163enum RuntimeKind {
164    TokioCurrentThread,
165    TokioMultiThread,
166}
167
168#[derive(Debug, Default)]
169struct StressBudgets {
170    ns_per_op: Option<f64>,
171    allocs_per_op: Option<f64>,
172    bytes_per_op: Option<f64>,
173    regression_pct: Option<f64>,
174    rsd_pct: Option<f64>,
175    peak_rss_mb: Option<f64>,
176}
177
178impl Default for StressAttrs {
179    fn default() -> Self {
180        Self {
181            name: None,
182            tier: 2,
183            ignore: false,
184            role: None,
185            budgets: StressBudgets::default(),
186            metadata: Vec::new(),
187            runtime: None,
188        }
189    }
190}
191
192impl StressAttrs {
193    fn parse(attr: proc_macro2::TokenStream) -> syn::Result<Self> {
194        let parser = syn::punctuated::Punctuated::<Meta, Token![,]>::parse_terminated;
195        let metas = parser.parse2(attr)?;
196        let mut attrs = Self::default();
197        let mut singleton_attributes = HashSet::new();
198
199        for meta in metas {
200            match meta {
201                Meta::Path(path) if path.is_ident("ignore") => {
202                    mark_singleton(&mut singleton_attributes, "ignore", &path)?;
203                    attrs.ignore = true;
204                }
205                Meta::NameValue(name_value) if name_value.path.is_ident("name") => {
206                    mark_singleton(&mut singleton_attributes, "name", &name_value)?;
207                    let name = string_value(&name_value)?;
208                    if name.trim().is_empty() {
209                        return Err(syn::Error::new_spanned(
210                            &name_value.value,
211                            "stress benchmark name must not be empty or whitespace",
212                        ));
213                    }
214                    attrs.name = Some(name);
215                }
216                Meta::NameValue(name_value) if name_value.path.is_ident("tier") => {
217                    mark_singleton(&mut singleton_attributes, "tier", &name_value)?;
218                    let tier = int_value(&name_value)?;
219                    if let Some(error) = tier_error(tier) {
220                        return Err(syn::Error::new_spanned(name_value, error));
221                    }
222                    attrs.tier = tier;
223                }
224                Meta::NameValue(name_value) if name_value.path.is_ident("role") => {
225                    mark_singleton(&mut singleton_attributes, "role", &name_value)?;
226                    attrs.role = Some(role_value(&name_value)?);
227                }
228                Meta::NameValue(name_value) if name_value.path.is_ident("runtime") => {
229                    mark_singleton(&mut singleton_attributes, "runtime", &name_value)?;
230                    attrs.runtime = Some(runtime_value(&name_value)?);
231                }
232                Meta::NameValue(name_value) if name_value.path.is_ident("mode") => {
233                    return Err(syn::Error::new_spanned(
234                        name_value,
235                        "mode is not a public stress attribute; choose tier = 1 through 6",
236                    ));
237                }
238                Meta::NameValue(name_value) if name_value.path.is_ident("max_ns_per_op") => {
239                    mark_singleton(&mut singleton_attributes, "max_ns_per_op", &name_value)?;
240                    attrs.budgets.ns_per_op = Some(nonnegative_budget_value(&name_value)?);
241                }
242                Meta::NameValue(name_value) if name_value.path.is_ident("max_allocs_per_op") => {
243                    mark_singleton(&mut singleton_attributes, "max_allocs_per_op", &name_value)?;
244                    attrs.budgets.allocs_per_op = Some(nonnegative_budget_value(&name_value)?);
245                }
246                Meta::NameValue(name_value) if name_value.path.is_ident("max_bytes_per_op") => {
247                    mark_singleton(&mut singleton_attributes, "max_bytes_per_op", &name_value)?;
248                    attrs.budgets.bytes_per_op = Some(nonnegative_budget_value(&name_value)?);
249                }
250                Meta::NameValue(name_value) if name_value.path.is_ident("max_regression_pct") => {
251                    mark_singleton(&mut singleton_attributes, "max_regression_pct", &name_value)?;
252                    attrs.budgets.regression_pct = Some(percentage_budget_value(&name_value)?);
253                }
254                Meta::NameValue(name_value) if name_value.path.is_ident("max_rsd_pct") => {
255                    mark_singleton(&mut singleton_attributes, "max_rsd_pct", &name_value)?;
256                    attrs.budgets.rsd_pct = Some(nonnegative_budget_value(&name_value)?);
257                }
258                Meta::NameValue(name_value) if name_value.path.is_ident("max_peak_rss_mb") => {
259                    mark_singleton(&mut singleton_attributes, "max_peak_rss_mb", &name_value)?;
260                    attrs.budgets.peak_rss_mb = Some(nonnegative_budget_value(&name_value)?);
261                }
262                Meta::List(list) if list.path.is_ident("metadata") => {
263                    let metadata = parse_metadata(list.tokens.clone())?;
264                    for (key, value) in metadata {
265                        if key == "trust_class" {
266                            return Err(syn::Error::new_spanned(
267                                &list,
268                                "trust_class is reserved; use role = \"diagnostic\" or role = \"experimental\"",
269                            ));
270                        }
271                        if attrs
272                            .metadata
273                            .iter()
274                            .any(|(existing_key, _)| existing_key == &key)
275                        {
276                            return Err(syn::Error::new_spanned(
277                                &list,
278                                format!("stress metadata key `{key}` may be specified only once"),
279                            ));
280                        }
281                        attrs.metadata.push((key, value));
282                    }
283                }
284                other => {
285                    return Err(syn::Error::new_spanned(
286                        other,
287                        "unsupported stress attribute; expected tier, name, ignore, role, runtime, a max_* budget, or metadata(...)"
288                    ));
289                }
290            }
291        }
292
293        Ok(attrs)
294    }
295}
296
297fn mark_singleton(
298    seen: &mut HashSet<&'static str>,
299    name: &'static str,
300    tokens: impl quote::ToTokens,
301) -> syn::Result<()> {
302    if seen.insert(name) {
303        Ok(())
304    } else {
305        Err(syn::Error::new_spanned(
306            tokens,
307            format!("stress attribute `{name}` may be specified only once"),
308        ))
309    }
310}
311
312fn role_value(name_value: &MetaNameValue) -> syn::Result<String> {
313    let value = string_value(name_value)?;
314    match value.as_str() {
315        "gate" | "diagnostic" | "experimental" => Ok(value),
316        _ => Err(syn::Error::new_spanned(
317            &name_value.value,
318            "stress role must be \"gate\", \"diagnostic\", or \"experimental\"",
319        )),
320    }
321}
322
323fn benchmark_invocation(
324    fn_name: &syn::Ident,
325    is_async: bool,
326    runtime: Option<RuntimeKind>,
327    stress_crate: &TokenStream2,
328) -> syn::Result<TokenStream2> {
329    Ok(match (is_async, runtime) {
330        (false, Some(_)) => {
331            return Err(syn::Error::new_spanned(
332                fn_name,
333                "#[stress(runtime = ...)] applies only to async benchmark functions",
334            ));
335        }
336        (false, None) => quote! { #fn_name(ctx) },
337        (true, None) => quote! {
338            #stress_crate::__private::block_on(#fn_name(ctx))
339        },
340        (true, Some(RuntimeKind::TokioCurrentThread)) => quote! {
341            #stress_crate::__stress_tokio_block_on!(current_thread, #fn_name(ctx))
342        },
343        (true, Some(RuntimeKind::TokioMultiThread)) => quote! {
344            #stress_crate::__stress_tokio_block_on!(multi_thread, #fn_name(ctx))
345        },
346    })
347}
348
349fn runtime_value(name_value: &MetaNameValue) -> syn::Result<RuntimeKind> {
350    match string_value(name_value)?.as_str() {
351        "tokio" => Ok(RuntimeKind::TokioCurrentThread),
352        "tokio-multi" => Ok(RuntimeKind::TokioMultiThread),
353        _ => Err(syn::Error::new_spanned(
354            &name_value.value,
355            "stress runtime must be \"tokio\" or \"tokio-multi\"",
356        )),
357    }
358}
359
360fn parse_metadata(tokens: proc_macro2::TokenStream) -> syn::Result<Vec<(String, String)>> {
361    let parser = syn::punctuated::Punctuated::<MetaNameValue, Token![,]>::parse_terminated;
362    parser
363        .parse2(tokens)?
364        .into_iter()
365        .map(|name_value| {
366            let key = name_value
367                .path
368                .get_ident()
369                .map(ToString::to_string)
370                .ok_or_else(|| {
371                    syn::Error::new_spanned(&name_value.path, "metadata keys must be identifiers")
372                })?;
373            let value = string_value(&name_value)?;
374            Ok((key, value))
375        })
376        .collect()
377}
378
379fn string_value(name_value: &MetaNameValue) -> syn::Result<String> {
380    match &name_value.value {
381        Expr::Lit(ExprLit {
382            lit: Lit::Str(value),
383            ..
384        }) => Ok(value.value()),
385        value => Err(syn::Error::new_spanned(value, "expected string literal")),
386    }
387}
388
389fn int_value(name_value: &MetaNameValue) -> syn::Result<u32> {
390    match &name_value.value {
391        Expr::Lit(ExprLit {
392            lit: Lit::Int(value),
393            ..
394        }) => value.base10_parse(),
395        value => Err(syn::Error::new_spanned(value, "expected integer literal")),
396    }
397}
398
399fn float_value(name_value: &MetaNameValue) -> syn::Result<f64> {
400    match &name_value.value {
401        Expr::Lit(ExprLit {
402            lit: Lit::Float(value),
403            ..
404        }) => value.base10_parse(),
405        Expr::Lit(ExprLit {
406            lit: Lit::Int(value),
407            ..
408        }) => value.base10_parse(),
409        value => Err(syn::Error::new_spanned(value, "expected numeric literal")),
410    }
411}
412
413fn nonnegative_budget_value(name_value: &MetaNameValue) -> syn::Result<f64> {
414    let value = float_value(name_value)?;
415    if value.is_finite() && value >= 0.0 {
416        Ok(value)
417    } else {
418        Err(syn::Error::new_spanned(
419            &name_value.value,
420            "stress budgets must be finite non-negative numbers",
421        ))
422    }
423}
424
425fn percentage_budget_value(name_value: &MetaNameValue) -> syn::Result<f64> {
426    let value = nonnegative_budget_value(name_value)?;
427    if value <= 100.0 {
428        Ok(value)
429    } else {
430        Err(syn::Error::new_spanned(
431            &name_value.value,
432            "max_regression_pct must be between 0 and 100",
433        ))
434    }
435}
436
437#[derive(Debug, Clone, Copy, PartialEq, Eq)]
438enum ModeKind {
439    Micro,
440    FixedOperations,
441    FixedDuration,
442}
443
444impl ModeKind {
445    fn tokens(self, stress_crate: &TokenStream2) -> TokenStream2 {
446        match self {
447            Self::Micro => quote! { #stress_crate::artifact::BenchmarkModeKind::Micro },
448            Self::FixedOperations => {
449                quote! { #stress_crate::artifact::BenchmarkModeKind::FixedOperations }
450            }
451            Self::FixedDuration => {
452                quote! { #stress_crate::artifact::BenchmarkModeKind::FixedDuration }
453            }
454        }
455    }
456}
457
458fn validate_stress_signature(signature: &Signature) -> syn::Result<()> {
459    if let Some(unsafety) = &signature.unsafety {
460        return Err(syn::Error::new_spanned(
461            unsafety,
462            "#[stress] benchmark functions cannot be unsafe",
463        ));
464    }
465    if let Some(abi) = &signature.abi {
466        return Err(syn::Error::new_spanned(
467            abi,
468            "#[stress] benchmark functions cannot use an extern ABI",
469        ));
470    }
471    if !signature.generics.params.is_empty() || signature.generics.where_clause.is_some() {
472        return Err(syn::Error::new_spanned(
473            &signature.generics,
474            "#[stress] benchmark functions cannot be generic",
475        ));
476    }
477    if let Some(variadic) = &signature.variadic {
478        return Err(syn::Error::new_spanned(
479            variadic,
480            "#[stress] benchmark functions cannot be variadic",
481        ));
482    }
483    if let Some(FnArg::Receiver(receiver)) = signature
484        .inputs
485        .iter()
486        .find(|parameter| matches!(parameter, FnArg::Receiver(_)))
487    {
488        return Err(syn::Error::new_spanned(
489            receiver,
490            "#[stress] benchmark functions cannot have a self receiver; use a free function with &mut StressContext",
491        ));
492    }
493    if signature.inputs.len() != 1 {
494        return Err(syn::Error::new_spanned(
495            &signature.inputs,
496            "#[stress] benchmark functions require exactly one parameter: &mut StressContext",
497        ));
498    }
499
500    let parameter = signature
501        .inputs
502        .first()
503        .expect("one benchmark parameter was checked above");
504    let FnArg::Typed(parameter) = parameter else {
505        unreachable!("self receivers were rejected above");
506    };
507    if !is_mut_stress_context(&parameter.ty) {
508        return Err(syn::Error::new_spanned(
509            &parameter.ty,
510            "#[stress] benchmark parameter must have type &mut StressContext",
511        ));
512    }
513    Ok(())
514}
515
516fn is_mut_stress_context(ty: &Type) -> bool {
517    let Type::Reference(reference) = ty else {
518        return false;
519    };
520    if reference.mutability.is_none() || reference.lifetime.is_some() {
521        return false;
522    }
523    let Type::Path(path) = reference.elem.as_ref() else {
524        return false;
525    };
526    path.qself.is_none()
527        && path
528            .path
529            .segments
530            .last()
531            .is_some_and(|segment| segment.ident == "StressContext" && segment.arguments.is_empty())
532}
533
534fn stress_crate_path() -> syn::Result<TokenStream2> {
535    match crate_name("cntryl-stress") {
536        Ok(FoundCrate::Itself) => Ok(quote! { ::cntryl_stress }),
537        Ok(FoundCrate::Name(name)) => {
538            let ident = syn::Ident::new(&name, Span::call_site());
539            Ok(quote! { ::#ident })
540        }
541        Err(error) => Err(syn::Error::new(
542            Span::call_site(),
543            format!(
544                "could not find the cntryl-stress dependency: {error}; add cntryl-stress to Cargo.toml"
545            ),
546        )),
547    }
548}
549
550const fn mode_for_tier(tier: u32) -> Option<ModeKind> {
551    match tier {
552        1 => Some(ModeKind::Micro),
553        2 => Some(ModeKind::FixedOperations),
554        3..=MAX_TIER => Some(ModeKind::FixedDuration),
555        _ => None,
556    }
557}
558
559fn tier_error(tier: u32) -> Option<String> {
560    if tier == 0 {
561        Some("cntryl-stress tiers start at 1".to_string())
562    } else if tier > MAX_TIER {
563        Some(format!("cntryl-stress tiers are 1 through {MAX_TIER}"))
564    } else {
565        None
566    }
567}
568
569/// Generate a `main` function for stress benchmark binaries.
570#[proc_macro]
571pub fn stress_main(input: TokenStream) -> TokenStream {
572    let input = TokenStream2::from(input);
573    if !input.is_empty() {
574        return syn::Error::new_spanned(input, "stress_main! takes no arguments")
575            .to_compile_error()
576            .into();
577    }
578    let stress_crate = match stress_crate_path() {
579        Ok(path) => path,
580        Err(error) => return error.to_compile_error().into(),
581    };
582    quote! {
583        fn main() {
584            #stress_crate::__private::stress_binary_main();
585        }
586    }
587    .into()
588}
589
590/// Returns the `#[cfg(...)]` attributes that must also gate the generated
591/// wrapper and registration items. `#[cfg_attr(...)]` is deliberately not
592/// copied: its payload (for example `inline`) may be invalid on a `static`.
593fn propagated_cfg_attrs(attrs: &[syn::Attribute]) -> Vec<&syn::Attribute> {
594    attrs
595        .iter()
596        .filter(|attr| attr.path().is_ident("cfg"))
597        .collect()
598}
599
600#[cfg(test)]
601mod tests {
602    use super::*;
603
604    #[test]
605    fn tier_validation_rejects_undefined_tiers() {
606        assert_eq!(
607            tier_error(0).as_deref(),
608            Some("cntryl-stress tiers start at 1")
609        );
610        assert_eq!(
611            tier_error(MAX_TIER + 1).as_deref(),
612            Some("cntryl-stress tiers are 1 through 6")
613        );
614        assert!(tier_error(1).is_none());
615        assert!(tier_error(MAX_TIER).is_none());
616    }
617
618    #[test]
619    fn mode_defaults_are_derived_from_tier() {
620        assert_eq!(mode_for_tier(1), Some(ModeKind::Micro));
621        assert_eq!(mode_for_tier(2), Some(ModeKind::FixedOperations));
622        for tier in 3..=MAX_TIER {
623            assert_eq!(mode_for_tier(tier), Some(ModeKind::FixedDuration));
624        }
625    }
626
627    #[test]
628    fn mode_attribute_is_rejected() {
629        let error = StressAttrs::parse(quote::quote! { tier = 1, mode = "micro" })
630            .expect_err("mode is not public");
631
632        assert!(error
633            .to_string()
634            .contains("mode is not a public stress attribute"));
635    }
636
637    #[test]
638    fn only_cfg_attributes_are_propagated_to_generated_items() {
639        let item: ItemFn = syn::parse_quote! {
640            #[cfg(all())]
641            #[cfg_attr(all(), inline)]
642            #[doc = "x"]
643            fn bench(ctx: &mut StressContext) {}
644        };
645        let propagated = propagated_cfg_attrs(&item.attrs);
646        assert_eq!(propagated.len(), 1);
647        assert!(propagated[0].path().is_ident("cfg"));
648    }
649}
cntryl_stress_macros/lib.rs

cntryl_stress_macros/
lib.rs
