DynamicPartitionControl: Add required_size to Prepare

Add out parameter required_size to PreparePartitionsForUpdate
to indicate the total size required on /userdata in order
to apply the update.

Bug: 138808058
Test: update_engine_unittests

Change-Id: I2768d13671e212fd24a1a22811b50c9738834459
diff --git a/common/dynamic_partition_control_interface.h b/common/dynamic_partition_control_interface.h
index c171fb5..19bb523 100644
--- a/common/dynamic_partition_control_interface.h
+++ b/common/dynamic_partition_control_interface.h
@@ -64,10 +64,15 @@
   // This is needed before calling MapPartitionOnDeviceMapper(), otherwise the
   // device would be mapped in an inconsistent way.
   // If |update| is set, create snapshots and writes super partition metadata.
+  // If |required_size| is not null and call fails due to insufficient space,
+  // |required_size| will be set to total free space required on userdata
+  // partition to apply the update. Otherwise (call succeeds, or fails
+  // due to other errors), |required_size| is set to zero.
   virtual bool PreparePartitionsForUpdate(uint32_t source_slot,
                                           uint32_t target_slot,
                                           const DeltaArchiveManifest& manifest,
-                                          bool update) = 0;
+                                          bool update,
+                                          uint64_t* required_size) = 0;
 
   // After writing to new partitions, before rebooting into the new slot, call
   // this function to indicate writes to new partitions are done.
diff --git a/common/dynamic_partition_control_stub.cc b/common/dynamic_partition_control_stub.cc
index bc792c8..974cd1b 100644
--- a/common/dynamic_partition_control_stub.cc
+++ b/common/dynamic_partition_control_stub.cc
@@ -43,7 +43,8 @@
     uint32_t source_slot,
     uint32_t target_slot,
     const DeltaArchiveManifest& manifest,
-    bool update) {
+    bool update,
+    uint64_t* required_size) {
   return true;
 }
 
diff --git a/common/dynamic_partition_control_stub.h b/common/dynamic_partition_control_stub.h
index 1704f05..09990a7 100644
--- a/common/dynamic_partition_control_stub.h
+++ b/common/dynamic_partition_control_stub.h
@@ -35,7 +35,8 @@
   bool PreparePartitionsForUpdate(uint32_t source_slot,
                                   uint32_t target_slot,
                                   const DeltaArchiveManifest& manifest,
-                                  bool update) override;
+                                  bool update,
+                                  uint64_t* required_size) override;
 
   bool FinishUpdate() override;
 };
diff --git a/dynamic_partition_control_android.cc b/dynamic_partition_control_android.cc
index 4ad02c7..bb85480 100644
--- a/dynamic_partition_control_android.cc
+++ b/dynamic_partition_control_android.cc
@@ -50,6 +50,7 @@
 using android::fs_mgr::Partition;
 using android::fs_mgr::PartitionOpener;
 using android::fs_mgr::SlotSuffixForSlotNumber;
+using android::snapshot::SnapshotManager;
 using android::snapshot::SourceCopyOperationIsClone;
 
 namespace chromeos_update_engine {
@@ -92,7 +93,7 @@
           GetFeatureFlag(kUseDynamicPartitions, kRetrfoitDynamicPartitions)),
       virtual_ab_(GetFeatureFlag(kVirtualAbEnabled, kVirtualAbRetrofit)) {
   if (GetVirtualAbFeatureFlag().IsEnabled()) {
-    snapshot_ = android::snapshot::SnapshotManager::New();
+    snapshot_ = SnapshotManager::New();
     CHECK(snapshot_ != nullptr) << "Cannot initialize SnapshotManager.";
   }
 }
@@ -372,9 +373,13 @@
     uint32_t source_slot,
     uint32_t target_slot,
     const DeltaArchiveManifest& manifest,
