Changed delta_generator to make use of --max_threads

Adding an option to configure max number of threads used during
ota_from_target_files

Bug: 277794445
Test: tested ota_from_target_files
(cherry picked from https://android-review.googlesource.com/q/commit:cbe8e633f29fa63a5734cd5b43dc7842b024efdc)
Merged-In: I9bae92ee43c21926434ab04563a2b4f0678b73eb
Change-Id: I9bae92ee43c21926434ab04563a2b4f0678b73eb
diff --git a/payload_generator/delta_diff_utils.cc b/payload_generator/delta_diff_utils.cc
index 8ee1436..152da4d 100644
--- a/payload_generator/delta_diff_utils.cc
+++ b/payload_generator/delta_diff_utils.cc
@@ -677,6 +677,10 @@
 
   size_t max_threads = GetMaxThreads();
 
+  if (config.max_threads > 0) {
+    max_threads = config.max_threads;
+  }
+
   // Sort the files in descending order based on number of new blocks to make
   // sure we start the largest ones first.
   if (file_delta_processors.size() > max_threads) {
diff --git a/payload_generator/generate_delta_main.cc b/payload_generator/generate_delta_main.cc
index 5a901d3..6616ee1 100644
--- a/payload_generator/generate_delta_main.cc
+++ b/payload_generator/generate_delta_main.cc
@@ -455,6 +455,11 @@
               "Compression parameter passed to mkfs.erofs's -z option. "
               "Example: lz4 lz4hc,9");
 
+DEFINE_int64(max_threads,
+             0,
+             "The maximum number of threads allowed for generating "
+             "ota.");
+
 void RoundDownPartitions(const ImageConfig& config) {
   for (const auto& part : config.partitions) {
     if (part.path.empty()) {
@@ -764,6 +769,8 @@
 
   payload_config.security_patch_level = FLAGS_security_patch_level;
 
+  payload_config.max_threads = FLAGS_max_threads;
+
   if (!FLAGS_partition_timestamps.empty()) {
     CHECK(ParsePerPartitionTimestamps(FLAGS_partition_timestamps,
                                       &payload_config));
diff --git a/payload_generator/payload_generation_config.h b/payload_generator/payload_generation_config.h
index fc56f56..225237a 100644
--- a/payload_generator/payload_generation_config.h
+++ b/payload_generator/payload_generation_config.h
@@ -263,6 +263,8 @@
 
   std::string security_patch_level;
 
+  uint32_t max_threads = 0;
+
   std::vector<bsdiff::CompressorType> compressors{
       bsdiff::CompressorType::kBZ2, bsdiff::CompressorType::kBrotli};
 
diff --git a/scripts/brillo_update_payload b/scripts/brillo_update_payload
index 6652b38..083bfc2 100755
--- a/scripts/brillo_update_payload
+++ b/scripts/brillo_update_payload
@@ -222,6 +222,8 @@
     "Compression parameter passed to mkfs.erofs's -z option."
   DEFINE_string security_patch_level "" \
     "Optional: security patch level of this OTA"
+  DEFINE_string max_threads "" \
+    "Optional: specifies max_threads used to generate OTA"
 fi
 if [[ "${COMMAND}" == "hash" || "${COMMAND}" == "sign" ]]; then
   DEFINE_string unsigned_payload "" "Path to the input unsigned payload."
@@ -761,6 +763,11 @@
       --disable_vabc="${FLAGS_disable_vabc}" )
   fi
 
+  if [[ -n "${FLAGS_max_threads}" ]]; then
+    GENERATOR_ARGS+=(
+      --max_threads="${FLAGS_max_threads}" )
+  fi
+
   # minor version is set only for delta or partial payload.
   if [[ -n "${FORCE_MINOR_VERSION}" ]]; then
     GENERATOR_ARGS+=( --minor_version="${FORCE_MINOR_VERSION}" )