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}
diff --git a/components/leveldb_proto/leveldb_database.cc b/components/leveldb_proto/leveldb_database.cc
index 0dfb5e3..5769710 100644
--- a/components/leveldb_proto/leveldb_database.cc
+++ b/components/leveldb_proto/leveldb_database.cc
@@ -138,6 +138,21 @@
   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_)
diff --git a/components/leveldb_proto/leveldb_database.h b/components/leveldb_proto/leveldb_database.h
index 4474a87..8c984b7 100644
--- a/components/leveldb_proto/leveldb_database.h
+++ b/components/leveldb_proto/leveldb_database.h
@@ -45,6 +45,7 @@
   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);
diff --git a/components/leveldb_proto/proto_database.h b/components/leveldb_proto/proto_database.h
index 20c630e..081233c 100644
--- a/components/leveldb_proto/proto_database.h
+++ b/components/leveldb_proto/proto_database.h
@@ -27,6 +27,9 @@
   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 @@
   // 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.
diff --git a/components/leveldb_proto/proto_database_impl.h b/components/leveldb_proto/proto_database_impl.h
index da02eed1..9997d38 100644
--- a/components/leveldb_proto/proto_database_impl.h
+++ b/components/leveldb_proto/proto_database_impl.h
@@ -53,6 +53,8 @@
       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,12 +96,20 @@
 
 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,
                     const bool* found,
@@ -171,6 +181,15 @@
   }
 }
 
+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,
@@ -286,6 +305,21 @@
 }
 
 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,
     const typename ProtoDatabase<T>::GetCallback& callback) {
diff --git a/components/leveldb_proto/proto_database_impl_unittest.cc b/components/leveldb_proto/proto_database_impl_unittest.cc
index bbbeab32..63f9f8836 100644
--- a/components/leveldb_proto/proto_database_impl_unittest.cc
+++ b/components/leveldb_proto/proto_database_impl_unittest.cc
@@ -29,6 +29,7 @@
 using base::ScopedTempDir;
 using testing::Invoke;
 using testing::Return;
+using testing::UnorderedElementsAre;
 using testing::_;
 
 namespace leveldb_proto {
@@ -248,6 +249,56 @@
   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"));
 
diff --git a/components/leveldb_proto/testing/fake_db.h b/components/leveldb_proto/testing/fake_db.h
index e67ceba..bd9efb57 100644
--- a/components/leveldb_proto/testing/fake_db.h
+++ b/components/leveldb_proto/testing/fake_db.h
@@ -39,6 +39,8 @@
       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 @@
 
   void LoadCallback(bool success);
 
+  void LoadKeysCallback(bool success);
+
   void GetCallback(bool success);
 
   void UpdateCallback(bool success);
@@ -63,6 +67,11 @@
       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 @@
 
   Callback init_callback_;
   Callback load_callback_;
+  Callback load_keys_callback_;
   Callback get_callback_;
   Callback update_callback_;
 };
@@ -118,6 +128,18 @@
 }
 
 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,
     const typename ProtoDatabase<T>::GetCallback& callback) {
@@ -152,6 +174,12 @@
 }
 
 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);
   get_callback_.Reset();
@@ -174,6 +202,15 @@
 
 // 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(
     const typename ProtoDatabase<T>::GetCallback& callback,
     std::unique_ptr<T> entry,