Extended the ProtoDatabase to provide LoadKeys() functionality.

Before this change, the only way to scan all keys was to call LoadEntries() which loads the whole database. In our use case, this means copying a lot of image data which is wasteful.
With LoadKeys(), we still need to iterate the whole database but we don't have to copy out (and parse) the values -- just the keys. LoadKeys() is useful when storing relatively large values and only access to the keys is required.

We'll use LoadKeys() for NTP content suggestions to garbage collect data. We do that at times when we have a list of all still alive elements and need to intersect that with the elements stored in the db.

BUG=649009

Review-Url: https://codereview.chromium.org/2379113002
Cr-Commit-Position: refs/heads/master@{#422772}

components/leveldb_proto/leveldb_database.cc

@@ -138,6 +138,21 @@ bool LevelDB::Load(std::vector<std::string>* entries) {
   return true;
 }
 
+bool LevelDB::LoadKeys(std::vector<std::string>* keys) {
+  DFAKE_SCOPED_LOCK(thread_checker_);
+  if (!db_)
+    return false;
+
+  leveldb::ReadOptions options;
+  options.fill_cache = false;
+  std::unique_ptr<leveldb::Iterator> db_iterator(db_->NewIterator(options));
+  for (db_iterator->SeekToFirst(); db_iterator->Valid(); db_iterator->Next()) {
+    leveldb::Slice key_slice = db_iterator->key();
+    keys->push_back(std::string(key_slice.data(), key_slice.size()));
+  }
+  return true;
+}
+
 bool LevelDB::Get(const std::string& key, bool* found, std::string* entry) {
   DFAKE_SCOPED_LOCK(thread_checker_);
   if (!db_)

components/leveldb_proto/leveldb_database.h

@@ -45,6 +45,7 @@ class LevelDB : base::trace_event::MemoryDumpProvider {
   virtual bool Save(const base::StringPairs& pairs_to_save,
                     const std::vector<std::string>& keys_to_remove);
   virtual bool Load(std::vector<std::string>* entries);
+  virtual bool LoadKeys(std::vector<std::string>* keys);
   virtual bool Get(const std::string& key, bool* found, std::string* entry);
 
   static bool Destroy(const base::FilePath& database_dir);

components/leveldb_proto/proto_database.h

@@ -27,6 +27,9 @@ class ProtoDatabase {
   using UpdateCallback = base::Callback<void(bool success)>;
   using LoadCallback =
       base::Callback<void(bool success, std::unique_ptr<std::vector<T>>)>;
+  using LoadKeysCallback =
+      base::Callback<void(bool success,
+                          std::unique_ptr<std::vector<std::string>>)>;
   using GetCallback = base::Callback<void(bool success, std::unique_ptr<T>)>;
   using DestroyCallback = base::Callback<void(bool success)>;
 
@@ -53,6 +56,10 @@ class ProtoDatabase {
   // when complete.
   virtual void LoadEntries(const LoadCallback& callback) = 0;
 
+  // Asynchronously loads all keys from the database and invokes |callback| with
+  // those keys when complete.
+  virtual void LoadKeys(const LoadKeysCallback& callback) = 0;
+
   // Asynchronously loads a single entry, identified by |key|, from the database
   // and invokes |callback| when complete. If no entry with |key| is found,
   // a nullptr is passed to the callback, but the success flag is still true.

components/leveldb_proto/proto_database_impl.h

@@ -53,6 +53,8 @@ class ProtoDatabaseImpl : public ProtoDatabase<T> {
       const typename ProtoDatabase<T>::UpdateCallback& callback) override;
   void LoadEntries(
       const typename ProtoDatabase<T>::LoadCallback& callback) override;
+  void LoadKeys(
+      const typename ProtoDatabase<T>::LoadKeysCallback& callback) override;
   void GetEntry(
       const std::string& key,
       const typename ProtoDatabase<T>::GetCallback& callback) override;
@@ -94,11 +96,19 @@ void RunUpdateCallback(
 
 template <typename T>
 void RunLoadCallback(const typename ProtoDatabase<T>::LoadCallback& callback,
-                     const bool* success,
+                     bool* success,
                      std::unique_ptr<std::vector<T>> entries) {
   callback.Run(*success, std::move(entries));
 }
 
+template <typename T>
+void RunLoadKeysCallback(
+    const typename ProtoDatabase<T>::LoadKeysCallback& callback,
+    std::unique_ptr<bool> success,
+    std::unique_ptr<std::vector<std::string>> keys) {
+  callback.Run(*success, std::move(keys));
+}
+
 template <typename T>
 void RunGetCallback(const typename ProtoDatabase<T>::GetCallback& callback,
                     const bool* success,
@@ -171,6 +181,15 @@ void LoadEntriesFromTaskRunner(LevelDB* database,
   }
 }
 
+void LoadKeysFromTaskRunner(LevelDB* database,
+                            std::vector<std::string>* keys,
+                            bool* success) {
+  DCHECK(success);
+  DCHECK(keys);
+  keys->clear();
+  *success = database->LoadKeys(keys);
+}
+
 template <typename T>
 void GetEntryFromTaskRunner(LevelDB* database,
                             const std::string& key,
@@ -285,6 +304,21 @@ void ProtoDatabaseImpl<T>::LoadEntries(
                  base::Passed(&entries)));
 }
 
+template <typename T>
+void ProtoDatabaseImpl<T>::LoadKeys(
+    const typename ProtoDatabase<T>::LoadKeysCallback& callback) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  auto success = base::MakeUnique<bool>(false);
+  auto keys = base::MakeUnique<std::vector<std::string>>();
+  auto load_task =
+      base::Bind(LoadKeysFromTaskRunner, base::Unretained(db_.get()),
+                 keys.get(), success.get());
+  task_runner_->PostTaskAndReply(
+      FROM_HERE, load_task,
+      base::Bind(RunLoadKeysCallback<T>, callback, base::Passed(&success),
+                 base::Passed(&keys)));
+}
+
 template <typename T>
 void ProtoDatabaseImpl<T>::GetEntry(
     const std::string& key,

components/leveldb_proto/proto_database_impl_unittest.cc

@@ -29,6 +29,7 @@ using base::MessageLoop;
 using base::ScopedTempDir;
 using testing::Invoke;
 using testing::Return;
+using testing::UnorderedElementsAre;
 using testing::_;
 
 namespace leveldb_proto {
@@ -248,6 +249,56 @@ TEST_F(ProtoDatabaseImplTest, TestDBGetSuccess) {
   base::RunLoop().RunUntilIdle();
 }
 
+TEST(ProtoDatabaseImplLevelDBTest, TestDBSaveAndLoadKeys) {
+  base::MessageLoop main_loop;
+
+  ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+  base::Thread db_thread("dbthread");
+  ASSERT_TRUE(db_thread.Start());
+  std::unique_ptr<ProtoDatabaseImpl<TestProto>> db(
+      new ProtoDatabaseImpl<TestProto>(db_thread.task_runner()));
+
+  auto expect_init_success =
+      base::Bind([](bool success) { EXPECT_TRUE(success); });
+  db->Init(kTestLevelDBClientName, temp_dir.GetPath(), expect_init_success);
+
+  base::RunLoop run_update_entries;
+  auto expect_update_success = base::Bind(
+      [](base::Closure signal, bool success) {
+        EXPECT_TRUE(success);
+        signal.Run();
+      },
+      run_update_entries.QuitClosure());
+  TestProto test_proto;
+  test_proto.set_data("some data");
+  ProtoDatabase<TestProto>::KeyEntryVector data_set(
+          {{"0", test_proto}, {"1", test_proto}, {"2", test_proto}});
+  db->UpdateEntries(
+      base::MakeUnique<ProtoDatabase<TestProto>::KeyEntryVector>(data_set),
+      base::MakeUnique<std::vector<std::string>>(), expect_update_success);
+  run_update_entries.Run();
+
+  base::RunLoop run_load_keys;
+  auto verify_loaded_keys = base::Bind(
+      [](base::Closure signal, bool success,
+         std::unique_ptr<std::vector<std::string>> keys) {
+        EXPECT_TRUE(success);
+        EXPECT_THAT(*keys, UnorderedElementsAre("0", "1", "2"));
+        signal.Run();
+      },
+      run_load_keys.QuitClosure());
+  db->LoadKeys(verify_loaded_keys);
+  run_load_keys.Run();
+
+  // Shutdown database.
+  db.reset();
+  base::RunLoop run_destruction;
+  db_thread.task_runner()->PostTaskAndReply(
+      FROM_HERE, base::Bind(base::DoNothing), run_destruction.QuitClosure());
+  run_destruction.Run();
+}
+
 TEST_F(ProtoDatabaseImplTest, TestDBGetNotFound) {
   base::FilePath path(FILE_PATH_LITERAL("/fake/path"));
 

components/leveldb_proto/testing/fake_db.h

@@ -39,6 +39,8 @@ class FakeDB : public ProtoDatabase<T> {
       const typename ProtoDatabase<T>::UpdateCallback& callback) override;
   void LoadEntries(
       const typename ProtoDatabase<T>::LoadCallback& callback) override;
+   void LoadKeys(
+      const typename ProtoDatabase<T>::LoadKeysCallback& callback) override;
   void GetEntry(
       const std::string& key,
       const typename ProtoDatabase<T>::GetCallback& callback) override;
@@ -51,6 +53,8 @@ class FakeDB : public ProtoDatabase<T> {
 
   void LoadCallback(bool success);
 
+  void LoadKeysCallback(bool success);
+
   void GetCallback(bool success);
 
   void UpdateCallback(bool success);
@@ -63,6 +67,11 @@ class FakeDB : public ProtoDatabase<T> {
       std::unique_ptr<typename std::vector<T>> entries,
       bool success);
 
+  static void RunLoadKeysCallback(
+      const typename ProtoDatabase<T>::LoadKeysCallback& callback,
+      std::unique_ptr<std::vector<std::string>> keys,
+      bool success);
+
   static void RunGetCallback(
       const typename ProtoDatabase<T>::GetCallback& callback,
       std::unique_ptr<T> entry,
@@ -73,6 +82,7 @@ class FakeDB : public ProtoDatabase<T> {
 
   Callback init_callback_;
   Callback load_callback_;
+  Callback load_keys_callback_;
   Callback get_callback_;
   Callback update_callback_;
 };
@@ -117,6 +127,18 @@ void FakeDB<T>::LoadEntries(
       base::Bind(RunLoadCallback, callback, base::Passed(&entries));
 }
 
+template <typename T>
+void FakeDB<T>::LoadKeys(
+    const typename ProtoDatabase<T>::LoadKeysCallback& callback) {
+  std::unique_ptr<std::vector<std::string>> keys(
+      new std::vector<std::string>());
+  for (const auto& pair : *db_)
+    keys->push_back(pair.first);
+
+  load_keys_callback_ =
+      base::Bind(RunLoadKeysCallback, callback, base::Passed(&keys));
+}
+
 template <typename T>
 void FakeDB<T>::GetEntry(
     const std::string& key,
@@ -151,6 +173,12 @@ void FakeDB<T>::LoadCallback(bool success) {
   load_callback_.Reset();
 }
 
+template <typename T>
+void FakeDB<T>::LoadKeysCallback(bool success) {
+  load_keys_callback_.Run(success);
+  load_keys_callback_.Reset();
+}
+
 template <typename T>
 void FakeDB<T>::GetCallback(bool success) {
   get_callback_.Run(success);
@@ -172,6 +200,15 @@ void FakeDB<T>::RunLoadCallback(
   callback.Run(success, std::move(entries));
 }
 
+// static
+template <typename T>
+void FakeDB<T>::RunLoadKeysCallback(
+    const typename ProtoDatabase<T>::LoadKeysCallback& callback,
+    std::unique_ptr<std::vector<std::string>> keys,
+    bool success) {
+  callback.Run(success, std::move(keys));
+}
+
 // static
 template <typename T>
 void FakeDB<T>::RunGetCallback(