Skip to content

Commit 91cf5ce

Browse files
helio-frotahelio-frota
committed
refactor: common code for criterion benches
Closes #1354
1 parent d94f5a0 commit 91cf5ce

File tree

3 files changed

+140
-249
lines changed

3 files changed

+140
-249
lines changed

‎modules/fundamental/benches/bench-otel.rs

Lines changed: 8 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -2,29 +2,16 @@ use std::time::Duration;
22

33
use criterion::{Criterion, criterion_group, criterion_main};
44

5+
mod common;
6+
57
#[allow(clippy::expect_used, clippy::unwrap_used)]
68
pub(crate) mod trustify_benches {
7-
use std::io::Error;
8-
use std::ops::Add;
9-
use std::str::FromStr;
10-
use std::sync::Arc;
119
use std::time::{Duration, Instant};
1210

13-
use bytes::Bytes;
14-
use cpe::cpe::Cpe;
15-
use cpe::uri::OwnedUri;
1611
use criterion::{Criterion, black_box};
17-
use csaf::Csaf;
18-
use csaf::definitions::{BranchesT, ProductIdentificationHelper};
19-
use csaf::product_tree::ProductTree;
20-
use csaf::vulnerability::Vulnerability;
21-
use packageurl::PackageUrl;
22-
use sea_orm::ConnectionTrait;
23-
use test_context::AsyncTestContext;
24-
use tokio::runtime::Runtime;
12+
use std::ops::Add;
2513
use trustify_entity::labels::Labels;
2614
use trustify_module_ingestor::service::Format;
27-
use trustify_test_context::{TrustifyContext, document};
2815

2916
use opentelemetry::trace::TracerProvider as _;
3017
use opentelemetry_sdk::{
@@ -35,6 +22,8 @@ pub(crate) mod trustify_benches {
3522
use tracing_opentelemetry::OpenTelemetryLayer;
3623
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
3724

25+
use crate::common;
26+
3827
fn resource() -> Resource {
3928
Resource::builder()
4029
.with_service_name("ingestion_csaf".to_string())
@@ -87,19 +76,19 @@ pub(crate) mod trustify_benches {
8776
static INIT: Once = Once::new();
8877

8978
pub fn ingestion(c: &mut Criterion) {
90-
let (runtime, ctx) = setup_runtime_and_ctx();
79+
let (runtime, ctx) = common::setup_runtime_and_ctx();
9180
c.bench_function("ingestion_csaf", |b| {
9281
b.to_async(&runtime).iter_custom(|count| {
9382
let ctx = ctx.clone();
9483
async move {
9584
INIT.call_once(init_tracing_subscriber);
9685
log::info!("db reset...");
97-
reset_db(&ctx).await;
86+
common::reset_db(&ctx).await;
9887

9988
let mut duration = Duration::default();
10089
for i in 0..count {
10190
log::info!("inserting document {}...", i);
102-
let data = document_generated_from("csaf/cve-2023-33201.json", i)
91+
let data = common::document_generated_from("csaf/cve-2023-33201.json", i)
10392
.await
10493
.expect("data ok");
10594

@@ -118,112 +107,6 @@ pub(crate) mod trustify_benches {
118107
});
119108
});
120109
}
121-
122-
async fn document_generated_from(path: &str, rev: u64) -> Result<Bytes, Error> {
123-
let (mut doc, _): (Csaf, _) = document(path).await.expect("load ok");
124-
doc.document.tracking.id = format!("{}-{}", doc.document.tracking.id, rev);
125-
126-
fn rev_branches(branches: &mut Option<BranchesT>, rev: u64) {
127-
if let Some(BranchesT(branches)) = branches {
128-
for branch in branches.iter_mut() {
129-
if let Some(product) = &mut branch.product {
130-
rev_product_helper(&mut product.product_identification_helper, rev);
131-
}
132-
rev_branches(&mut branch.branches, rev);
133-
}
134-
}
135-
}
136-
137-
fn rev_product_helper(helper: &mut Option<ProductIdentificationHelper>, rev: u64) {
138-
if let Some(helper) = helper {
139-
if let Some(cpe) = &mut helper.cpe {
140-
helper.cpe = Some(rev_cpe(cpe, rev))
141-
}
142-
if let Some(purl) = &mut helper.purl {
143-
helper.purl = Some(rev_purl(purl.clone(), rev));
144-
}
145-
}
146-
}
147-
148-
fn rev_purl(from: PackageUrl, rev: u64) -> PackageUrl {
149-
let name = from.name();
150-
let new_name = format!("{}-{}", name, rev);
151-
PackageUrl::from_str(
152-
from.to_string()
153-
.replacen(name, new_name.as_str(), 1)
154-
.as_str(),
155-
)
156-
.unwrap()
157-
}
158-
fn rev_cpe(cpe: &OwnedUri, rev: u64) -> OwnedUri {
159-
let uri = format!("{:0}", cpe);
160-
let mut uri = cpe::uri::Uri::parse(uri.as_str()).unwrap();
161-
let x = format!("{}_{}", uri.product(), rev);
162-
uri.set_product(x.as_str()).unwrap();
163-
uri.to_owned()
164-
}
165-
fn rev_product_tree(product_tree: &mut ProductTree, rev: u64) {
166-
rev_branches(&mut product_tree.branches, rev);
167-
for relationships in product_tree.relationships.iter_mut() {
168-
for relationship in relationships.iter_mut() {
169-
// rev_product_id(&mut relationship.full_product_name.product_id, rev);
170-
rev_product_helper(
171-
&mut relationship.full_product_name.product_identification_helper,
172-
rev,
173-
);
174-
}
175-
}
176-
}
177-
fn rev_vulnerability(vulnerability: &mut Vulnerability, rev: u64) {
178-
for cve in vulnerability.cve.iter_mut() {
179-
*cve = format!("{}-{}", cve, rev);
180-
}
181-
}
182-
183-
for product_tree in doc.product_tree.iter_mut() {
184-
rev_product_tree(product_tree, rev);
185-
}
186-
for vulnerabilities in doc.vulnerabilities.iter_mut() {
187-
for vulnerability in vulnerabilities.iter_mut() {
188-
rev_vulnerability(vulnerability, rev);
189-
}
190-
}
191-
192-
let data = serde_json::to_vec_pretty(&doc).expect("serialize ok");
193-
Ok(Bytes::from(data))
194-
}
195-
196-
fn setup_runtime_and_ctx() -> (Runtime, Arc<TrustifyContext>) {
197-
let runtime = tokio::runtime::Builder::new_multi_thread()
198-
.enable_all()
199-
.build()
200-
.unwrap();
201-
let ctx = runtime.block_on(async { TrustifyContext::setup().await });
202-
(runtime, Arc::new(ctx))
203-
}
204-
205-
async fn reset_db(ctx: &Arc<TrustifyContext>) {
206-
// reset DB tables to a clean state...
207-
for table in [
208-
"advisory",
209-
"base_purl",
210-
"versioned_purl",
211-
"qualified_purl",
212-
"cvss3",
213-
"cpe",
214-
"version_range",
215-
"vulnerability",
216-
] {
217-
ctx.db
218-
.execute_unprepared(format!("DELETE FROM {table}").as_str())
219-
.await
220-
.expect("DELETE ok");
221-
}
222-
ctx.db
223-
.execute_unprepared("VACUUM ANALYZE")
224-
.await
225-
.expect("vacuum analyze ok");
226-
}
227110
}
228111

229112
criterion_group! {

‎modules/fundamental/benches/bench.rs

Lines changed: 6 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -2,43 +2,31 @@ use std::time::Duration;
22

33
use criterion::{Criterion, criterion_group, criterion_main};
44

5+
mod common;
6+
57
#[allow(clippy::expect_used, clippy::unwrap_used)]
68
pub(crate) mod trustify_benches {
7-
use std::io::Error;
89
use std::ops::Add;
9-
use std::str::FromStr;
10-
use std::sync::Arc;
1110
use std::time::{Duration, Instant};
1211

13-
use bytes::Bytes;
14-
use cpe::cpe::Cpe;
15-
use cpe::uri::OwnedUri;
12+
use crate::common;
1613
use criterion::{Criterion, black_box};
17-
use csaf::Csaf;
18-
use csaf::definitions::{BranchesT, ProductIdentificationHelper};
19-
use csaf::product_tree::ProductTree;
20-
use csaf::vulnerability::Vulnerability;
21-
use packageurl::PackageUrl;
22-
use sea_orm::ConnectionTrait;
23-
use test_context::AsyncTestContext;
24-
use tokio::runtime::Runtime;
2514
use trustify_entity::labels::Labels;
2615
use trustify_module_ingestor::service::Format;
27-
use trustify_test_context::{TrustifyContext, document};
2816

2917
pub fn ingestion(c: &mut Criterion) {
30-
let (runtime, ctx) = setup_runtime_and_ctx();
18+
let (runtime, ctx) = common::setup_runtime_and_ctx();
3119
c.bench_function("ingestion_csaf", |b| {
3220
b.to_async(&runtime).iter_custom(|count| {
3321
let ctx = ctx.clone();
3422
async move {
3523
log::info!("db reset...");
36-
reset_db(&ctx).await;
24+
common::reset_db(&ctx).await;
3725

3826
let mut duration = Duration::default();
3927
for i in 0..count {
4028
log::info!("inserting document {}...", i);
41-
let data = document_generated_from("csaf/cve-2023-33201.json", i)
29+
let data = common::document_generated_from("csaf/cve-2023-33201.json", i)
4230
.await
4331
.expect("data ok");
4432

@@ -57,112 +45,6 @@ pub(crate) mod trustify_benches {
5745
});
5846
});
5947
}
60-
61-
async fn document_generated_from(path: &str, rev: u64) -> Result<Bytes, Error> {
62-
let (mut doc, _): (Csaf, _) = document(path).await.expect("load ok");
63-
doc.document.tracking.id = format!("{}-{}", doc.document.tracking.id, rev);
64-
65-
fn rev_branches(branches: &mut Option<BranchesT>, rev: u64) {
66-
if let Some(BranchesT(branches)) = branches {
67-
for branch in branches.iter_mut() {
68-
if let Some(product) = &mut branch.product {
69-
rev_product_helper(&mut product.product_identification_helper, rev);
70-
}
71-
rev_branches(&mut branch.branches, rev);
72-
}
73-
}
74-
}
75-
76-
fn rev_product_helper(helper: &mut Option<ProductIdentificationHelper>, rev: u64) {
77-
if let Some(helper) = helper {
78-
if let Some(cpe) = &mut helper.cpe {
79-
helper.cpe = Some(rev_cpe(cpe, rev))
80-
}
81-
if let Some(purl) = &mut helper.purl {
82-
helper.purl = Some(rev_purl(purl.clone(), rev));
83-
}
84-
}
85-
}
86-
87-
fn rev_purl(from: PackageUrl, rev: u64) -> PackageUrl {
88-
let name = from.name();
89-
let new_name = format!("{}-{}", name, rev);
90-
PackageUrl::from_str(
91-
from.to_string()
92-
.replacen(name, new_name.as_str(), 1)
93-
.as_str(),
94-
)
95-
.unwrap()
96-
}
97-
fn rev_cpe(cpe: &OwnedUri, rev: u64) -> OwnedUri {
98-
let uri = format!("{:0}", cpe);
99-
let mut uri = cpe::uri::Uri::parse(uri.as_str()).unwrap();
100-
let x = format!("{}_{}", uri.product(), rev);
101-
uri.set_product(x.as_str()).unwrap();
102-
uri.to_owned()
103-
}
104-
fn rev_product_tree(product_tree: &mut ProductTree, rev: u64) {
105-
rev_branches(&mut product_tree.branches, rev);
106-
for relationships in product_tree.relationships.iter_mut() {
107-
for relationship in relationships.iter_mut() {
108-
// rev_product_id(&mut relationship.full_product_name.product_id, rev);
109-
rev_product_helper(
110-
&mut relationship.full_product_name.product_identification_helper,
111-
rev,
112-
);
113-
}
114-
}
115-
}
116-
fn rev_vulnerability(vulnerability: &mut Vulnerability, rev: u64) {
117-
for cve in vulnerability.cve.iter_mut() {
118-
*cve = format!("{}-{}", cve, rev);
119-
}
120-
}
121-
122-
for product_tree in doc.product_tree.iter_mut() {
123-
rev_product_tree(product_tree, rev);
124-
}
125-
for vulnerabilities in doc.vulnerabilities.iter_mut() {
126-
for vulnerability in vulnerabilities.iter_mut() {
127-
rev_vulnerability(vulnerability, rev);
128-
}
129-
}
130-
131-
let data = serde_json::to_vec_pretty(&doc).expect("serialize ok");
132-
Ok(Bytes::from(data))
133-
}
134-
135-
fn setup_runtime_and_ctx() -> (Runtime, Arc<TrustifyContext>) {
136-
let runtime = tokio::runtime::Builder::new_multi_thread()
137-
.enable_all()
138-
.build()
139-
.unwrap();
140-
let ctx = runtime.block_on(async { TrustifyContext::setup().await });
141-
(runtime, Arc::new(ctx))
142-
}
143-
144-
async fn reset_db(ctx: &Arc<TrustifyContext>) {
145-
// reset DB tables to a clean state...
146-
for table in [
147-
"advisory",
148-
"base_purl",
149-
"versioned_purl",
150-
"qualified_purl",
151-
"cvss3",
152-
"cpe",
153-
"version_range",
154-
"vulnerability",
155-
] {
156-
ctx.db
157-
.execute_unprepared(format!("DELETE FROM {table}").as_str())
158-
.await
159-
.expect("DELETE ok");
160-
}
161-
ctx.db
162-
.execute_unprepared("VACUUM ANALYZE")
163-
.await
164-
.expect("vacuum analyze ok");
165-
}
16648
}
16749

16850
criterion_group! {

0 commit comments

Comments
 (0)