From 9f7bf7eff316b2775ef16835a443aab3fa08f00c Mon Sep 17 00:00:00 2001
From: Daniel Zheng <zhengdaniel@google.com>
Date: Fri, 3 Mar 2023 07:16:46 +0000
Subject: [PATCH] Added support for Resize Task

Test: tested flashall on raven
Bug: 194686221
Change-Id: I60586756b2d5f99848a664def8204c0bd47d7e67
---
 fastboot/fastboot.cpp | 16 +++++++++-------
 fastboot/task.cpp     | 13 +++++++++++++
 fastboot/task.h       | 13 +++++++++++++
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 22d15da19..ad9b4246d 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1631,13 +1631,13 @@ void FlashAllTool::Flash() {
         update_super_task->Run();
         // Resize any logical partition to 0, so each partition is reset to 0
         // extents, and will achieve more optimal allocation.
+        std::vector<std::unique_ptr<ResizeTask>> resize_tasks;
         for (const auto& [image, slot] : os_images_) {
-            auto resize_partition = [](const std::string& partition) -> void {
-                if (is_logical(partition)) {
-                    fb->ResizePartition(partition, "0");
-                }
-            };
-            do_for_partitions(image->part_name, slot, resize_partition, false);
+            resize_tasks.emplace_back(
+                    std::make_unique<ResizeTask>(fp_, image->part_name, "0", slot));
+        }
+        for (auto& i : resize_tasks) {
+            i->Run();
         }
     }
     FlashImages(os_images_);
@@ -2356,7 +2356,9 @@ int FastBootTool::Main(int argc, char* argv[]) {
         } else if (command == FB_CMD_RESIZE_PARTITION) {
             std::string partition = next_arg(&args);
             std::string size = next_arg(&args);
-            fb->ResizePartition(partition, size);
+            std::unique_ptr<ResizeTask> resize_task =
+                    std::make_unique<ResizeTask>(fp.get(), partition, size, slot_override);
+            resize_task->Run();
         } else if (command == "gsi") {
             std::string arg = next_arg(&args);
             if (arg == "wipe") {
diff --git a/fastboot/task.cpp b/fastboot/task.cpp
index 59abf831a..799b5c536 100644
--- a/fastboot/task.cpp
+++ b/fastboot/task.cpp
@@ -167,3 +167,16 @@ void UpdateSuperTask::Run() {
     }
     fp_->fb->RawCommand(command, "Updating super partition");
 }
+
+ResizeTask::ResizeTask(FlashingPlan* fp, const std::string& pname, const std::string& size,
+                       const std::string& slot)
+    : fp_(fp), pname_(pname), size_(size), slot_(slot) {}
+
+void ResizeTask::Run() {
+    auto resize_partition = [this](const std::string& partition) -> void {
+        if (is_logical(partition)) {
+            fp_->fb->ResizePartition(partition, size_);
+        }
+    };
+    do_for_partitions(pname_, slot_, resize_partition, false);
+}
diff --git a/fastboot/task.h b/fastboot/task.h
index 630278ae8..7aa19a3ad 100644
--- a/fastboot/task.h
+++ b/fastboot/task.h
@@ -79,3 +79,16 @@ class UpdateSuperTask : public Task {
   private:
     FlashingPlan* fp_;
 };
+
+class ResizeTask : public Task {
+  public:
+    ResizeTask(FlashingPlan* fp, const std::string& pname, const std::string& size,
+               const std::string& slot);
+    void Run() override;
+
+  private:
+    FlashingPlan* fp_;
+    const std::string pname_;
+    const std::string size_;
+    const std::string slot_;
+};