Add disk fallback to LRU cache

Test: OTA on reference hardware. Reduces path time for a 40MB puffin
patch from 786to 135s

Overall OTA update time reduced from 1915s to 1053s (~45% reduction)

Bug: 320443894

Change-Id: Iabf0f0bece04bd8561d3b9fb3a0fb30882fce593

3 files changed, 104 insertions, 11 deletions

diff --git a/Android.bp b/Android.bp
index fdaab47..3ecc9b3 100644
--- a/Android.bp
+++ b/Android.bp
@@ -62,6 +62,7 @@ cc_library_static {
     ],
     static_libs: [
         "libbspatch",
+        "libc++fs",
     ],
     whole_static_libs: [
         "libzucchini",
@@ -84,6 +85,7 @@ cc_library_static {
         "libbsdiff",
         "libzucchini",
         "libpuffpatch",
+        "libc++fs",
     ],
 }
 
@@ -105,6 +107,7 @@ cc_binary {
         "libdivsufsort64",
         "libpuffdiff",
         "libpuffpatch",
+        "libc++fs",
     ],
 }
 
@@ -138,5 +141,6 @@ cc_test {
         "libdivsufsort64",
         "libpuffdiff",
         "libpuffpatch",
+        "libc++fs",
     ],
 }
diff --git a/src/puffin_stream.cc b/src/puffin_stream.cc
index c50ad93..4eec797 100644
--- a/src/puffin_stream.cc
+++ b/src/puffin_stream.cc
@@ -3,8 +3,11 @@
 // found in the LICENSE file.
 
 #include "puffin/src/puffin_stream.h"
+#include <fcntl.h>
+#include <sys/stat.h>
 
 #include <algorithm>
+#include <filesystem>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -448,19 +451,80 @@ bool PuffinStream::GetPuffCache(int puff_id,
     lru_cache_.put(puff_id, cache);
   }
 
-  *buffer = cache;
+  *buffer = std::move(cache);
   return found;
 }
 
