Merge 5d3613e99ec327d2ba487cc00de13b334618886f on remote branch

Change-Id: I2a577c1058cf2e4b65e1c13cb3bc9eccb1b2bea4

30 files changed, 597 insertions, 110 deletions

diff --git a/Android.bp b/Android.bp
index a4b7978e..dab97110 100644
--- a/Android.bp
+++ b/Android.bp
@@ -66,7 +66,7 @@ cc_defaults {
     ],
     include_dirs: ["system"],
     local_include_dirs: ["client_library/include"],
-    static_libs: ["libgtest_prod"],
+    header_libs: ["libgtest_prod_headers"],
     shared_libs: [
         "libbrillo-stream",
         "libbrillo",
@@ -426,6 +426,7 @@ cc_binary {
     // TODO(deymo): Remove external/cros/system_api/dbus once the strings are moved
     // out of the DBus interface.
     include_dirs: ["external/cros/system_api/dbus"],
+    header_libs: ["libgtest_prod_headers"],
 
     srcs: [
         "aosp/hardware_android.cc",
@@ -460,7 +461,6 @@ cc_binary {
         "gkiprops",
         "libevent",
         "libmodpb64",
-        "libgtest_prod",
         "libprotobuf-cpp-lite",
         "libbrillo-stream",
         "libbrillo",
@@ -763,6 +763,7 @@ cc_test {
 
     srcs: [
         "aosp/apex_handler_android_unittest.cc",
+        "aosp/cleanup_previous_update_action_unittest.cc",
         "aosp/dynamic_partition_control_android_unittest.cc",
         "aosp/update_attempter_android_unittest.cc",
         "certificate_checker_unittest.cc",
@@ -895,4 +896,4 @@ cc_binary_host {
         "libz",
         "update_metadata-protos",
     ],
-}
\ No newline at end of file
+}
diff --git a/aosp/cleanup_previous_update_action.cc b/aosp/cleanup_previous_update_action.cc
index b4475098..68954f6f 100644
--- a/aosp/cleanup_previous_update_action.cc
+++ b/aosp/cleanup_previous_update_action.cc
@@ -207,6 +207,7 @@ void CleanupPreviousUpdateAction::CheckSlotMarkedSuccessfulOrSchedule() {
   if (!kIsRecovery &&
       !boot_control_->IsSlotMarkedSuccessful(boot_control_->GetCurrentSlot())) {
     ScheduleWaitMarkBootSuccessful();
+    return;
   }
 
   if (metadata_device_ == nullptr) {
@@ -487,7 +488,11 @@ void CleanupPreviousUpdateAction::ReportMergeStats() {
                              vab_retrofit,
                              static_cast<int64_t>(report.cow_file_size()),
                              vab_compression_enabled,
-                             vab_compression_used);
+                             vab_compression_used,
+                             report.total_cow_size_bytes(),
+                             report.estimated_cow_size_bytes(),
+                             report.boot_complete_time_ms(),
+                             report.boot_complete_to_merge_start_time_ms());
 #endif
 }
 
diff --git a/aosp/cleanup_previous_update_action_unittest.cc b/aosp/cleanup_previous_update_action_unittest.cc
new file mode 100644
index 00000000..0d2b4e6a
--- /dev/null
+++ b/aosp/cleanup_previous_update_action_unittest.cc
@@ -0,0 +1,174 @@
+//
+// 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.
+//
+
+#include <algorithm>
+
+#include <brillo/message_loops/fake_message_loop.h>
+#include <gtest/gtest.h>
+#include <libsnapshot/snapshot.h>
+#include <libsnapshot/mock_snapshot.h>
+#include <libsnapshot/mock_snapshot_merge_stats.h>
+
+#include "update_engine/aosp/cleanup_previous_update_action.h"
+#include "update_engine/common/mock_boot_control.h"
+#include "update_engine/common/mock_dynamic_partition_control.h"
+#include "update_engine/common/mock_prefs.h"
+
+namespace chromeos_update_engine {
+
+using android::snapshot::AutoDevice;
+using android::snapshot::MockSnapshotManager;
+using android::snapshot::MockSnapshotMergeStats;
+using android::snapshot::UpdateState;
+using testing::_;
+using testing::AtLeast;
+using testing::Return;
+
+class MockCleanupPreviousUpdateActionDelegate final
+    : public CleanupPreviousUpdateActionDelegateInterface {
+  MOCK_METHOD(void, OnCleanupProgressUpdate, (double), (override));
+};
+
+class MockActionProcessor : public ActionProcessor {
+ public:
+  MOCK_METHOD(void, ActionComplete, (AbstractAction*, ErrorCode), (override));
+};
+
+class MockAutoDevice : public AutoDevice {
+ public:
+  explicit MockAutoDevice(std::string name) : AutoDevice(name) {}
+  ~MockAutoDevice() = default;
+};
+
+class CleanupPreviousUpdateActionTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    ON_CALL(boot_control_, GetDynamicPartitionControl())
+        .WillByDefault(Return(&dynamic_control_));
+    ON_CALL(boot_control_, GetCurrentSlot()).WillByDefault(Return(0));
+    ON_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+        .WillByDefault(Return(&mock_stats_));
+    action_.SetProcessor(&mock_processor_);
+    loop_.SetAsCurrent();
+  }
+
+  constexpr static FeatureFlag LAUNCH{FeatureFlag::Value::LAUNCH};
+  constexpr static FeatureFlag NONE{FeatureFlag::Value::NONE};
+  MockSnapshotManager mock_snapshot_;
+  MockPrefs mock_prefs_;
+  MockBootControl boot_control_;
+  MockDynamicPartitionControl dynamic_control_{};
+  MockCleanupPreviousUpdateActionDelegate mock_delegate_;
+  MockSnapshotMergeStats mock_stats_;
+  MockActionProcessor mock_processor_;
+  brillo::FakeMessageLoop loop_{nullptr};
+  CleanupPreviousUpdateAction action_{
+      &mock_prefs_, &boot_control_, &mock_snapshot_, &mock_delegate_};
+};
+
+TEST_F(CleanupPreviousUpdateActionTest, NonVabTest) {
+  // Since VAB isn't even enabled, |GetSnapshotMergeStatsInstance| shouldn't be
+  // called at all
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance()).Times(0);
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(NONE));
+  action_.PerformAction();
+}
+
+TEST_F(CleanupPreviousUpdateActionTest, VABSlotSuccessful) {
+  // Expectaion: if VABC is enabled, Clenup action should call
+  // |SnapshotMergeStats::Start()| to start merge, and wait for it to finish
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+      .Times(AtLeast(1));
+  EXPECT_CALL(mock_snapshot_, EnsureMetadataMounted())
+      .Times(AtLeast(1))
+      .WillRepeatedly(
+          []() { return std::make_unique<MockAutoDevice>("mock_device"); });
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(LAUNCH));
+  // CleanupPreviousUpdateAction should use whatever slot returned by
+  // |GetCurrentSlot()|
+  EXPECT_CALL(boot_control_, GetCurrentSlot())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(1));
+  EXPECT_CALL(boot_control_, IsSlotMarkedSuccessful(1))
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(true));
+  EXPECT_CALL(mock_snapshot_, ProcessUpdateState(_, _))
+      .Times(AtLeast(2))
+      .WillOnce(Return(UpdateState::Merging))
+      .WillRepeatedly(Return(UpdateState::MergeCompleted));
+  EXPECT_CALL(mock_stats_, Start())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(true));
+  EXPECT_CALL(mock_processor_, ActionComplete(&action_, ErrorCode::kSuccess))
+      .Times(1);
+  action_.PerformAction();
+  while (loop_.PendingTasks()) {
+    ASSERT_TRUE(loop_.RunOnce(true));
+  }
+}
+
+TEST_F(CleanupPreviousUpdateActionTest, VabSlotNotReady) {
+  // Cleanup action should repeatly query boot control until the slot is marked
+  // successful.
+  static constexpr auto MAX_TIMEPOINT =
+      std::chrono::steady_clock::time_point::max();
+  EXPECT_CALL(mock_snapshot_, GetSnapshotMergeStatsInstance())
+      .Times(AtLeast(1));
+  EXPECT_CALL(mock_snapshot_, EnsureMetadataMounted())
+      .Times(AtLeast(1))
+      .WillRepeatedly(
+          []() { return std::make_unique<MockAutoDevice>("mock_device"); });
+  EXPECT_CALL(dynamic_control_, GetVirtualAbFeatureFlag())
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(LAUNCH));
+  auto slot_success_time = MAX_TIMEPOINT;
+  auto merge_start_time = MAX_TIMEPOINT;
+  EXPECT_CALL(boot_control_, IsSlotMarkedSuccessful(_))
+      .Times(AtLeast(3))
+      .WillOnce(Return(false))
+      .WillOnce(Return(false))
+      .WillOnce([&slot_success_time]() {
+        slot_success_time =
+            std::min(slot_success_time, std::chrono::steady_clock::now());
+        return true;
+      });
+
+  EXPECT_CALL(mock_stats_, Start())
+      .Times(1)
+      .WillRepeatedly([&merge_start_time]() {
+        merge_start_time =
+            std::min(merge_start_time, std::chrono::steady_clock::now());
+        return true;
+      });
+
+  EXPECT_CALL(mock_snapshot_, ProcessUpdateState(_, _))
+      .Times(AtLeast(1))
+      .WillRepeatedly(Return(UpdateState::MergeCompleted));
+  EXPECT_CALL(mock_processor_, ActionComplete(&action_, ErrorCode::kSuccess))
+      .Times(1);
+  action_.PerformAction();
+  while (loop_.PendingTasks()) {
+    ASSERT_TRUE(loop_.RunOnce(true));
+  }
+  ASSERT_LT(slot_success_time, merge_start_time)
+      << "Merge should not be started until slot is marked successful";
+}
+
+}  // namespace chromeos_update_engine
diff --git a/aosp/dynamic_partition_control_android.cc b/aosp/dynamic_partition_control_android.cc
index 444fe42c..ab349a82 100644
--- a/aosp/dynamic_partition_control_android.cc
+++ b/aosp/dynamic_partition_control_android.cc
@@ -32,6 +32,7 @@
 #include <base/files/file_util.h>
 #include <base/logging.h>
 #include <base/strings/string_util.h>
