ota_extractor: parallelize partition extractions

Each sub-partition will be extracted in a separate thread.

Big partitions will still be the bottleneck point, but this still gives
meaningful improvements.

A simple test of Nothing Phone 2's series of OTA extractions (2.0.1 ->
2.0.1a -> 2.0.2 -> 2.0.2a -> 2.0.3 -> 2.0.4 -> 2.5.1 -> 2.5.1a ->
2.5.2 -> 2.5.3 -> 2.5.5 -> 2.5.6 -> 2.6.0, 5.1 GiB in total) in
AMD Ryzen Threadripper 7980X gives a speed boost of 2.55x.

[Before]
	Percent of CPU this job got: 102%
	Elapsed (wall clock) time (h:mm:ss or m:ss): 15:57.80

[After]
	Percent of CPU this job got: 258%
	Elapsed (wall clock) time (h:mm:ss or m:ss): 6:15.34

Users can use `-single_thread` to fallback to the previous behavior.

Test: m ota_extractor
Change-Id: Ia0b587949eede2ae6920f5c19fa716b6b8c246f1
Signed-off-by: Juhyung Park <qkrwngud825@gmail.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;