Handle resume of VABC updates by emitting labels am: 24599af599

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

Change-Id: I8a412fdb3078595160ee5ab030c6774053c44ae8
diff --git a/common/constants.cc b/common/constants.cc
index 8883668..1648596 100644
--- a/common/constants.cc
+++ b/common/constants.cc
@@ -98,6 +98,8 @@
 const char kPrefsUpdateStateNextDataLength[] = "update-state-next-data-length";
 const char kPrefsUpdateStateNextDataOffset[] = "update-state-next-data-offset";
 const char kPrefsUpdateStateNextOperation[] = "update-state-next-operation";
+const char kPrefsUpdateStatePartitionNextOperation[] =
+    "update-state-partition-next-operation";
 const char kPrefsUpdateStatePayloadIndex[] = "update-state-payload-index";
 const char kPrefsUpdateStateSHA256Context[] = "update-state-sha-256-context";
 const char kPrefsUpdateStateSignatureBlob[] = "update-state-signature-blob";
diff --git a/common/constants.h b/common/constants.h
index 3685102..2a2a62a 100644
--- a/common/constants.h
+++ b/common/constants.h
@@ -95,6 +95,7 @@
 extern const char kPrefsUpdateStateNextDataLength[];
 extern const char kPrefsUpdateStateNextDataOffset[];
 extern const char kPrefsUpdateStateNextOperation[];
+extern const char kPrefsUpdateStatePartitionNextOperation[];
 extern const char kPrefsUpdateStatePayloadIndex[];
 extern const char kPrefsUpdateStateSHA256Context[];
 extern const char kPrefsUpdateStateSignatureBlob[];
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index 30bd1ef..b75c8cf 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -48,6 +48,7 @@
 #include "update_engine/common/prefs_interface.h"
 #include "update_engine/common/subprocess.h"
 #include "update_engine/common/terminator.h"
+#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/bzip_extent_writer.h"
 #include "update_engine/payload_consumer/cached_file_descriptor.h"
 #include "update_engine/payload_consumer/certificate_parser_interface.h"
@@ -247,6 +248,7 @@
       install_part,
       dynamic_control,
       block_size_,
+      prefs_,
       interactive_,
       IsDynamicPartition(install_part.name));
   // Open source fds if we have a delta payload, or for partitions in the
@@ -1335,6 +1337,13 @@
         next_operation != kUpdateStateOperationInvalid && next_operation > 0))
     return false;
 
+  int64_t partition_next_op = -1;
+  if (!(prefs->GetInt64(kPrefsUpdateStatePartitionNextOperation,
+                        &partition_next_op) &&
+        partition_next_op >= 0)) {
+    return false;
+  }
+
   string interrupted_hash;
   if (!(prefs->GetString(kPrefsUpdateCheckResponseHash, &interrupted_hash) &&
         !interrupted_hash.empty() &&
@@ -1387,6 +1396,7 @@
     prefs->SetString(kPrefsUpdateStateSignatureBlob, "");
     prefs->SetInt64(kPrefsManifestMetadataSize, -1);
     prefs->SetInt64(kPrefsManifestSignatureSize, -1);
+    prefs->SetInt64(kPrefsUpdateStatePartitionNextOperation, -1);
     prefs->SetInt64(kPrefsResumedUpdateFailures, 0);
     prefs->Delete(kPrefsPostInstallSucceeded);
     prefs->Delete(kPrefsVerityWritten);
@@ -1431,6 +1441,7 @@
           partitions_[partition_index].operations(partition_operation_num);
       TEST_AND_RETURN_FALSE(
           prefs_->SetInt64(kPrefsUpdateStateNextDataLength, op.data_length()));
+      partition_writer_->CheckpointUpdateProgress(partition_operation_num);
     } else {
       TEST_AND_RETURN_FALSE(
           prefs_->SetInt64(kPrefsUpdateStateNextDataLength, 0));
diff --git a/payload_consumer/delta_performer_integration_test.cc b/payload_consumer/delta_performer_integration_test.cc
index 374131e..e83a20d 100644
--- a/payload_consumer/delta_performer_integration_test.cc
+++ b/payload_consumer/delta_performer_integration_test.cc
@@ -718,6 +718,8 @@
       .WillOnce(Return(true));
   EXPECT_CALL(prefs, SetInt64(kPrefsUpdateStateNextOperation, _))
       .WillRepeatedly(Return(true));
+  EXPECT_CALL(prefs, SetInt64(kPrefsUpdateStatePartitionNextOperation, _))
+      .WillRepeatedly(Return(true));
   EXPECT_CALL(prefs, GetInt64(kPrefsUpdateStateNextOperation, _))
       .WillOnce(Return(false));
   EXPECT_CALL(prefs, SetInt64(kPrefsUpdateStateNextDataOffset, _))
diff --git a/payload_consumer/partition_writer.cc b/payload_consumer/partition_writer.cc
index b4b869c..bdc6015 100644
--- a/payload_consumer/partition_writer.cc
+++ b/payload_consumer/partition_writer.cc
@@ -242,12 +242,14 @@
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
     size_t block_size,
+    PrefsInterface* prefs,
     bool is_interactive)
     : partition_update_(partition_update),
       install_part_(install_part),
       dynamic_control_(dynamic_control),
       interactive_(is_interactive),
-      block_size_(block_size) {}
+      block_size_(block_size),
+      prefs_(prefs) {}
 
 PartitionWriter::~PartitionWriter() {
   Close();
diff --git a/payload_consumer/partition_writer.h b/payload_consumer/partition_writer.h
index 1acbddc..a67339e 100644
--- a/payload_consumer/partition_writer.h
+++ b/payload_consumer/partition_writer.h
@@ -25,6 +25,7 @@
 #include <gtest/gtest_prod.h>
 
 #include "update_engine/common/dynamic_partition_control_interface.h"
+#include "update_engine/common/prefs_interface.h"
 #include "update_engine/payload_consumer/extent_writer.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
@@ -36,6 +37,7 @@
                   const InstallPlan::Partition& install_part,
                   DynamicPartitionControlInterface* dynamic_control,
                   size_t block_size,
+                  PrefsInterface* prefs,
                   bool is_interactive);
   virtual ~PartitionWriter();
   static bool ValidateSourceHash(const brillo::Blob& calculated_hash,
@@ -48,6 +50,11 @@
   [[nodiscard]] virtual bool Init(const InstallPlan* install_plan,
                                   bool source_may_exist);
 
+  // This will be called by DeltaPerformer after applying an InstallOp.
+  // |next_op_index| is index of next operation that should be applied.
+  // |next_op_index-1| is the last operation that is already applied.
+  virtual void CheckpointUpdateProgress(size_t next_op_index) {}
+
   int Close();
 
   // These perform a specific type of operation and return true on success.
@@ -111,6 +118,8 @@
   // Used to avoid re-opening the same source partition if it is not actually
   // error corrected.
   bool source_ecc_open_failure_{false};
+
+  PrefsInterface* prefs_;
 };
 
 namespace partition_writer {
@@ -121,6 +130,7 @@
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
     size_t block_size,
+    PrefsInterface* prefs,
     bool is_interactive,
     bool is_dynamic_partition);
 }  // namespace partition_writer
