[AF][Traffic Light] Cache for StrikeDatabase

Implemented a cache in form of a [key, StrikeData] map for set or
retrieved values to avoid hitting the disk repeatedly for the same
card.

Bug: 884717
Change-Id: I5e8a5936a54f30a64a919a97314cd04d3874b2a2
Reviewed-on: https://chromium-review.googlesource.com/c/1249130
Commit-Queue: Anne Lim <annelim@google.com>
Reviewed-by: Mathieu Perreault <mathp@chromium.org>
Reviewed-by: Sebastien Seguin-Gagnon <sebsg@chromium.org>
Reviewed-by: Jared Saul <jsaul@google.com>
Cr-Commit-Position: refs/heads/master@{#597235}

components/autofill/core/browser/BUILD.gn

@@ -365,6 +365,7 @@ jumbo_static_library("test_support") {
     "//base/test:test_support",
     "//components/autofill/core/browser",
     "//components/autofill/core/common",
+    "//components/leveldb_proto:test_support",
     "//components/os_crypt",
     "//components/os_crypt:test_support",
     "//components/pref_registry",

components/autofill/core/browser/strike_database.cc

@@ -5,6 +5,7 @@
 #include "components/autofill/core/browser/strike_database.h"
 
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -35,14 +36,14 @@ StrikeDatabase::StrikeDatabase(const base::FilePath& database_dir)
 
 StrikeDatabase::~StrikeDatabase() {}
 
-void StrikeDatabase::GetStrikes(const std::string key,
+void StrikeDatabase::GetStrikes(const std::string& key,
                                 const StrikesCallback& outer_callback) {
   GetStrikeData(key, base::BindRepeating(&StrikeDatabase::OnGetStrikes,
                                          base::Unretained(this),
-                                         std::move(outer_callback)));
+                                         std::move(outer_callback), key));
 }
 
-void StrikeDatabase::AddStrike(const std::string key,
+void StrikeDatabase::AddStrike(const std::string& key,
                                const StrikesCallback& outer_callback) {
   GetStrikeData(key, base::BindRepeating(&StrikeDatabase::OnAddStrike,
                                          base::Unretained(this),
@@ -60,7 +61,11 @@ void StrikeDatabase::ClearAllStrikesForKey(
           leveldb_proto::ProtoDatabase<StrikeData>::KeyEntryVector>(),
       /*keys_to_remove=*/std::move(keys_to_remove),
       base::BindRepeating(&StrikeDatabase::OnClearAllStrikesForKey,
-                          base::Unretained(this), outer_callback));
+                          base::Unretained(this), outer_callback, key));
+}
+
+void StrikeDatabase::ClearCache() {
+  strike_map_cache_.clear();
 }
 
 std::string StrikeDatabase::GetKeyForCreditCardSave(
@@ -70,64 +75,85 @@ std::string StrikeDatabase::GetKeyForCreditCardSave(
 
 void StrikeDatabase::OnDatabaseInit(bool success) {}
 
-void StrikeDatabase::GetStrikeData(const std::string key,
-                                   const GetValueCallback& callback) {
-  db_->GetEntry(key, callback);
+void StrikeDatabase::GetStrikeData(const std::string& key,
+                                   const GetValueCallback& inner_callback) {
+  std::unordered_map<std::string, StrikeData>::iterator it =
+      strike_map_cache_.find(key);
+  if (it != strike_map_cache_.end()) {  // key is in cache
+    StrikeData data_copy(it->second);
+    inner_callback.Run(/*success=*/true,
+                       std::make_unique<StrikeData>(data_copy));
+  } else {
+    db_->GetEntry(key, inner_callback);
+  }
 }
 
 void StrikeDatabase::SetStrikeData(const std::string& key,
                                    const StrikeData& data,
-                                   const SetValueCallback& callback) {
+                                   const SetValueCallback& inner_callback) {
   std::unique_ptr<StrikeDataProto::KeyEntryVector> entries(
       new StrikeDataProto::KeyEntryVector());
   entries->push_back(std::make_pair(key, data));
   db_->UpdateEntries(
       /*entries_to_save=*/std::move(entries),
       /*keys_to_remove=*/std::make_unique<std::vector<std::string>>(),
-      callback);
+      inner_callback);
 }
 
-void StrikeDatabase::OnGetStrikes(StrikesCallback callback,
+void StrikeDatabase::OnGetStrikes(StrikesCallback outer_callback,
+                                  const std::string& key,
                                   bool success,
                                   std::unique_ptr<StrikeData> strike_data) {
-  if (success && strike_data)
-    callback.Run(strike_data->num_strikes());
-  else
-    callback.Run(0);
+  if (success && strike_data) {
+    outer_callback.Run(strike_data->num_strikes());
+    UpdateCache(key, *strike_data.get());
+  } else {
+    outer_callback.Run(0);
+  }
 }
 
-void StrikeDatabase::OnAddStrike(StrikesCallback callback,
-                                 std::string key,
+void StrikeDatabase::OnAddStrike(StrikesCallback outer_callback,
+                                 const std::string& key,
                                  bool success,
                                  std::unique_ptr<StrikeData> strike_data) {
   if (!success) {
-    // Failed to get strike data; abort adding strike.
-    callback.Run(0);
+    outer_callback.Run(0);
     return;
   }
   int num_strikes = strike_data ? strike_data->num_strikes() + 1 : 1;
-  StrikeData data;
-  data.set_num_strikes(num_strikes);
-  data.set_last_update_timestamp(
+  std::unique_ptr<StrikeData> data = std::make_unique<StrikeData>(StrikeData());
+  data->set_num_strikes(num_strikes);
+  data->set_last_update_timestamp(
       base::Time::Now().ToDeltaSinceWindowsEpoch().InMicroseconds());
-  SetStrikeData(
-      key, data,
-      base::BindRepeating(&StrikeDatabase::OnAddStrikeComplete,
-                          base::Unretained(this), callback, num_strikes));
+  SetStrikeData(key, *data.get(),
+                base::BindRepeating(&StrikeDatabase::OnAddStrikeComplete,
+                                    base::Unretained(this), outer_callback, key,
+                                    *data.get()));
 }
 
-void StrikeDatabase::OnAddStrikeComplete(StrikesCallback callback,
-                                         int num_strikes,
+void StrikeDatabase::OnAddStrikeComplete(StrikesCallback outer_callback,
+                                         const std::string& key,
+                                         const StrikeData& data,
                                          bool success) {
-  if (success)
-    callback.Run(num_strikes);
-  else
-    callback.Run(0);
+  if (success) {
+    outer_callback.Run(data.num_strikes());
+    UpdateCache(key, data);
+  } else {
+    outer_callback.Run(0);
+  }
+}
+
+void StrikeDatabase::OnClearAllStrikesForKey(
+    ClearStrikesCallback outer_callback,
+    const std::string& key,
+    bool success) {
+  strike_map_cache_.erase(key);
+  outer_callback.Run(success);
 }
 
-void StrikeDatabase::OnClearAllStrikesForKey(ClearStrikesCallback callback,
-                                             bool success) {
-  callback.Run(success);
+void StrikeDatabase::UpdateCache(const std::string& key,
+                                 const StrikeData& data) {
+  strike_map_cache_[key] = data;
 }
 
 std::string StrikeDatabase::CreateKey(const std::string& type_prefix,

components/autofill/core/browser/strike_database.h

@@ -7,6 +7,7 @@
 
 #include <memory>
 #include <string>
+#include <unordered_map>
 
 #include "base/callback_forward.h"
 #include "base/memory/weak_ptr.h"
@@ -21,6 +22,8 @@ class StrikeData;
 // fails.
 class StrikeDatabase {
  public:
+  using AddToCacheCallback = base::RepeatingCallback<void(bool success)>;
+
   using ClearStrikesCallback = base::RepeatingCallback<void(bool success)>;
 
   using StrikesCallback = base::RepeatingCallback<void(int num_strikes)>;
@@ -39,13 +42,14 @@ class StrikeDatabase {
   // Passes the number of strikes for |key| to |outer_callback|. In the case
   // that the database fails to retrieve the strike update or if no entry is
   // found for |key|, 0 is passed.
-  void GetStrikes(const std::string key, const StrikesCallback& outer_callback);
+  void GetStrikes(const std::string& key,
+                  const StrikesCallback& outer_callback);
 
   // Increments strike count by 1 and passes the updated strike count to the
   // callback. In the case of |key| has no entry, a StrikeData entry with strike
   // count of 1 is added to the database. If the database fails to save or
   // retrieve the strike update, 0 is passed to |outer_callback|.
-  void AddStrike(const std::string key, const StrikesCallback& outer_callback);
+  void AddStrike(const std::string& key, const StrikesCallback& outer_callback);
 
   // Removes database entry for |key|, which implicitly sets strike count to 0.
   void ClearAllStrikesForKey(const std::string& key,
@@ -56,13 +60,23 @@ class StrikeDatabase {
   std::string GetKeyForCreditCardSave(const std::string& card_last_four_digits);
 
  protected:
+  void ClearCache();
+
+  void OnDatabaseInit(bool success);
+
+  // The database for storing strike information.
   std::unique_ptr<leveldb_proto::ProtoDatabase<StrikeData>> db_;
 
+  // Cached data for the keys which have been loaded or recently saved.
+  std::unordered_map<std::string, StrikeData> strike_map_cache_;
+
+  base::WeakPtrFactory<StrikeDatabase> weak_ptr_factory_;
+
  private:
-  void OnDatabaseInit(bool success);
+  friend class StrikeDatabaseTest;
 
   // Passes success status and StrikeData entry for |key| to |inner_callback|.
-  void GetStrikeData(const std::string key,
+  void GetStrikeData(const std::string& key,
                      const GetValueCallback& inner_callback);
 
   // Sets the entry for |key| to |strike_data|. Success status is passed to the
@@ -71,23 +85,30 @@ class StrikeDatabase {
                      const StrikeData& strike_data,
                      const SetValueCallback& inner_callback);
 
-  // Passes number of strikes to |outer_callback|.
+  // Passes number of strikes to |outer_callback| and updates
+  // |strike_map_cache_| with the entry of |key| to |strike_data|.
   void OnGetStrikes(StrikesCallback outer_callback,
+                    const std::string& key,
                     bool success,
                     std::unique_ptr<StrikeData> strike_data);
 
   // Updates database entry for |key| to increment num_strikes by 1, then passes
   // the updated strike count to |outer_callback|.
   void OnAddStrike(StrikesCallback outer_callback,
-                   std::string key,
+                   const std::string& key,
                    bool success,
                    std::unique_ptr<StrikeData> strike_data);
 
+  // Runs |outer_callback| with number of strikes as input, and updates
+  // |strike_map_cache_|.
   void OnAddStrikeComplete(StrikesCallback outer_callback,
-                           int num_strikes,
+                           const std::string& key,
+                           const StrikeData& data,
                            bool success);
 
+  // Removes the entry for |key| in |strike_map_cache_|.
   void OnClearAllStrikesForKey(ClearStrikesCallback outer_callback,
+                               const std::string& key,
                                bool success);
 
   // Concatenates type prefix and identifier suffix to create a key.
@@ -96,7 +117,11 @@ class StrikeDatabase {
 
   std::string GetKeyPrefixForCreditCardSave();
 
-  base::WeakPtrFactory<StrikeDatabase> weak_ptr_factory_;
+  // Sets the entry for |key| in |strike_map_cache_| to |data|.
+  void UpdateCache(const std::string& key, const StrikeData& data);
+
+  FRIEND_TEST_ALL_PREFIXES(StrikeDatabaseTest,
+                           NoDatabaseCallsWhenEntryIsCachedTest);
 };
 
 }  // namespace autofill

components/autofill/core/browser/strike_database_unittest.cc

@@ -11,6 +11,7 @@
 #include "base/run_loop.h"
 #include "base/test/scoped_task_environment.h"
 #include "base/threading/thread_task_runner_handle.h"
+#include "build/build_config.h"
 #include "components/autofill/core/browser/proto/strike_data.pb.h"
 #include "components/autofill/core/browser/test_strike_database.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -19,12 +20,12 @@ namespace autofill {
 
 class StrikeDatabaseTest : public ::testing::Test {
  public:
-  StrikeDatabaseTest() : db_(InitFilePath()) {}
+  StrikeDatabaseTest() : strike_database_(InitFilePath()) {}
 
   void AddEntries(
       std::vector<std::pair<std::string, StrikeData>> entries_to_add) {
     base::RunLoop run_loop;
-    db_.AddEntries(
+    strike_database_.AddEntries(
         entries_to_add,
         base::BindRepeating(&StrikeDatabaseTest::OnAddEntries,
                             base::Unretained(this), run_loop.QuitClosure()));
@@ -42,9 +43,10 @@ class StrikeDatabaseTest : public ::testing::Test {
 
   int GetStrikes(std::string key) {
     base::RunLoop run_loop;
-    db_.GetStrikes(key, base::BindRepeating(&StrikeDatabaseTest::OnGetStrikes,
-                                            base::Unretained(this),
-                                            run_loop.QuitClosure()));
+    strike_database_.GetStrikes(
+        key,
+        base::BindRepeating(&StrikeDatabaseTest::OnGetStrikes,
+                            base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
     return num_strikes_;
   }
@@ -56,9 +58,10 @@ class StrikeDatabaseTest : public ::testing::Test {
 
   int AddStrike(std::string key) {
     base::RunLoop run_loop;
-    db_.AddStrike(key, base::BindRepeating(&StrikeDatabaseTest::OnAddStrike,
-                                           base::Unretained(this),
-                                           run_loop.QuitClosure()));
+    strike_database_.AddStrike(
+        key,
+        base::BindRepeating(&StrikeDatabaseTest::OnAddStrike,
+                            base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
     return num_strikes_;
   }
@@ -70,16 +73,22 @@ class StrikeDatabaseTest : public ::testing::Test {
 
   void ClearAllStrikesForKey(const std::string key) {
     base::RunLoop run_loop;
-    db_.ClearAllStrikesForKey(
+    strike_database_.ClearAllStrikesForKey(
         key,
         base::BindRepeating(&StrikeDatabaseTest::OnClearAllStrikesForKey,
                             base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
   }
 
+  void ClearCache() { strike_database_.ClearCache(); }
+
+  int GetNumberOfDatabaseCalls() {
+    return strike_database_.GetNumberOfDatabaseCalls();
+  }
+
  protected:
   base::test::ScopedTaskEnvironment scoped_task_environment_;
-  TestStrikeDatabase db_;
+  TestStrikeDatabase strike_database_;
 
  private:
   static const base::FilePath InitFilePath() {
@@ -94,18 +103,19 @@ class StrikeDatabaseTest : public ::testing::Test {
   std::unique_ptr<StrikeData> strike_data_;
 };
 
+#if !defined(OS_IOS)
+
 TEST_F(StrikeDatabaseTest, AddStrikeTest) {
   const std::string key = "12345";
-  int strikes = AddStrike(key);
-  EXPECT_EQ(1, strikes);
-  strikes = AddStrike(key);
-  EXPECT_EQ(2, strikes);
+  EXPECT_EQ(1, AddStrike(key));
+  EXPECT_EQ(2, AddStrike(key));
+  EXPECT_EQ(2, GetStrikes(key));
+  EXPECT_EQ(3, AddStrike(key));
 }
 
 TEST_F(StrikeDatabaseTest, GetStrikeForZeroStrikesTest) {
   const std::string key = "12345";
-  int strikes = GetStrikes(key);
-  EXPECT_EQ(0, strikes);
+  EXPECT_EQ(0, GetStrikes(key));
 }
 
 TEST_F(StrikeDatabaseTest, GetStrikeForNonZeroStrikesTest) {
@@ -117,15 +127,13 @@ TEST_F(StrikeDatabaseTest, GetStrikeForNonZeroStrikesTest) {
   entries.push_back(std::make_pair(key, data));
   AddEntries(entries);
 
-  int strikes = GetStrikes(key);
-  EXPECT_EQ(3, strikes);
+  EXPECT_EQ(3, GetStrikes(key));
 }
 
 TEST_F(StrikeDatabaseTest, ClearStrikesForZeroStrikesTest) {
   const std::string key = "12345";
   ClearAllStrikesForKey(key);
-  int strikes = GetStrikes(key);
-  EXPECT_EQ(0, strikes);
+  EXPECT_EQ(0, GetStrikes(key));
 }
 
 TEST_F(StrikeDatabaseTest, ClearStrikesForNonZeroStrikesTest) {
@@ -137,11 +145,9 @@ TEST_F(StrikeDatabaseTest, ClearStrikesForNonZeroStrikesTest) {
   entries.push_back(std::make_pair(key, data));
   AddEntries(entries);
 
-  int strikes = GetStrikes(key);
-  EXPECT_EQ(3, strikes);
+  EXPECT_EQ(3, GetStrikes(key));
   ClearAllStrikesForKey(key);
-  strikes = GetStrikes(key);
-  EXPECT_EQ(0, strikes);
+  EXPECT_EQ(0, GetStrikes(key));
 }
 
 TEST_F(StrikeDatabaseTest, ClearStrikesForMultipleNonZeroStrikesEntriesTest) {
@@ -158,20 +164,44 @@ TEST_F(StrikeDatabaseTest, ClearStrikesForMultipleNonZeroStrikesEntriesTest) {
   entries.push_back(std::make_pair(key2, data2));
   AddEntries(entries);
 
-  int strikes = GetStrikes(key1);
-  EXPECT_EQ(3, strikes);
-  strikes = GetStrikes(key2);
-  EXPECT_EQ(5, strikes);
+  EXPECT_EQ(3, GetStrikes(key1));
+  EXPECT_EQ(5, GetStrikes(key2));
   ClearAllStrikesForKey(key1);
-  strikes = GetStrikes(key1);
-  EXPECT_EQ(0, strikes);
-  strikes = GetStrikes(key2);
-  EXPECT_EQ(5, strikes);
+  EXPECT_EQ(0, GetStrikes(key1));
+  EXPECT_EQ(5, GetStrikes(key2));
+}
+
+TEST_F(StrikeDatabaseTest, NoDatabaseCallsWhenEntryIsCachedTest) {
+  ClearCache();
+  int database_calls = GetNumberOfDatabaseCalls();
+  EXPECT_EQ(0, database_calls);
+
+  // Set up database with 3 pre-existing strikes at |key|.
+  const std::string key = "12345";
+  std::vector<std::pair<std::string, StrikeData>> entries;
+  StrikeData data;
+  data.set_num_strikes(3);
+  entries.push_back(std::make_pair(key, data));
+  AddEntries(entries);
+  EXPECT_EQ(1, GetNumberOfDatabaseCalls());
+  // Calling GetStrikes(~) should add its returned value to the cache.
+  EXPECT_EQ(3, GetStrikes(key));
+  EXPECT_EQ(2, GetNumberOfDatabaseCalls());
+  // GetStrikes(~) should not hit the db because it's already cached.
+  GetStrikes(key);
+  EXPECT_EQ(2, GetNumberOfDatabaseCalls());
+  ClearCache();
+  // GetStrikes(~) will hit the db because the cache was cleared.
+  GetStrikes(key);
+  EXPECT_EQ(3, GetNumberOfDatabaseCalls());
 }
 
 TEST_F(StrikeDatabaseTest, GetKeyForCreditCardSave) {
   const std::string last_four = "1234";
-  EXPECT_EQ("creditCardSave__1234", db_.GetKeyForCreditCardSave(last_four));
+  EXPECT_EQ("creditCardSave__1234",
+            strike_database_.GetKeyForCreditCardSave(last_four));
 }
 
+#endif
+
 }  // namespace autofill

components/autofill/core/browser/test_strike_database.cc

@@ -4,12 +4,27 @@
 
 #include "components/autofill/core/browser/test_strike_database.h"
 
+#include "base/task/post_task.h"
 #include "components/autofill/core/browser/proto/strike_data.pb.h"
+#include "components/leveldb_proto/testing/test_proto_database_impl.h"
 
 namespace autofill {
 
+namespace {
+const char kDatabaseClientName[] = "TestStrikeService";
+}  // namespace
+
 TestStrikeDatabase::TestStrikeDatabase(const base::FilePath& database_dir)
-    : StrikeDatabase(database_dir) {}
+    : StrikeDatabase(database_dir) {
+  db_ = std::make_unique<leveldb_proto::TestProtoDatabaseImpl<StrikeData>>(
+      base::CreateSequencedTaskRunnerWithTraits(
+          {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
+           base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN}));
+  db_->Init(kDatabaseClientName, database_dir,
+            leveldb_proto::CreateSimpleOptions(),
+            base::BindRepeating(&TestStrikeDatabase::OnDatabaseInit,
+                                weak_ptr_factory_.GetWeakPtr()));
+}
 
 void TestStrikeDatabase::AddEntries(
     std::vector<std::pair<std::string, StrikeData>> entries_to_add,
@@ -25,4 +40,10 @@ void TestStrikeDatabase::AddEntries(
       callback);
 }
 
+int TestStrikeDatabase::GetNumberOfDatabaseCalls() {
+  leveldb_proto::TestProtoDatabaseImpl<StrikeData>* test_db =
+      static_cast<leveldb_proto::TestProtoDatabaseImpl<StrikeData>*>(db_.get());
+  return test_db->number_of_db_calls();
+}
+
 }  // namespace autofill

components/autofill/core/browser/test_strike_database.h

@@ -11,8 +11,10 @@
 #include <vector>
 
 #include "components/autofill/core/browser/strike_database.h"
+#include "components/leveldb_proto/testing/test_proto_database_impl.h"
 
 namespace autofill {
+class StrikeData;
 
 class TestStrikeDatabase : public StrikeDatabase {
  public:
@@ -21,6 +23,8 @@ class TestStrikeDatabase : public StrikeDatabase {
   void AddEntries(
       std::vector<std::pair<std::string, StrikeData>> entries_to_add,
       const SetValueCallback& callback);
+
+  int GetNumberOfDatabaseCalls();
 };
 
 }  // namespace autofill

components/leveldb_proto/BUILD.gn

@@ -21,6 +21,7 @@ source_set("test_support") {
   testonly = true
   sources = [
     "testing/fake_db.h",
+    "testing/test_proto_database_impl.h",
   ]
 
   public_deps = [

components/leveldb_proto/testing/test_proto_database_impl.h

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef COMPONENTS_LEVELDB_PROTO_TESTING_TEST_PROTO_DATABASE_IMPL_H_
+#define COMPONENTS_LEVELDB_PROTO_TESTING_TEST_PROTO_DATABASE_IMPL_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "base/sequenced_task_runner.h"
+#include "components/leveldb_proto/proto_database_impl.h"
+
+namespace leveldb_proto {
+
+// TestProtoDatabaseImpl is a wrapper of ProtoDataBaseImpl, with the additional
+// functionality of counting the number of times the database is hit, which can
+// be accessed with GetNumberOfDatabaseCalls().
+template <typename T>
+class TestProtoDatabaseImpl : public ProtoDatabaseImpl<T> {
+ public:
+  explicit TestProtoDatabaseImpl(
+      const scoped_refptr<base::SequencedTaskRunner>& task_runner);
+
+  void UpdateEntries(
+      std::unique_ptr<typename ProtoDatabase<T>::KeyEntryVector>
+          entries_to_save,
+      std::unique_ptr<KeyVector> keys_to_remove,
+      typename ProtoDatabase<T>::UpdateCallback callback) override;
+  void UpdateEntriesWithRemoveFilter(
+      std::unique_ptr<typename ProtoDatabase<T>::KeyEntryVector>
+          entries_to_save,
+      const LevelDB::KeyFilter& delete_key_filter,
+      typename ProtoDatabase<T>::UpdateCallback callback) override;
+  void LoadEntries(typename ProtoDatabase<T>::LoadCallback callback) override;
+  void LoadEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      typename ProtoDatabase<T>::LoadCallback callback) override;
+  void LoadEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      const leveldb::ReadOptions& options,
+      const std::string& target_prefix,
+      typename ProtoDatabase<T>::LoadCallback callback) override;
+  void LoadKeys(typename ProtoDatabase<T>::LoadKeysCallback callback) override;
+  void GetEntry(const std::string& key,
+                typename ProtoDatabase<T>::GetCallback callback) override;
+  void Destroy(typename ProtoDatabase<T>::DestroyCallback callback) override;
+
+  int number_of_db_calls();
+
+ private:
+  int IncrementDatabaseCalls();
+
+  // The number of times the database is called.
+  int db_calls_ = 0;
+};
+
+template <typename T>
+TestProtoDatabaseImpl<T>::TestProtoDatabaseImpl(
+    const scoped_refptr<base::SequencedTaskRunner>& task_runner)
+    : ProtoDatabaseImpl<T>(task_runner) {}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::UpdateEntries(
+    std::unique_ptr<typename ProtoDatabase<T>::KeyEntryVector> entries_to_save,
+    std::unique_ptr<KeyVector> keys_to_remove,
+    typename ProtoDatabase<T>::UpdateCallback callback) {
+  IncrementDatabaseCalls();
+  leveldb_proto::ProtoDatabaseImpl<T>::UpdateEntries(std::move(entries_to_save),
+                                                     std::move(keys_to_remove),
+                                                     std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::UpdateEntriesWithRemoveFilter(
+    std::unique_ptr<typename ProtoDatabase<T>::KeyEntryVector> entries_to_save,
+    const LevelDB::KeyFilter& delete_key_filter,
+    typename ProtoDatabase<T>::UpdateCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::UpdateEntriesWithRemoveFilter(
+      std::move(entries_to_save), delete_key_filter, std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::LoadEntries(
+    typename ProtoDatabase<T>::LoadCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::LoadEntries(std::move(callback));
+}
+
+template <typename T>
+void leveldb_proto::TestProtoDatabaseImpl<T>::LoadEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    typename ProtoDatabase<T>::LoadCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::LoadEntriesWithFilter(key_filter, std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::LoadEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    const leveldb::ReadOptions& options,
+    const std::string& target_prefix,
+    typename ProtoDatabase<T>::LoadCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::LoadEntriesWithFilter(
+      key_filter, options, target_prefix, std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::LoadKeys(
+    typename ProtoDatabase<T>::LoadKeysCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::LoadKeys(std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::GetEntry(
+    const std::string& key,
+    typename ProtoDatabase<T>::GetCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::GetEntry(key, std::move(callback));
+}
+
+template <typename T>
+void TestProtoDatabaseImpl<T>::Destroy(
+    typename ProtoDatabase<T>::DestroyCallback callback) {
+  IncrementDatabaseCalls();
+  ProtoDatabaseImpl<T>::Destroy(std::move(callback));
+}
+
+template <typename T>
+int TestProtoDatabaseImpl<T>::number_of_db_calls() {
+  return db_calls_;
+}
+
+template <typename T>
+int TestProtoDatabaseImpl<T>::IncrementDatabaseCalls() {
+  return ++db_calls_;
+}
+
+}  // namespace leveldb_proto
+
+#endif  // COMPONENTS_LEVELDB_PROTO_TESTING_TEST_PROTO_DATABASE_IMPL_H_