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 = 0; let executor = K6Executor::new()?
649 .with_local_ips(local_ips.to_string())
650 .with_discard_response_bodies(true);
651 let results = executor
652 .execute_with_port(&script_path, Some(&output_dir), verbose, Some(api_port))
653 .await;
654
655 match results {
656 Ok(k6_results) => Ok(TargetResult {
657 target_url: target.url.clone(),
658 target_index,
659 results: k6_results,
660 output_dir,
661 success: true,
662 error: None,
663 }),
664 Err(e) => Ok(TargetResult {
665 target_url: target.url.clone(),
666 target_index,
667 results: K6Results::default(),
668 output_dir,
669 success: false,
670 error: Some(e.to_string()),
671 }),
672 }
673 }
674}
675
676#[cfg(test)]
677mod tests {
678 use super::*;
679
680 #[test]
681 fn test_aggregated_metrics_from_results() {
682 let results = vec![
683 TargetResult {
684 target_url: "http://api1.com".to_string(),
685 target_index: 0,
686 results: K6Results {
687 total_requests: 100,
688 failed_requests: 5,
689 avg_duration_ms: 100.0,
690 p95_duration_ms: 200.0,
691 p99_duration_ms: 300.0,
692 ..Default::default()
693 },
694 output_dir: PathBuf::from("output1"),
695 success: true,
696 error: None,
697 },
698 TargetResult {
699 target_url: "http://api2.com".to_string(),
700 target_index: 1,
701 results: K6Results {
702 total_requests: 200,
703 failed_requests: 10,
704 avg_duration_ms: 150.0,
705 p95_duration_ms: 250.0,
706 p99_duration_ms: 350.0,
707 ..Default::default()
708 },
709 output_dir: PathBuf::from("output2"),
710 success: true,
711 error: None,
712 },
713 ];
714
715 let metrics = AggregatedMetrics::from_results(&results);
716 assert_eq!(metrics.total_requests, 300);
717 assert_eq!(metrics.total_failed_requests, 15);
718 assert_eq!(metrics.avg_duration_ms, 125.0); }
720
721 #[test]
722 fn test_aggregated_metrics_with_failed_targets() {
723 let results = vec![
724 TargetResult {
725 target_url: "http://api1.com".to_string(),
726 target_index: 0,
727 results: K6Results {
728 total_requests: 100,
729 failed_requests: 5,
730 avg_duration_ms: 100.0,
731 p95_duration_ms: 200.0,
732 p99_duration_ms: 300.0,
733 ..Default::default()
734 },
735 output_dir: PathBuf::from("output1"),
736 success: true,
737 error: None,
738 },
739 TargetResult {
740 target_url: "http://api2.com".to_string(),
741 target_index: 1,
742 results: K6Results::default(),
743 output_dir: PathBuf::from("output2"),
744 success: false,
745 error: Some("Network error".to_string()),
746 },
747 ];
748
749 let metrics = AggregatedMetrics::from_results(&results);
750 assert_eq!(metrics.total_requests, 100);
752 assert_eq!(metrics.total_failed_requests, 5);
753 assert_eq!(metrics.avg_duration_ms, 100.0);
754 }
755
756 #[test]
757 fn test_aggregated_metrics_empty_results() {
758 let results = vec![];
759 let metrics = AggregatedMetrics::from_results(&results);
760 assert_eq!(metrics.total_requests, 0);
761 assert_eq!(metrics.total_failed_requests, 0);
762 assert_eq!(metrics.avg_duration_ms, 0.0);
763 assert_eq!(metrics.error_rate, 0.0);
764 }
765
766 #[test]
767 fn test_aggregated_metrics_error_rate_calculation() {
768 let results = vec![TargetResult {
769 target_url: "http://api1.com".to_string(),
770 target_index: 0,
771 results: K6Results {
772 total_requests: 1000,
773 failed_requests: 50,
774 avg_duration_ms: 100.0,
775 p95_duration_ms: 200.0,
776 p99_duration_ms: 300.0,
777 ..Default::default()
778 },
779 output_dir: PathBuf::from("output1"),
780 success: true,
781 error: None,
782 }];
783
784 let metrics = AggregatedMetrics::from_results(&results);
785 assert_eq!(metrics.error_rate, 5.0); }
787
788 #[test]
789 fn test_aggregated_metrics_p95_p99_calculation() {
790 let results = vec![
791 TargetResult {
792 target_url: "http://api1.com".to_string(),
793 target_index: 0,
794 results: K6Results {
795 total_requests: 100,
796 failed_requests: 0,
797 avg_duration_ms: 100.0,
798 p95_duration_ms: 150.0,
799 p99_duration_ms: 200.0,
800 ..Default::default()
801 },
802 output_dir: PathBuf::from("output1"),
803 success: true,
804 error: None,
805 },
806 TargetResult {
807 target_url: "http://api2.com".to_string(),
808 target_index: 1,
809 results: K6Results {
810 total_requests: 100,
811 failed_requests: 0,
812 avg_duration_ms: 200.0,
813 p95_duration_ms: 250.0,
814 p99_duration_ms: 300.0,
815 ..Default::default()
816 },
817 output_dir: PathBuf::from("output2"),
818 success: true,
819 error: None,
820 },
821 TargetResult {
822 target_url: "http://api3.com".to_string(),
823 target_index: 2,
824 results: K6Results {
825 total_requests: 100,
826 failed_requests: 0,
827 avg_duration_ms: 300.0,
828 p95_duration_ms: 350.0,
829 p99_duration_ms: 400.0,
830 ..Default::default()
831 },
832 output_dir: PathBuf::from("output3"),
833 success: true,
834 error: None,
835 },
836 ];
837
838 let metrics = AggregatedMetrics::from_results(&results);
839 assert_eq!(metrics.p95_duration_ms, 350.0);
842 assert_eq!(metrics.p99_duration_ms, 400.0);
843 }
844}