-    bool update) {
+    bool update,
+    uint64_t* required_size) {
   source_slot_ = source_slot;
   target_slot_ = target_slot;
+  if (required_size != nullptr) {
+    *required_size = 0;
+  }
 
   if (fs_mgr_overlayfs_is_setup()) {
     // Non DAP devices can use overlayfs as well.
@@ -420,7 +425,7 @@
     // - If !target_supports_snapshot_, explicitly CancelUpdate().
     if (target_supports_snapshot_) {
       return PrepareSnapshotPartitionsForUpdate(
-          source_slot, target_slot, manifest);
+          source_slot, target_slot, manifest, required_size);
     }
     if (!snapshot_->CancelUpdate()) {
       LOG(ERROR) << "Cannot cancel previous update.";
@@ -473,13 +478,19 @@
 bool DynamicPartitionControlAndroid::PrepareSnapshotPartitionsForUpdate(
     uint32_t source_slot,
     uint32_t target_slot,
-    const DeltaArchiveManifest& manifest) {
+    const DeltaArchiveManifest& manifest,
+    uint64_t* required_size) {
   if (!snapshot_->BeginUpdate()) {
     LOG(ERROR) << "Cannot begin new update.";
     return false;
   }
-  if (!snapshot_->CreateUpdateSnapshots(manifest)) {
-    LOG(ERROR) << "Cannot create update snapshots.";
+  auto ret = snapshot_->CreateUpdateSnapshots(manifest);
+  if (!ret) {
+    LOG(ERROR) << "Cannot create update snapshots: " << ret.string();
+    if (required_size != nullptr &&
+        ret.error_code() == SnapshotManager::Return::ErrorCode::NO_SPACE) {
+      *required_size = ret.required_size();
+    }
     return false;
   }
   return true;
diff --git a/dynamic_partition_control_android.h b/dynamic_partition_control_android.h
index 13fbb1a..a79f41a 100644
--- a/dynamic_partition_control_android.h
+++ b/dynamic_partition_control_android.h
@@ -42,7 +42,8 @@
   bool PreparePartitionsForUpdate(uint32_t source_slot,
                                   uint32_t target_slot,
                                   const DeltaArchiveManifest& manifest,
-                                  bool update) override;
+                                  bool update,
+                                  uint64_t* required_size) override;
   bool FinishUpdate() override;
 
   // Return the device for partition |partition_name| at slot |slot|.
@@ -151,7 +152,8 @@
   // Virtual A/B update.
   bool PrepareSnapshotPartitionsForUpdate(uint32_t source_slot,
                                           uint32_t target_slot,
-                                          const DeltaArchiveManifest& manifest);
+                                          const DeltaArchiveManifest& manifest,
+                                          uint64_t* required_size);
 
   enum class DynamicPartitionDeviceStatus {
     SUCCESS,
diff --git a/dynamic_partition_control_android_unittest.cc b/dynamic_partition_control_android_unittest.cc
index 207a97e..3e8375c 100644
--- a/dynamic_partition_control_android_unittest.cc
+++ b/dynamic_partition_control_android_unittest.cc
@@ -120,7 +120,11 @@
   }
   bool PreparePartitionsForUpdate(const PartitionSizes& partition_sizes) {
     return dynamicControl().PreparePartitionsForUpdate(
-        source(), target(), PartitionSizesToManifest(partition_sizes), true);
+        source(),
+        target(),
+        PartitionSizesToManifest(partition_sizes),
+        true,
+        nullptr);
   }
   void SetSlots(const TestParam& slots) { slots_ = slots; }
 
@@ -317,7 +321,7 @@
   // DynamicPartitionControlAndroidTest::PreparePartitionsForUpdate(), since we
   // don't want any default group in the PartitionMetadata.
   EXPECT_TRUE(dynamicControl().PreparePartitionsForUpdate(
-      source(), target(), {}, true));
+      source(), target(), {}, true, nullptr));
 
   // Should use dynamic source partitions.
   EXPECT_CALL(dynamicControl(), GetState(S("system")))
@@ -371,7 +375,8 @@
       source(),
       target(),
       PartitionSizesToManifest({{"system", 2_GiB}, {"vendor", 1_GiB}}),
-      false));
+      false,
+      nullptr));
 
   // Dynamic partition "system".
   EXPECT_CALL(dynamicControl(), GetState(S("system")))
@@ -622,7 +627,11 @@
 
 TEST_P(DynamicPartitionControlAndroidTestP, ShouldSkipOperationTest) {
   ASSERT_TRUE(dynamicControl().PreparePartitionsForUpdate(
-      source(), target(), PartitionSizesToManifest({{"foo", 4_MiB}}), false));
+      source(),
+      target(),
+      PartitionSizesToManifest({{"foo", 4_MiB}}),
+      false,
+      nullptr));
   dynamicControl().set_fake_mapped_devices({T("foo")});
 
   InstallOperation iop;
diff --git a/mock_dynamic_partition_control.h b/mock_dynamic_partition_control.h
index 09b825d..ffabac7 100644
--- a/mock_dynamic_partition_control.h
+++ b/mock_dynamic_partition_control.h
@@ -38,8 +38,9 @@
                     std::string*));
   MOCK_METHOD0(Cleanup, void());
   MOCK_METHOD0(GetDynamicPartitionsFeatureFlag, FeatureFlag());
-  MOCK_METHOD4(PreparePartitionsForUpdate,
-               bool(uint32_t, uint32_t, const DeltaArchiveManifest&, bool));
+  MOCK_METHOD5(
+      PreparePartitionsForUpdate,
+      bool(uint32_t, uint32_t, const DeltaArchiveManifest&, bool, uint64_t*));
   MOCK_METHOD0(GetVirtualAbFeatureFlag, FeatureFlag());
   MOCK_METHOD0(FinishUpdate, bool());
 };
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index c49474c..8cec076 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -951,7 +951,8 @@
           boot_control_->GetCurrentSlot(),
           install_plan_->target_slot,
           manifest_,
-          !metadata_updated)) {
+          !metadata_updated,
+          nullptr /* required_size */)) {
     LOG(ERROR) << "Unable to initialize partition metadata for slot "
                << BootControlInterface::SlotName(install_plan_->target_slot);
     return false;