-const LRUCache::Value LRUCache::get(const LRUCache::Key& key) {
+const LRUCache::Value LRUCache::get(const Key& key) {
   auto it = items_map_.find(key);
   if (it == items_map_.end()) {
+    if (ondisk_items_.count(key) > 0) {
+      auto&& data = ReadFromDisk(key);
+      put(key, data);
+      return data;
+    }
     return nullptr;
   }
   items_list_.splice(items_list_.begin(), items_list_, it->second);
   return it->second->second;
 }
 
+LRUCache::Value LRUCache::ReadFromDisk(const LRUCache::Key& key) {
+  const auto path = tmpdir_ / std::to_string(key);
+  int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC);
+  if (fd < 0) {
+    PLOG(ERROR) << "Failed to open " << path;
+    return {};
+  }
+  auto fd_delete = [](int* fd) { close(*fd); };
+  std::unique_ptr<int, decltype(fd_delete)> guard(&fd, fd_delete);
+  struct stat st {};
+  const int ret = stat(path.c_str(), &st);
+  if (ret < 0) {
+    PLOG(ERROR) << "Failed to stat " << path << " ret: " << ret;
+    return {};
+  }
+  LRUCache::Value data = std::make_shared<std::vector<uint8_t>>(st.st_size);
+  const auto bytes_read =
+      TEMP_FAILURE_RETRY(pread(fd, data->data(), data->size(), 0));
+  if (static_cast<size_t>(bytes_read) != data->size()) {
+    PLOG(ERROR) << "Failed to read " << data->size() << " bytes data from "
+                << path;
+    return {};
+  }
+  return data;
+}
+
+LRUCache::~LRUCache() {
+  std::filesystem::remove_all(tmpdir_);
+}
+
+bool LRUCache::WriteToDisk(const LRUCache::Key& key,
+                           const LRUCache::Value& value) {
+  if (tmpdir_.empty()) {
+    return false;
+  }
+  const auto path = tmpdir_ / std::to_string(key);
+  int fd = open(path.c_str(), O_RDWR | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
+  if (fd < 0) {
+    PLOG(ERROR) << "Failed to open " << path;
+    return false;
+  }
+  auto fd_delete = [](int* fd) { close(*fd); };
+  std::unique_ptr<int, decltype(fd_delete)> guard(&fd, fd_delete);
+  const auto written =
+      TEMP_FAILURE_RETRY(write(fd, value->data(), value->size()));
+  if (static_cast<size_t>(written) != value->size()) {
+    PLOG(ERROR) << "Failed to write " << value->size() << " bytes data to "
+                << path;
+    return false;
+  }
+  close(fd);
+  guard.release();
+  ondisk_items_.insert(key);
+  return true;
+}
+
 void LRUCache::put(const LRUCache::Key& key, const LRUCache::Value& value) {
   auto it = items_map_.find(key);
   if (it != items_map_.end()) {
@@ -474,26 +538,43 @@ void LRUCache::put(const LRUCache::Key& key, const LRUCache::Value& value) {
   items_map_[key] = items_list_.begin();
 }
 
-void LRUCache::EvictLRUItem() {
+bool LRUCache::EvictLRUItem() {
   if (items_list_.empty()) {
-    return;
+    return false;
   }
-  auto last = items_list_.back();
+  const auto last = items_list_.back();
   cache_size_ -= last.second->capacity();
+  // Only write puffs large enough to disk, as disk writes have latency.
+  if (last.second->size() > 16 * 1024) {
+    WriteToDisk(last.first, last.second);
+  }
   items_map_.erase(last.first);
   items_list_.pop_back();
+  return true;
 }
 
 void LRUCache::EnsureSpace(size_t size) {
   if (size > max_size_) {
-    items_list_.clear();
-    items_map_.clear();
-    cache_size_ = 0;
+    while (EvictLRUItem())
+      ;
     return;
   }
   while (cache_size_ + size > max_size_) {
-    EvictLRUItem();
+    if (!EvictLRUItem()) {
+      return;
+    }
+  }
+}
+
+LRUCache::LRUCache(size_t max_size) : max_size_(max_size) {
+  std::string buffer = std::filesystem::temp_directory_path() / "lru.XXXXXX";
+  const char* dirname = mkdtemp(buffer.data());
+  if (dirname == nullptr) {
+    LOG(ERROR) << "Failed to create a unique temporary dir";
+    return;
   }
+  tmpdir_ = dirname;
+  std::filesystem::create_directory(tmpdir_);
 }
 
 }  // namespace puffin
diff --git a/src/puffin_stream.h b/src/puffin_stream.h
index 1f8ec4f..c3ccf81 100644
--- a/src/puffin_stream.h
+++ b/src/puffin_stream.h
@@ -5,8 +5,10 @@
 #ifndef SRC_PUFFIN_STREAM_H_
 #define SRC_PUFFIN_STREAM_H_
 
+#include <filesystem>
 #include <list>
 #include <memory>
+#include <set>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -25,11 +27,13 @@ class LRUCache {
   typedef typename std::pair<Key, Value> key_value_pair_t;
   typedef typename std::list<key_value_pair_t>::iterator iterator;
 
-  LRUCache(size_t max_size) : max_size_(max_size) {}
+  LRUCache(size_t max_size);
+  ~LRUCache();
+  LRUCache(const LRUCache&) = delete;
 
   void put(const Key& key, const Value& value);
 
-  void EvictLRUItem();
+  bool EvictLRUItem();
 
   // Ensure that cache has sufficient space for a new item of |size| bytes
   void EnsureSpace(size_t size);
@@ -44,10 +48,14 @@ class LRUCache {
   size_t capacity() const { return max_size_; }
 
  private:
+  bool WriteToDisk(const Key& key, const Value& value);
+  Value ReadFromDisk(const Key& key);
   std::list<key_value_pair_t> items_list_;
   std::unordered_map<Key, iterator> items_map_;
   size_t cache_size_ = 0;
   size_t max_size_;
+  std::filesystem::path tmpdir_;
+  std::set<Key> ondisk_items_;
 };
 
 // A class for puffing a deflate stream and huffing into a deflate stream. The