libsnapshot: reorder COW ops vector

Reorder COW ops vector based on merge sequence. We don't
need additional vector to be stored in memory.

Memory usage for a full OTA on Pixel:

Without Patch:
RssAnon:       61020 kB

With Patch:
RssAnon:	   51112 kB

Bug: 237490659
Test: OTA on Pixel
Signed-off-by: Akilesh Kailash <akailash@google.com>
Change-Id: I543dd73acfa7cf4e57379e82bc184e943072e7c8

fs_mgr/libsnapshot/cow_reader.cpp

@@ -34,12 +34,13 @@
 namespace android {
 namespace snapshot {
 
-CowReader::CowReader(ReaderFlags reader_flag)
+CowReader::CowReader(ReaderFlags reader_flag, bool is_merge)
     : fd_(-1),
       header_(),
       fd_size_(0),
       block_pos_index_(std::make_shared<std::vector<int>>()),
-      reader_flag_(reader_flag) {}
+      reader_flag_(reader_flag),
+      is_merge_(is_merge) {}
 
 static void SHA256(const void*, size_t, uint8_t[]) {
 #if 0
@@ -64,6 +65,7 @@ std::unique_ptr<CowReader> CowReader::CloneCowReader() {
     cow->has_seq_ops_ = has_seq_ops_;
     cow->data_loc_ = data_loc_;
     cow->block_pos_index_ = block_pos_index_;
+    cow->is_merge_ = is_merge_;
     return cow;
 }
 
@@ -476,15 +478,28 @@ bool CowReader::PrepMergeOps() {
 
     merge_op_blocks->insert(merge_op_blocks->end(), other_ops.begin(), other_ops.end());
 
-    for (auto block : *merge_op_blocks) {
-        block_pos_index_->push_back(block_map->at(block));
-    }
-
     num_total_data_ops_ = merge_op_blocks->size();
     if (header_.num_merge_ops > 0) {
         merge_op_start_ = header_.num_merge_ops;
     }
 
+    if (is_merge_) {
+        // Metadata ops are not required for merge. Thus, just re-arrange
+        // the ops vector as required for merge operations.
+        auto merge_ops_buffer = std::make_shared<std::vector<CowOperation>>();
+        merge_ops_buffer->reserve(num_total_data_ops_);
+        for (auto block : *merge_op_blocks) {
+            merge_ops_buffer->emplace_back(ops_->data()[block_map->at(block)]);
+        }
+        ops_->clear();
+        ops_ = merge_ops_buffer;
+        ops_->shrink_to_fit();
+    } else {
+        for (auto block : *merge_op_blocks) {
+            block_pos_index_->push_back(block_map->at(block));
+        }
+    }
+
     block_map->clear();
     merge_op_blocks->clear();
 
@@ -548,7 +563,7 @@ bool CowReader::GetLastLabel(uint64_t* label) {
 
 class CowOpIter final : public ICowOpIter {
   public:
-    CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops);
+    CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops, uint64_t start);
 
     bool Done() override;
     const CowOperation& Get() override;
@@ -562,9 +577,9 @@ class CowOpIter final : public ICowOpIter {
     std::vector<CowOperation>::iterator op_iter_;
 };
 
-CowOpIter::CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops) {
+CowOpIter::CowOpIter(std::shared_ptr<std::vector<CowOperation>>& ops, uint64_t start) {
     ops_ = ops;
-    op_iter_ = ops_->begin();
+    op_iter_ = ops_->begin() + start;
 }
 
 bool CowOpIter::RDone() {
@@ -691,8 +706,8 @@ const CowOperation& CowRevMergeOpIter::Get() {
     return ops_->data()[*block_riter_];
 }
 
-std::unique_ptr<ICowOpIter> CowReader::GetOpIter() {
-    return std::make_unique<CowOpIter>(ops_);
+std::unique_ptr<ICowOpIter> CowReader::GetOpIter(bool merge_progress) {
+    return std::make_unique<CowOpIter>(ops_, merge_progress ? merge_op_start_ : 0);
 }
 
 std::unique_ptr<ICowOpIter> CowReader::GetRevMergeOpIter(bool ignore_progress) {

fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h

@@ -74,7 +74,7 @@ class ICowReader {
     virtual bool GetLastLabel(uint64_t* label) = 0;
 
     // Return an iterator for retrieving CowOperation entries.
-    virtual std::unique_ptr<ICowOpIter> GetOpIter() = 0;
+    virtual std::unique_ptr<ICowOpIter> GetOpIter(bool merge_progress) = 0;
 
     // Return an iterator for retrieving CowOperation entries in reverse merge order
     virtual std::unique_ptr<ICowOpIter> GetRevMergeOpIter(bool ignore_progress) = 0;
@@ -115,7 +115,7 @@ class CowReader final : public ICowReader {
         USERSPACE_MERGE = 1,
     };
 
-    CowReader(ReaderFlags reader_flag = ReaderFlags::DEFAULT);
+    CowReader(ReaderFlags reader_flag = ReaderFlags::DEFAULT, bool is_merge = false);
     ~CowReader() { owned_fd_ = {}; }
 
     // Parse the COW, optionally, up to the given label. If no label is
@@ -135,7 +135,7 @@ class CowReader final : public ICowReader {
     // CowOperation objects. Get() returns a unique CowOperation object
     // whose lifetime depends on the CowOpIter object; the return
     // value of these will never be null.
-    std::unique_ptr<ICowOpIter> GetOpIter() override;
+    std::unique_ptr<ICowOpIter> GetOpIter(bool merge_progress = false) override;
     std::unique_ptr<ICowOpIter> GetRevMergeOpIter(bool ignore_progress = false) override;
     std::unique_ptr<ICowOpIter> GetMergeOpIter(bool ignore_progress = false) override;
 
@@ -177,6 +177,7 @@ class CowReader final : public ICowReader {
     bool has_seq_ops_{};
     std::shared_ptr<std::unordered_map<uint64_t, uint64_t>> data_loc_;
     ReaderFlags reader_flag_;
+    bool is_merge_{};
 };
 
 }  // namespace snapshot

fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.cpp

@@ -162,7 +162,7 @@ bool SnapshotHandler::CheckMergeCompletionStatus() {
 }
 
 bool SnapshotHandler::ReadMetadata() {
-    reader_ = std::make_unique<CowReader>(CowReader::ReaderFlags::USERSPACE_MERGE);
+    reader_ = std::make_unique<CowReader>(CowReader::ReaderFlags::USERSPACE_MERGE, true);
     CowHeader header;
     CowOptions options;
 
@@ -193,7 +193,7 @@ bool SnapshotHandler::ReadMetadata() {
     UpdateMergeCompletionPercentage();
 
     // Initialize the iterator for reading metadata
-    std::unique_ptr<ICowOpIter> cowop_iter = reader_->GetMergeOpIter();
+    std::unique_ptr<ICowOpIter> cowop_iter = reader_->GetOpIter(true);
 
     int num_ra_ops_per_iter = ((GetBufferDataSize()) / BLOCK_SZ);
     int ra_index = 0;

fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_merge.cpp

@@ -466,7 +466,7 @@ bool Worker::SyncMerge() {
 }
 
 bool Worker::Merge() {
-    cowop_iter_ = reader_->GetMergeOpIter();
+    cowop_iter_ = reader_->GetOpIter(true);
 
     bool retry = false;
     bool ordered_ops_merge_status;

fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp

@@ -772,7 +772,7 @@ bool ReadAhead::InitReader() {
 }
 
 void ReadAhead::InitializeRAIter() {
-    cowop_iter_ = reader_->GetMergeOpIter();
+    cowop_iter_ = reader_->GetOpIter(true);
 }
 
 bool ReadAhead::RAIterDone() {