1 use criterion::{black_box, criterion_group, criterion_main, Benchmark, BenchmarkId, Criterion};
2 use serde_json::Value;
3 use std::{collections::HashMap, fs, io::Read, path::Path};
4 use vector::{
5     transforms::{wasm::Wasm, Transform},
6     Event,
7 };
8 
parse_event_artifact(path: impl AsRef<Path>) -> vector::Result<Event>9 fn parse_event_artifact(path: impl AsRef<Path>) -> vector::Result<Event> {
10     let mut event = Event::new_empty_log();
11     let mut test_file = fs::File::open(path)?;
12 
13     let mut buf = String::new();
14     test_file.read_to_string(&mut buf)?;
15     let test_json: HashMap<String, Value> = serde_json::from_str(&buf)?;
16 
17     for (key, value) in test_json {
18         event.as_mut_log().insert(key, value.clone());
19     }
20     Ok(event)
21 }
22 
protobuf(c: &mut Criterion)23 pub fn protobuf(c: &mut Criterion) {
24     let input = parse_event_artifact("tests/data/wasm/protobuf/demo.json").unwrap();
25     let cloned_input = input.clone();
26     c.bench(
27         "wasm/protobuf",
28         Benchmark::new("wasm", move |b| {
29             let input = cloned_input.clone();
30             let mut transform = Wasm::new(
31                 toml::from_str(
32                     r#"
33                 module = "target/wasm32-wasi/release/protobuf.wasm"
34                 artifact_cache = "target/artifacts/"
35                 "#,
36                 )
37                 .unwrap(),
38             )
39             .unwrap();
40             b.iter_with_setup(
41                 || input.clone(),
42                 |input| {
43                     let output = transform.transform(input);
44                     black_box(output)
45                 },
46             )
47         }),
48     );
49 }
50 
drop(criterion: &mut Criterion)51 pub fn drop(criterion: &mut Criterion) {
52     let transforms: Vec<(&str, Box<dyn Transform>)> = vec![
53         (
54             "lua",
55             Box::new(
56                 vector::transforms::lua::v2::Lua::new(
57                     &toml::from_str(
58                         r#"
59                 hooks.process = """
60                 function (event, emit)
61                 end
62                 """
63                 "#,
64                     )
65                     .unwrap(),
66                 )
67                 .unwrap(),
68             ),
69         ),
70         (
71             "wasm",
72             Box::new(
73                 Wasm::new(
74                     toml::from_str(
75                         r#"
76               module = "target/wasm32-wasi/release/drop.wasm"
77               artifact_cache = "target/artifacts/"
78               "#,
79                     )
80                     .unwrap(),
81                 )
82                 .unwrap(),
83             ),
84         ),
85     ];
86     let parameters = vec![0, 2, 8, 16];
87 
88     bench_group_transforms_over_parameterized_event_sizes(
89         criterion,
90         "wasm/drop",
91         transforms,
92         parameters,
93     );
94 }
95 
add_fields(criterion: &mut Criterion)96 pub fn add_fields(criterion: &mut Criterion) {
97     let transforms: Vec<(&str, Box<dyn Transform>)> = vec![
98         (
99             "lua",
100             Box::new(
101                 vector::transforms::lua::v2::Lua::new(
102                     &toml::from_str(
103                         r#"
104                 hooks.process = """
105                 function (event, emit)
106                     event.log.test_key = "test_value"
107                     event.log.test_key2 = "test_value2"
108                     emit(event)
109                 end
110                 """
111                 "#,
112                     )
113                     .unwrap(),
114                 )
115                 .unwrap(),
116             ),
117         ),
118         (
119             "wasm",
120             Box::new(
121                 Wasm::new(
122                     toml::from_str(
123                         r#"
124               module = "target/wasm32-wasi/release/add_fields.wasm"
125               artifact_cache = "target/artifacts/"
126               "#,
127                     )
128                     .unwrap(),
129                 )
130                 .unwrap(),
131             ),
132         ),
133         (
134             "native",
135             Box::new({
136                 let mut fields = indexmap::IndexMap::default();
137                 fields.insert("test_key".into(), "test_value".into());
138                 fields.insert("test_key2".into(), "test_value2".into());
139                 vector::transforms::add_fields::AddFields::new(fields, false)
140             }),
141         ),
142     ];
143     let parameters = vec![0, 2, 8, 16];
144 
145     bench_group_transforms_over_parameterized_event_sizes(
146         criterion,
147         "wasm/add_fields",
148         transforms,
149         parameters,
150     );
151 }
152 
bench_group_transforms_over_parameterized_event_sizes( criterion: &mut Criterion, group: &str, transforms: Vec<(&str, Box<dyn Transform>)>, parameters: Vec<usize>, )153 fn bench_group_transforms_over_parameterized_event_sizes(
154     criterion: &mut Criterion,
155     group: &str,
156     transforms: Vec<(&str, Box<dyn Transform>)>,
157     parameters: Vec<usize>,
158 ) {
159     let mut group = criterion.benchmark_group(group);
160     for (name, mut transform) in transforms {
161         for &parameter in &parameters {
162             let mut input = Event::new_empty_log();
163             for key in 0..parameter {
164                 input
165                     .as_mut_log()
166                     .insert(format!("key-{}", key), format!("value-{}", key));
167             }
168 
169             let id = BenchmarkId::new(name.clone(), parameter);
170 
171             group.bench_with_input(id, &input, |bencher, input| {
172                 bencher.iter_with_setup(
173                     || input.clone(),
174                     |input| {
175                         let output = transform.transform(input);
176                         black_box(output)
177                     },
178                 )
179             });
180         }
181     }
182     group.finish();
183 }
184 
185 criterion_group!(wasm, protobuf, drop, add_fields);
186 criterion_main!(wasm);
187