Add a new flag --extract_super_image_fragments to cvd fetch.

Databean · Databean · commit 1f60a7d1a70b · 2025-10-23T16:14:26.000-07:00
This enables a new flow to download the super image components out of the target_files zip, using a path that only downloads the fragments of the zip necessary. Downloading the super image fragments is comparable to downloading the image zip, but these files can be used to build a super image ourselves with any last-minute modifications we need. Some performance and correctness improvements in the partial zip download are included in this change: increasing the chunk size for partial downloads, and working around libzip doing read-all behavior from its read calls. The partial zip download takes twice the time of the full image zip download, so hopefully more performance improvements are possible afterward. This change also supersedes #1558 . Bug: b/439877215
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel b/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel
@@ -138,11 +138,16 @@ cf_cc_library(
         "//cuttlefish/host/libs/web:android_build",
         "//cuttlefish/host/libs/web:android_build_api",
         "//cuttlefish/host/libs/web:android_build_string",
+        "//cuttlefish/host/libs/web:build_api_zip",
         "//cuttlefish/host/libs/web:build_zip_name",
         "//cuttlefish/host/libs/web:chrome_os_build_string",
         "//cuttlefish/host/libs/web:luci_build_api",
         "//cuttlefish/host/libs/web/http_client:curl_global_init",
+        "//cuttlefish/host/libs/zip:zip_cc",
+        "//cuttlefish/host/libs/zip:zip_file",
+        "//cuttlefish/host/libs/zip:zip_string",
         "//libbase",
+        "@abseil-cpp//absl/strings",
     ],
 )
 
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/download_flags.cc b/base/cvd/cuttlefish/host/commands/cvd/fetch/download_flags.cc
@@ -27,6 +27,9 @@ DownloadFlags DownloadFlags::Create(const VectorFlags& flags, const int index) {
       .download_target_files_zip =
           GetOptional(flags.download_target_files_zip, index)
               .value_or(kDefaultDownloadTargetFilesZip),
+      .extract_super_image_fragments =
+          GetOptional(flags.extract_super_image_fragments, index)
+              .value_or(kDefaultExtractSuperImageFragments),
   };
 }
 
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/download_flags.h b/base/cvd/cuttlefish/host/commands/cvd/fetch/download_flags.h
@@ -24,6 +24,7 @@ struct DownloadFlags {
 
   bool download_img_zip;
   bool download_target_files_zip;
+  bool extract_super_image_fragments;
 };
 
 }  // namespace cuttlefish
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc b/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc
@@ -28,6 +28,7 @@
 #include <vector>
 
 #include <android-base/strings.h>
+#include "absl/strings/str_split.h"
 
 #include "cuttlefish/common/libs/utils/archive.h"
 #include "cuttlefish/common/libs/utils/contains.h"
@@ -48,10 +49,14 @@
 #include "cuttlefish/host/libs/web/android_build.h"
 #include "cuttlefish/host/libs/web/android_build_api.h"
 #include "cuttlefish/host/libs/web/android_build_string.h"
+#include "cuttlefish/host/libs/web/build_api_zip.h"
 #include "cuttlefish/host/libs/web/build_zip_name.h"
 #include "cuttlefish/host/libs/web/chrome_os_build_string.h"
 #include "cuttlefish/host/libs/web/http_client/curl_global_init.h"
 #include "cuttlefish/host/libs/web/luci_build_api.h"
