ota_extractor: parallelize partition extractions am: c8c0057a7d

Original change: https://android-review.googlesource.com/c/platform/system/update_engine/+/3164519

Change-Id: I3ba5f45ea75d0b1c799c8ecc99d31c631f1f3234
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/aosp/ota_extractor.cc b/aosp/ota_extractor.cc
index a61a430..d3b2705 100644
--- a/aosp/ota_extractor.cc
+++ b/aosp/ota_extractor.cc
@@ -17,6 +17,7 @@
 #include <array>
 #include <cstdint>
 #include <cstdio>
+#include <future>
 #include <iterator>
 #include <memory>
 
@@ -53,6 +54,7 @@
               "",
               "Comma separated list of partitions to extract, leave empty for "
               "extracting all partitions");
+DEFINE_bool(single_thread, false, "Limit extraction to a single thread");
 
 using chromeos_update_engine::DeltaArchiveManifest;
 using chromeos_update_engine::PayloadMetadata;
@@ -195,19 +197,33 @@
                             metadata.GetMetadataSignatureSize() +
                             payload_offset;
 
-  for (const auto& partition : manifest.partitions()) {
-    if (!partitions.empty() &&
-        partitions.count(partition.partition_name()) == 0) {
-      continue;
+  if (FLAGS_single_thread) {
+    for (const auto& partition : manifest.partitions()) {
+      if (!ExtractImageFromPartition(manifest,
+                                     partition,
+                                     data_begin,
+                                     payload_fd,
+                                     input_dir,
+                                     output_dir)) {
+        return false;
+      }
     }
-
-    if (!ExtractImageFromPartition(manifest,
+  } else {
+    std::vector<std::future<bool>> futures;
+    for (const auto& partition : manifest.partitions()) {
+      futures.push_back(std::async(std::launch::async,
+                                   ExtractImageFromPartition,
+                                   manifest,
                                    partition,
                                    data_begin,
                                    payload_fd,
                                    input_dir,
-                                   output_dir)) {
-      return false;
+                                   output_dir));
+    }
+    for (auto& future : futures) {
+      if (!future.get()) {
+        return false;
+      }
     }
   }
   return true;