diff --git a/payload_consumer/partition_writer_factory_android.cc b/payload_consumer/partition_writer_factory_android.cc
index 0c9f7ea..5960d9b 100644
--- a/payload_consumer/partition_writer_factory_android.cc
+++ b/payload_consumer/partition_writer_factory_android.cc
@@ -19,6 +19,7 @@
 
 #include <base/logging.h>
 
+#include "update_engine/common/prefs_interface.h"
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
 
 namespace chromeos_update_engine::partition_writer {
@@ -28,6 +29,7 @@
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
     size_t block_size,
+    PrefsInterface* prefs,
     bool is_interactive,
     bool is_dynamic_partition) {
   if (dynamic_control &&
@@ -40,6 +42,7 @@
                                                  install_part,
                                                  dynamic_control,
                                                  block_size,
+                                                 prefs,
                                                  is_interactive);
   } else {
     LOG(INFO) << "Virtual AB Compression disabled, using Partition Writer for `"
@@ -48,6 +51,7 @@
                                              install_part,
                                              dynamic_control,
                                              block_size,
+                                             prefs,
                                              is_interactive);
   }
 }
diff --git a/payload_consumer/partition_writer_factory_chromeos.cc b/payload_consumer/partition_writer_factory_chromeos.cc
index 609f043..d89beb7 100644
--- a/payload_consumer/partition_writer_factory_chromeos.cc
+++ b/payload_consumer/partition_writer_factory_chromeos.cc
@@ -27,12 +27,14 @@
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
     size_t block_size,
+    PrefsInterface* prefs,
     bool is_interactive,
     bool is_dynamic_partition) {
   return std::make_unique<PartitionWriter>(partition_update,
                                            install_part,
                                            dynamic_control,
                                            block_size,
+                                           prefs,
                                            is_interactive);
 }
 }  // namespace chromeos_update_engine::partition_writer
diff --git a/payload_consumer/partition_writer_unittest.cc b/payload_consumer/partition_writer_unittest.cc
index 1ef4783..ca2ef41 100644
--- a/payload_consumer/partition_writer_unittest.cc
+++ b/payload_consumer/partition_writer_unittest.cc
@@ -112,8 +112,12 @@
   DeltaArchiveManifest manifest_{};
   PartitionUpdate partition_update_{};
   InstallPlan::Partition install_part_{};