+#include "cuttlefish/host/libs/zip/zip_cc.h"
+#include "cuttlefish/host/libs/zip/zip_file.h"
+#include "cuttlefish/host/libs/zip/zip_string.h"
 
 namespace cuttlefish {
 namespace {
@@ -237,19 +242,58 @@ Result<void> FetchDefaultTarget(BuildApi& build_api, const Builds& builds,
     trace.CompletePhase("Desparse image files");
   }
 
+  std::string target_files_name =
+      GetBuildZipName(*builds.default_build, "target_files");
+  std::optional<std::string> target_files;
   if (builds.system || flags.download_target_files_zip) {
     LOG(INFO) << "Downloading target files zip for " << *builds.default_build;
-    std::string target_files_name =
-        GetBuildZipName(*builds.default_build, "target_files");
-    std::string target_files = CF_EXPECT(build_api.DownloadFile(
+    target_files = CF_EXPECT(build_api.DownloadFile(
         *builds.default_build, target_directories.default_target_files,
         target_files_name));
     trace.CompletePhase("Download Target Files");
     LOG(INFO) << "Adding target files for default build";
-    CF_EXPECT(config.AddFilesToConfig(FileSource::DEFAULT_BUILD,
-                                      default_build_id, default_build_target,
-                                      {target_files}, target_directories.root));
+    CF_EXPECT(config.AddFilesToConfig(
+        FileSource::DEFAULT_BUILD, default_build_id, default_build_target,
+        {*target_files}, target_directories.root));
+  }
+
+  if (flags.extract_super_image_fragments) {
+    ReadableZip target_files_zip =
+        target_files.has_value()
+            ? CF_EXPECT(ZipOpenRead(*target_files))
+            : CF_EXPECT(
+                  OpenZip(build_api, *builds.default_build, target_files_name));
+    ReadableZipSource ab_partitions_source =
+        CF_EXPECT(target_files_zip.GetFile("META/ab_partitions.txt"));
+    std::string ab_partitions_contents =
+        CF_EXPECT(ReadToString(ab_partitions_source));
+
+    std::string ab_partitions_file =
+        target_directories.default_target_files + "/ab_partitions.txt";
+    CF_EXPECT(android::base::WriteStringToFile(ab_partitions_contents,
+                                               ab_partitions_file));
+    CF_EXPECT(config.AddFilesToConfig(
+        FileSource::DEFAULT_BUILD, default_build_id, default_build_target,
+        {ab_partitions_file}, target_directories.default_target_files));
+
+    std::vector<std::string> ab_files =
+        absl::StrSplit(ab_partitions_contents, '\n');
+    ab_files.emplace_back("super_empty");
+    for (const std::string& ab_file : ab_files) {
+      if (ab_file.empty()) {
+        continue;
+      }
+      std::string output = fmt::format(
+          "{}/{}.img", target_directories.default_target_files, ab_file);
+      CF_EXPECTF(ExtractFile(target_files_zip,
+                             fmt::format("IMAGES/{}.img", ab_file), output),
+                 "Failed to extract {}", output);
+      CF_EXPECT(config.AddFilesToConfig(
+          FileSource::DEFAULT_BUILD, default_build_id, default_build_target,
+          {output}, target_directories.default_target_files));
+    }
   }
+
   return {};
 }
 
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/vector_flags.cc b/base/cvd/cuttlefish/host/commands/cvd/fetch/vector_flags.cc
@@ -68,6 +68,11 @@ std::vector<Flag> VectorFlags::Flags() {
                        this->download_target_files_zip,
                        kDefaultDownloadTargetFilesZip)
           .Help("Whether to fetch the -target_files-*.zip file."));
+  flags.emplace_back(
+      GflagsCompatFlag("extract_super_image_fragments",
+                       this->extract_super_image_fragments,
+                       kDefaultExtractSuperImageFragments)
+          .Help("Whether to fetch the -target_files-*.zip file."));
 
   return flags;
 }
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/vector_flags.h b/base/cvd/cuttlefish/host/commands/cvd/fetch/vector_flags.h
@@ -27,6 +27,8 @@ namespace cuttlefish {
 
 inline constexpr bool kDefaultDownloadImgZip = true;
 inline constexpr bool kDefaultDownloadTargetFilesZip = false;
+// TODO: schuffelen - Enable this by default.
+inline constexpr bool kDefaultExtractSuperImageFragments = false;
 
 struct VectorFlags {
   std::vector<Flag> Flags();
@@ -44,6 +46,7 @@ struct VectorFlags {
   std::vector<bool> download_img_zip;
   std::vector<bool> download_target_files_zip;
   std::vector<std::string> boot_artifact;
+  std::vector<bool> extract_super_image_fragments;
 };
 
 }  // namespace cuttlefish
diff --git a/base/cvd/cuttlefish/host/libs/web/build_api_zip.cc b/base/cvd/cuttlefish/host/libs/web/build_api_zip.cc
@@ -31,7 +31,7 @@ Result<ReadableZip> OpenZip(BuildApi& build_api, const Build& build,
   SeekableZipSource source = CF_EXPECT(build_api.FileReader(build, name));
 
   SeekableZipSource buffered =
-      CF_EXPECT(BufferZipSource(std::move(source), 1 << 16));
+      CF_EXPECT(BufferZipSource(std::move(source), 1 << 26));
 
   return CF_EXPECT(ReadableZip::FromSource(std::move(buffered)));
 }
diff --git a/base/cvd/cuttlefish/host/libs/zip/BUILD.bazel b/base/cvd/cuttlefish/host/libs/zip/BUILD.bazel
@@ -13,6 +13,7 @@ cf_cc_library(
     deps = [
         "//cuttlefish/common/libs/utils:result",
         "//cuttlefish/host/libs/zip:zip_cc",
+        "//libbase",
     ],
 )
 
@@ -107,6 +108,7 @@ cf_cc_library(
         "//cuttlefish/common/libs/utils:result",
         "//cuttlefish/host/libs/web/http_client",
         "//cuttlefish/host/libs/zip:zip_cc",
+        "//libbase",
         "@abseil-cpp//absl/strings",
         "@fmt",
     ],
diff --git a/base/cvd/cuttlefish/host/libs/zip/buffered_zip_source.cc b/base/cvd/cuttlefish/host/libs/zip/buffered_zip_source.cc
@@ -25,6 +25,8 @@
 #include <utility>
 #include <vector>
 
+#include "android-base/logging.h"
+
 #include "cuttlefish/common/libs/utils/result.h"
 #include "cuttlefish/host/libs/zip/zip_cc.h"
 
@@ -71,6 +73,7 @@ class BufferedZipSourceCallbacks : public SeekableZipSourceCallback {
       if (!reader_->SeekFromStart(offset_).ok()) {
         return false;
       }
+      LOG(VERBOSE) << "Bypassing buffer, reading " << len;
       Result<uint64_t> data_read = reader_->Read(data, len);
       if (data_read.ok()) {
         offset_ += *data_read;
@@ -97,6 +100,7 @@ class BufferedZipSourceCallbacks : public SeekableZipSourceCallback {
     if (!reader_->SeekFromStart(offset_).ok()) {
       return -1;
     }
+    LOG(VERBOSE) << "Filling buffer with " << buffer_fill;
     Result<size_t> inner_read = reader_->Read(buffer_.data(), buffer_fill);
     if (!inner_read.ok()) {
       return -1;
@@ -106,9 +110,14 @@ class BufferedZipSourceCallbacks : public SeekableZipSourceCallback {
     return Read(data, len);
   }
   uint64_t Size() override { return size_; }
-  bool SetOffset(int64_t offset) override {
-    offset_ = offset;
-    buffer_remaining_ = 0;
+  bool SetOffset(int64_t new_offset) override {
+    if (new_offset >= offset_ && new_offset < offset_ + buffer_remaining_) {
+      offset_in_buffer_ += new_offset - offset_;
+      buffer_remaining_ -= new_offset - offset_;
+    } else {
+      buffer_remaining_ = 0;
+    }
+    offset_ = new_offset;
     return true;
   }
   int64_t Offset() override { return offset_; }
diff --git a/base/cvd/cuttlefish/host/libs/zip/cached_zip_source.cc b/base/cvd/cuttlefish/host/libs/zip/cached_zip_source.cc
@@ -47,6 +47,7 @@ class CachedZipSourceCallbacks : public SeekableZipSourceCallback {
   }
 
   int64_t Read(char* data, uint64_t len) override {
+    LOG(VERBOSE) << "Reading " << len;
     if (Result<void> res = source_.Seek(offset_); !res.ok()) {
       LOG(ERROR) << res.error().FormatForEnv();
       return -1;
@@ -64,6 +65,7 @@ class CachedZipSourceCallbacks : public SeekableZipSourceCallback {
   uint64_t Size() override { return size_; }
 
   bool SetOffset(int64_t offset) override {
+    LOG(VERBOSE) << "Setting offset to " << offset;
     offset_ = offset;
     return true;
   }
@@ -113,7 +115,7 @@ Result<SeekableZipSource> CacheZipSource(SeekableZipSource inner,
                                                   std::move(reader));
 
   LazilyLoadedFile file = CF_EXPECT(LazilyLoadedFile::Create(
-      std::move(file_path), std::move(file_callbacks)));
+      std::move(file_path), size, std::move(file_callbacks)));
 
   CachedZipSourceCallbacks callbacks(std::move(file), size);
 
diff --git a/base/cvd/cuttlefish/host/libs/zip/lazily_loaded_file.cc b/base/cvd/cuttlefish/host/libs/zip/lazily_loaded_file.cc
@@ -17,6 +17,7 @@
 #include <fcntl.h>
 #include <stddef.h>
 #include <stdint.h>
+#include <string.h>
 #include <sys/types.h>
 #include <unistd.h>  // NOLINT(misc-include-cleaner): SEEK_SET
 
@@ -25,6 +26,7 @@
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include <android-base/logging.h>
 
@@ -37,6 +39,8 @@
 
 namespace cuttlefish {
 
+static constexpr size_t kMinReadSize = 1 << 26;
+
 LazilyLoadedFileReadCallback::~LazilyLoadedFileReadCallback() = default;
 
 struct LazilyLoadedFile::Impl {
@@ -51,10 +55,12 @@ struct LazilyLoadedFile::Impl {
   std::unique_ptr<LazilyLoadedFileReadCallback> callback_;
   DisjointRangeSet already_downloaded_;
   size_t seek_pos_;
+  size_t size_;
+  std::vector<char> extended_read_buffer_;
 };
 
 Result<LazilyLoadedFile> LazilyLoadedFile::Create(
-    std::string filename,
+    std::string filename, size_t size,
     std::unique_ptr<LazilyLoadedFileReadCallback> callback) {
   std::unique_ptr<Impl> impl = std::make_unique<Impl>();
   CF_EXPECT(impl.get());
@@ -63,6 +69,8 @@ Result<LazilyLoadedFile> LazilyLoadedFile::Create(
   impl->filename_ = std::move(filename);
   impl->callback_ = std::move(callback);
   impl->seek_pos_ = 0;
+  impl->size_ = size;
+  impl->extended_read_buffer_ = std::vector<char>(kMinReadSize);
 
   CF_EXPECT(impl->ReadMetadata());
 
@@ -99,6 +107,7 @@ Result<size_t> LazilyLoadedFile::Read(char* data, size_t size) {
 
 Result<void> LazilyLoadedFile::Seek(size_t location) {
   CF_EXPECT(impl_.get());
+  LOG(VERBOSE) << "Seeking to " << location;
   impl_->seek_pos_ = location;
   return {};
 }
@@ -141,9 +150,15 @@ Result<void> LazilyLoadedFile::Impl::WriteMetadata() {
 }
 
 Result<size_t> LazilyLoadedFile::Impl::Read(char* data, size_t size) {
+  LOG(VERBOSE) << "Reading " << size << ", seek pos " << seek_pos_;
   // NOLINTNEXTLINE(misc-include-cleaner): SEEK_SET
   CF_EXPECT_EQ(contents_file_->LSeek(seek_pos_, SEEK_SET), seek_pos_,
                contents_file_->StrError());
+  auto all_ranges = already_downloaded_.AllRanges();
+  for (const auto& range : all_ranges) {
+    LOG(VERBOSE) << "Already downloaded: [" << range.first << ", "
+                 << range.second << ")";
+  }
   std::optional<uint64_t> end_of_present_data =
       already_downloaded_.EndOfContainingRange(seek_pos_);
   // In terms of IO performance, this aims to minimize round trips over
@@ -152,16 +167,39 @@ Result<size_t> LazilyLoadedFile::Impl::Read(char* data, size_t size) {
     size_t read_request = std::min(*end_of_present_data - seek_pos_, size);
     ssize_t data_read = contents_file_->Read(data, read_request);
     CF_EXPECT_GE(data_read, 0, contents_file_->StrError());
+    LOG(VERBOSE) << "Read " << data_read << " from local storage, seek pos was "
+                 << seek_pos_;
     seek_pos_ += data_read;
     return data_read;
   }
   CF_EXPECT(callback_->Seek(seek_pos_));
-  size_t data_read = CF_EXPECT(callback_->Read(data, size));
-  CF_EXPECT_EQ(WriteAll(contents_file_, data, data_read), data_read,
-               contents_file_->StrError());
-  already_downloaded_.InsertRange(seek_pos_, seek_pos_ + data_read);
-  seek_pos_ += data_read;
-  return data_read;
+  if (size < kMinReadSize) {
+    size_t extended_read_size = std::min(kMinReadSize, size_ - seek_pos_);
+    LOG(VERBOSE) << "Extending read request from " << size << " to "
+                 << extended_read_size;
+    size_t data_read = CF_EXPECT(
+        callback_->Read(extended_read_buffer_.data(), extended_read_size));
+    CF_EXPECT_EQ(
+        WriteAll(contents_file_, extended_read_buffer_.data(), data_read),
+        data_read, contents_file_->StrError());
+    already_downloaded_.InsertRange(seek_pos_, seek_pos_ + data_read);
+    LOG(VERBOSE) << "Read " << data_read << " from source, seek pos was "
+                 << seek_pos_;
+    size_t reported_size = std::min(data_read, size);
+    memcpy(data, extended_read_buffer_.data(), reported_size);
+    seek_pos_ += reported_size;
+    return reported_size;
+  } else {
+    LOG(VERBOSE) << "Passing down read request of " << size;
+    size_t data_read = CF_EXPECT(callback_->Read(data, size));
+    CF_EXPECT_EQ(WriteAll(contents_file_, data, data_read), data_read,
+                 contents_file_->StrError());
+    already_downloaded_.InsertRange(seek_pos_, seek_pos_ + data_read);
+    LOG(VERBOSE) << "Read " << data_read << " from source, seek pos was "
+                 << seek_pos_;
+    seek_pos_ += data_read;
+    return data_read;
+  }
 }
 
 }  // namespace cuttlefish
diff --git a/base/cvd/cuttlefish/host/libs/zip/lazily_loaded_file.h b/base/cvd/cuttlefish/host/libs/zip/lazily_loaded_file.h
@@ -34,7 +34,8 @@ class LazilyLoadedFileReadCallback {
 class LazilyLoadedFile {
  public:
   static Result<LazilyLoadedFile> Create(
-      std::string filename, std::unique_ptr<LazilyLoadedFileReadCallback>);
+      std::string filename, size_t size,
+      std::unique_ptr<LazilyLoadedFileReadCallback>);
 
   LazilyLoadedFile(LazilyLoadedFile&&);
   ~LazilyLoadedFile();
diff --git a/base/cvd/cuttlefish/host/libs/zip/remote_zip.cc b/base/cvd/cuttlefish/host/libs/zip/remote_zip.cc
diff --git a/base/cvd/cuttlefish/host/libs/zip/zip_copy.cc b/base/cvd/cuttlefish/host/libs/zip/zip_copy.cc

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,9 @@ DownloadFlags DownloadFlags::Create(const VectorFlags& flags, const int index) {`
`27`	`27`	`.download_target_files_zip =`
`28`	`28`	`GetOptional(flags.download_target_files_zip, index)`
`29`	`29`	`.value_or(kDefaultDownloadTargetFilesZip),`
	`30`	`+ .extract_super_image_fragments =`
	`31`	`+ GetOptional(flags.extract_super_image_fragments, index)`
	`32`	`+ .value_or(kDefaultExtractSuperImageFragments),`
`30`	`33`	`};`
`31`	`34`	`}`
`32`	`35`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ Result<ReadableZip> OpenZip(BuildApi& build_api, const Build& build,`
`31`	`31`	`SeekableZipSource source = CF_EXPECT(build_api.FileReader(build, name));`
`32`	`32`
`33`	`33`	`SeekableZipSource buffered =`
`34`		`- CF_EXPECT(BufferZipSource(std::move(source), 1 << 16));`
	`34`	`+ CF_EXPECT(BufferZipSource(std::move(source), 1 << 26));`
`35`	`35`
`36`	`36`	`return CF_EXPECT(ReadableZip::FromSource(std::move(buffered)));`
`37`	`37`	`}`