+#include <base/strings/stringprintf.h>
 #include <bootloader_message/bootloader_message.h>
 #include <fs_mgr.h>
 #include <fs_mgr_dm_linear.h>
@@ -70,6 +71,7 @@ using android::snapshot::Return;
 using android::snapshot::SnapshotManager;
 using android::snapshot::SnapshotManagerStub;
 using android::snapshot::UpdateState;
+using base::StringPrintf;
 
 namespace chromeos_update_engine {
 
@@ -830,33 +832,90 @@ bool DynamicPartitionControlAndroid::PrepareDynamicPartitionsForUpdate(
   return StoreMetadata(target_device, builder.get(), target_slot);
 }
 
+DynamicPartitionControlAndroid::SpaceLimit
+DynamicPartitionControlAndroid::GetSpaceLimit(bool use_snapshot) {
+  // On device retrofitting dynamic partitions, allocatable_space = "super",
+  // where "super" is the sum of all block devices for that slot. Since block
+  // devices are dedicated for the corresponding slot, there's no need to halve
+  // the allocatable space.
+  if (GetDynamicPartitionsFeatureFlag().IsRetrofit())
+    return SpaceLimit::ERROR_IF_EXCEEDED_SUPER;
+
+  // On device launching dynamic partitions w/o VAB, regardless of recovery
+  // sideload, super partition must be big enough to hold both A and B slots of
+  // groups. Hence,
+  // allocatable_space = super / 2
+  if (!GetVirtualAbFeatureFlag().IsEnabled())
+    return SpaceLimit::ERROR_IF_EXCEEDED_HALF_OF_SUPER;
+
+  // Source build supports VAB. Super partition must be big enough to hold
+  // one slot of groups (ERROR_IF_EXCEEDED_SUPER). However, there are cases
+  // where additional warning messages needs to be written.
+
+  // If using snapshot updates, implying that target build also uses VAB,
+  // allocatable_space = super
+  if (use_snapshot)
+    return SpaceLimit::ERROR_IF_EXCEEDED_SUPER;
+
+  // Source build supports VAB but not using snapshot updates. There are
+  // several cases, as listed below.
+  // Sideloading: allocatable_space = super.
+  if (IsRecovery())
+    return SpaceLimit::ERROR_IF_EXCEEDED_SUPER;
+
+  // On launch VAB device, this implies secondary payload.
+  // Technically, we don't have to check anything, but sum(groups) < super
+  // still applies.
+  if (!GetVirtualAbFeatureFlag().IsRetrofit())
+    return SpaceLimit::ERROR_IF_EXCEEDED_SUPER;
+
+  // On retrofit VAB device, either of the following:
+  // - downgrading: allocatable_space = super / 2
+  // - secondary payload: don't check anything
+  // These two cases are indistinguishable,
+  // hence emit warning if sum(groups) > super / 2
+  return SpaceLimit::WARN_IF_EXCEEDED_HALF_OF_SUPER;
+}
+
 bool DynamicPartitionControlAndroid::CheckSuperPartitionAllocatableSpace(
     android::fs_mgr::MetadataBuilder* builder,
     const DeltaArchiveManifest& manifest,
     bool use_snapshot) {
-  uint64_t total_size = 0;
+  uint64_t sum_groups = 0;
   for (const auto& group : manifest.dynamic_partition_metadata().groups()) {
-    total_size += group.size();
-  }
-
-  std::string expr;
-  uint64_t allocatable_space = builder->AllocatableSpace();
-  // On device retrofitting dynamic partitions, allocatable_space = super.
-  // On device launching dynamic partitions w/o VAB,
-  //   allocatable_space = super / 2.
-  // On device launching dynamic partitions with VAB, allocatable_space = super.
-  // For recovery sideload, allocatable_space = super.
-  if (!GetDynamicPartitionsFeatureFlag().IsRetrofit() && !use_snapshot &&
-      !IsRecovery()) {
-    allocatable_space /= 2;
-    expr = "half of ";
-  }
-  if (total_size > allocatable_space) {
-    LOG(ERROR) << "The maximum size of all groups for the target slot"
-               << " (" << total_size << ") has exceeded " << expr
-               << "allocatable space for dynamic partitions "
-               << allocatable_space << ".";
-    return false;
+    sum_groups += group.size();
+  }
+
+  uint64_t full_space = builder->AllocatableSpace();
+  uint64_t half_space = full_space / 2;
+  constexpr const char* fmt =
+      "The maximum size of all groups for the target slot (%" PRIu64
+      ") has exceeded %sallocatable space for dynamic partitions %" PRIu64 ".";
+  switch (GetSpaceLimit(use_snapshot)) {
+    case SpaceLimit::ERROR_IF_EXCEEDED_HALF_OF_SUPER: {
+      if (sum_groups > half_space) {
+        LOG(ERROR) << StringPrintf(fmt, sum_groups, "HALF OF ", half_space);
+        return false;
+      }
+      // If test passes, it implies that the following two conditions also pass.
+      break;
+    }
+    case SpaceLimit::WARN_IF_EXCEEDED_HALF_OF_SUPER: {
+      if (sum_groups > half_space) {
+        LOG(WARNING) << StringPrintf(fmt, sum_groups, "HALF OF ", half_space)
+                     << " This is allowed for downgrade or secondary OTA on "
+                        "retrofit VAB device.";
+      }
+      // still check sum(groups) < super
+      [[fallthrough]];
+    }
+    case SpaceLimit::ERROR_IF_EXCEEDED_SUPER: {
+      if (sum_groups > full_space) {
+        LOG(ERROR) << base::StringPrintf(fmt, sum_groups, "", full_space);
+        return false;
+      }
+      break;
+    }
   }
 
   return true;
@@ -910,9 +969,16 @@ bool DynamicPartitionControlAndroid::UpdatePartitionMetadata(
     uint32_t target_slot,
     const DeltaArchiveManifest& manifest) {
   // Check preconditions.
-  LOG_IF(WARNING, !GetVirtualAbFeatureFlag().IsEnabled() || IsRecovery())
-      << "UpdatePartitionMetadata is called on a Virtual A/B device "
-         "but source partitions is not deleted. This is not allowed.";
+  if (GetVirtualAbFeatureFlag().IsEnabled()) {
+    CHECK(!target_supports_snapshot_ || IsRecovery())
+        << "Must use snapshot on VAB device when target build supports VAB and "
+           "not sideloading.";
+    LOG_IF(INFO, !target_supports_snapshot_)
+        << "Not using snapshot on VAB device because target build does not "
+           "support snapshot. Secondary or downgrade OTA?";
+    LOG_IF(INFO, IsRecovery())
+        << "Not using snapshot on VAB device because sideloading.";
+  }
 
   // If applying downgrade from Virtual A/B to non-Virtual A/B, the left-over
   // COW group needs to be deleted to ensure there are enough space to create
diff --git a/aosp/dynamic_partition_control_android.h b/aosp/dynamic_partition_control_android.h
index b7aa7eaa..df914018 100644
--- a/aosp/dynamic_partition_control_android.h
+++ b/aosp/dynamic_partition_control_android.h
@@ -258,6 +258,18 @@ class DynamicPartitionControlAndroid : public DynamicPartitionControlInterface {
                                           const DeltaArchiveManifest& manifest,
                                           uint64_t* required_size);
 
+  enum SpaceLimit {
+    // Most restricted: if sum(groups) > super / 2, error
+    ERROR_IF_EXCEEDED_HALF_OF_SUPER,
+    // Implies ERROR_IF_EXCEEDED_SUPER; then, if sum(groups) > super / 2, warn
+    WARN_IF_EXCEEDED_HALF_OF_SUPER,
+    // Least restricted: if sum(groups) > super, error
+    ERROR_IF_EXCEEDED_SUPER,
+  };
+  // Helper of CheckSuperPartitionAllocatableSpace. Determine limit for groups
+  // and partitions.
+  SpaceLimit GetSpaceLimit(bool use_snapshot);
+
   // Returns true if the allocatable space in super partition is larger than
   // the size of dynamic partition groups in the manifest.
   bool CheckSuperPartitionAllocatableSpace(
diff --git a/aosp/metrics_reporter_android.cc b/aosp/metrics_reporter_android.cc
index 1f8f45a6..a324fab1 100644
--- a/aosp/metrics_reporter_android.cc
+++ b/aosp/metrics_reporter_android.cc
@@ -18,6 +18,8 @@
 
 #include <stdint.h>
 
+#include <algorithm>
+#include <any>
 #include <memory>
 #include <string>
 
@@ -30,6 +32,7 @@
 #include <statslog.h>
 
 #include "update_engine/common/constants.h"
+#include "update_engine/payload_consumer/install_plan.h"
 
 using android::fs_mgr::GetPartitionGroupName;
 using android::fs_mgr::LpMetadata;
@@ -48,6 +51,22 @@ constexpr auto kMetricsReporterEnumOffset = 10000;
 int32_t GetStatsdEnumValue(int32_t value) {
   return kMetricsReporterEnumOffset + value;
 }
+
+bool IsHashTreeEnabled(
+    const chromeos_update_engine::InstallPlan* install_plan) {
+  return std::any_of(
+      install_plan->partitions.begin(),
+      install_plan->partitions.end(),
+      [](const auto& partition) { return partition.hash_tree_size > 0; });
+}
+
+bool IsFECEnabled(const chromeos_update_engine::InstallPlan* install_plan) {
+  return std::any_of(
+      install_plan->partitions.begin(),
+      install_plan->partitions.end(),
+      [](const auto& partition) { return partition.fec_size > 0; });
+}
+
 }  // namespace
 
 namespace chromeos_update_engine {
@@ -55,8 +74,10 @@ namespace chromeos_update_engine {
 namespace metrics {
 
 std::unique_ptr<MetricsReporterInterface> CreateMetricsReporter(
-    DynamicPartitionControlInterface* dynamic_partition_control) {
-  return std::make_unique<MetricsReporterAndroid>(dynamic_partition_control);
+    DynamicPartitionControlInterface* dynamic_partition_control,
+    const InstallPlan* install_plan) {
+  return std::make_unique<MetricsReporterAndroid>(dynamic_partition_control,
+                                                  install_plan);
 }
 
 }  // namespace metrics
@@ -164,7 +185,9 @@ void MetricsReporterAndroid::ReportSuccessfulUpdateMetrics(
       static_cast<int32_t>(total_bytes_downloaded),
       static_cast<int32_t>(download_overhead_percentage),
       static_cast<int32_t>(total_duration.InMinutes()),
-      static_cast<int32_t>(reboot_count));
+      static_cast<int32_t>(reboot_count),
+      IsHashTreeEnabled(install_plan_),
+      IsFECEnabled(install_plan_));
 }
 
 void MetricsReporterAndroid::ReportAbnormallyTerminatedUpdateAttemptMetrics() {
diff --git a/aosp/metrics_reporter_android.h b/aosp/metrics_reporter_android.h
index abe7c277..aeb579a1 100644
--- a/aosp/metrics_reporter_android.h
+++ b/aosp/metrics_reporter_android.h
@@ -22,14 +22,17 @@
 #include "update_engine/common/error_code.h"
 #include "update_engine/common/metrics_constants.h"
 #include "update_engine/common/metrics_reporter_interface.h"
+#include "update_engine/payload_consumer/install_plan.h"
 
 namespace chromeos_update_engine {
 
 class MetricsReporterAndroid : public MetricsReporterInterface {
  public:
   explicit MetricsReporterAndroid(
-      DynamicPartitionControlInterface* dynamic_partition_control)
-      : dynamic_partition_control_(dynamic_partition_control) {}
+      DynamicPartitionControlInterface* dynamic_partition_control,
+      const InstallPlan* install_plan)
+      : dynamic_partition_control_(dynamic_partition_control),
+        install_plan_(install_plan) {}
 
   ~MetricsReporterAndroid() override = default;
 
@@ -93,7 +96,8 @@ class MetricsReporterAndroid : public MetricsReporterInterface {
       bool has_time_restriction_policy, int time_to_update_days) override {}
 
  private:
-  DynamicPartitionControlInterface* dynamic_partition_control_;
+  DynamicPartitionControlInterface* dynamic_partition_control_{};
+  const InstallPlan* install_plan_{};
 
   DISALLOW_COPY_AND_ASSIGN(MetricsReporterAndroid);
 };
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index 1080acba..ba61f255 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -64,7 +64,6 @@ using base::Bind;
 using base::Time;
 using base::TimeDelta;
 using base::TimeTicks;
-using std::shared_ptr;
 using std::string;
 using std::vector;
 using update_engine::UpdateEngineStatus;
@@ -143,7 +142,7 @@ UpdateAttempterAndroid::UpdateAttempterAndroid(
       processor_(new ActionProcessor()),
       clock_(new Clock()) {
   metrics_reporter_ = metrics::CreateMetricsReporter(
-      boot_control_->GetDynamicPartitionControl());
+      boot_control_->GetDynamicPartitionControl(), &install_plan_);
   network_selector_ = network::CreateNetworkSelector();
 }
 
diff --git a/common/metrics_reporter_interface.h b/common/metrics_reporter_interface.h
index 29d13faf..a7a91a54 100644
--- a/common/metrics_reporter_interface.h
+++ b/common/metrics_reporter_interface.h
@@ -26,6 +26,7 @@
 #include "update_engine/common/dynamic_partition_control_interface.h"
 #include "update_engine/common/error_code.h"
 #include "update_engine/common/metrics_constants.h"
+#include "update_engine/payload_consumer/install_plan.h"
 
 namespace chromeos_update_engine {
 
@@ -237,7 +238,8 @@ class MetricsReporterInterface {
 namespace metrics {
 
 std::unique_ptr<MetricsReporterInterface> CreateMetricsReporter(
-    DynamicPartitionControlInterface* dynamic_partition_control);
+    DynamicPartitionControlInterface* dynamic_partition_control,
+    const InstallPlan* install_plan);
 
 }  // namespace metrics
 
diff --git a/common/metrics_reporter_stub.cc b/common/metrics_reporter_stub.cc
index 96b519bf..61559d95 100644
--- a/common/metrics_reporter_stub.cc
+++ b/common/metrics_reporter_stub.cc
@@ -23,7 +23,8 @@ namespace chromeos_update_engine {
 namespace metrics {
 
 std::unique_ptr<MetricsReporterInterface> CreateMetricsReporter(
-    DynamicPartitionControlInterface* dynamic_partition_control) {
+    DynamicPartitionControlInterface* dynamic_partition_control,
+    const InstallPlan* install_plan) {
   return std::make_unique<MetricsReporterStub>();
 }
 
diff --git a/common/mock_boot_control.h b/common/mock_boot_control.h
new file mode 100644
index 00000000..f75ce5ec
--- /dev/null
+++ b/common/mock_boot_control.h
@@ -0,0 +1,74 @@
+//
+// 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_COMMON_MOCK_BOOT_CONTROL_H_
+#define UPDATE_ENGINE_COMMON_MOCK_BOOT_CONTROL_H_
+
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+
+#include "update_engine/common/boot_control_stub.h"
+
+namespace chromeos_update_engine {
+
+class MockBootControl final : public BootControlStub {
+ public:
+  MOCK_METHOD(bool,
+              IsSlotMarkedSuccessful,
+              (BootControlInterface::Slot),
+              (const override));
+  MOCK_METHOD(unsigned int, GetNumSlots, (), (const override));
+  MOCK_METHOD(BootControlInterface::Slot, GetCurrentSlot, (), (const override));
+  MOCK_METHOD(bool,
+              GetPartitionDevice,
+              (const std::string&, Slot, bool, std::string*, bool*),
+              (const override));
+  MOCK_METHOD(bool,
+              GetPartitionDevice,
+              (const std::string&, BootControlInterface::Slot, std::string*),
+              (const override));
+  MOCK_METHOD(std::optional<PartitionDevice>,
+              GetPartitionDevice,
+              (const std::string&, uint32_t, uint32_t, bool),
+              (const override));
+
+  MOCK_METHOD(bool,
+              IsSlotBootable,
+              (BootControlInterface::Slot),
+              (const override));
+  MOCK_METHOD(bool,
+              MarkSlotUnbootable,
+              (BootControlInterface::Slot),
+              (override));
+  MOCK_METHOD(bool,
+              SetActiveBootSlot,
+              (BootControlInterface::Slot),
+              (override));
+  MOCK_METHOD(bool,
+              MarkBootSuccessfulAsync,
+              (base::Callback<void(bool)>),
+              (override));
+  MOCK_METHOD(DynamicPartitionControlInterface*,
+              GetDynamicPartitionControl,
+              (),
+              (override));
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // UPDATE_ENGINE_COMMON_MOCK_BOOT_CONTROL_H_
diff --git a/payload_consumer/bzip_extent_writer.cc b/payload_consumer/bzip_extent_writer.cc
index 0c25c71f..26fdc5f4 100644
--- a/payload_consumer/bzip_extent_writer.cc
+++ b/payload_consumer/bzip_extent_writer.cc
@@ -29,8 +29,7 @@ BzipExtentWriter::~BzipExtentWriter() {
   TEST_AND_RETURN(input_buffer_.empty());
 }
 
-bool BzipExtentWriter::Init(FileDescriptorPtr fd,
-                            const RepeatedPtrField<Extent>& extents,
+bool BzipExtentWriter::Init(const RepeatedPtrField<Extent>& extents,
                             uint32_t block_size) {
   // Init bzip2 stream
   int rc = BZ2_bzDecompressInit(&stream_,
@@ -39,7 +38,7 @@ bool BzipExtentWriter::Init(FileDescriptorPtr fd,
 
   TEST_AND_RETURN_FALSE(rc == BZ_OK);
 
-  return next_->Init(fd, extents, block_size);
+  return next_->Init(extents, block_size);
 }
 
 bool BzipExtentWriter::Write(const void* bytes, size_t count) {
diff --git a/payload_consumer/bzip_extent_writer.h b/payload_consumer/bzip_extent_writer.h
index ec181a78..38c041a8 100644
--- a/payload_consumer/bzip_extent_writer.h
+++ b/payload_consumer/bzip_extent_writer.h
@@ -40,8 +40,7 @@ class BzipExtentWriter : public ExtentWriter {
   }
   ~BzipExtentWriter() override;
 
-  bool Init(FileDescriptorPtr fd,
-            const google::protobuf::RepeatedPtrField<Extent>& extents,
+  bool Init(const google::protobuf::RepeatedPtrField<Extent>& extents,
             uint32_t block_size) override;
   bool Write(const void* bytes, size_t count) override;
 
diff --git a/payload_consumer/bzip_extent_writer_unittest.cc b/payload_consumer/bzip_extent_writer_unittest.cc
index b5870401..c93545ab 100644
--- a/payload_consumer/bzip_extent_writer_unittest.cc
+++ b/payload_consumer/bzip_extent_writer_unittest.cc
@@ -29,7 +29,6 @@
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 
-using google::protobuf::RepeatedPtrField;
 using std::min;
 using std::string;
 using std::vector;
@@ -64,9 +63,8 @@ TEST_F(BzipExtentWriterTest, SimpleTest) {
       0x22, 0x9c, 0x28, 0x48, 0x66, 0x61, 0xb8, 0xea, 0x00,
   };
 
-  BzipExtentWriter bzip_writer(std::make_unique<DirectExtentWriter>());
-  EXPECT_TRUE(
-      bzip_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  BzipExtentWriter bzip_writer(std::make_unique<DirectExtentWriter>(fd_));
+  EXPECT_TRUE(bzip_writer.Init({extents.begin(), extents.end()}, kBlockSize));
   EXPECT_TRUE(bzip_writer.Write(test, sizeof(test)));
 
   brillo::Blob buf;
@@ -97,9 +95,8 @@ TEST_F(BzipExtentWriterTest, ChunkedTest) {
 
   vector<Extent> extents = {ExtentForBytes(kBlockSize, 0, kDecompressedLength)};
 
-  BzipExtentWriter bzip_writer(std::make_unique<DirectExtentWriter>());
-  EXPECT_TRUE(
-      bzip_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  BzipExtentWriter bzip_writer(std::make_unique<DirectExtentWriter>(fd_));
+  EXPECT_TRUE(bzip_writer.Init({extents.begin(), extents.end()}, kBlockSize));
 
   brillo::Blob original_compressed_data = compressed_data;
   for (brillo::Blob::size_type i = 0; i < compressed_data.size();
diff --git a/payload_consumer/extent_writer.h b/payload_consumer/extent_writer.h
index 9e53561e..8b1b5326 100644
--- a/payload_consumer/extent_writer.h
+++ b/payload_consumer/extent_writer.h
@@ -38,8 +38,7 @@ class ExtentWriter {
   virtual ~ExtentWriter() = default;
 
   // Returns true on success.
-  virtual bool Init(FileDescriptorPtr fd,
-                    const google::protobuf::RepeatedPtrField<Extent>& extents,
+  virtual bool Init(const google::protobuf::RepeatedPtrField<Extent>& extents,
                     uint32_t block_size) = 0;
 
   // Returns true on success.
@@ -51,13 +50,11 @@ class ExtentWriter {
 
 class DirectExtentWriter : public ExtentWriter {
  public:
-  DirectExtentWriter() = default;
+  explicit DirectExtentWriter(FileDescriptorPtr fd) : fd_(fd) {}
   ~DirectExtentWriter() override = default;
 
-  bool Init(FileDescriptorPtr fd,
-            const google::protobuf::RepeatedPtrField<Extent>& extents,
+  bool Init(const google::protobuf::RepeatedPtrField<Extent>& extents,
             uint32_t block_size) override {
-    fd_ = fd;
     block_size_ = block_size;
     extents_ = extents;
     cur_extent_ = extents_.begin();
diff --git a/payload_consumer/extent_writer_unittest.cc b/payload_consumer/extent_writer_unittest.cc
index afebb1a9..5c67d3e8 100644
--- a/payload_consumer/extent_writer_unittest.cc
+++ b/payload_consumer/extent_writer_unittest.cc
@@ -65,9 +65,8 @@ class ExtentWriterTest : public ::testing::Test {
 TEST_F(ExtentWriterTest, SimpleTest) {
   vector<Extent> extents = {ExtentForRange(1, 1)};
   const string bytes = "1234";
-  DirectExtentWriter direct_writer;
-  EXPECT_TRUE(
-      direct_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  DirectExtentWriter direct_writer{fd_};
+  EXPECT_TRUE(direct_writer.Init({extents.begin(), extents.end()}, kBlockSize));
   EXPECT_TRUE(direct_writer.Write(bytes.data(), bytes.size()));
 
   EXPECT_EQ(static_cast<off_t>(kBlockSize + bytes.size()),
@@ -84,9 +83,8 @@ TEST_F(ExtentWriterTest, SimpleTest) {
 
 TEST_F(ExtentWriterTest, ZeroLengthTest) {
   vector<Extent> extents = {ExtentForRange(1, 1)};
-  DirectExtentWriter direct_writer;
-  EXPECT_TRUE(
-      direct_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  DirectExtentWriter direct_writer{fd_};
+  EXPECT_TRUE(direct_writer.Init({extents.begin(), extents.end()}, kBlockSize));
   EXPECT_TRUE(direct_writer.Write(nullptr, 0));
 }
 
@@ -109,9 +107,8 @@ void ExtentWriterTest::WriteAlignedExtents(size_t chunk_size,
   brillo::Blob data(kBlockSize * 3);
   test_utils::FillWithData(&data);
 
-  DirectExtentWriter direct_writer;
-  EXPECT_TRUE(
-      direct_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  DirectExtentWriter direct_writer{fd_};
+  EXPECT_TRUE(direct_writer.Init({extents.begin(), extents.end()}, kBlockSize));
 
   size_t bytes_written = 0;
   while (bytes_written < data.size()) {
@@ -150,9 +147,8 @@ TEST_F(ExtentWriterTest, SparseFileTest) {
   brillo::Blob data(17);
   test_utils::FillWithData(&data);
 
-  DirectExtentWriter direct_writer;
-  EXPECT_TRUE(
-      direct_writer.Init(fd_, {extents.begin(), extents.end()}, kBlockSize));
+  DirectExtentWriter direct_writer{fd_};
+  EXPECT_TRUE(direct_writer.Init({extents.begin(), extents.end()}, kBlockSize));
 
   size_t bytes_written = 0;
   while (bytes_written < (block_count * kBlockSize)) {
diff --git a/payload_consumer/fake_extent_writer.h b/payload_consumer/fake_extent_writer.h
index 7b2b7ac3..680b1b3b 100644
--- a/payload_consumer/fake_extent_writer.h
+++ b/payload_consumer/fake_extent_writer.h
@@ -33,8 +33,7 @@ class FakeExtentWriter : public ExtentWriter {
   ~FakeExtentWriter() override = default;
 
   // ExtentWriter overrides.
-  bool Init(FileDescriptorPtr /* fd */,
-            const google::protobuf::RepeatedPtrField<Extent>& /* extents */,
+  bool Init(const google::protobuf::RepeatedPtrField<Extent>& /* extents */,
             uint32_t /* block_size */) override {
     init_called_ = true;
     return true;
diff --git a/payload_consumer/file_descriptor_utils.cc b/payload_consumer/file_descriptor_utils.cc
index 846cbd72..9a6a601e 100644
--- a/payload_consumer/file_descriptor_utils.cc
+++ b/payload_consumer/file_descriptor_utils.cc
@@ -82,8 +82,8 @@ bool CopyAndHashExtents(FileDescriptorPtr source,
                         const RepeatedPtrField<Extent>& tgt_extents,
                         uint64_t block_size,
                         brillo::Blob* hash_out) {
-  DirectExtentWriter writer;
-  TEST_AND_RETURN_FALSE(writer.Init(target, tgt_extents, block_size));
+  DirectExtentWriter writer{target};
+  TEST_AND_RETURN_FALSE(writer.Init(tgt_extents, block_size));
   TEST_AND_RETURN_FALSE(utils::BlocksInExtents(src_extents) ==
                         utils::BlocksInExtents(tgt_extents));
   TEST_AND_RETURN_FALSE(
diff --git a/payload_consumer/filesystem_verifier_action_unittest.cc b/payload_consumer/filesystem_verifier_action_unittest.cc
index 2cad5232..c1006842 100644
--- a/payload_consumer/filesystem_verifier_action_unittest.cc
+++ b/payload_consumer/filesystem_verifier_action_unittest.cc
@@ -26,8 +26,8 @@
 #include <brillo/message_loops/message_loop_utils.h>
 #include <brillo/secure_blob.h>
 #include <gtest/gtest.h>
+#include <libsnapshot/snapshot_writer.h>
 
-#include "gmock/gmock-actions.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"
@@ -463,4 +463,26 @@ TEST_F(FilesystemVerifierActionTest, RunWithVABCNoVerity) {
   ASSERT_EQ(actual_read_size, part.target_size);
 }
 
+TEST_F(FilesystemVerifierActionTest, ReadAfterWrite) {
+  constexpr auto BLOCK_SIZE = 4096;
+  ScopedTempFile cow_device_file("cow_device.XXXXXX", true);
+  android::snapshot::CompressedSnapshotWriter snapshot_writer{
+      {.block_size = BLOCK_SIZE}};
+  snapshot_writer.SetCowDevice(android::base::unique_fd{cow_device_file.fd()});
+  snapshot_writer.Initialize();
+  std::vector<unsigned char> buffer;
+  buffer.resize(BLOCK_SIZE);
+  std::fill(buffer.begin(), buffer.end(), 123);
+
+  ASSERT_TRUE(snapshot_writer.AddRawBlocks(0, buffer.data(), buffer.size()));
+  ASSERT_TRUE(snapshot_writer.Finalize());
+  auto cow_reader = snapshot_writer.OpenReader();
+  ASSERT_NE(cow_reader, nullptr);
+  ASSERT_TRUE(snapshot_writer.AddRawBlocks(1, buffer.data(), buffer.size()));
+  ASSERT_TRUE(snapshot_writer.AddRawBlocks(2, buffer.data(), buffer.size()));
+  ASSERT_TRUE(snapshot_writer.Finalize());
+  cow_reader = snapshot_writer.OpenReader();
+  ASSERT_NE(cow_reader, nullptr);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_writer.cc b/payload_consumer/partition_writer.cc
index 6f06dd2f..6f98ba36 100644
--- a/payload_consumer/partition_writer.cc
+++ b/payload_consumer/partition_writer.cc
@@ -326,8 +326,7 @@ bool PartitionWriter::PerformReplaceOperation(const InstallOperation& operation,
     writer.reset(new XzExtentWriter(std::move(writer)));
   }
 
-  TEST_AND_RETURN_FALSE(
-      writer->Init(target_fd_, operation.dst_extents(), block_size_));
+  TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   TEST_AND_RETURN_FALSE(writer->Write(data, operation.data_length()));
 
   return true;
@@ -377,7 +376,7 @@ bool PartitionWriter::PerformSourceCopyOperation(
   const auto& partition_control = dynamic_control_;
 
   InstallOperation buf;
-  bool should_optimize = partition_control->OptimizeOperation(
+  const bool should_optimize = partition_control->OptimizeOperation(
       partition.partition_name(), operation, &buf);
   const InstallOperation& optimized = should_optimize ? buf : operation;
 
@@ -493,8 +492,7 @@ bool PartitionWriter::PerformSourceBsdiffOperation(
       utils::BlocksInExtents(operation.src_extents()) * block_size_);
 
   auto writer = CreateBaseExtentWriter();
-  TEST_AND_RETURN_FALSE(
-      writer->Init(target_fd_, operation.dst_extents(), block_size_));
+  TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   auto dst_file = std::make_unique<BsdiffExtentFile>(
       std::move(writer),
       utils::BlocksInExtents(operation.dst_extents()) * block_size_);
@@ -522,8 +520,7 @@ bool PartitionWriter::PerformPuffDiffOperation(
       utils::BlocksInExtents(operation.src_extents()) * block_size_));
 
   auto writer = CreateBaseExtentWriter();
-  TEST_AND_RETURN_FALSE(
-      writer->Init(target_fd_, operation.dst_extents(), block_size_));
+  TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   puffin::UniqueStreamPtr dst_stream(new PuffinExtentStream(
       std::move(writer),
       utils::BlocksInExtents(operation.dst_extents()) * block_size_));
@@ -658,7 +655,7 @@ void PartitionWriter::CheckpointUpdateProgress(size_t next_op_index) {
 }
 
 std::unique_ptr<ExtentWriter> PartitionWriter::CreateBaseExtentWriter() {
-  return std::make_unique<DirectExtentWriter>();
+  return std::make_unique<DirectExtentWriter>(target_fd_);
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_writer_unittest.cc b/payload_consumer/partition_writer_unittest.cc
index 91e5e26e..263f338e 100644
--- a/payload_consumer/partition_writer_unittest.cc
+++ b/payload_consumer/partition_writer_unittest.cc
@@ -82,10 +82,10 @@ class PartitionWriterTest : public testing::Test {
   brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
                                    const brillo::Blob blob_data) {
     ScopedTempFile source_partition("Blob-XXXXXX");
-    DirectExtentWriter extent_writer;
     FileDescriptorPtr fd(new EintrSafeFileDescriptor());
+    DirectExtentWriter extent_writer{fd};
     EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
-    EXPECT_TRUE(extent_writer.Init(fd, op.src_extents(), kBlockSize));
+    EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
     EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
 
     ScopedTempFile target_partition("Blob-XXXXXX");
diff --git a/payload_consumer/snapshot_extent_writer.cc b/payload_consumer/snapshot_extent_writer.cc
index c9e6f312..242e7260 100644
--- a/payload_consumer/snapshot_extent_writer.cc
+++ b/payload_consumer/snapshot_extent_writer.cc
@@ -36,7 +36,6 @@ SnapshotExtentWriter::~SnapshotExtentWriter() {
 }
 
 bool SnapshotExtentWriter::Init(
-    FileDescriptorPtr /*fd*/,
     const google::protobuf::RepeatedPtrField<Extent>& extents,
     uint32_t block_size) {
   extents_ = extents;
diff --git a/payload_consumer/snapshot_extent_writer.h b/payload_consumer/snapshot_extent_writer.h
index 6d9fe7d6..c3c64cda 100644
--- a/payload_consumer/snapshot_extent_writer.h
+++ b/payload_consumer/snapshot_extent_writer.h
@@ -29,8 +29,7 @@ class SnapshotExtentWriter : public chromeos_update_engine::ExtentWriter {
   explicit SnapshotExtentWriter(android::snapshot::ICowWriter* cow_writer);
   ~SnapshotExtentWriter();
   // Returns true on success.
-  bool Init(FileDescriptorPtr fd,
-            const google::protobuf::RepeatedPtrField<Extent>& extents,
+  bool Init(const google::protobuf::RepeatedPtrField<Extent>& extents,
             uint32_t block_size) override;
   // Returns true on success.
   // This will construct a COW_REPLACE operation and forward it to CowWriter. It
diff --git a/payload_consumer/snapshot_extent_writer_unittest.cc b/payload_consumer/snapshot_extent_writer_unittest.cc
index 0e22482b..22010433 100644
--- a/payload_consumer/snapshot_extent_writer_unittest.cc
+++ b/payload_consumer/snapshot_extent_writer_unittest.cc
@@ -107,7 +107,7 @@ void AddExtent(google::protobuf::RepeatedPtrField<Extent>* extents,
 TEST_F(SnapshotExtentWriterTest, BufferWrites) {
   google::protobuf::RepeatedPtrField<Extent> extents;
   AddExtent(&extents, 123, 1);
-  writer_.Init(nullptr, extents, kBlockSize);
+  writer_.Init(extents, kBlockSize);
 
   std::vector<uint8_t> buf(kBlockSize, 0);
   buf[123] = 231;
@@ -130,7 +130,7 @@ TEST_F(SnapshotExtentWriterTest, NonBufferedWrites) {
   google::protobuf::RepeatedPtrField<Extent> extents;
   AddExtent(&extents, 123, 1);
   AddExtent(&extents, 125, 1);
-  writer_.Init(nullptr, extents, kBlockSize);
+  writer_.Init(extents, kBlockSize);
 
   std::vector<uint8_t> buf(kBlockSize * 2, 0);
   buf[123] = 231;
@@ -153,7 +153,7 @@ TEST_F(SnapshotExtentWriterTest, WriteAcrossBlockBoundary) {
   google::protobuf::RepeatedPtrField<Extent> extents;
   AddExtent(&extents, 123, 1);
   AddExtent(&extents, 125, 2);
-  writer_.Init(nullptr, extents, kBlockSize);
+  writer_.Init(extents, kBlockSize);
 
   std::vector<uint8_t> buf(kBlockSize * 3);
   std::memset(buf.data(), 0, buf.size());
diff --git a/payload_consumer/xz_extent_writer.cc b/payload_consumer/xz_extent_writer.cc
index a5b939db..a6483519 100644
--- a/payload_consumer/xz_extent_writer.cc
+++ b/payload_consumer/xz_extent_writer.cc
@@ -57,12 +57,11 @@ XzExtentWriter::~XzExtentWriter() {
   TEST_AND_RETURN(input_buffer_.empty());
 }
 
-bool XzExtentWriter::Init(FileDescriptorPtr fd,
-                          const RepeatedPtrField<Extent>& extents,
+bool XzExtentWriter::Init(const RepeatedPtrField<Extent>& extents,
                           uint32_t block_size) {
   stream_ = xz_dec_init(XZ_DYNALLOC, kXzMaxDictSize);
   TEST_AND_RETURN_FALSE(stream_ != nullptr);
-  return underlying_writer_->Init(fd, extents, block_size);
+  return underlying_writer_->Init(extents, block_size);
 }
 
 bool XzExtentWriter::Write(const void* bytes, size_t count) {
diff --git a/payload_consumer/xz_extent_writer.h b/payload_consumer/xz_extent_writer.h
index e022274e..70338f20 100644
--- a/payload_consumer/xz_extent_writer.h
+++ b/payload_consumer/xz_extent_writer.h
@@ -39,8 +39,7 @@ class XzExtentWriter : public ExtentWriter {
       : underlying_writer_(std::move(underlying_writer)) {}
   ~XzExtentWriter() override;
 
-  bool Init(FileDescriptorPtr fd,
-            const google::protobuf::RepeatedPtrField<Extent>& extents,
+  bool Init(const google::protobuf::RepeatedPtrField<Extent>& extents,
             uint32_t block_size) override;
   bool Write(const void* bytes, size_t count) override;
 
diff --git a/payload_consumer/xz_extent_writer_unittest.cc b/payload_consumer/xz_extent_writer_unittest.cc
index 34980a9a..5269dbc9 100644
--- a/payload_consumer/xz_extent_writer_unittest.cc
+++ b/payload_consumer/xz_extent_writer_unittest.cc
@@ -87,7 +87,7 @@ class XzExtentWriterTest : public ::testing::Test {
   }
 
   void WriteAll(const brillo::Blob& compressed) {
-    EXPECT_TRUE(xz_writer_->Init(fd_, {}, 1024));
+    EXPECT_TRUE(xz_writer_->Init({}, 1024));
     EXPECT_TRUE(xz_writer_->Write(compressed.data(), compressed.size()));
 
     EXPECT_TRUE(fake_extent_writer_->InitCalled());
@@ -130,7 +130,7 @@ TEST_F(XzExtentWriterTest, CompressedDataBiggerThanTheBuffer) {
 }
 
 TEST_F(XzExtentWriterTest, GarbageDataRejected) {
-  EXPECT_TRUE(xz_writer_->Init(fd_, {}, 1024));
+  EXPECT_TRUE(xz_writer_->Init({}, 1024));
   // The sample_data_ is an uncompressed string.
   EXPECT_FALSE(xz_writer_->Write(sample_data_.data(), sample_data_.size()));
 }
@@ -138,7 +138,7 @@ TEST_F(XzExtentWriterTest, GarbageDataRejected) {
 TEST_F(XzExtentWriterTest, PartialDataIsKept) {
   brillo::Blob compressed(std::begin(kCompressed30KiBofA),
                           std::end(kCompressed30KiBofA));
-  EXPECT_TRUE(xz_writer_->Init(fd_, {}, 1024));
+  EXPECT_TRUE(xz_writer_->Init({}, 1024));
   for (uint8_t byte : compressed) {
     EXPECT_TRUE(xz_writer_->Write(&byte, 1));
   }
diff --git a/payload_generator/zip_unittest.cc b/payload_generator/zip_unittest.cc
index e357b156..10e899b0 100644
--- a/payload_generator/zip_unittest.cc
+++ b/payload_generator/zip_unittest.cc
@@ -33,7 +33,6 @@
 using chromeos_update_engine::test_utils::kRandomString;
 using google::protobuf::RepeatedPtrField;
 using std::string;
-using std::vector;
 
 namespace chromeos_update_engine {
 
@@ -50,8 +49,7 @@ class MemoryExtentWriter : public ExtentWriter {
   }
   ~MemoryExtentWriter() override = default;
 
-  bool Init(FileDescriptorPtr fd,
-            const RepeatedPtrField<Extent>& extents,
+  bool Init(const RepeatedPtrField<Extent>& extents,
             uint32_t block_size) override {
     return true;
   }
@@ -72,7 +70,7 @@ bool DecompressWithWriter(const brillo::Blob& in, brillo::Blob* out) {
   std::unique_ptr<ExtentWriter> writer(
       new W(std::make_unique<MemoryExtentWriter>(out)));
   // Init() parameters are ignored by the testing MemoryExtentWriter.
-  bool ok = writer->Init(nullptr, {}, 1);
+  bool ok = writer->Init({}, 1);
   ok = writer->Write(in.data(), in.size()) && ok;
   return ok;
 }
diff --git a/scripts/ota_stress_test.py b/scripts/ota_stress_test.py
new file mode 100644
index 00000000..55aa4b1e
--- /dev/null
+++ b/scripts/ota_stress_test.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+#
+# 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.
+#
+
+"""Repeatedly install an A/B update to an Android device over adb."""
+
+import argparse
+import sys
+from pathlib import Path
+import subprocess
+import signal
+
+
+def CleanupLoopDevices():
+  # b/184716804 clean up unused loop devices
+  subprocess.check_call(["adb", "shell", "su", "0", "losetup", '-D'])
+
+
+def CancelOTA():
+  subprocess.call(["adb", "shell", "su", "0",
+                  "update_engine_client", "--cancel"])
+
+
+def PerformOTAThenPause(otafile: Path, update_device_script: Path):
+  python = sys.executable
+  ota_cmd = [python, str(update_device_script), str(otafile),
+             "--no-postinstall", "--no-slot-switch"]
+  p = subprocess.Popen(ota_cmd)
+  pid = p.pid
+  try:
+    ret = p.wait(10)
+    if ret is not None and ret != 0:
+      raise RuntimeError("OTA failed to apply")
+    if ret == 0:
+      print("OTA finished early? Surprise.")
+      return
+  except subprocess.TimeoutExpired:
+    pass
+  print(f"Killing {pid}")
+  subprocess.check_call(["pkill", "-INT", "-P", str(pid)])
+  p.send_signal(signal.SIGINT)
+  p.wait()
+
+
+def PerformTest(otafile: Path, resumes: int, timeout: int):
+  """Install an OTA to device, raising exceptions on failure
+
+  Args:
+    otafile: Path to the ota.zip to install
+
+  Return:
+    None if no error, if there's an error exception will be thrown
+  """
+  assert otafile.exists()
+  print("Applying", otafile)
+  script_dir = Path(__file__).parent.absolute()
+  update_device_script = script_dir / "update_device.py"
+  assert update_device_script.exists()
+  print(update_device_script)
+  python = sys.executable
+
+  for i in range(resumes):
+    print("Pause/Resume for the", i+1, "th time")
+    PerformOTAThenPause(otafile, update_device_script)
+  CancelOTA()
+  CleanupLoopDevices()
+
+  ota_cmd = [python, str(update_device_script),
+             str(otafile), "--no-postinstall"]
+  print("Finishing OTA Update", ota_cmd)
+  output = subprocess.check_output(
+      ota_cmd, stderr=subprocess.STDOUT, timeout=timeout).decode()
+  print(output)
+  if "onPayloadApplicationComplete(ErrorCode::kSuccess" not in output:
+    raise RuntimeError("Failed to finish OTA")
+  subprocess.call(
+      ["adb", "shell", "su", "0", "update_engine_client", "--cancel"])
+  subprocess.check_call(
+      ["adb", "shell", "su", "0", "update_engine_client", "--reset_status"])
+  CleanupLoopDevices()
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Android A/B OTA stress test helper.')
+  parser.add_argument('otafile', metavar='PAYLOAD', type=Path,
+                      help='the OTA package file (a .zip file) or raw payload \
+                      if device uses Omaha.')
+  parser.add_argument('-n', "--iterations", type=int, default=10,
+                      metavar='ITERATIONS',
+                      help='The number of iterations to run the stress test, or\
+                       -1 to keep running until CTRL+C')
+  parser.add_argument('-r', "--resumes", type=int, default=5, metavar='RESUMES',
+                      help='The number of iterations to pause the update when \
+                        installing')
+  parser.add_argument('-t', "--timeout", type=int, default=60*60,
+                      metavar='TIMEOUTS',
+                      help='Timeout, in seconds, when waiting for OTA to \
+                        finish')
+  args = parser.parse_args()
+  print(args)
+  n = args.iterations
+  while n != 0:
+    PerformTest(args.otafile, args.resumes, args.timeout)
+    n -= 1
+
+
+if __name__ == "__main__":
+  main()
diff --git a/scripts/update_device.py b/scripts/update_device.py
index 074b91d6..f672cda0 100755
--- a/scripts/update_device.py
+++ b/scripts/update_device.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 #
 # Copyright (C) 2017 The Android Open Source Project
 #
@@ -350,7 +350,7 @@ class AdbHost(object):
     if self._device_serial:
       self._command_prefix += ['-s', self._device_serial]
 
-  def adb(self, command):
+  def adb(self, command, timeout_seconds: float = None):
     """Run an ADB command like "adb push".
 
     Args:
@@ -365,7 +365,7 @@ class AdbHost(object):
     command = self._command_prefix + command
     logging.info('Running: %s', ' '.join(str(x) for x in command))
     p = subprocess.Popen(command, universal_newlines=True)
-    p.wait()
+    p.wait(timeout_seconds)
     return p.returncode
 
   def adb_output(self, command):
@@ -428,12 +428,14 @@ def main():
                       help='Update with the secondary payload in the package.')
   parser.add_argument('--no-slot-switch', action='store_true',
                       help='Do not perform slot switch after the update.')
-  parser.add_argument('--allocate_only', action='store_true',
+  parser.add_argument('--no-postinstall', action='store_true',
+                      help='Do not execute postinstall scripts after the update.')
+  parser.add_argument('--allocate-only', action='store_true',
                       help='Allocate space for this OTA, instead of actually \
                         applying the OTA.')
-  parser.add_argument('--verify_only', action='store_true',
+  parser.add_argument('--verify-only', action='store_true',
                       help='Verify metadata then exit, instead of applying the OTA.')
-  parser.add_argument('--no_care_map', action='store_true',
+  parser.add_argument('--no-care-map', action='store_true',
                       help='Do not push care_map.pb to device.')
   args = parser.parse_args()
   logging.basicConfig(
@@ -472,6 +474,8 @@ def main():
 
   if args.no_slot_switch:
     args.extra_headers += "\nSWITCH_SLOT_ON_REBOOT=0"
+  if args.no_postinstall:
+    args.extra_headers += "\nRUN_POST_INSTALL=0"
 
   with zipfile.ZipFile(args.otafile) as zfp:
     CARE_MAP_ENTRY_NAME = "care_map.pb"
@@ -482,7 +486,7 @@ def main():
         care_map_fp.write(zfp.read(CARE_MAP_ENTRY_NAME))
         care_map_fp.flush()
         dut.adb(["push", care_map_fp.name,
-                 "/data/ota_package/" + CARE_MAP_ENTRY_NAME])
+                "/data/ota_package/" + CARE_MAP_ENTRY_NAME])
 
   if args.file:
     # Update via pushing a file to /data.
@@ -542,7 +546,7 @@ def main():
     if server_thread:
       server_thread.StopServer()
     for cmd in finalize_cmds:
-      dut.adb(cmd)
+      dut.adb(cmd, 5)
 
   return 0