fix: bound snapshot manifest cache

Co-authored-by: Codex <codex@openai.com>

Author: manuzhang

src/iceberg/file_io.cc

@@ -253,6 +253,8 @@ Status FileIO::ConfigureMetadataCache(
   std::lock_guard lock(metadata_cache_state_->mutex);
   if (metadata_cache_state_->cache == nullptr) {
     metadata_cache_state_->cache = std::move(cache);
+    metadata_cache_state_->snapshot_cache =
+        internal::MakeSnapshotCacheData(metadata_cache_state_->cache);
     return {};
   }
   if (metadata_cache_state_->cache->options() == options) {
@@ -288,7 +290,8 @@ std::shared_ptr<MetadataCache> FileIO::GetMetadataCache() const {
 std::shared_ptr<internal::SnapshotCacheData> FileIO::GetSnapshotCacheData() const {
   std::lock_guard lock(metadata_cache_state_->mutex);
   if (metadata_cache_state_->snapshot_cache == nullptr) {
-    metadata_cache_state_->snapshot_cache = internal::MakeSnapshotCacheData();
+    metadata_cache_state_->snapshot_cache =
+        internal::MakeSnapshotCacheData(metadata_cache_state_->cache);
   }
   return metadata_cache_state_->snapshot_cache;
 }

src/iceberg/snapshot.cc

@@ -30,6 +30,7 @@
 #include "iceberg/manifest/manifest_entry.h"
 #include "iceberg/manifest/manifest_list.h"
 #include "iceberg/manifest/manifest_reader.h"
+#include "iceberg/metadata_cache.h"
 #include "iceberg/util/macros.h"
 #include "iceberg/util/string_util.h"
 
@@ -39,59 +40,107 @@ namespace internal {
 
 class SnapshotCacheData {
  public:
-  Result<std::reference_wrapper<SnapshotCache::ManifestsCache>> Get(
+  explicit SnapshotCacheData(std::shared_ptr<MetadataCache> metadata_cache)
+      : metadata_cache_(std::move(metadata_cache)) {}
+
+  Result<std::shared_ptr<const SnapshotCache::ManifestsCache>> Get(
       const Snapshot* snapshot, std::shared_ptr<FileIO> file_io) {
+    if (metadata_cache_ == nullptr || !metadata_cache_->options().enabled) {
+      ICEBERG_ASSIGN_OR_RAISE(
+          auto loaded, SnapshotCache::LoadManifestsCache(snapshot, std::move(file_io)));
+      return std::make_shared<const SnapshotCache::ManifestsCache>(std::move(loaded));
+    }
+
+    const auto& location = snapshot->manifest_list;
     std::shared_ptr<Entry> entry;
-    {
+    while (true) {
+      // Checking the content cache also applies its expiration policy and refreshes its
+      // LRU position. Parsed entries are reusable only while the matching bytes remain
+      // cached.
+      auto content = metadata_cache_->GetIfPresent(location);
       std::unique_lock lock(mutex_);
-      while (true) {
-        auto [it, inserted] = entries_.try_emplace(snapshot->manifest_list);
-        if (inserted) {
-          it->second = std::make_shared<Entry>();
+      PruneExpired();
+      auto [it, inserted] = entries_.try_emplace(location);
+      if (inserted) {
+        it->second = std::make_shared<Entry>();
+      }
+      entry = it->second;
+      if (entry->value != nullptr) {
+        if (content != nullptr && entry->content.lock() == content) {
+          return entry->value;
         }
-        entry = it->second;
-        if (entry->value.has_value()) {
-          return std::ref(*entry->value);
+        entries_.erase(it);
+        continue;
+      }
+      if (entry->loading) {
+        entry->loaded.wait(lock, [&entry] { return !entry->loading; });
+        if (entry->error.has_value()) {
+          return std::unexpected<Error>(*entry->error);
         }
-        if (!entry->loading) {
-          entry->loading = true;
-          break;
+        if (entry->value != nullptr) {
+          return entry->value;
         }
-        entry->loaded.wait(lock, [&entry] { return !entry->loading; });
+        continue;
       }
+      entry->loading = true;
+      break;
     }
 
-    auto loaded = SnapshotCache::InitManifestsCache(snapshot, std::move(file_io));
+    auto loaded = SnapshotCache::LoadManifestsCache(snapshot, std::move(file_io));
+    auto content = metadata_cache_->GetIfPresent(location);
     std::lock_guard lock(mutex_);
-    auto it = entries_.find(snapshot->manifest_list);
+    auto it = entries_.find(location);
     if (!loaded.has_value()) {
       if (it != entries_.end() && it->second == entry) {
         entries_.erase(it);
       }
+      entry->error = loaded.error();
       entry->loading = false;
       entry->loaded.notify_all();
       return std::unexpected<Error>(std::move(loaded).error());
     }
 
-    entry->value = std::move(loaded).value();
+    entry->value =
+        std::make_shared<const SnapshotCache::ManifestsCache>(std::move(loaded).value());
+    if (content != nullptr) {
+      entry->content = content;
+    } else if (it != entries_.end() && it->second == entry) {
+      // Oversized content is not retained by MetadataCache, so do not retain its parsed
+      // representation in the cross-query cache either.
+      entries_.erase(it);
+    }
     entry->loading = false;
     entry->loaded.notify_all();
-    return std::ref(*entry->value);
+    return entry->value;
   }
 
  private:
   struct Entry {
     bool loading = false;
-    std::optional<SnapshotCache::ManifestsCache> value;
+    std::shared_ptr<const SnapshotCache::ManifestsCache> value;
+    std::weak_ptr<const std::string> content;
+    std::optional<Error> error;
     std::condition_variable loaded;
   };
 
+  void PruneExpired() {
+    for (auto it = entries_.begin(); it != entries_.end();) {
+      if (!it->second->loading && it->second->content.expired()) {
+        it = entries_.erase(it);
+      } else {
+        ++it;
+      }
+    }
+  }
+
+  std::shared_ptr<MetadataCache> metadata_cache_;
   std::mutex mutex_;
   std::unordered_map<std::string, std::shared_ptr<Entry>> entries_;
 };
 
-std::shared_ptr<SnapshotCacheData> MakeSnapshotCacheData() {
-  return std::make_shared<SnapshotCacheData>();
+std::shared_ptr<SnapshotCacheData> MakeSnapshotCacheData(
+    std::shared_ptr<MetadataCache> metadata_cache) {
+  return std::make_shared<SnapshotCacheData>(std::move(metadata_cache));
 }
 
 }  // namespace internal
@@ -267,7 +316,14 @@ Result<std::unique_ptr<Snapshot>> Snapshot::Make(
   });
 }
 
-Result<SnapshotCache::ManifestsCache> SnapshotCache::InitManifestsCache(
+Result<std::shared_ptr<const SnapshotCache::ManifestsCache>>
+SnapshotCache::InitManifestsCache(const Snapshot* snapshot,
+                                  std::shared_ptr<FileIO> file_io) {
+  auto cache_data = file_io->GetSnapshotCacheData();
+  return cache_data->Get(snapshot, std::move(file_io));
+}
+
+Result<SnapshotCache::ManifestsCache> SnapshotCache::LoadManifestsCache(
     const Snapshot* snapshot, std::shared_ptr<FileIO> file_io) {
   if (file_io == nullptr) {
     return InvalidArgument("Cannot cache manifests: FileIO is null");
@@ -304,29 +360,29 @@ Result<std::span<const ManifestFile>> SnapshotCache::Manifests(
     std::shared_ptr<FileIO> file_io) const {
   ICEBERG_PRECHECK(snapshot_ != nullptr, "Cannot cache manifests for a null snapshot");
   ICEBERG_PRECHECK(file_io != nullptr, "Cannot cache manifests: FileIO is null");
-  auto cache_data = file_io->GetSnapshotCacheData();
-  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref, cache_data->Get(snapshot_, std::move(file_io)));
-  auto& cache = cache_ref.get();
+  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref,
+                          manifests_cache_.Get(snapshot_, std::move(file_io)));
+  const auto& cache = *cache_ref.get();
   return std::span<const ManifestFile>(cache.first.data(), cache.first.size());
 }
 
 Result<std::span<const ManifestFile>> SnapshotCache::DataManifests(
     std::shared_ptr<FileIO> file_io) const {
   ICEBERG_PRECHECK(snapshot_ != nullptr, "Cannot cache manifests for a null snapshot");
   ICEBERG_PRECHECK(file_io != nullptr, "Cannot cache manifests: FileIO is null");
-  auto cache_data = file_io->GetSnapshotCacheData();
-  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref, cache_data->Get(snapshot_, std::move(file_io)));
-  auto& cache = cache_ref.get();
+  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref,
+                          manifests_cache_.Get(snapshot_, std::move(file_io)));
+  const auto& cache = *cache_ref.get();
   return std::span<const ManifestFile>(cache.first.data(), cache.second);
 }
 
 Result<std::span<const ManifestFile>> SnapshotCache::DeleteManifests(
     std::shared_ptr<FileIO> file_io) const {
   ICEBERG_PRECHECK(snapshot_ != nullptr, "Cannot cache manifests for a null snapshot");
   ICEBERG_PRECHECK(file_io != nullptr, "Cannot cache manifests: FileIO is null");
-  auto cache_data = file_io->GetSnapshotCacheData();
-  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref, cache_data->Get(snapshot_, std::move(file_io)));
-  auto& cache = cache_ref.get();
+  ICEBERG_ASSIGN_OR_RAISE(auto cache_ref,
+                          manifests_cache_.Get(snapshot_, std::move(file_io)));
+  const auto& cache = *cache_ref.get();
   const size_t delete_start = cache.second;
   const size_t delete_count = cache.first.size() - delete_start;
   return std::span<const ManifestFile>(cache.first.data() + delete_start, delete_count);

src/iceberg/snapshot.h

@@ -34,13 +34,17 @@
 #include "iceberg/manifest/manifest_list.h"
 #include "iceberg/result.h"
 #include "iceberg/type_fwd.h"
+#include "iceberg/util/lazy.h"
 #include "iceberg/util/timepoint.h"
 
 namespace iceberg {
 
+class MetadataCache;
+
 namespace internal {
 class SnapshotCacheData;
-ICEBERG_EXPORT std::shared_ptr<SnapshotCacheData> MakeSnapshotCacheData();
+ICEBERG_EXPORT std::shared_ptr<SnapshotCacheData> MakeSnapshotCacheData(
+    std::shared_ptr<MetadataCache> metadata_cache);
 }  // namespace internal
 
 /// \brief The type of snapshot reference
@@ -510,12 +514,18 @@ class ICEBERG_EXPORT SnapshotCache {
   /// \param snapshot The snapshot to initialize the manifests cache for
   /// \param file_io The FileIO instance to use for reading the manifest list
   /// \return A result containing the manifests cache
-  static Result<ManifestsCache> InitManifestsCache(const Snapshot* snapshot,
+  static Result<std::shared_ptr<const ManifestsCache>> InitManifestsCache(
+      const Snapshot* snapshot, std::shared_ptr<FileIO> file_io);
+
+  static Result<ManifestsCache> LoadManifestsCache(const Snapshot* snapshot,
                                                    std::shared_ptr<FileIO> file_io);
 
   /// The underlying snapshot data
   const Snapshot* snapshot_;
 
+  /// Keep the selected shared cache entry alive while spans from this object are used.
+  Lazy<InitManifestsCache> manifests_cache_;
+
   friend class internal::SnapshotCacheData;
 };
 

src/iceberg/test/manifest_writer_versions_test.cc

@@ -17,6 +17,7 @@
  * under the License.
  */
 
+#include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -33,13 +34,16 @@
 #include "iceberg/manifest/manifest_list.h"
 #include "iceberg/manifest/manifest_reader.h"
 #include "iceberg/manifest/manifest_writer.h"
+#include "iceberg/metadata_cache.h"
 #include "iceberg/metrics.h"
 #include "iceberg/partition_spec.h"
 #include "iceberg/row/partition_values.h"
 #include "iceberg/schema.h"
 #include "iceberg/schema_field.h"
+#include "iceberg/snapshot.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/test/matchers.h"
+#include "iceberg/test/mock_io.h"
 #include "iceberg/transform.h"
 #include "iceberg/type.h"
 
@@ -127,6 +131,36 @@ std::unique_ptr<DataFile> CreateDeleteFile() {
   return delete_file;
 }
 
+class CountingInputFile : public InputFile {
+ public:
+  CountingInputFile(std::unique_ptr<InputFile> input_file, int* open_count)
+      : input_file_(std::move(input_file)), open_count_(open_count) {}
+
+  std::string_view location() const override { return input_file_->location(); }
+
+  Result<int64_t> Size() const override { return input_file_->Size(); }
+
+  Result<std::unique_ptr<SeekableInputStream>> Open() override {
+    ++*open_count_;
+    return input_file_->Open();
+  }
+
+ private:
+  std::unique_ptr<InputFile> input_file_;
+  int* open_count_;
+};
+
+class CountingMockFileIO : public MockFileIO {
+ public:
+  Result<std::unique_ptr<InputFile>> NewInputFile(std::string file_location) override {
+    ICEBERG_ASSIGN_OR_RAISE(auto input_file,
+                            MockFileIO::NewInputFile(std::move(file_location)));
+    return std::make_unique<CountingInputFile>(std::move(input_file), &open_count);
+  }
+
+  int open_count = 0;
+};
+
 }  // namespace
 
 class ManifestWriterVersionsTest : public ::testing::Test {
@@ -410,6 +444,82 @@ class ManifestWriterVersionsTest : public ::testing::Test {
   std::shared_ptr<FileIO> file_io_{nullptr};
 };
 
+TEST_F(ManifestWriterVersionsTest, SnapshotEntriesFollowContentCacheEviction) {
+  auto file_io = std::make_shared<CountingMockFileIO>();
+  auto write_manifest_list = [&](const std::string& location, int64_t snapshot_id) {
+    auto writer_result = ManifestListWriter::MakeWriter(
+        /*format_version=*/2, snapshot_id, /*parent_snapshot_id=*/std::nullopt, location,
+        file_io, /*sequence_number=*/1);
+    ASSERT_THAT(writer_result, IsOk());
+    auto writer = std::move(writer_result).value();
+    ManifestFile manifest{
+        .manifest_path = location + ".manifest",
+        .manifest_length = 10,
+        .partition_spec_id = PartitionSpec::kInitialSpecId,
+        .content = ManifestContent::kData,
+        .sequence_number = 1,
+        .min_sequence_number = 1,
+        .added_snapshot_id = snapshot_id,
+        .added_files_count = 1,
+        .existing_files_count = 0,
+        .deleted_files_count = 0,
+        .added_rows_count = 1,
+        .existing_rows_count = 0,
+        .deleted_rows_count = 0,
+    };
+    ASSERT_THAT(writer->AddAll({manifest}), IsOk());
+    ASSERT_THAT(writer->Close(), IsOk());
+  };
+
+  const std::string first_location = "first-manifest-list.avro";
+  const std::string second_location = "second-manifest-list.avro";
+  write_manifest_list(first_location, /*snapshot_id=*/1);
+  write_manifest_list(second_location, /*snapshot_id=*/2);
+  const auto max_list_size = std::max(file_io->FileData(first_location).size(),
+                                      file_io->FileData(second_location).size());
+  ASSERT_THAT(file_io->ConfigureMetadataCache(
+                  {{std::string(MetadataCacheOptions::kEnabled), "true"},
+                   {std::string(MetadataCacheOptions::kMaxTotalBytes),
+                    std::to_string(max_list_size)},
+                   {std::string(MetadataCacheOptions::kMaxContentLength),
+                    std::to_string(max_list_size)}}),
+              IsOk());
+
+  Snapshot first_snapshot{
+      .snapshot_id = 1,
+      .parent_snapshot_id = std::nullopt,
+      .sequence_number = 1,
+      .timestamp_ms = TimePointMs{},
+      .manifest_list = first_location,
+  };
+  Snapshot second_snapshot{
+      .snapshot_id = 2,
+      .parent_snapshot_id = 1,
+      .sequence_number = 2,
+      .timestamp_ms = TimePointMs{},
+      .manifest_list = second_location,
+  };
+
+  SnapshotCache first_cache(&first_snapshot);
+  ICEBERG_UNWRAP_OR_FAIL(auto first_manifests, first_cache.Manifests(file_io));
+  ASSERT_EQ(first_manifests.size(), 1);
+  EXPECT_EQ(file_io->open_count, 1);
+
+  SnapshotCache shared_first_cache(&first_snapshot);
+  EXPECT_THAT(shared_first_cache.Manifests(file_io), IsOk());
+  EXPECT_EQ(file_io->open_count, 1);
+
+  SnapshotCache second_cache(&second_snapshot);
+  EXPECT_THAT(second_cache.Manifests(file_io), IsOk());
+  EXPECT_EQ(file_io->open_count, 2);
+  // Evicting the shared entry must not invalidate spans owned by an existing wrapper.
+  EXPECT_EQ(first_manifests.front().manifest_path, first_location + ".manifest");
+
+  SnapshotCache reloaded_first_cache(&first_snapshot);
+  EXPECT_THAT(reloaded_first_cache.Manifests(file_io), IsOk());
+  EXPECT_EQ(file_io->open_count, 3);
+}
+
 TEST_F(ManifestWriterVersionsTest, TestV1Write) {
   auto manifest = WriteManifest(/*format_version=*/1, {data_file_});
   CheckManifest(manifest, kInvalidSequenceNumber, kInvalidSequenceNumber);