Trim AutofillAiModelCache

This CL introduces a TrimEntries method to AutofillAiModelCacheImpl.
TrimEntries does two things:
- It removes entries (starting with the oldest) if the cache exceeds
  the target maximum size.
- It removes entries once they exceed the maximum target age.

The TrimEntries method is triggered whenever the DB is loaded (i.e.,
on start) and when a cache entry is updated/added.

Bug: 389631477
Change-Id: If9e1fee273e6596121ba45514a9114fe97b2be04
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6317566
Reviewed-by: Christoph Schwering <schwering@google.com>
Commit-Queue: Christoph Schwering <schwering@google.com>
Cr-Commit-Position: refs/heads/main@{#1427183}

chrome/browser/autofill/autofill_ai_model_cache_factory.cc

@@ -49,7 +49,9 @@ AutofillAiModelCacheFactory::BuildServiceInstanceForBrowserContext(
   Profile* profile = Profile::FromBrowserContext(context);
   return std::make_unique<AutofillAiModelCacheImpl>(
       profile->GetDefaultStoragePartition()->GetProtoDatabaseProvider(),
-      profile->GetPath());
+      profile->GetPath(),
+      autofill::features::kAutofillAiServerModelCacheSize.Get(),
+      autofill::features::kAutofillAiServerModelCacheAge.Get());
 }
 
 bool AutofillAiModelCacheFactory::ServiceIsCreatedWithBrowserContext() const {

components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache_impl.cc

@@ -4,12 +4,13 @@
 
 #include "components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache_impl.h"
 
-#include <iterator>
+#include <algorithm>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "base/containers/to_vector.h"
 #include "base/files/file_path.h"
 #include "base/functional/bind.h"
 #include "base/strings/string_number_conversions.h"
@@ -20,9 +21,21 @@
 
 namespace autofill {
 
+namespace {
+
+// A helper for serializing time used in this DB.
+constexpr int64_t SerializeTime(base::Time time) {
+  return time.ToDeltaSinceWindowsEpoch().InMicroseconds();
+}
+
+}  // namespace
+
 AutofillAiModelCacheImpl::AutofillAiModelCacheImpl(
     leveldb_proto::ProtoDatabaseProvider* db_provider,
-    const base::FilePath& profile_path) {
+    const base::FilePath& profile_path,
+    size_t max_cache_size,
+    base::TimeDelta max_cache_age)
+    : max_cache_size_(max_cache_size), max_cache_age_(max_cache_age) {
   auto database_task_runner = base::ThreadPool::CreateSequencedTaskRunner(
       {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
        base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN});
@@ -41,14 +54,59 @@ void AutofillAiModelCacheImpl::Update(FormSignature form_signature,
                                       CacheEntry entry) {
   CacheEntryWithMetadata entry_with_metadata;
   *entry_with_metadata.mutable_server_response() = std::move(entry);
-  entry_with_metadata.set_creation_time(
-      base::Time::Now().ToDeltaSinceWindowsEpoch().InMicroseconds());
+  entry_with_metadata.set_creation_time(SerializeTime(base::Time::Now()));
   UpdateInDatabase(form_signature, entry_with_metadata);
   in_memory_cache_[form_signature] = std::move(entry_with_metadata);
+  TrimEntries();
 }
 
 bool AutofillAiModelCacheImpl::Contains(FormSignature form_signature) const {
-  return in_memory_cache_.contains(form_signature);
+  auto it = in_memory_cache_.find(form_signature);
+  return it != in_memory_cache_.end() &&
+         it->second.creation_time() >=
+             SerializeTime(base::Time::Now() - max_cache_age_);
+}
+
+void AutofillAiModelCacheImpl::TrimEntries() {
+  // Transform the cache into an ordered list of (creation_time, FormSignature)
+  // pairs.
+  std::vector<std::pair<int64_t, FormSignature>> entries_by_creation_time =
+      base::ToVector(in_memory_cache_,
+                     [](const std::pair<FormSignature, CacheEntryWithMetadata>&
+                            map_entry) {
+                       const auto& [signature, cache_entry] = map_entry;
+                       return std::pair(cache_entry.creation_time(), signature);
+                     });
+  std::ranges::sort(entries_by_creation_time);
+
+  // Entries with a time stamp smaller than `cut_off_time` should be removed.
+  const int64_t cut_off_time =
+      SerializeTime(base::Time::Now() - max_cache_age_);
+  const size_t deletions_due_to_time_cutoff = std::distance(
+      entries_by_creation_time.begin(),
+      std::ranges::lower_bound(entries_by_creation_time, cut_off_time, {},
+                               &std::pair<int64_t, FormSignature>::first));
+  // The remaining cache should at most have size `max_cache_size_`.
+  size_t deletions_overall = deletions_due_to_time_cutoff;
+  if (entries_by_creation_time.size() - deletions_overall > max_cache_size_) {
+    deletions_overall = entries_by_creation_time.size() - max_cache_size_;
+  }
+
+  if (deletions_overall == 0) {
+    return;
+  }
+
+  auto entries_to_save = std::make_unique<
+      leveldb_proto::ProtoDatabase<CacheEntryWithMetadata>::KeyEntryVector>();
+  auto keys_to_remove = std::make_unique<std::vector<std::string>>();
+  keys_to_remove->reserve(deletions_overall);
+  for (const auto [_, signature_to_delete] :
+       base::span(entries_by_creation_time).first(deletions_overall)) {
+    in_memory_cache_.erase(signature_to_delete);
+    keys_to_remove->emplace_back(base::NumberToString(*signature_to_delete));
+  }
+  db_->UpdateEntries(std::move(entries_to_save), std::move(keys_to_remove),
+                     /*callback=*/base::DoNothing());
 }
 
 void AutofillAiModelCacheImpl::UpdateInDatabase(
@@ -88,10 +146,12 @@ void AutofillAiModelCacheImpl::OnDatabaseLoadKeysAndEntries(
   for (auto& [signature_str, entry] : *entries) {
     uint64_t signature_uint64;
     if (!base::StringToUint64(signature_str, &signature_uint64)) {
+      // TODO(crbug.com/389631477): Clean up malformed entries.
       continue;
     }
     in_memory_cache_[FormSignature(signature_uint64)] = std::move(entry);
   }
+  TrimEntries();
 }
 
 }  // namespace autofill

components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache_impl.h

@@ -11,6 +11,7 @@
 
 #include "base/files/file_path.h"
 #include "base/memory/weak_ptr.h"
+#include "base/time/time.h"
 #include "components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache.h"
 #include "components/autofill/core/browser/proto/autofill_ai_model_cache.pb.h"
 #include "components/autofill/core/common/signatures.h"
@@ -32,12 +33,12 @@ inline constexpr base::FilePath::StringViewType
 // `AutofillAiModelCacheImpl` implements `AutofillAiModelCache` using a
 // LevelDB as a persistence layer and a map as an in-memory cache to allow
 // synchronous access.
-// TODO(crbug.com/389631477): Implement a maximum size of the cache.
-// TODO(crbug.com/389631477): Implement a maximum lifetime for cache entries.
 class AutofillAiModelCacheImpl : public AutofillAiModelCache {
  public:
   AutofillAiModelCacheImpl(leveldb_proto::ProtoDatabaseProvider* db_provider,
-                           const base::FilePath& profile_path);
+                           const base::FilePath& profile_path,
+                           size_t max_cache_size,
+                           base::TimeDelta max_cache_age);
   AutofillAiModelCacheImpl(const AutofillAiModelCacheImpl&) = delete;
   AutofillAiModelCacheImpl& operator=(const AutofillAiModelCacheImpl&) = delete;
   AutofillAiModelCacheImpl(AutofillAiModelCacheImpl&&) = delete;
@@ -51,6 +52,9 @@ class AutofillAiModelCacheImpl : public AutofillAiModelCache {
  private:
   using CacheEntryWithMetadata = AutofillAiModelCacheEntryWithMetadata;
 
+  // Removes expired cache entries and limits the cache size to `max_cache_size`
+  // by removing the oldest entries.
+  void TrimEntries();
   void UpdateInDatabase(FormSignature form_signature,
                         const CacheEntryWithMetadata& entry);
 
@@ -63,6 +67,9 @@ class AutofillAiModelCacheImpl : public AutofillAiModelCache {
     return weak_ptr_factory_.GetWeakPtr();
   }
 
+  const size_t max_cache_size_;
+  const base::TimeDelta max_cache_age_;
+
   // An in-memory cache that allows for synchronous access. Should contain the
   // same content as the database.
   std::map<FormSignature, CacheEntryWithMetadata> in_memory_cache_;

components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache_impl_unittest.cc

@@ -8,7 +8,9 @@
 
 #include "base/files/scoped_temp_dir.h"
 #include "base/test/task_environment.h"
+#include "base/time/time.h"
 #include "components/autofill/core/browser/ml_model/autofill_ai/autofill_ai_model_cache.h"
+#include "components/autofill/core/browser/proto/autofill_ai_model_cache.pb.h"
 #include "components/autofill/core/common/signatures.h"
 #include "components/leveldb_proto/public/proto_database_provider.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -33,12 +35,17 @@ class AutofillAiModelCacheImplTest : public ::testing::Test {
     task_environment_.RunUntilIdle();
   }
 
+  void AdvanceClock(base::TimeDelta delta) {
+    task_environment_.AdvanceClock(delta);
+  }
+
   // Simulates restart of the browser by recreating the cache.
-  void RecreateCache() {
+  void RecreateCache(size_t max_cache_size = 50,
+                     base::TimeDelta max_cache_age = base::Days(7)) {
     // Process remaining operations.
     task_environment_.RunUntilIdle();
-    cache_ = std::make_unique<AutofillAiModelCacheImpl>(db_provider_.get(),
-                                                        temp_dir_.GetPath());
+    cache_ = std::make_unique<AutofillAiModelCacheImpl>(
+        db_provider_.get(), temp_dir_.GetPath(), max_cache_size, max_cache_age);
     // Wait until database has loaded.
     task_environment_.RunUntilIdle();
   }
@@ -46,7 +53,8 @@ class AutofillAiModelCacheImplTest : public ::testing::Test {
   AutofillAiModelCache& cache() { return *cache_; }
 
  private:
-  base::test::TaskEnvironment task_environment_;
+  base::test::TaskEnvironment task_environment_{
+      base::test::TaskEnvironment::TimeSource::MOCK_TIME};
   base::ScopedTempDir temp_dir_;
   std::unique_ptr<leveldb_proto::ProtoDatabaseProvider> db_provider_;
   std::unique_ptr<AutofillAiModelCache> cache_;
@@ -77,6 +85,77 @@ TEST_F(AutofillAiModelCacheImplTest, CacheSurvivesRestart) {
   EXPECT_TRUE(cache().Contains(signature));
 }
 
+// Tests that the maximum cache size is enforced by removing the oldest entries
+// that exceed the cache size.
+TEST_F(AutofillAiModelCacheImplTest, MaxCacheSize) {
+  constexpr auto signature1 = FormSignature(123);
+  constexpr auto signature2 = FormSignature(1234);
+  constexpr auto signature3 = FormSignature(12345);
+  constexpr auto signature4 = FormSignature(123456);
+
+  RecreateCache(/*max_cache_size=*/3);
+  cache().Update(signature1, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  cache().Update(signature2, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  cache().Update(signature3, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  EXPECT_TRUE(cache().Contains(signature1));
+  EXPECT_TRUE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+  EXPECT_FALSE(cache().Contains(signature4));
+
+  // Adding a fourth entry removes the first one.
+  cache().Update(signature4, AutofillAiModelCache::CacheEntry());
+  EXPECT_FALSE(cache().Contains(signature1));
+  EXPECT_TRUE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+  EXPECT_TRUE(cache().Contains(signature4));
+
+  // This remains true after a restart.
+  RecreateCache();
+  EXPECT_FALSE(cache().Contains(signature1));
+  EXPECT_TRUE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+  EXPECT_TRUE(cache().Contains(signature4));
+}
+
+// Tests that the maximum cache age is enforced.
+TEST_F(AutofillAiModelCacheImplTest, MaxCacheAge) {
+  constexpr auto signature1 = FormSignature(123);
+  constexpr auto signature2 = FormSignature(1234);
+  constexpr auto signature3 = FormSignature(12345);
+
+  RecreateCache(/*max_cache_size=*/10, /*max_cache_age=*/base::Days(3));
+  cache().Update(signature1, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  cache().Update(signature2, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  cache().Update(signature3, AutofillAiModelCache::CacheEntry());
+  AdvanceClock(base::Days(1));
+  EXPECT_TRUE(cache().Contains(signature1));
+  EXPECT_TRUE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+
+  // If we advance the clock further, the first entry expires.
+  AdvanceClock(base::Hours(1));
+  EXPECT_FALSE(cache().Contains(signature1));
+  EXPECT_TRUE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+
+  // A day later, the second entry expires as well.
+  AdvanceClock(base::Days(1));
+  EXPECT_FALSE(cache().Contains(signature1));
+  EXPECT_FALSE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+
+  // This is still true after a restart.
+  RecreateCache(/*max_cache_size=*/10, /*max_cache_age=*/base::Days(3));
+  EXPECT_FALSE(cache().Contains(signature1));
+  EXPECT_FALSE(cache().Contains(signature2));
+  EXPECT_TRUE(cache().Contains(signature3));
+}
+
 }  // namespace
 
 }  // namespace autofill

components/autofill/core/common/autofill_features.cc

@@ -57,6 +57,17 @@ BASE_FEATURE(kAutofillAiServerModel,
              "AutofillAiServerModel",
              base::FEATURE_DISABLED_BY_DEFAULT);
 
+// The maximum duration for which an AutofillAI server model response is kept in
+// the local cache. NOTE: It is advisable to choose a value that is at least as
+// large as the cache duration for Autofill server responses to limit cases in
+// which the model is run multiple times for the same form.
+const base::FeatureParam<base::TimeDelta> kAutofillAiServerModelCacheAge{
+    &kAutofillAiServerModel, "autofill_ai_model_cache_age", base::Days(7)};
+
+// The maximum size of the AutofillAI server model cache.
+const base::FeatureParam<int> kAutofillAiServerModelCacheSize{
+    &kAutofillAiServerModel, "autofill_ai_model_cache_size", 100};
+
 // Enables the second iteration AutofillAI.
 // This feature is independent of `autofill_ai::kAutofillAi`.
 BASE_FEATURE(kAutofillAiWithDataSchema,

components/autofill/core/common/autofill_features.h

@@ -26,6 +26,10 @@ BASE_DECLARE_FEATURE(kAutofillAddressUserPerceptionSurvey);
 COMPONENT_EXPORT(AUTOFILL)
 BASE_DECLARE_FEATURE(kAutofillAiServerModel);
 COMPONENT_EXPORT(AUTOFILL)
+extern const base::FeatureParam<base::TimeDelta> kAutofillAiServerModelCacheAge;
+COMPONENT_EXPORT(AUTOFILL)
+extern const base::FeatureParam<int> kAutofillAiServerModelCacheSize;
+COMPONENT_EXPORT(AUTOFILL)
 BASE_DECLARE_FEATURE(kAutofillAiWithDataSchema);
 COMPONENT_EXPORT(AUTOFILL)
 extern const base::FeatureParam<int>