Move FinishedSnapshotWrites call back to PostinstallRunnerAction

After FinishedSnapshotWrites is called, if device reboots w/o slot
switch, libsnapshot will discard all update state. This makes slot
switching after reboot difficult. For better UX, move back.

This reverts commit 5b00dc5386e67db5609e79f9e5c603d4ef1098a9.

Reason for revert: b/318986391

Change-Id: If33d7661a907d779f1a860b439a707885a8882bf
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index 1adaabc..e41531c 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -1325,21 +1325,22 @@
   CHECK_NE(install_plan_.source_slot, UINT32_MAX);
   CHECK_NE(install_plan_.target_slot, UINT32_MAX);
 
-  auto install_plan_action = std::make_unique<InstallPlanAction>(install_plan_);
   auto postinstall_runner_action =
       std::make_unique<PostinstallRunnerAction>(boot_control_, hardware_);
-  SetStatusAndNotify(UpdateStatus::VERIFYING);
   postinstall_runner_action->set_delegate(this);
-  ErrorCode error_code{};
 
   // If last error code is kUpdatedButNotActive, we know that we reached this
   // state by calling applyPayload() with switch_slot=false. That applyPayload()
   // call would have already performed filesystem verification, therefore, we
   // can safely skip the verification to save time.
   if (last_error_ == ErrorCode::kUpdatedButNotActive) {
+    auto install_plan_action =
+        std::make_unique<InstallPlanAction>(install_plan_);
     BondActions(install_plan_action.get(), postinstall_runner_action.get());
     processor_->EnqueueAction(std::move(install_plan_action));
+    SetStatusAndNotify(UpdateStatus::FINALIZING);
   } else {
+    ErrorCode error_code{};
     if (!boot_control_->GetDynamicPartitionControl()
              ->PreparePartitionsForUpdate(GetCurrentSlot(),
                                           GetTargetSlot(),
@@ -1361,7 +1362,8 @@
                                 utils::ErrorCodeToString(error_code),
                             error_code);
     }
-
+    auto install_plan_action =
+        std::make_unique<InstallPlanAction>(install_plan_);
     auto filesystem_verifier_action =
         std::make_unique<FilesystemVerifierAction>(
             boot_control_->GetDynamicPartitionControl());
@@ -1371,6 +1373,7 @@
                 postinstall_runner_action.get());
     processor_->EnqueueAction(std::move(install_plan_action));
     processor_->EnqueueAction(std::move(filesystem_verifier_action));
+    SetStatusAndNotify(UpdateStatus::VERIFYING);
   }
 
   processor_->EnqueueAction(std::move(postinstall_runner_action));
diff --git a/payload_consumer/filesystem_verifier_action.cc b/payload_consumer/filesystem_verifier_action.cc
index 5342e6f..5345085 100644
--- a/payload_consumer/filesystem_verifier_action.cc
+++ b/payload_consumer/filesystem_verifier_action.cc
@@ -136,13 +136,6 @@
   partition_fd_.reset();
   // This memory is not used anymore.
   buffer_.clear();
-  if (code == ErrorCode::kSuccess && !cancelled_) {
-    if (!dynamic_control_->FinishUpdate(install_plan_.powerwash_required)) {
-      LOG(ERROR) << "Failed to FinishUpdate("
-                 << install_plan_.powerwash_required << ")";
-      code = ErrorCode::kFilesystemVerifierError;
-    }
-  }
 
   // If we didn't write verity, partitions were maped. Releaase resource now.
   if (!install_plan_.write_verity &&
diff --git a/payload_consumer/filesystem_verifier_action_unittest.cc b/payload_consumer/filesystem_verifier_action_unittest.cc
index 565976a..f2967db 100644
--- a/payload_consumer/filesystem_verifier_action_unittest.cc
+++ b/payload_consumer/filesystem_verifier_action_unittest.cc
@@ -32,7 +32,6 @@
 #include <libsnapshot/cow_writer.h>
 #include <sys/stat.h>
 
-#include "gmock/gmock-spec-builders.h"
 #include "update_engine/common/dynamic_partition_control_stub.h"
 #include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/mock_dynamic_partition_control.h"
@@ -557,13 +556,12 @@
     install_plan_.write_verity = true;
     ASSERT_NO_FATAL_FAILURE(SetHashWithVerity(&part));
   }
-  NiceMock<MockDynamicPartitionControl> dynamic_control;
   if (clear_target_hash) {
     part.target_hash.clear();
-  } else {
-    EXPECT_CALL(dynamic_control, FinishUpdate(0)).WillOnce(Return(true));
   }
 
