Refactor both PartitionWriter and VABCPartitionWriter

Previously, VABCPartitionWriter is a subclass of PartitionWriter, and
these 2 classes share some data via protected members. Now both classes
inherit from PartitionWriterInterface, no protected members.

Test: th
Change-Id: Ib7759df6d8895c20efacd0b038467feb4d81e422
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index 7c07f1d..31fa90e 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -1491,7 +1491,7 @@
       part_name, slot);
 }
 
-std::unique_ptr<PartitionWriter> DeltaPerformer::CreatePartitionWriter(
+std::unique_ptr<PartitionWriterInterface> DeltaPerformer::CreatePartitionWriter(
     const PartitionUpdate& partition_update,
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index 1d95e1e..30f5e9c 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -196,7 +196,7 @@
 
  protected:
   // Exposed as virtual for testing purposes.
-  virtual std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+  virtual std::unique_ptr<PartitionWriterInterface> CreatePartitionWriter(
       const PartitionUpdate& partition_update,
       const InstallPlan::Partition& install_part,
       DynamicPartitionControlInterface* dynamic_control,
@@ -426,7 +426,7 @@
       base::TimeDelta::FromSeconds(kCheckpointFrequencySeconds)};
   base::TimeTicks update_checkpoint_time_;
 
-  std::unique_ptr<PartitionWriter> partition_writer_;
+  std::unique_ptr<PartitionWriterInterface> partition_writer_;
 
   DISALLOW_COPY_AND_ASSIGN(DeltaPerformer);
 };
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index 27f846e..ed89f89 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -1116,7 +1116,7 @@
  public:
   using DeltaPerformer::DeltaPerformer;
 
-  std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+  std::unique_ptr<PartitionWriterInterface> CreatePartitionWriter(
       const PartitionUpdate& partition_update,
       const InstallPlan::Partition& install_part,
       DynamicPartitionControlInterface* dynamic_control,
diff --git a/payload_consumer/partition_writer.cc b/payload_consumer/partition_writer.cc
index 9db7ae0..7922eb2 100644
--- a/payload_consumer/partition_writer.cc
+++ b/payload_consumer/partition_writer.cc
@@ -39,6 +39,7 @@
 #include "update_engine/payload_consumer/extent_reader.h"
 #include "update_engine/payload_consumer/extent_writer.h"
 #include "update_engine/payload_consumer/file_descriptor_utils.h"
+#include "update_engine/payload_consumer/install_operation_executor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/mount_history.h"
 #include "update_engine/payload_consumer/payload_constants.h"
@@ -228,23 +229,6 @@
   return true;
 }
 
-std::ostream& operator<<(std::ostream& out,
-                         const RepeatedPtrField<Extent>& extents) {
-  if (extents.size() == 0) {
-    out << "[]";
-    return out;
-  }
-  out << "[";
-  auto begin = extents.begin();
-  out << *begin;
-  for (int i = 1; i < extents.size(); i++) {
-    ++begin;
-    out << ", " << *begin;
-  }
-  out << "]";
-  return out;
-}
-
 bool PartitionWriter::PerformSourceCopyOperation(
     const InstallOperation& operation, ErrorCode* error) {
   // The device may optimize the SOURCE_COPY operation.
@@ -264,7 +248,7 @@
   if (source_fd == nullptr) {
     LOG(ERROR) << "Unrecoverable source hash mismatch found on partition "
                << partition.partition_name()
-               << " extents: " << operation.src_extents();
+               << " extents: " << ExtentsToString(operation.src_extents());
     return false;
   }
 
diff --git a/payload_consumer/partition_writer.h b/payload_consumer/partition_writer.h
index 14bd18a..e11d987 100644
--- a/payload_consumer/partition_writer.h
+++ b/payload_consumer/partition_writer.h
@@ -29,19 +29,19 @@
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_operation_executor.h"
 #include "update_engine/payload_consumer/install_plan.h"
+#include "update_engine/payload_consumer/partition_writer_interface.h"
 #include "update_engine/payload_consumer/verified_source_fd.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
-
-class PartitionWriter {
+class PartitionWriter : public PartitionWriterInterface {
  public:
   PartitionWriter(const PartitionUpdate& partition_update,
                   const InstallPlan::Partition& install_part,
                   DynamicPartitionControlInterface* dynamic_control,
                   size_t block_size,
                   bool is_interactive);
-  virtual ~PartitionWriter();
+  ~PartitionWriter();
   static bool ValidateSourceHash(const brillo::Blob& calculated_hash,
                                  const InstallOperation& operation,
                                  const FileDescriptorPtr source_fd,
@@ -49,63 +49,58 @@
 
   // Perform necessary initialization work before InstallOperation can be
   // applied to this partition
-  [[nodiscard]] virtual bool Init(const InstallPlan* install_plan,
-                                  bool source_may_exist,
-                                  size_t next_op_index);
+  [[nodiscard]] bool Init(const InstallPlan* install_plan,
+                          bool source_may_exist,
+                          size_t next_op_index) override;
 
   // |CheckpointUpdateProgress| will be called after SetNextOpIndex(), but it's
   // optional. DeltaPerformer may or may not call this everytime an operation is
   // applied.
   //   |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);
+  void CheckpointUpdateProgress(size_t next_op_index) override;
 
   // Close partition writer, when calling this function there's no guarantee
   // that all |InstallOperations| are sent to |PartitionWriter|. This function
   // will be called even if we are pausing/aborting the update.
-  int Close();
+  int Close() override;
 
   // These perform a specific type of operation and return true on success.
   // |error| will be set if source hash mismatch, otherwise |error| might not be
   // set even if it fails.
-  [[nodiscard]] virtual bool PerformReplaceOperation(
-      const InstallOperation& operation, const void* data, size_t count);
-  [[nodiscard]] virtual bool PerformZeroOrDiscardOperation(
-      const InstallOperation& operation);
+  [[nodiscard]] bool PerformReplaceOperation(const InstallOperation& operation,
+                                             const void* data,
+                                             size_t count) override;
+  [[nodiscard]] bool PerformZeroOrDiscardOperation(
+      const InstallOperation& operation) override;
 
-  [[nodiscard]] virtual bool PerformSourceCopyOperation(
-      const InstallOperation& operation, ErrorCode* error);
-  [[nodiscard]] virtual bool PerformSourceBsdiffOperation(
+  [[nodiscard]] bool PerformSourceCopyOperation(
+      const InstallOperation& operation, ErrorCode* error) override;
+  [[nodiscard]] bool PerformSourceBsdiffOperation(
       const InstallOperation& operation,
       ErrorCode* error,
       const void* data,
-      size_t count);
-  [[nodiscard]] virtual bool PerformPuffDiffOperation(
-      const InstallOperation& operation,
-      ErrorCode* error,
-      const void* data,
-      size_t count);
+      size_t count) override;
+  [[nodiscard]] bool PerformPuffDiffOperation(const InstallOperation& operation,
+                                              ErrorCode* error,
+                                              const void* data,
+                                              size_t count) override;
 
   // |DeltaPerformer| calls this when all Install Ops are sent to partition
   // writer. No |Perform*Operation| methods will be called in the future, and
   // the partition writer is expected to be closed soon.
-  [[nodiscard]] virtual bool FinishedInstallOps() { return true; }
+  [[nodiscard]] bool FinishedInstallOps() override { return true; }
 
- protected:
+ private:
   friend class PartitionWriterTest;
   FRIEND_TEST(PartitionWriterTest, ChooseSourceFDTest);
 
   [[nodiscard]] bool OpenSourcePartition(uint32_t source_slot,
                                          bool source_may_exist);
-
-  bool OpenCurrentECCPartition();
-  // For a given operation, choose the source fd to be used (raw device or error
-  // correction device) based on the source operation hash.
-  // Returns nullptr if the source hash mismatch cannot be corrected, and set
-  // the |error| accordingly.
-  FileDescriptorPtr ChooseSourceFD(const InstallOperation& operation,
+  FileDescriptorPtr ChooseSourceFD(const InstallOperation& op,
                                    ErrorCode* error);
-  [[nodiscard]] virtual std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
+
+  [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
 
   const PartitionUpdate& partition_update_;
   const InstallPlan::Partition& install_part_;
@@ -128,7 +123,7 @@
 namespace partition_writer {
 // Return a PartitionWriter instance for perform InstallOps on this partition.
 // Uses VABCPartitionWriter for Virtual AB Compression
-std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+std::unique_ptr<PartitionWriterInterface> CreatePartitionWriter(
     const PartitionUpdate& partition_update,
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
diff --git a/payload_consumer/partition_writer_factory_android.cc b/payload_consumer/partition_writer_factory_android.cc
index 184e2d5..0a9a3fb 100644
--- a/payload_consumer/partition_writer_factory_android.cc
+++ b/payload_consumer/partition_writer_factory_android.cc
@@ -23,7 +23,7 @@
 
 namespace chromeos_update_engine::partition_writer {
 
-std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+std::unique_ptr<PartitionWriterInterface> CreatePartitionWriter(
     const PartitionUpdate& partition_update,
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
diff --git a/payload_consumer/partition_writer_factory_chromeos.cc b/payload_consumer/partition_writer_factory_chromeos.cc
index 609f043..0f7a11c 100644
--- a/payload_consumer/partition_writer_factory_chromeos.cc
+++ b/payload_consumer/partition_writer_factory_chromeos.cc
@@ -22,7 +22,7 @@
 #include "update_engine/payload_consumer/partition_writer.h"
 
 namespace chromeos_update_engine::partition_writer {
-std::unique_ptr<PartitionWriter> CreatePartitionWriter(
+std::unique_ptr<PartitionWriterInterface> CreatePartitionWriter(
     const PartitionUpdate& partition_update,
     const InstallPlan::Partition& install_part,
     DynamicPartitionControlInterface* dynamic_control,
diff --git a/payload_consumer/partition_writer_interface.h b/payload_consumer/partition_writer_interface.h
new file mode 100644
index 0000000..f8d6b9c
--- /dev/null
+++ b/payload_consumer/partition_writer_interface.h
@@ -0,0 +1,83 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef UPDATE_ENGINE_PARTITION_WRITER_INTERFACE_H_
+#define UPDATE_ENGINE_PARTITION_WRITER_INTERFACE_H_
+
+#include <cstdint>
+#include <string>
+
+#include <brillo/secure_blob.h>
+#include <gtest/gtest_prod.h>
+
+#include "update_engine/common/dynamic_partition_control_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"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+class PartitionWriterInterface {
+ public:
+  virtual ~PartitionWriterInterface() = default;
+
+  // Perform necessary initialization work before InstallOperation can be
+  // applied to this partition
+  [[nodiscard]] virtual bool Init(const InstallPlan* install_plan,
+                                  bool source_may_exist,
+                                  size_t next_op_index) = 0;
+
+  // |CheckpointUpdateProgress| will be called after SetNextOpIndex(), but it's
+  // optional. DeltaPerformer may or may not call this everytime an operation is
+  // applied.
+  //   |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) = 0;
+
+  // Close partition writer, when calling this function there's no guarantee
+  // that all |InstallOperations| are sent to |PartitionWriter|. This function
+  // will be called even if we are pausing/aborting the update.
+  virtual int Close() = 0;
+
+  // These perform a specific type of operation and return true on success.
+  // |error| will be set if source hash mismatch, otherwise |error| might not be
+  // set even if it fails.
+  [[nodiscard]] virtual bool PerformReplaceOperation(
+      const InstallOperation& operation, const void* data, size_t count) = 0;
+  [[nodiscard]] virtual bool PerformZeroOrDiscardOperation(
+      const InstallOperation& operation) = 0;
+
+  [[nodiscard]] virtual bool PerformSourceCopyOperation(
+      const InstallOperation& operation, ErrorCode* error) = 0;
+  [[nodiscard]] virtual bool PerformSourceBsdiffOperation(
+      const InstallOperation& operation,
+      ErrorCode* error,
+      const void* data,
+      size_t count) = 0;
+  [[nodiscard]] virtual bool PerformPuffDiffOperation(
+      const InstallOperation& operation,
+      ErrorCode* error,
+      const void* data,
+      size_t count) = 0;
+
+  // |DeltaPerformer| calls this when all Install Ops are sent to partition
+  // writer. No |Perform*Operation| methods will be called in the future, and
+  // the partition writer is expected to be closed soon.
+  [[nodiscard]] virtual bool FinishedInstallOps() = 0;
+};
+}  // namespace chromeos_update_engine
+
+#endif
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index a27847c..d181c75 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -18,6 +18,7 @@
 
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <libsnapshot/cow_writer.h>
@@ -55,12 +56,27 @@
 // label 3, Which contains all operation 2's data, but none of operation 3's
 // data.
 
+VABCPartitionWriter::VABCPartitionWriter(
+    const PartitionUpdate& partition_update,
+    const InstallPlan::Partition& install_part,
+    DynamicPartitionControlInterface* dynamic_control,
+    size_t block_size,
+    bool is_interactive)
+    : partition_update_(partition_update),
+      install_part_(install_part),
+      dynamic_control_(dynamic_control),
+      interactive_(is_interactive),
+      block_size_(block_size),
+      executor_(block_size),
+      verified_source_fd_(block_size, install_part.source_path) {}
+
 bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
                                bool source_may_exist,
                                size_t next_op_index) {
   TEST_AND_RETURN_FALSE(install_plan != nullptr);
-  TEST_AND_RETURN_FALSE(
-      OpenSourcePartition(install_plan->source_slot, source_may_exist));
+  if (source_may_exist) {
+    TEST_AND_RETURN_FALSE(verified_source_fd_.Open());
+  }
   std::optional<std::string> source_path;
   if (!install_part_.source_path.empty()) {
     // TODO(zhangkelvin) Make |source_path| a std::optional<std::string>
@@ -158,6 +174,43 @@
   return true;
 }
 
+bool VABCPartitionWriter::PerformReplaceOperation(const InstallOperation& op,
+                                                  const void* data,
+                                                  size_t count) {
+  // Setup the ExtentWriter stack based on the operation type.
+  std::unique_ptr<ExtentWriter> writer = CreateBaseExtentWriter();
+
+  return executor_.ExecuteReplaceOperation(op, std::move(writer), data, count);
+}
+
+bool VABCPartitionWriter::PerformSourceBsdiffOperation(
+    const InstallOperation& operation,
+    ErrorCode* error,
+    const void* data,
+    size_t count) {
+  FileDescriptorPtr source_fd =
+      verified_source_fd_.ChooseSourceFD(operation, error);
+  TEST_AND_RETURN_FALSE(source_fd != nullptr);
+
+  auto writer = CreateBaseExtentWriter();
+  return executor_.ExecuteSourceBsdiffOperation(
+      operation, std::move(writer), source_fd, data, count);
+}
+
+bool VABCPartitionWriter::PerformPuffDiffOperation(
+    const InstallOperation& operation,
+    ErrorCode* error,
+    const void* data,
+    size_t count) {
+  FileDescriptorPtr source_fd =
+      verified_source_fd_.ChooseSourceFD(operation, error);
+  TEST_AND_RETURN_FALSE(source_fd != nullptr);
+
+  auto writer = CreateBaseExtentWriter();
+  return executor_.ExecutePuffDiffOperation(
+      operation, std::move(writer), source_fd, data, count);
+}
+
 void VABCPartitionWriter::CheckpointUpdateProgress(size_t next_op_index) {
   // No need to call fsync/sync, as CowWriter flushes after a label is added
   // added.
@@ -175,9 +228,15 @@
 }
 
 VABCPartitionWriter::~VABCPartitionWriter() {
+  Close();
+}
+
+int VABCPartitionWriter::Close() {
   if (cow_writer_) {
     cow_writer_->Finalize();
+    cow_writer_ = nullptr;
   }
+  return 0;
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index 7fb2a2c..e8601a9 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -18,25 +18,29 @@
 #define UPDATE_ENGINE_VABC_PARTITION_WRITER_H_
 
 #include <memory>
+#include <string>
 #include <vector>
 
 #include <libsnapshot/snapshot_writer.h>
 
 #include "update_engine/common/cow_operation_convert.h"
+#include "update_engine/payload_consumer/install_operation_executor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
 
 namespace chromeos_update_engine {
-class VABCPartitionWriter final : public PartitionWriter {
+class VABCPartitionWriter final : public PartitionWriterInterface {
  public:
-  using PartitionWriter::PartitionWriter;
+  VABCPartitionWriter(const PartitionUpdate& partition_update,
+                      const InstallPlan::Partition& install_part,
+                      DynamicPartitionControlInterface* dynamic_control,
+                      size_t block_size,
+                      bool is_interactive);
   [[nodiscard]] bool Init(const InstallPlan* install_plan,
                           bool source_may_exist,
                           size_t next_op_index) override;
   ~VABCPartitionWriter() override;
 
-  [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter() override;
-
   // Only ZERO and SOURCE_COPY InstallOperations are treated special by VABC
   // Partition Writer. These operations correspond to COW_ZERO and COW_COPY. All
   // other operations just get converted to COW_REPLACE.
@@ -45,6 +49,20 @@
   [[nodiscard]] bool PerformSourceCopyOperation(
       const InstallOperation& operation, ErrorCode* error) override;
 
+  [[nodiscard]] bool PerformReplaceOperation(const InstallOperation& operation,
+                                             const void* data,
+                                             size_t count) override;
+
+  [[nodiscard]] bool PerformSourceBsdiffOperation(
+      const InstallOperation& operation,
+      ErrorCode* error,
+      const void* data,
+      size_t count) override;
+  [[nodiscard]] bool PerformPuffDiffOperation(const InstallOperation& operation,
+                                              ErrorCode* error,
+                                              const void* data,
+                                              size_t count) override;
+
   void CheckpointUpdateProgress(size_t next_op_index) override;
 
   static bool WriteAllCowOps(size_t block_size,
@@ -53,9 +71,24 @@
                              FileDescriptorPtr source_fd);
 
   [[nodiscard]] bool FinishedInstallOps() override;
+  int Close() override;
 
  private:
   std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
+
+  bool OpenCurrentECCPartition();
+  [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
+
+  const PartitionUpdate& partition_update_;
+  const InstallPlan::Partition& install_part_;
+  DynamicPartitionControlInterface* dynamic_control_;
+  // Path to source partition
+  std::string source_path_;
+
+  const bool interactive_;
+  const size_t block_size_;
+  InstallOperationExecutor executor_;
+  VerifiedSourceFd verified_source_fd_;
 };
 
 }  // namespace chromeos_update_engine