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/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()) {