-  PartitionWriter writer_{
-      partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
+  PartitionWriter writer_{partition_update_,
+                          install_part_,
+                          &dynamic_control_,
+                          kBlockSize,
+                          &prefs_,
+                          false};
 };
 // Test that the error-corrected file descriptor is used to read a partition
 // when no hash is available for SOURCE_COPY but it falls back to the normal
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index 1578f29..9e4d9b8 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -29,6 +29,30 @@
 #include "update_engine/payload_consumer/snapshot_extent_writer.h"
 
 namespace chromeos_update_engine {
+// Expected layout of COW file:
+// === Beginning of Cow Image ===
+// All Source Copy Operations
+// ========== Label 0 ==========
+// Operation 0 in PartitionUpdate
+// ========== Label 1 ==========
+// Operation 1 in PartitionUpdate
+// ========== label 2 ==========
+// Operation 2 in PartitionUpdate
+// ========== label 3 ==========
+// .
+// .
+// .
+
+// When resuming, pass |kPrefsUpdateStatePartitionNextOperation| as label to
+// |InitializeWithAppend|.
+// For example, suppose we finished writing SOURCE_COPY, and we finished writing
+// operation 2 completely. Update is suspended when we are half way through
+// operation 3.
+// |kPrefsUpdateStatePartitionNextOperation| would be 3, so we pass 3 as
+// label to |InitializeWithAppend|. The CowWriter will retain all data before
+// label 3, Which contains all operation 2's data, but none of operation 3's
+// data.
+
 bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
                                bool source_may_exist) {
   TEST_AND_RETURN_FALSE(install_plan != nullptr);
@@ -37,18 +61,37 @@
       install_part_.name, install_part_.source_path, install_plan->is_resume);
   TEST_AND_RETURN_FALSE(cow_writer_ != nullptr);
 
-  // TODO(zhangkelvin) Emit a label before writing SOURCE_COPY. When resuming,
+  // Emit a label before writing SOURCE_COPY. When resuming,
   // use pref or CowWriter::GetLastLabel to determine if the SOURCE_COPY ops are
   // written. No need to handle SOURCE_COPY operations when resuming.
 
   // ===== Resume case handling code goes here ====
+  if (install_plan->is_resume) {
+    int64_t next_op = 0;
+    if (!prefs_->GetInt64(kPrefsUpdateStatePartitionNextOperation, &next_op)) {
+      LOG(ERROR)
+          << "Resuming an update but can't fetch |next_op| from saved prefs.";
+      return false;
+    }
+    if (next_op < 0) {
+      TEST_AND_RETURN_FALSE(cow_writer_->Initialize());
+    } else {
+      TEST_AND_RETURN_FALSE(cow_writer_->InitializeAppend(next_op));
+      return true;
+    }
+  } else {
+    TEST_AND_RETURN_FALSE(cow_writer_->Initialize());
+  }
 
   // ==============================================
+  TEST_AND_RETURN_FALSE(
+      prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, -1));
 
   // TODO(zhangkelvin) Rewrite this in C++20 coroutine once that's available.
   auto converted = ConvertToCowOperations(partition_update_.operations(),
                                           partition_update_.merge_operations());
   std::vector<uint8_t> buffer(block_size_);
+
   for (const auto& cow_op : converted) {
     switch (cow_op.op) {
       case CowOperation::CowCopy:
@@ -71,6 +114,11 @@
         break;
     }
   }
+
+  // Emit label 0 to mark end of all SOURCE_COPY operations
+  cow_writer_->AddLabel(0);
+  TEST_AND_RETURN_FALSE(
+      prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, 0));
   return true;
 }
 
@@ -95,11 +143,27 @@
 }
 
 bool VABCPartitionWriter::Flush() {
-  // No need to do anything, as CowWriter automatically flushes every OP added.
+  // No need to call fsync/sync, as CowWriter flushes after a label is added
+  // added.
+  int64_t next_op = 0;
+  // |kPrefsUpdateStatePartitionNextOperation| will be maintained and set by
+  // CheckpointUpdateProgress()
+  TEST_AND_RETURN_FALSE(
+      prefs_->GetInt64(kPrefsUpdateStatePartitionNextOperation, &next_op));
+  // +1 because label 0 is reserved for SOURCE_COPY. See beginning of this
+  // file for explanation for cow format.
+  cow_writer_->AddLabel(next_op + 1);
   return true;
 }
 
+void VABCPartitionWriter::CheckpointUpdateProgress(size_t next_op_index) {
+  prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, next_op_index);
+}
+
 VABCPartitionWriter::~VABCPartitionWriter() {
+  // Reset |kPrefsUpdateStatePartitionNextOperation| once we finished a
+  // partition.
+  prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, -1);
   cow_writer_->Finalize();
 }
 
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index d65ac4a..3fc97ce 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -42,6 +42,7 @@
   [[nodiscard]] bool PerformSourceCopyOperation(
       const InstallOperation& operation, ErrorCode* error) override;
   [[nodiscard]] bool Flush() override;
+  void CheckpointUpdateProgress(size_t next_op_index) override;
 
  private:
   std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;