libsnapshot: Export COW image size

Enable libsnapshot to export the size of the COW image, representing the
fraction of the COW device that is temporarily created in the /data
partition, computed just before initializing a merge operation.
Thise additional information can be used by other components (i.e.,
update engine) to enrich Virtual A/B metrics.

Bug: 154016862
Test: manual OTA
Signed-off-by: Alessio Balsini <balsini@google.com>
Change-Id: I4faa0785f23884e26161e0d51268dd7a305d86bf
Merged-In: I4faa0785f23884e26161e0d51268dd7a305d86bf
diff --git a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
index 2ac0c44..587089e 100644
--- a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
+++ b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
@@ -140,4 +140,7 @@
     // Number of reboots that occurred after issuing and before completeing the
     // merge of all the snapshot devices.
     int32 resume_count = 2;
+
+    // Total size of all the COW images before the update.
+    uint64 cow_file_size = 4;
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
index cf0b085..4457de3 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
@@ -25,7 +25,7 @@
     MOCK_METHOD(bool, BeginUpdate, (), (override));
     MOCK_METHOD(bool, CancelUpdate, (), (override));
     MOCK_METHOD(bool, FinishedSnapshotWrites, (bool wipe), (override));
-    MOCK_METHOD(bool, InitiateMerge, (), (override));
+    MOCK_METHOD(bool, InitiateMerge, (uint64_t * cow_file_size), (override));
 
     MOCK_METHOD(UpdateState, ProcessUpdateState,
                 (const std::function<bool()>& callback, const std::function<bool()>& before_cancel),
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 2d6071f..3c2c776 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -124,7 +124,7 @@
 
     // Initiate a merge on all snapshot devices. This should only be used after an
     // update has been marked successful after booting.
-    virtual bool InitiateMerge() = 0;
+    virtual bool InitiateMerge(uint64_t* cow_file_size = nullptr) = 0;
 
     // Perform any necessary post-boot actions. This should be run soon after
     // /data is mounted.
@@ -281,7 +281,7 @@
     bool BeginUpdate() override;
     bool CancelUpdate() override;
     bool FinishedSnapshotWrites(bool wipe) override;
-    bool InitiateMerge() override;
+    bool InitiateMerge(uint64_t* cow_file_size = nullptr) override;
     UpdateState ProcessUpdateState(const std::function<bool()>& callback = {},
                                    const std::function<bool()>& before_cancel = {}) override;
     UpdateState GetUpdateState(double* progress = nullptr) override;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
index 4caf632..d691d4f 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
@@ -29,6 +29,8 @@
     // Called when merge starts or resumes.
     virtual bool Start() = 0;
     virtual void set_state(android::snapshot::UpdateState state) = 0;
+    virtual void set_cow_file_size(uint64_t cow_file_size) = 0;
+    virtual uint64_t cow_file_size() = 0;
 
     // Called when merge ends. Properly clean up permanent storage.
     class Result {
@@ -50,6 +52,8 @@
     // ISnapshotMergeStats overrides
     bool Start() override;
     void set_state(android::snapshot::UpdateState state) override;
+    void set_cow_file_size(uint64_t cow_file_size) override;
+    uint64_t cow_file_size() override;
     std::unique_ptr<Result> Finish() override;
 
   private:
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
index 9c82906..7a27fad 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
@@ -27,7 +27,7 @@
     bool BeginUpdate() override;
     bool CancelUpdate() override;
     bool FinishedSnapshotWrites(bool wipe) override;
-    bool InitiateMerge() override;
+    bool InitiateMerge(uint64_t* cow_file_size = nullptr) override;
     UpdateState ProcessUpdateState(const std::function<bool()>& callback = {},
                                    const std::function<bool()>& before_cancel = {}) override;
     UpdateState GetUpdateState(double* progress = nullptr) override;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 488009a..5909cff 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -555,7 +555,7 @@
     return true;
 }
 
-bool SnapshotManager::InitiateMerge() {
+bool SnapshotManager::InitiateMerge(uint64_t* cow_file_size) {
     auto lock = LockExclusive();
     if (!lock) return false;
 
@@ -618,6 +618,7 @@
         }
     }
 
+    uint64_t total_cow_file_size = 0;
     DmTargetSnapshot::Status initial_target_values = {};
     for (const auto& snapshot : snapshots) {
         DmTargetSnapshot::Status current_status;
@@ -627,6 +628,16 @@
         initial_target_values.sectors_allocated += current_status.sectors_allocated;
         initial_target_values.total_sectors += current_status.total_sectors;
         initial_target_values.metadata_sectors += current_status.metadata_sectors;
+
+        SnapshotStatus snapshot_status;
+        if (!ReadSnapshotStatus(lock.get(), snapshot, &snapshot_status)) {
+            return false;
+        }
+        total_cow_file_size += snapshot_status.cow_file_size();
+    }
+
+    if (cow_file_size) {
+        *cow_file_size = total_cow_file_size;
     }
 
     SnapshotUpdateStatus initial_status;
diff --git a/fs_mgr/libsnapshot/snapshot_stats.cpp b/fs_mgr/libsnapshot/snapshot_stats.cpp
index 5da7b98..3723730 100644
--- a/fs_mgr/libsnapshot/snapshot_stats.cpp
+++ b/fs_mgr/libsnapshot/snapshot_stats.cpp
@@ -88,6 +88,15 @@
     report_.set_state(state);
 }
 
+void SnapshotMergeStats::set_cow_file_size(uint64_t cow_file_size) {
+    report_.set_cow_file_size(cow_file_size);
+    WriteState();
+}
+
+uint64_t SnapshotMergeStats::cow_file_size() {
+    return report_.cow_file_size();
+}
+
 class SnapshotMergeStatsResultImpl : public SnapshotMergeStats::Result {
   public:
     SnapshotMergeStatsResultImpl(const SnapshotMergeReport& report,
diff --git a/fs_mgr/libsnapshot/snapshot_stub.cpp b/fs_mgr/libsnapshot/snapshot_stub.cpp
index 2aaa78c..9b6f758 100644
--- a/fs_mgr/libsnapshot/snapshot_stub.cpp
+++ b/fs_mgr/libsnapshot/snapshot_stub.cpp
@@ -42,7 +42,7 @@
     return false;
 }
 
-bool SnapshotManagerStub::InitiateMerge() {
+bool SnapshotManagerStub::InitiateMerge(uint64_t*) {
     LOG(ERROR) << __FUNCTION__ << " should never be called.";
     return false;
 }
@@ -118,6 +118,8 @@
 class SnapshotMergeStatsStub : public ISnapshotMergeStats {
     bool Start() override { return false; }
     void set_state(android::snapshot::UpdateState) override {}
+    void set_cow_file_size(uint64_t) override {}
+    uint64_t cow_file_size() override { return 0; }
     std::unique_ptr<Result> Finish() override { return nullptr; }
 };