Files
rustytorch/crates/data/rtx-etl/benches/etl_benchmarks.rs
T
2026-03-04 00:08:42 +00:00

613 lines
20 KiB
Rust

//! Performance benchmarks for rtx-etl components
//!
//! Comprehensive benchmarks covering:
//! - DAG scheduling and execution performance
//! - Stream processing throughput
//! - Incremental processing efficiency
//! - Data transformation speed
//! - State management operations
//! - Memory usage optimization
use chrono::Utc;
use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use std::collections::HashMap;
use std::time::Duration;
use tokio::runtime::Runtime;
use uuid::Uuid;
use rtx_etl::{
DataPayload, DataRecord, DataValue, EtlConfigBuilder, EtlEngine, RecordMetadata, Result, Task,
TaskBuilder, TaskGraph,
dag::{DagConfig, ResourceRequirements, TaskType},
monitoring::EtlMetrics,
state::{StateBackendType, StateConfig},
};
/// Create a test data record for benchmarking
fn create_benchmark_record(id: u64) -> DataRecord {
let mut data = HashMap::new();
data.insert("id".to_string(), DataValue::Int(id as i64));
data.insert(
"name".to_string(),
DataValue::String(format!("user_{}", id)),
);
data.insert("score".to_string(), DataValue::Float(id as f64 * 0.1));
data.insert("active".to_string(), DataValue::Bool(id % 2 == 0));
DataRecord {
id: Uuid::new_v4(),
event_time: Utc::now(),
process_time: Utc::now(),
partition_key: Some(format!("partition_{}", id % 10)),
data: DataPayload::Structured(data),
metadata: RecordMetadata {
source: "benchmark".to_string(),
lineage: Vec::new(),
quality_scores: None,
attributes: HashMap::new(),
schema_version: Some("1.0".to_string()),
checksum: None,
},
}
}
/// Create a simple DAG for benchmarking
fn create_benchmark_dag(task_count: usize) -> TaskGraph {
let mut dag = TaskGraph::new();
// Create a chain of dependent tasks
for i in 0..task_count {
let mut task_builder = TaskBuilder::new(&format!("task_{}", i))
.name(&format!("Benchmark Task {}", i))
.task_type(TaskType::Transform);
if i > 0 {
task_builder = task_builder.depends_on(&format!("task_{}", i - 1));
}
dag.add_task(task_builder.build());
}
dag
}
/// Create a parallel DAG for benchmarking concurrent execution
fn create_parallel_dag(task_count: usize) -> TaskGraph {
let mut dag = TaskGraph::new();
// Create independent tasks that can run in parallel
for i in 0..task_count {
let task = TaskBuilder::new(&format!("parallel_task_{}", i))
.name(&format!("Parallel Task {}", i))
.task_type(TaskType::Extract)
.priority(i as i32)
.build();
dag.add_task(task);
}
dag
}
/// Benchmark ETL engine initialization
fn bench_engine_initialization(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
c.bench_function("etl_engine_init", |b| {
b.iter(|| {
rt.block_on(async {
let config = EtlConfigBuilder::new()
.with_dag_scheduling(DagConfig::default())
.build()
.unwrap();
let engine = EtlEngine::new(black_box(config)).await.unwrap();
black_box(engine)
})
})
});
}
/// Benchmark DAG validation performance
fn bench_dag_validation(c: &mut Criterion) {
let mut group = c.benchmark_group("dag_validation");
for task_count in [10, 50, 100, 500].iter() {
group.bench_with_input(
BenchmarkId::new("sequential", task_count),
task_count,
|b, &task_count| {
let dag = create_benchmark_dag(task_count);
b.iter(|| {
let mut dag_clone = dag.clone();
black_box(dag_clone.validate()).unwrap()
})
},
);
group.bench_with_input(
BenchmarkId::new("parallel", task_count),
task_count,
|b, &task_count| {
let dag = create_parallel_dag(task_count);
b.iter(|| {
let mut dag_clone = dag.clone();
black_box(dag_clone.validate()).unwrap()
})
},
);
}
group.finish();
}
/// Benchmark DAG execution performance
fn bench_dag_execution(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("dag_execution");
group.sample_size(10); // Reduce sample size for longer operations
for task_count in [5, 10, 20, 50].iter() {
group.bench_with_input(
BenchmarkId::new("execution", task_count),
task_count,
|b, &task_count| {
b.iter(|| {
rt.block_on(async {
let config = EtlConfigBuilder::new()
.with_dag_scheduling(DagConfig::default())
.with_max_concurrent_tasks(8)
.build()
.unwrap();
let engine = EtlEngine::new(config).await.unwrap();
let dag = create_parallel_dag(task_count);
let result = engine.execute_dag(black_box(dag)).await.unwrap();
black_box(result)
})
})
},
);
}
group.finish();
}
/// Benchmark concurrent task execution
fn bench_concurrent_execution(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("concurrent_execution");
group.sample_size(10);
for concurrency in [1, 2, 4, 8, 16].iter() {
group.bench_with_input(
BenchmarkId::new("concurrency", concurrency),
concurrency,
|b, &concurrency| {
b.iter(|| {
rt.block_on(async {
let config = EtlConfigBuilder::new()
.with_dag_scheduling(DagConfig::default())
.with_max_concurrent_tasks(concurrency)
.build()
.unwrap();
let engine = EtlEngine::new(config).await.unwrap();
let dag = create_parallel_dag(20); // Fixed task count
let result = engine.execute_dag(black_box(dag)).await.unwrap();
black_box(result)
})
})
},
);
}
group.finish();
}
/// Benchmark data record processing
fn bench_record_processing(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("record_processing");
for record_count in [100, 1000, 10000, 100000].iter() {
group.bench_with_input(
BenchmarkId::new("create_records", record_count),
record_count,
|b, &record_count| {
b.iter(|| {
let records: Vec<DataRecord> = (0..record_count)
.map(|i| create_benchmark_record(i as u64))
.collect();
black_box(records)
})
},
);
group.bench_with_input(
BenchmarkId::new("serialize_records", record_count),
record_count,
|b, &record_count| {
let records: Vec<DataRecord> = (0..record_count)
.map(|i| create_benchmark_record(i as u64))
.collect();
b.iter(|| {
let serialized: Vec<String> = records
.iter()
.map(|r| serde_json::to_string(r).unwrap())
.collect();
black_box(serialized)
})
},
);
}
group.finish();
}
/// Benchmark metrics collection performance
fn bench_metrics_collection(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("metrics_collection");
group.bench_function("counter_increment", |b| {
b.iter(|| {
rt.block_on(async {
let metrics = EtlMetrics::new();
for i in 0..1000 {
metrics
.increment_counter("test_counter", black_box(i))
.await;
}
black_box(metrics)
})
})
});
group.bench_function("gauge_update", |b| {
b.iter(|| {
rt.block_on(async {
let metrics = EtlMetrics::new();
for i in 0..1000 {
metrics.set_gauge("test_gauge", black_box(i as f64)).await;
}
black_box(metrics)
})
})
});
group.bench_function("histogram_record", |b| {
b.iter(|| {
rt.block_on(async {
let metrics = EtlMetrics::new();
for i in 0..1000 {
metrics
.record_histogram("test_histogram", black_box(i as f64))
.await;
}
black_box(metrics)
})
})
});
group.finish();
}
/// Benchmark state management operations
fn bench_state_management(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("state_management");
group.bench_function("state_operations", |b| {
b.iter(|| {
rt.block_on(async {
let config = rtx_etl::state::StateConfig {
backend: StateBackendType::InMemory,
..Default::default()
};
let state_manager = rtx_etl::state::StateManager::new(config).await.unwrap();
// Benchmark state operations
for i in 0..100 {
let key = format!("key_{}", i);
let value = serde_json::json!({"value": i, "timestamp": Utc::now()});
state_manager
.set_state(&key, black_box(value))
.await
.unwrap();
let retrieved = state_manager.get_state(&key).await.unwrap();
black_box(retrieved);
}
black_box(state_manager)
})
})
});
group.finish();
}
/// Benchmark memory usage and optimization
fn bench_memory_usage(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("memory_usage");
for data_size in [1024, 10240, 102400, 1048576].iter() {
group.bench_with_input(
BenchmarkId::new("large_payload", data_size),
data_size,
|b, &data_size| {
b.iter(|| {
rt.block_on(async {
// Create records with large payloads
let large_data = vec![0u8; data_size];
let mut data = HashMap::new();
data.insert("large_field".to_string(), DataValue::Binary(large_data));
let record = DataRecord {
id: Uuid::new_v4(),
event_time: Utc::now(),
process_time: Utc::now(),
partition_key: Some("test".to_string()),
data: DataPayload::Structured(data),
metadata: RecordMetadata {
source: "benchmark".to_string(),
lineage: Vec::new(),
quality_scores: None,
attributes: HashMap::new(),
schema_version: Some("1.0".to_string()),
checksum: None,
},
};
black_box(record)
})
})
},
);
}
group.finish();
}
/// Benchmark data transformation operations
fn bench_transformations(c: &mut Criterion) {
let mut group = c.benchmark_group("transformations");
for record_count in [100, 1000, 10000].iter() {
group.bench_with_input(
BenchmarkId::new("filter_transform", record_count),
record_count,
|b, &record_count| {
let records: Vec<DataRecord> = (0..record_count)
.map(|i| create_benchmark_record(i as u64))
.collect();
b.iter(|| {
let filtered: Vec<&DataRecord> = records
.iter()
.filter(|r| {
if let DataPayload::Structured(data) = &r.data {
if let Some(DataValue::Bool(active)) = data.get("active") {
return *active;
}
}
false
})
.collect();
black_box(filtered)
})
},
);
group.bench_with_input(
BenchmarkId::new("map_transform", record_count),
record_count,
|b, &record_count| {
let records: Vec<DataRecord> = (0..record_count)
.map(|i| create_benchmark_record(i as u64))
.collect();
b.iter(|| {
let transformed: Vec<HashMap<String, DataValue>> = records
.iter()
.map(|r| {
if let DataPayload::Structured(data) = &r.data {
let mut new_data = data.clone();
if let Some(DataValue::Float(score)) = data.get("score") {
new_data.insert(
"double_score".to_string(),
DataValue::Float(score * 2.0),
);
}
new_data
} else {
HashMap::new()
}
})
.collect();
black_box(transformed)
})
},
);
}
group.finish();
}
/// Benchmark resource allocation and task scheduling
fn bench_resource_allocation(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("resource_allocation");
group.sample_size(10);
group.bench_function("priority_scheduling", |b| {
b.iter(|| {
rt.block_on(async {
let config = EtlConfigBuilder::new()
.with_dag_scheduling(DagConfig {
priority_scheduling: true,
..Default::default()
})
.with_max_concurrent_tasks(4)
.build()
.unwrap();
let engine = EtlEngine::new(config).await.unwrap();
// Create DAG with varied priorities and resource requirements
let mut dag = TaskGraph::new();
for i in 0..20 {
let task = TaskBuilder::new(&format!("task_{}", i))
.name(&format!("Task {}", i))
.task_type(TaskType::Transform)
.priority((i % 5) as i32) // Varied priorities
.resources(ResourceRequirements {
cpu_cores: (i % 3) as f32 + 1.0, // 1-3 CPU cores
memory_mb: (i % 4 + 1) * 256, // 256MB-1GB memory
disk_mb: 1024,
..Default::default()
})
.build();
dag.add_task(task);
}
let result = engine.execute_dag(black_box(dag)).await.unwrap();
black_box(result)
})
})
});
group.finish();
}
/// Comprehensive end-to-end benchmark
fn bench_end_to_end_pipeline(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("end_to_end_pipeline");
group.sample_size(5); // Very small sample size for comprehensive test
group.measurement_time(Duration::from_secs(60)); // Longer measurement time
group.bench_function("complete_etl_pipeline", |b| {
b.iter(|| {
rt.block_on(async {
let config = EtlConfigBuilder::new()
.with_dag_scheduling(DagConfig::default())
.with_max_concurrent_tasks(8)
.build()
.unwrap();
let engine = EtlEngine::new(config).await.unwrap();
// Create a realistic ETL pipeline
let mut dag = TaskGraph::new();
// Data extraction
dag.add_task(
TaskBuilder::new("extract_users")
.task_type(TaskType::Extract)
.build(),
);
dag.add_task(
TaskBuilder::new("extract_orders")
.task_type(TaskType::Extract)
.build(),
);
dag.add_task(
TaskBuilder::new("extract_products")
.task_type(TaskType::Extract)
.build(),
);
// Data transformation
dag.add_task(
TaskBuilder::new("clean_users")
.depends_on("extract_users")
.task_type(TaskType::Transform)
.build(),
);
dag.add_task(
TaskBuilder::new("clean_orders")
.depends_on("extract_orders")
.task_type(TaskType::Transform)
.build(),
);
dag.add_task(
TaskBuilder::new("clean_products")
.depends_on("extract_products")
.task_type(TaskType::Transform)
.build(),
);
// Data aggregation
dag.add_task(
TaskBuilder::new("aggregate_orders")
.depends_on("clean_orders")
.task_type(TaskType::Transform)
.build(),
);
// Data joining
dag.add_task(
TaskBuilder::new("join_user_orders")
.depends_on("clean_users")
.depends_on("aggregate_orders")
.task_type(TaskType::Transform)
.build(),
);
// Quality checks
dag.add_task(
TaskBuilder::new("quality_check")
.depends_on("join_user_orders")
.task_type(TaskType::QualityCheck)
.build(),
);
// Data loading
dag.add_task(
TaskBuilder::new("load_warehouse")
.depends_on("quality_check")
.task_type(TaskType::Load)
.build(),
);
let result = engine.execute_dag(black_box(dag)).await.unwrap();
black_box(result)
})
})
});
group.finish();
}
criterion_group!(
benches,
bench_engine_initialization,
bench_dag_validation,
bench_dag_execution,
bench_concurrent_execution,
bench_record_processing,
bench_metrics_collection,
bench_state_management,
bench_memory_usage,
bench_transformations,
bench_resource_allocation,
bench_end_to_end_pipeline
);
criterion_main!(benches);