+  NiceMock<MockDynamicPartitionControl> dynamic_control;
+
   EnableVABC(&dynamic_control, part.name);
   auto open_cow = [part]() {
     auto cow_fd = std::make_unique<EintrSafeFileDescriptor>();
@@ -663,7 +661,6 @@
   part.readonly_target_path = target_part_.path();
   NiceMock<MockDynamicPartitionControl> dynamic_control;
   EnableVABC(&dynamic_control, part.name);
-  ON_CALL(dynamic_control, FinishUpdate(_)).WillByDefault(Return(true));
 
   // b/186196758 is only visible if we repeatedely run FS verification w/o
   // writing verity
diff --git a/payload_consumer/postinstall_runner_action.cc b/payload_consumer/postinstall_runner_action.cc
index fa4654d..2cfd3c6 100644
--- a/payload_consumer/postinstall_runner_action.cc
+++ b/payload_consumer/postinstall_runner_action.cc
@@ -109,12 +109,11 @@
   auto dynamic_control = boot_control_->GetDynamicPartitionControl();
   CHECK(dynamic_control);
 
-  // Mount snapshot partitions for Virtual AB Compression
+  // Mount snapshot partitions for Virtual AB Compression Compression.
   if (dynamic_control->UpdateUsesSnapshotCompression()) {
     // Before calling MapAllPartitions to map snapshot devices, all CowWriters
     // must be closed, and MapAllPartitions() should be called.
     if (!install_plan_.partitions.empty()) {
-      dynamic_control->UnmapAllPartitions();
       if (!dynamic_control->MapAllPartitions()) {
         return CompletePostinstall(ErrorCode::kPostInstallMountError);
       }
@@ -429,13 +428,39 @@
   PerformPartitionPostinstall();
 }
 
+PostinstallRunnerAction::~PostinstallRunnerAction() {
+  if (!install_plan_.partitions.empty()) {
+    auto dynamic_control = boot_control_->GetDynamicPartitionControl();
+    CHECK(dynamic_control);
+    dynamic_control->UnmapAllPartitions();
+    LOG(INFO) << "Unmapped all partitions.";
+  }
+}
+
 void PostinstallRunnerAction::CompletePostinstall(ErrorCode error_code) {
   // We only attempt to mark the new slot as active if all the postinstall
   // steps succeeded.
+  DEFER {
+    if (error_code != ErrorCode::kSuccess &&
+        error_code != ErrorCode::kUpdatedButNotActive) {
+      LOG(ERROR) << "Postinstall action failed.";
+
+      // Undo any changes done to trigger Powerwash.
+      if (powerwash_scheduled_)
+        hardware_->CancelPowerwash();
+    }
+    processor_->ActionComplete(this, error_code);
+  };
   if (error_code == ErrorCode::kSuccess) {
     if (install_plan_.switch_slot_on_reboot) {
-      if (!boot_control_->SetActiveBootSlot(install_plan_.target_slot)) {
-        LOG(ERROR) << "Failed to switch slot to " << install_plan_.target_slot;
+      if (!boot_control_->GetDynamicPartitionControl()->MapAllPartitions()) {
+        LOG(ERROR) << "Failed to map all partitions before marking snapshot as "
+                      "ready for slot switch.";
+        return;
+      }
+      if (!boot_control_->GetDynamicPartitionControl()->FinishUpdate(
+              install_plan_.powerwash_required) ||
+          !boot_control_->SetActiveBootSlot(install_plan_.target_slot)) {
         error_code = ErrorCode::kPostinstallRunnerError;
       } else {
         // Schedules warm reset on next reboot, ignores the error.
@@ -447,26 +472,6 @@
       error_code = ErrorCode::kUpdatedButNotActive;
     }
   }
-  if (!install_plan_.partitions.empty()) {
-    auto dynamic_control = boot_control_->GetDynamicPartitionControl();
-    CHECK(dynamic_control);
-    dynamic_control->UnmapAllPartitions();
-    LOG(INFO) << "Unmapped all partitions.";
-  }
-
-  ScopedActionCompleter completer(processor_, this);
-  completer.set_code(error_code);
-
-  if (error_code != ErrorCode::kSuccess &&
-      error_code != ErrorCode::kUpdatedButNotActive) {
-    LOG(ERROR) << "Postinstall action failed.";
-
-    // Undo any changes done to trigger Powerwash.
-    if (powerwash_scheduled_)
-      hardware_->CancelPowerwash();
-
-    return;
-  }
 
   LOG(INFO) << "All post-install commands succeeded";
   if (HasOutputPipe()) {
diff --git a/payload_consumer/postinstall_runner_action.h b/payload_consumer/postinstall_runner_action.h
index 66721af..6017069 100644
--- a/payload_consumer/postinstall_runner_action.h
+++ b/payload_consumer/postinstall_runner_action.h
@@ -42,6 +42,7 @@
  public:
   PostinstallRunnerAction(BootControlInterface* boot_control,
                           HardwareInterface* hardware);
+  ~PostinstallRunnerAction();
 
   // InstallPlanAction overrides.
   void PerformAction() override;