1use crate::command::BenchCommand;
8use crate::error::{BenchError, Result};
9use crate::executor::{K6Executor, K6Results};
10use crate::k6_gen::{K6Config, K6ScriptGenerator};
11use crate::reporter::TerminalReporter;
12use crate::request_gen::RequestGenerator;
13use crate::scenarios::LoadScenario;
14use crate::spec_parser::SpecParser;
15use crate::target_parser::TargetConfig;
16use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
17use std::collections::HashMap;
18use std::path::{Path, PathBuf};
19use std::str::FromStr;
20use std::sync::Arc;
21use tokio::sync::Semaphore;
22use tokio::task::JoinHandle;
23
24#[derive(Debug, Clone)]
26pub struct TargetResult {
27 pub target_url: String,
29 pub target_index: usize,
31 pub results: K6Results,
33 pub output_dir: PathBuf,
35 pub success: bool,
37 pub error: Option<String>,
39}
40
41#[derive(Debug, Clone)]
43pub struct AggregatedResults {
44 pub target_results: Vec<TargetResult>,
46 pub total_targets: usize,
48 pub successful_targets: usize,
49 pub failed_targets: usize,
50 pub aggregated_metrics: AggregatedMetrics,
52}
53
54#[derive(Debug, Clone)]
56pub struct AggregatedMetrics {
57 pub total_requests: u64,
59 pub total_failed_requests: u64,
61 pub avg_duration_ms: f64,
63 pub p95_duration_ms: f64,
65 pub p99_duration_ms: f64,
67 pub error_rate: f64,
69 pub total_rps: f64,
71 pub avg_rps: f64,
73 pub total_vus_max: u32,
75 pub total_connections_opened: u64,
79 pub total_iterations_completed: u64,
81}
82
83impl AggregatedMetrics {
84 fn from_results(results: &[TargetResult]) -> Self {
86 let mut total_requests = 0u64;
87 let mut total_failed_requests = 0u64;
88 let mut durations = Vec::new();
89 let mut p95_values = Vec::new();
90 let mut p99_values = Vec::new();
91 let mut total_rps = 0.0f64;
92 let mut total_vus_max = 0u32;
93 let mut total_connections_opened = 0u64;
94 let mut total_iterations_completed = 0u64;
95 let mut successful_count = 0usize;
96
97 for result in results {
98 if result.success {
99 total_requests += result.results.total_requests;
100 total_failed_requests += result.results.failed_requests;
101 durations.push(result.results.avg_duration_ms);
102 p95_values.push(result.results.p95_duration_ms);
103 p99_values.push(result.results.p99_duration_ms);
104 total_rps += result.results.rps;
105 total_vus_max += result.results.vus_max;
106 total_connections_opened += result.results.tcp_connect_samples;
107 total_iterations_completed += result.results.iterations_completed;
108 successful_count += 1;
109 }
110 }
111
112 let avg_duration_ms = if !durations.is_empty() {
113 durations.iter().sum::<f64>() / durations.len() as f64
114 } else {
115 0.0
116 };
117
118 let p95_duration_ms = if !p95_values.is_empty() {
119 p95_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
120 let index = (p95_values.len() as f64 * 0.95).ceil() as usize - 1;
121 p95_values[index.min(p95_values.len() - 1)]
122 } else {
123 0.0
124 };
125
126 let p99_duration_ms = if !p99_values.is_empty() {
127 p99_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
128 let index = (p99_values.len() as f64 * 0.99).ceil() as usize - 1;
129 p99_values[index.min(p99_values.len() - 1)]
130 } else {
131 0.0
132 };
133
134 let error_rate = if total_requests > 0 {
135 (total_failed_requests as f64 / total_requests as f64) * 100.0
136 } else {
137 0.0
138 };
139
140 let avg_rps = if successful_count > 0 {
141 total_rps / successful_count as f64
142 } else {
143 0.0
144 };
145
146 Self {
147 total_requests,
148 total_failed_requests,
149 avg_duration_ms,
150 p95_duration_ms,
151 p99_duration_ms,
152 error_rate,
153 total_rps,
154 avg_rps,
155 total_vus_max,
156 total_connections_opened,
157 total_iterations_completed,
158 }
159 }
160}
161
162pub struct ParallelExecutor {
164 base_command: BenchCommand,
166 targets: Vec<TargetConfig>,
168 max_concurrency: usize,
170 base_output: PathBuf,
172}
173
174impl ParallelExecutor {
175 pub fn new(
177 base_command: BenchCommand,
178 targets: Vec<TargetConfig>,
179 max_concurrency: usize,
180 ) -> Self {
181 let base_output = base_command.output.clone();
182 Self {
183 base_command,
184 targets,
185 max_concurrency,
186 base_output,
187 }
188 }
189
190 pub async fn execute_all(&self) -> Result<AggregatedResults> {
192 let total_targets = self.targets.len();
193 TerminalReporter::print_progress(&format!(
194 "Starting parallel execution for {} targets (max concurrency: {})",
195 total_targets, self.max_concurrency
196 ));
197
198 if !K6Executor::is_k6_installed() {
200 TerminalReporter::print_error("k6 is not installed");
201 TerminalReporter::print_warning(
202 "Install k6 from: https://k6.io/docs/get-started/installation/",
203 );
204 return Err(BenchError::K6NotFound);
205 }
206
207 TerminalReporter::print_progress("Loading OpenAPI specification(s)...");
209 let merged_spec = self.base_command.load_and_merge_specs().await?;
210 let parser = SpecParser::from_spec(merged_spec);
211 TerminalReporter::print_success("Specification(s) loaded");
212
213 let operations = if let Some(filter) = &self.base_command.operations {
215 parser.filter_operations(filter)?
216 } else {
217 parser.get_operations()
218 };
219
220 if operations.is_empty() {
221 return Err(BenchError::Other("No operations found in spec".to_string()));
222 }
223
224 TerminalReporter::print_success(&format!("Found {} operations", operations.len()));
225
226 TerminalReporter::print_progress("Generating request templates...");
228 let templates: Vec<_> = operations
229 .iter()
230 .map(RequestGenerator::generate_template)
231 .collect::<Result<Vec<_>>>()?;
232 TerminalReporter::print_success("Request templates generated");
233
234 let mut per_target_data: HashMap<
236 PathBuf,
237 (Vec<crate::request_gen::RequestTemplate>, Option<String>),
238 > = HashMap::new();
239 {
240 let mut unique_specs: Vec<PathBuf> = Vec::new();
241 for t in &self.targets {
242 if let Some(spec_path) = &t.spec {
243 if !unique_specs.contains(spec_path) {
244 unique_specs.push(spec_path.clone());
245 }
246 }
247 }
248 for spec_path in &unique_specs {
249 TerminalReporter::print_progress(&format!(
250 "Loading per-target spec: {}",
251 spec_path.display()
252 ));
253 match SpecParser::from_file(spec_path).await {
254 Ok(target_parser) => {
255 let target_ops = if let Some(filter) = &self.base_command.operations {
256 match target_parser.filter_operations(filter) {
257 Ok(ops) => ops,
258 Err(e) => {
259 TerminalReporter::print_warning(&format!(
260 "Failed to filter operations from {}: {}. Using shared spec.",
261 spec_path.display(),
262 e
263 ));
264 continue;
265 }
266 }
267 } else {
268 target_parser.get_operations()
269 };
270 let target_templates: Vec<_> = match target_ops
271 .iter()
272 .map(RequestGenerator::generate_template)
273 .collect::<Result<Vec<_>>>()
274 {
275 Ok(t) => t,
276 Err(e) => {
277 TerminalReporter::print_warning(&format!(
278 "Failed to generate templates from {}: {}. Using shared spec.",
279 spec_path.display(),
280 e
281 ));
282 continue;
283 }
284 };
285 let target_base_path = if let Some(cli_bp) = &self.base_command.base_path {
286 if cli_bp.is_empty() {
287 None
288 } else {
289 Some(cli_bp.clone())
290 }
291 } else {
292 target_parser.get_base_path()
293 };
294 TerminalReporter::print_success(&format!(
295 "Loaded {} operations from {}",
296 target_templates.len(),
297 spec_path.display()
298 ));
299 per_target_data
300 .insert(spec_path.clone(), (target_templates, target_base_path));
301 }
302 Err(e) => {
303 TerminalReporter::print_warning(&format!(
304 "Failed to load per-target spec {}: {}. Targets using this spec will use the shared spec.",
305 spec_path.display(),
306 e
307 ));
308 }
309 }
310 }
311 }
312
313 let base_headers = self.base_command.parse_headers()?;
315
316 let base_path = self.resolve_base_path(&parser);
318 if let Some(ref bp) = base_path {
319 TerminalReporter::print_progress(&format!("Using base path: {}", bp));
320 }
321
322 let scenario = LoadScenario::from_str(&self.base_command.scenario)
324 .map_err(BenchError::InvalidScenario)?;
325
326 let duration_secs_val = BenchCommand::parse_duration(&self.base_command.duration)?;
327
328 let security_testing_enabled_val =
330 self.base_command.security_test || self.base_command.wafbench_dir.is_some();
331
332 let has_advanced_features = self.base_command.data_file.is_some()
334 || self.base_command.error_rate.is_some()
335 || self.base_command.security_test
336 || self.base_command.parallel_create.is_some()
337 || self.base_command.wafbench_dir.is_some();
338
339 let enhancement_code = if has_advanced_features {
340 let dummy_script = "export const options = {};";
341 let enhanced = self.base_command.generate_enhanced_script(dummy_script)?;
342 if let Some(pos) = enhanced.find("export const options") {
343 enhanced[..pos].to_string()
344 } else {
345 String::new()
346 }
347 } else {
348 String::new()
349 };
350
351 let semaphore = Arc::new(Semaphore::new(self.max_concurrency));
353 let multi_progress = MultiProgress::new();
354
355 let progress_bars: Vec<ProgressBar> = (0..total_targets)
357 .map(|i| {
358 let pb = multi_progress.add(ProgressBar::new(1));
359 pb.set_style(
360 ProgressStyle::default_bar()
361 .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} {msg}")
362 .unwrap(),
363 );
364 pb.set_message(format!("Target {}", i + 1));
365 pb
366 })
367 .collect();
368
369 let mut handles: Vec<JoinHandle<Result<TargetResult>>> = Vec::new();
371
372 for (index, target) in self.targets.iter().enumerate() {
373 let target = target.clone();
374 let duration = self.base_command.duration.clone();
376 let vus = self.base_command.vus;
377 let scenario_str = self.base_command.scenario.clone();
378 let operations = self.base_command.operations.clone();
379 let auth = self.base_command.auth.clone();
380 let headers = self.base_command.headers.clone();
381 let threshold_percentile = self.base_command.threshold_percentile.clone();
382 let threshold_ms = self.base_command.threshold_ms;
383 let max_error_rate = self.base_command.max_error_rate;
384 let verbose = self.base_command.verbose;
385 let skip_tls_verify = self.base_command.skip_tls_verify;
386 let chunked_request_bodies = self.base_command.chunked_request_bodies;
387 let target_rps = self.base_command.target_rps;
388 let no_keep_alive = self.base_command.no_keep_alive;
389 let local_ips = self.base_command.source_ips.join(",");
396 let geo_source_ips = self.base_command.geo_source_ips.clone();
397 let geo_source_headers = self.base_command.geo_source_headers.clone();
398
399 let (templates, base_path) = if let Some(spec_path) = &target.spec {
401 if let Some((t, bp)) = per_target_data.get(spec_path) {
402 (t.clone(), bp.clone())
403 } else {
404 (templates.clone(), base_path.clone())
405 }
406 } else {
407 (templates.clone(), base_path.clone())
408 };
409
410 let base_headers = base_headers.clone();
411 let scenario = scenario.clone();
412 let duration_secs = duration_secs_val;
413 let base_output = self.base_output.clone();
414 let semaphore = semaphore.clone();
415 let progress_bar = progress_bars[index].clone();
416 let target_index = index;
417 let security_testing_enabled = security_testing_enabled_val;
418 let enhancement_code = enhancement_code.clone();
419
420 let handle = tokio::spawn(async move {
421 let _permit = semaphore.acquire().await.map_err(|e| {
423 BenchError::Other(format!("Failed to acquire semaphore: {}", e))
424 })?;
425
426 progress_bar.set_message(format!("Testing {}", target.url));
427
428 let result = Self::execute_single_target_internal(
430 &duration,
431 vus,
432 &scenario_str,
433 &operations,
434 &auth,
435 &headers,
436 &threshold_percentile,
437 threshold_ms,
438 max_error_rate,
439 verbose,
440 skip_tls_verify,
441 base_path.as_ref(),
442 &target,
443 target_index,
444 &templates,
445 &base_headers,
446 &scenario,
447 duration_secs,
448 &base_output,
449 security_testing_enabled,
450 chunked_request_bodies,
451 target_rps,
452 no_keep_alive,
453 &enhancement_code,
454 &local_ips,
455 &geo_source_ips,
456 &geo_source_headers,
457 )
458 .await;
459
460 progress_bar.inc(1);
461 progress_bar.finish_with_message(format!("Completed {}", target.url));
462
463 result
464 });
465
466 handles.push(handle);
467 }
468
469 let mut target_results = Vec::new();
471 for (index, handle) in handles.into_iter().enumerate() {
472 match handle.await {
473 Ok(Ok(result)) => {
474 target_results.push(result);
475 }
476 Ok(Err(e)) => {
477 let target_url = self.targets[index].url.clone();
479 target_results.push(TargetResult {
480 target_url: target_url.clone(),
481 target_index: index,
482 results: K6Results::default(),
483 output_dir: self.base_output.join(format!("target_{}", index + 1)),
484 success: false,
485 error: Some(e.to_string()),
486 });
487 }
488 Err(e) => {
489 let target_url = self.targets[index].url.clone();
491 target_results.push(TargetResult {
492 target_url: target_url.clone(),
493 target_index: index,
494 results: K6Results::default(),
495 output_dir: self.base_output.join(format!("target_{}", index + 1)),
496 success: false,
497 error: Some(format!("Task join error: {}", e)),
498 });
499 }
500 }
501 }
502
503 target_results.sort_by_key(|r| r.target_index);
505
506 let aggregated_metrics = AggregatedMetrics::from_results(&target_results);
508
509 let successful_targets = target_results.iter().filter(|r| r.success).count();
510 let failed_targets = total_targets - successful_targets;
511
512 Ok(AggregatedResults {
513 target_results,
514 total_targets,
515 successful_targets,
516 failed_targets,
517 aggregated_metrics,
518 })
519 }
520
521 fn resolve_base_path(&self, parser: &SpecParser) -> Option<String> {
523 if let Some(cli_base_path) = &self.base_command.base_path {
525 if cli_base_path.is_empty() {
526 return None;
527 }
528 return Some(cli_base_path.clone());
529 }
530 parser.get_base_path()
532 }
533
534 #[allow(clippy::too_many_arguments)]
536 async fn execute_single_target_internal(
537 _duration: &str,
538 vus: u32,
539 _scenario_str: &str,
540 _operations: &Option<String>,
541 auth: &Option<String>,
542 _headers: &[String],
543 threshold_percentile: &str,
544 threshold_ms: u64,
545 max_error_rate: f64,
546 verbose: bool,
547 skip_tls_verify: bool,
548 base_path: Option<&String>,
549 target: &TargetConfig,
550 target_index: usize,
551 templates: &[crate::request_gen::RequestTemplate],
552 base_headers: &HashMap<String, String>,
553 scenario: &LoadScenario,
554 duration_secs: u64,
555 base_output: &Path,
556 security_testing_enabled: bool,
557 chunked_request_bodies: bool,
558 target_rps: Option<u32>,
559 no_keep_alive: bool,
560 enhancement_code: &str,
561 local_ips: &str,
562 geo_source_ips: &[String],
563 geo_source_headers: &[String],
564 ) -> Result<TargetResult> {
565 let mut custom_headers = base_headers.clone();
567 if let Some(target_headers) = &target.headers {
568 custom_headers.extend(target_headers.clone());
569 }
570
571 let auth_header = target.auth.as_ref().or(auth.as_ref()).cloned();
573
574 let k6_config = K6Config {
576 target_url: target.url.clone(),
577 base_path: base_path.cloned(),
578 scenario: scenario.clone(),
579 duration_secs,
580 max_vus: vus,
581 threshold_percentile: threshold_percentile.to_string(),
582 threshold_ms,
583 max_error_rate,
584 auth_header,
585 custom_headers,
586 skip_tls_verify,
587 security_testing_enabled,
588 chunked_request_bodies,
589 target_rps,
590 no_keep_alive,
591 geo_source_ips: geo_source_ips.to_vec(),
592 geo_source_headers: geo_source_headers.to_vec(),
593 };
594
595 let generator = K6ScriptGenerator::new(k6_config, templates.to_vec());
597 let mut script = generator.generate()?;
598
599 if !enhancement_code.is_empty() {
601 if let Some(pos) = script.find("export const options") {
602 script.insert_str(pos, enhancement_code);
603 }
604 }
605
606 let validation_errors = K6ScriptGenerator::validate_script(&script);
608 if !validation_errors.is_empty() {
609 return Err(BenchError::Other(format!(
610 "Script validation failed for target {}: {}",
611 target.url,
612 validation_errors.join(", ")
613 )));
614 }
615
616 let output_dir = base_output.join(format!("target_{}", target_index + 1));
618 std::fs::create_dir_all(&output_dir)?;
619
620 let script_path = output_dir.join("k6-script.js");
622 std::fs::write(&script_path, script)?;
623
624 let api_port = 6565 + (target_index as u16) + 1; let executor = K6Executor::new()?
635 .with_local_ips(local_ips.to_string())
636 .with_discard_response_bodies(true);
637 let results = executor
638 .execute_with_port(&script_path, Some(&output_dir), verbose, Some(api_port))
639 .await;
640
641 match results {
642 Ok(k6_results) => Ok(TargetResult {
643 target_url: target.url.clone(),
644 target_index,
645 results: k6_results,
646 output_dir,
647 success: true,
648 error: None,
649 }),
650 Err(e) => Ok(TargetResult {
651 target_url: target.url.clone(),
652 target_index,
653 results: K6Results::default(),
654 output_dir,
655 success: false,
656 error: Some(e.to_string()),
657 }),
658 }
659 }
660}
661
662#[cfg(test)]
663mod tests {
664 use super::*;
665
666 #[test]
667 fn test_aggregated_metrics_from_results() {
668 let results = vec![
669 TargetResult {
670 target_url: "http://api1.com".to_string(),
671 target_index: 0,
672 results: K6Results {
673 total_requests: 100,
674 failed_requests: 5,
675 avg_duration_ms: 100.0,
676 p95_duration_ms: 200.0,
677 p99_duration_ms: 300.0,
678 ..Default::default()
679 },
680 output_dir: PathBuf::from("output1"),
681 success: true,
682 error: None,
683 },
684 TargetResult {
685 target_url: "http://api2.com".to_string(),
686 target_index: 1,
687 results: K6Results {
688 total_requests: 200,
689 failed_requests: 10,
690 avg_duration_ms: 150.0,
691 p95_duration_ms: 250.0,
692 p99_duration_ms: 350.0,
693 ..Default::default()
694 },
695 output_dir: PathBuf::from("output2"),
696 success: true,
697 error: None,
698 },
699 ];
700
701 let metrics = AggregatedMetrics::from_results(&results);
702 assert_eq!(metrics.total_requests, 300);
703 assert_eq!(metrics.total_failed_requests, 15);
704 assert_eq!(metrics.avg_duration_ms, 125.0); }
706
707 #[test]
708 fn test_aggregated_metrics_with_failed_targets() {
709 let results = vec![
710 TargetResult {
711 target_url: "http://api1.com".to_string(),
712 target_index: 0,
713 results: K6Results {
714 total_requests: 100,
715 failed_requests: 5,
716 avg_duration_ms: 100.0,
717 p95_duration_ms: 200.0,
718 p99_duration_ms: 300.0,
719 ..Default::default()
720 },
721 output_dir: PathBuf::from("output1"),
722 success: true,
723 error: None,
724 },
725 TargetResult {
726 target_url: "http://api2.com".to_string(),
727 target_index: 1,
728 results: K6Results::default(),
729 output_dir: PathBuf::from("output2"),
730 success: false,
731 error: Some("Network error".to_string()),
732 },
733 ];
734
735 let metrics = AggregatedMetrics::from_results(&results);
736 assert_eq!(metrics.total_requests, 100);
738 assert_eq!(metrics.total_failed_requests, 5);
739 assert_eq!(metrics.avg_duration_ms, 100.0);
740 }
741
742 #[test]
743 fn test_aggregated_metrics_empty_results() {
744 let results = vec![];
745 let metrics = AggregatedMetrics::from_results(&results);
746 assert_eq!(metrics.total_requests, 0);
747 assert_eq!(metrics.total_failed_requests, 0);
748 assert_eq!(metrics.avg_duration_ms, 0.0);
749 assert_eq!(metrics.error_rate, 0.0);
750 }
751
752 #[test]
753 fn test_aggregated_metrics_error_rate_calculation() {
754 let results = vec![TargetResult {
755 target_url: "http://api1.com".to_string(),
756 target_index: 0,
757 results: K6Results {
758 total_requests: 1000,
759 failed_requests: 50,
760 avg_duration_ms: 100.0,
761 p95_duration_ms: 200.0,
762 p99_duration_ms: 300.0,
763 ..Default::default()
764 },
765 output_dir: PathBuf::from("output1"),
766 success: true,
767 error: None,
768 }];
769
770 let metrics = AggregatedMetrics::from_results(&results);
771 assert_eq!(metrics.error_rate, 5.0); }
773
774 #[test]
775 fn test_aggregated_metrics_p95_p99_calculation() {
776 let results = vec![
777 TargetResult {
778 target_url: "http://api1.com".to_string(),
779 target_index: 0,
780 results: K6Results {
781 total_requests: 100,
782 failed_requests: 0,
783 avg_duration_ms: 100.0,
784 p95_duration_ms: 150.0,
785 p99_duration_ms: 200.0,
786 ..Default::default()
787 },
788 output_dir: PathBuf::from("output1"),
789 success: true,
790 error: None,
791 },
792 TargetResult {
793 target_url: "http://api2.com".to_string(),
794 target_index: 1,
795 results: K6Results {
796 total_requests: 100,
797 failed_requests: 0,
798 avg_duration_ms: 200.0,
799 p95_duration_ms: 250.0,
800 p99_duration_ms: 300.0,
801 ..Default::default()
802 },
803 output_dir: PathBuf::from("output2"),
804 success: true,
805 error: None,
806 },
807 TargetResult {
808 target_url: "http://api3.com".to_string(),
809 target_index: 2,
810 results: K6Results {
811 total_requests: 100,
812 failed_requests: 0,
813 avg_duration_ms: 300.0,
814 p95_duration_ms: 350.0,
815 p99_duration_ms: 400.0,
816 ..Default::default()
817 },
818 output_dir: PathBuf::from("output3"),
819 success: true,
820 error: None,
821 },
822 ];
823
824 let metrics = AggregatedMetrics::from_results(&results);
825 assert_eq!(metrics.p95_duration_ms, 350.0);
828 assert_eq!(metrics.p99_duration_ms, 400.0);
829 }
830}