|
1 | 1 | #![cfg_attr(not(target_os = "windows"), allow(unused))] |
| 2 | +#![allow(clippy::test_attr_in_doctest)] |
2 | 3 |
|
3 | 4 | use proc_macro::TokenStream; |
4 | | -use quote::quote; |
5 | | -use syn::{LitStr, parse_macro_input}; |
| 5 | +use quote::{ToTokens, quote}; |
| 6 | +use syn::{ItemFn, LitStr, parse_macro_input, parse_quote}; |
6 | 7 |
|
7 | 8 | /// A macro used in tests for cross-platform path string literals in tests. On Windows it replaces |
8 | 9 | /// `/` with `\\` and adds `C:` to the beginning of absolute paths. On other platforms, the path is |
@@ -87,3 +88,145 @@ pub fn line_endings(input: TokenStream) -> TokenStream { |
87 | 88 | #text |
88 | 89 | }) |
89 | 90 | } |
| 91 | + |
| 92 | +/// Inner data for the perf macro. |
| 93 | +struct PerfArgs { |
| 94 | + /// How many times to loop a test before rerunning the test binary. |
| 95 | + /// If left empty, the test harness will auto-determine this value. |
| 96 | + iterations: Option<syn::Expr>, |
| 97 | +} |
| 98 | + |
| 99 | +impl syn::parse::Parse for PerfArgs { |
| 100 | + fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> { |
| 101 | + if input.is_empty() { |
| 102 | + return Ok(PerfArgs { iterations: None }); |
| 103 | + } |
| 104 | + |
| 105 | + let mut iterations = None; |
| 106 | + // In principle we only have one possible argument, but leave this as |
| 107 | + // a loop in case we expand this in the future. |
| 108 | + for meta in |
| 109 | + syn::punctuated::Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated(input)? |
| 110 | + { |
| 111 | + match &meta { |
| 112 | + syn::Meta::NameValue(meta_name_value) => { |
| 113 | + if meta_name_value.path.is_ident("iterations") { |
| 114 | + iterations = Some(meta_name_value.value.clone()); |
| 115 | + } else { |
| 116 | + return Err(syn::Error::new_spanned( |
| 117 | + &meta_name_value.path, |
| 118 | + "unexpected argument, expected 'iterations'", |
| 119 | + )); |
| 120 | + } |
| 121 | + } |
| 122 | + _ => { |
| 123 | + return Err(syn::Error::new_spanned( |
| 124 | + meta, |
| 125 | + "expected name-value argument like 'iterations = 1'", |
| 126 | + )); |
| 127 | + } |
| 128 | + } |
| 129 | + } |
| 130 | + |
| 131 | + Ok(PerfArgs { iterations }) |
| 132 | + } |
| 133 | +} |
| 134 | + |
| 135 | +/// Marks a test as perf-sensitive, to be triaged when checking the performance |
| 136 | +/// of a build. This also automatically applies `#[test]`. |
| 137 | +/// |
| 138 | +/// By default, the number of iterations when profiling this test is auto-determined. |
| 139 | +/// If this needs to be overwritten, pass the desired iteration count to the macro |
| 140 | +/// as a parameter (`#[perf(iterations = n)]`). Note that the actual profiler may still |
| 141 | +/// run the test an arbitrary number times; this flag just sets the number of executions |
| 142 | +/// before the process is restarted and global state is reset. |
| 143 | +/// |
| 144 | +/// # Usage notes |
| 145 | +/// This should probably not be applied to tests that do any significant fs IO, as |
| 146 | +/// locks on files may not be released in time when repeating a test many times. This |
| 147 | +/// might lead to spurious failures. |
| 148 | +/// |
| 149 | +/// # Examples |
| 150 | +/// ```rust |
| 151 | +/// use util_macros::perf; |
| 152 | +/// |
| 153 | +/// #[perf] |
| 154 | +/// fn expensive_computation_test() { |
| 155 | +/// // Test goes here. |
| 156 | +/// } |
| 157 | +/// ``` |
| 158 | +/// |
| 159 | +/// This also works with `#[gpui::test]`s, though in most cases it shouldn't |
| 160 | +/// be used with automatic iterations. |
| 161 | +/// ```rust,ignore |
| 162 | +/// use util_macros::perf; |
| 163 | +/// |
| 164 | +/// #[perf(iterations = 1)] |
| 165 | +/// #[gpui::test] |
| 166 | +/// fn oneshot_test(_cx: &mut gpui::TestAppContext) { |
| 167 | +/// // Test goes here. |
| 168 | +/// } |
| 169 | +/// ``` |
| 170 | +#[proc_macro_attribute] |
| 171 | +pub fn perf(our_attr: TokenStream, input: TokenStream) -> TokenStream { |
| 172 | + // If any of the below constants are changed, make sure to also update the perf |
| 173 | + // profiler to match! |
| 174 | + |
| 175 | + /// The suffix on tests marked with `#[perf]`. |
| 176 | + const SUF_NORMAL: &str = "__ZED_PERF"; |
| 177 | + /// The suffix on tests marked with `#[perf(iterations = n)]`. |
| 178 | + const SUF_FIXED: &str = "__ZED_PERF_FIXEDITER"; |
| 179 | + /// The env var in which we pass the iteration count to our tests. |
| 180 | + const ITER_ENV_VAR: &str = "ZED_PERF_ITER"; |
| 181 | + |
| 182 | + let iter_count = parse_macro_input!(our_attr as PerfArgs).iterations; |
| 183 | + |
| 184 | + let ItemFn { |
| 185 | + mut attrs, |
| 186 | + vis, |
| 187 | + mut sig, |
| 188 | + block, |
| 189 | + } = parse_macro_input!(input as ItemFn); |
| 190 | + attrs.push(parse_quote!(#[test])); |
| 191 | + attrs.push(parse_quote!(#[allow(non_snake_case)])); |
| 192 | + |
| 193 | + let block: Box<syn::Block> = if cfg!(perf_enabled) { |
| 194 | + // Make the ident obvious when calling, for the test parser. |
| 195 | + let mut new_ident = sig.ident.to_string(); |
| 196 | + if iter_count.is_some() { |
| 197 | + new_ident.push_str(SUF_FIXED); |
| 198 | + } else { |
| 199 | + new_ident.push_str(SUF_NORMAL); |
| 200 | + } |
| 201 | + |
| 202 | + let new_ident = syn::Ident::new(&new_ident, sig.ident.span()); |
| 203 | + sig.ident = new_ident; |
| 204 | + // If we have a preset iteration count, just use that. |
| 205 | + if let Some(iter_count) = iter_count { |
| 206 | + parse_quote!({ |
| 207 | + for _ in 0..#iter_count { |
| 208 | + #block |
| 209 | + } |
| 210 | + }) |
| 211 | + } else { |
| 212 | + // Otherwise, the perf harness will pass us the value in an env var. |
| 213 | + parse_quote!({ |
| 214 | + let iter_count = std::env::var(#ITER_ENV_VAR).unwrap().parse::<usize>().unwrap(); |
| 215 | + for _ in 0..iter_count { |
| 216 | + #block |
| 217 | + } |
| 218 | + }) |
| 219 | + } |
| 220 | + } else { |
| 221 | + block |
| 222 | + }; |
| 223 | + |
| 224 | + ItemFn { |
| 225 | + attrs, |
| 226 | + vis, |
| 227 | + sig, |
| 228 | + block, |
| 229 | + } |
| 230 | + .into_token_stream() |
| 231 | + .into() |
| 232 | +} |
0 commit comments