From d45c50911f989caacc6c46f79ee2e4fd66615c65 Mon Sep 17 00:00:00 2001
From: Kelvin Zhang <zhangkelvin@google.com>
Date: Tue, 5 Dec 2023 15:21:46 -0800
Subject: [PATCH] Fix cow v3 size estimation errors

Current cow size estimation is computed by taking the offset of next
data write. However, during COW size estimation we repeatdly increment
op_count_max. Since data section is placed right after op section,
incrementing  op_count_max has the effect of moving the beginning of
data section. next_data_pos_ is never updated in this process, causing
the final estiamte to be off.

Test: th
Bug: 313962438
Change-Id: I250dff54c470c9c20d6db33d91bac898358dee31
---
 .../libsnapshot/libsnapshot_cow/test_v3.cpp   | 21 +++++++++++++++----
 .../libsnapshot/libsnapshot_cow/writer_v3.cpp |  7 ++++++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/test_v3.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/test_v3.cpp
index aceed243c..27accdc7b 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/test_v3.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/test_v3.cpp
@@ -15,17 +15,15 @@
 #include <sys/stat.h>
 
 #include <cstdio>
-#include <iostream>
 #include <memory>
-#include <string_view>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/unique_fd.h>
 #include <gtest/gtest.h>
+#include <libsnapshot/cow_format.h>
 #include <libsnapshot/cow_reader.h>
 #include <libsnapshot/cow_writer.h>
-#include "cow_decompress.h"
-#include "libsnapshot/cow_format.h"
 #include "writer_v2.h"
 #include "writer_v3.h"
 
@@ -658,5 +656,20 @@ TEST_F(CowTestV3, SetTypeSourceInverleave) {
     ASSERT_EQ(op.type(), kCowReplaceOp);
 }
 
+TEST_F(CowTestV3, CowSizeEstimate) {
+    CowOptions options{};
+    options.compression = "none";
+    auto estimator = android::snapshot::CreateCowEstimator(3, options);
+    ASSERT_TRUE(estimator->AddZeroBlocks(0, 1024 * 1024));
+    const auto cow_size = estimator->GetCowSize();
+    options.op_count_max = 1024 * 1024;
+    options.max_blocks = 1024 * 1024;
+    CowWriterV3 writer(options, GetCowFd());
+    ASSERT_TRUE(writer.Initialize());
+    ASSERT_TRUE(writer.AddZeroBlocks(0, 1024 * 1024));
+
+    ASSERT_LE(writer.GetCowSize(), cow_size);
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
index 5bdc5fd2c..07f6f002f 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
@@ -327,7 +327,12 @@ bool CowWriterV3::EmitSequenceData(size_t num_ops, const uint32_t* data) {
 bool CowWriterV3::WriteOperation(const CowOperationV3& op, const void* data, size_t size) {
     if (IsEstimating()) {
         header_.op_count++;
-        header_.op_count_max++;
+        if (header_.op_count > header_.op_count_max) {
+            // If we increment op_count_max, the offset of data section would
+            // change. So need to update |next_data_pos_|
+            next_data_pos_ += (header_.op_count - header_.op_count_max) * sizeof(CowOperationV3);
+            header_.op_count_max = header_.op_count;
+        }
         next_data_pos_ += op.data_length;
         return true;
     }