From 3f3162c21717003ff6724ab3121da45307ac053e Mon Sep 17 00:00:00 2001
From: Daniel Zheng <zhengdaniel@google.com>
Date: Sun, 17 Dec 2023 21:06:47 -0800
Subject: [PATCH] libsnapshot: get options from protobuf fields

Update snapshot.cpp to grab estimate_op_buffer_size &
estimate_sequence_buffer_size from update_engine. Update v3 writer to
use these options to size the buffers appropriately.

we probably don't need the fields for merge metrics yet but will leave
it here for now

Test: th
Bug: 313962438
Change-Id: I08252ff66174de9bafaf8dbe9115d9d049084c4c
---
 fs_mgr/libsnapshot/android/snapshot/snapshot.proto  | 6 ++++++
 fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h | 2 +-
 fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp    | 1 -
 fs_mgr/libsnapshot/partition_cow_creator.cpp        | 1 +
 fs_mgr/libsnapshot/snapshot.cpp                     | 7 ++++---
 5 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
index fa04c437d..7e97dc0fb 100644
--- a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
+++ b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
@@ -114,6 +114,9 @@ message SnapshotStatus {
 
     // Enable batching for COW writes
     bool batched_writes = 14;
+
+    // Size of v3 operation buffer. Needs to be determined during writer initialization
+    uint64 estimated_ops_buffer_size = 15;
 }
 
 // Next: 8
@@ -250,4 +253,7 @@ message SnapshotMergeReport {
 
     // Whether this update attempt used io_uring.
     bool iouring_used = 13;
+
+    // Size of v3 operation buffer. Needs to be determined during writer initialization
+    uint64 estimated_op_count_max = 14;
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
index ca0ebe182..7df976d85 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_writer.h
@@ -59,7 +59,7 @@ struct CowOptions {
     bool batch_write = false;
 
     // Size of the cow operation buffer; used in v3 only.
-    uint32_t op_count_max = 0;
+    uint64_t op_count_max = 0;
 };
 
 // Interface for writing to a snapuserd COW. All operations are ordered; merges
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
index 824fa39c7..bd69ac685 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/writer_v3.cpp
@@ -392,7 +392,6 @@ bool CowWriterV3::EmitLabel(uint64_t label) {
 }
 
 bool CowWriterV3::EmitSequenceData(size_t num_ops, const uint32_t* data) {
-    // TODO: size sequence buffer based on options
     if (header_.op_count > 0 || !cached_ops_.empty()) {
         LOG(ERROR) << "There's " << header_.op_count << " operations written to disk and "
                    << cached_ops_.size()
diff --git a/fs_mgr/libsnapshot/partition_cow_creator.cpp b/fs_mgr/libsnapshot/partition_cow_creator.cpp
index 5bc7e65f7..c0d20734e 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator.cpp
+++ b/fs_mgr/libsnapshot/partition_cow_creator.cpp
@@ -217,6 +217,7 @@ std::optional<PartitionCowCreator::Return> PartitionCowCreator::Run() {
 
     if (update && update->has_estimate_cow_size()) {
         ret.snapshot_status.set_estimated_cow_size(update->estimate_cow_size());
+        ret.snapshot_status.set_estimated_ops_buffer_size(update->estimate_op_count_max());
     }
 
     if (ret.snapshot_status.snapshot_size() == 0) {
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index e33bdfff8..9eb41b2cc 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -3551,6 +3551,9 @@ Return SnapshotManager::InitializeUpdateSnapshots(
                 options.scratch_space = false;
             }
             options.compression = it->second.compression_algorithm();
+            if (cow_version >= 3) {
+                options.op_count_max = it->second.estimated_ops_buffer_size();
+            }
 
             auto writer = CreateCowWriter(cow_version, options, std::move(fd));
             if (!writer->Finalize()) {
@@ -3662,9 +3665,7 @@ std::unique_ptr<ICowWriter> SnapshotManager::OpenCompressedSnapshotWriter(
     cow_options.max_blocks = {status.device_size() / cow_options.block_size};
     cow_options.batch_write = status.batched_writes();
     cow_options.num_compress_threads = status.enable_threading() ? 2 : 1;
-    // TODO(b/313962438) Improve op_count estimate. For now, use number of
-    // blocks as an upper bound.
-    cow_options.op_count_max = status.device_size() / cow_options.block_size;
+    cow_options.op_count_max = status.estimated_ops_buffer_size();
     // Disable scratch space for vts tests
     if (device()->IsTestDevice()) {
         cow_options.scratch_space = false;