From d54ec406df47ca49ab0f191650f7062d367122c0 Mon Sep 17 00:00:00 2001
From: Jon Cinque <jon.cinque@gmail.com>
Date: Fri, 15 Apr 2022 13:49:43 +0200
Subject: [PATCH] sdk: Add --jobs parameter in build/test bpf (#24359)

---
 sdk/cargo-build-bpf/src/main.rs | 16 ++++++++++++++++
 sdk/cargo-test-bpf/src/main.rs  | 16 ++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/sdk/cargo-build-bpf/src/main.rs b/sdk/cargo-build-bpf/src/main.rs
index e2f330f078a075..3274cd73e82c23 100644
--- a/sdk/cargo-build-bpf/src/main.rs
+++ b/sdk/cargo-build-bpf/src/main.rs
@@ -29,6 +29,7 @@ struct Config<'a> {
     offline: bool,
     verbose: bool,
     workspace: bool,
+    jobs: Option<String>,
 }
 
 impl Default for Config<'_> {
@@ -51,6 +52,7 @@ impl Default for Config<'_> {
             offline: false,
             verbose: false,
             workspace: false,
+            jobs: None,
         }
     }
 }
@@ -545,6 +547,10 @@ fn build_bpf_package(config: &Config, target_directory: &Path, package: &cargo_m
     if config.verbose {
         cargo_build_args.push("--verbose");
     }
+    if let Some(jobs) = &config.jobs {
+        cargo_build_args.push("--jobs");
+        cargo_build_args.push(jobs);
+    }
     if let Some(args) = &config.cargo_args {
         for arg in args {
             cargo_build_args.push(arg);
@@ -793,6 +799,15 @@ fn main() {
                 .alias("all")
                 .help("Build all BPF packages in the workspace"),
         )
+        .arg(
+            Arg::new("jobs")
+                .short('j')
+                .long("jobs")
+                .takes_value(true)
+                .value_name("N")
+                .validator(|val| val.parse::<usize>().map_err(|e| e.to_string()))
+                .help("Number of parallel jobs, defaults to # of CPUs"),
+        )
         .get_matches_from(args);
 
     let bpf_sdk: PathBuf = matches.value_of_t_or_exit("bpf_sdk");
@@ -827,6 +842,7 @@ fn main() {
         offline: matches.is_present("offline"),
         verbose: matches.is_present("verbose"),
         workspace: matches.is_present("workspace"),
+        jobs: matches.value_of_t("jobs").ok(),
     };
     let manifest_path: Option<PathBuf> = matches.value_of_t("manifest_path").ok();
     build_bpf(config, manifest_path);
diff --git a/sdk/cargo-test-bpf/src/main.rs b/sdk/cargo-test-bpf/src/main.rs
index 37230b7bd87f4f..d0a820cfccf83d 100644
--- a/sdk/cargo-test-bpf/src/main.rs
+++ b/sdk/cargo-test-bpf/src/main.rs
@@ -24,6 +24,7 @@ struct Config {
     offline: bool,
     verbose: bool,
     workspace: bool,
+    jobs: Option<String>,
 }
 
 impl Default for Config {
@@ -42,6 +43,7 @@ impl Default for Config {
             offline: false,
             verbose: false,
             workspace: false,
+            jobs: None,
         }
     }
 }
@@ -116,6 +118,10 @@ fn test_bpf_package(config: &Config, target_directory: &Path, package: &cargo_me
     if config.verbose {
         cargo_args.push("--verbose");
     }
+    if let Some(jobs) = &config.jobs {
+        cargo_args.push("--jobs");
+        cargo_args.push(jobs);
+    }
 
     let mut build_bpf_args = cargo_args.clone();
     if let Some(bpf_sdk) = config.bpf_sdk.as_ref() {
@@ -294,6 +300,15 @@ fn main() {
                 .alias("all")
                 .help("Test all BPF packages in the workspace"),
         )
+        .arg(
+            Arg::new("jobs")
+                .short('j')
+                .long("jobs")
+                .takes_value(true)
+                .value_name("N")
+                .validator(|val| val.parse::<usize>().map_err(|e| e.to_string()))
+                .help("Number of parallel jobs, defaults to # of CPUs"),
+        )
         .arg(
             Arg::new("extra_cargo_test_args")
                 .value_name("extra args for cargo test and the test binary")
@@ -319,6 +334,7 @@ fn main() {
         offline: matches.is_present("offline"),
         verbose: matches.is_present("verbose"),
         workspace: matches.is_present("workspace"),
+        jobs: matches.value_of_t("jobs").ok(),
         ..Config::default()
     };