Add ability to load keys/entries at same time in LevelDB/ProtoDB.

Keys and entries can currently be loaded separately, but never at the
same time, meaning it requires separate calls to LoadKeys/LoadEntries
to get keys and entries. Some users of ProtoDatabase have resorted to
storing the keys in the entries to get around this.

This CL adds LoadKeysAndEntries(WithFilter) that returns both at the
same time. ProtoDatabase's new LoadKeysAndEntriesCallback gives back a
vector of string/proto pairs.

Bug: 883409
Change-Id: Id624bbdeb356395c6ce679528d85c1f934a9772e
Reviewed-on: https://chromium-review.googlesource.com/1228305
Commit-Queue: Troy Hildebrandt <thildebr@chromium.org>
Reviewed-by: Tommy Nyquist <nyquist@chromium.org>
Reviewed-by: David Trainor <dtrainor@chromium.org>
Cr-Commit-Position: refs/heads/master@{#595175}

components/leveldb_proto/leveldb_database.cc

@@ -4,6 +4,7 @@
 
 #include "components/leveldb_proto/leveldb_database.h"
 
+#include <map>
 #include <string>
 #include <vector>
 
@@ -178,6 +179,34 @@ bool LevelDB::LoadWithFilter(const KeyFilter& filter,
                              std::vector<std::string>* entries,
                              const leveldb::ReadOptions& options,
                              const std::string& target_prefix) {
+  std::map<std::string, std::string> keys_entries;
+  bool result = LoadKeysAndEntriesWithFilter(filter, &keys_entries, options,
+                                             target_prefix);
+  if (!result)
+    return false;
+
+  for (const auto& pair : keys_entries)
+    entries->push_back(pair.second);
+  return true;
+}
+
+bool LevelDB::LoadKeysAndEntries(
+    std::map<std::string, std::string>* keys_entries) {
+  return LoadKeysAndEntriesWithFilter(KeyFilter(), keys_entries);
+}
+
+bool LevelDB::LoadKeysAndEntriesWithFilter(
+    const KeyFilter& filter,
+    std::map<std::string, std::string>* keys_entries) {
+  return LoadKeysAndEntriesWithFilter(filter, keys_entries,
+                                      leveldb::ReadOptions(), std::string());
+}
+
+bool LevelDB::LoadKeysAndEntriesWithFilter(
+    const KeyFilter& filter,
+    std::map<std::string, std::string>* keys_entries,
+    const leveldb::ReadOptions& options,
+    const std::string& target_prefix) {
   DFAKE_SCOPED_LOCK(thread_checker_);
   if (!db_)
     return false;
@@ -187,30 +216,30 @@ bool LevelDB::LoadWithFilter(const KeyFilter& filter,
   for (db_iterator->Seek(target);
        db_iterator->Valid() && db_iterator->key().starts_with(target);
        db_iterator->Next()) {
-    if (!filter.is_null()) {
-      leveldb::Slice key_slice = db_iterator->key();
-      if (!filter.Run(std::string(key_slice.data(), key_slice.size())))
-        continue;
+    leveldb::Slice key_slice = db_iterator->key();
+    std::string key_slice_str(key_slice.data(), key_slice.size());
+    if (!filter.is_null() && !filter.Run(key_slice_str)) {
+      continue;
     }
+
     leveldb::Slice value_slice = db_iterator->value();
-    std::string entry(value_slice.data(), value_slice.size());
-    entries->push_back(entry);
+    keys_entries->insert(std::make_pair(
+        key_slice_str, std::string(value_slice.data(), value_slice.size())));
   }
   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()));
-  }
+  std::map<std::string, std::string> keys_entries;
+  bool result = LoadKeysAndEntriesWithFilter(KeyFilter(), &keys_entries,
+                                             options, std::string());
+  if (!result)
+    return false;
+
+  for (const auto& pair : keys_entries)
+    keys->push_back(pair.first);
   return true;
 }
 

components/leveldb_proto/leveldb_database.h

@@ -5,6 +5,7 @@
 #ifndef COMPONENTS_LEVELDB_PROTO_LEVELDB_DATABASE_H_
 #define COMPONENTS_LEVELDB_PROTO_LEVELDB_DATABASE_H_
 
+#include <map>
 #include <memory>
 #include <string>
 #include <vector>
@@ -55,6 +56,7 @@ class LevelDB {
                     const std::vector<std::string>& keys_to_remove);
   virtual bool UpdateWithRemoveFilter(const base::StringPairs& entries_to_save,
                                       const KeyFilter& delete_key_filter);
+
   virtual bool Load(std::vector<std::string>* entries);
   virtual bool LoadWithFilter(const KeyFilter& filter,
                               std::vector<std::string>* entries);
@@ -62,6 +64,18 @@ class LevelDB {
                               std::vector<std::string>* entries,
                               const leveldb::ReadOptions& options,
                               const std::string& target_prefix);
+
+  virtual bool LoadKeysAndEntries(
+      std::map<std::string, std::string>* keys_entries);
+  virtual bool LoadKeysAndEntriesWithFilter(
+      const KeyFilter& filter,
+      std::map<std::string, std::string>* keys_entries);
+  virtual bool LoadKeysAndEntriesWithFilter(
+      const KeyFilter& filter,
+      std::map<std::string, std::string>* keys_entries,
+      const leveldb::ReadOptions& options,
+      const std::string& target_prefix);
+
   virtual bool LoadKeys(std::vector<std::string>* keys);
   virtual bool Get(const std::string& key, bool* found, std::string* entry);
   // Close (if currently open) and then destroy (i.e. delete) the database

components/leveldb_proto/proto_database.h

@@ -5,6 +5,7 @@
 #ifndef COMPONENTS_LEVELDB_PROTO_PROTO_DATABASE_H_
 #define COMPONENTS_LEVELDB_PROTO_PROTO_DATABASE_H_
 
+#include <map>
 #include <memory>
 #include <string>
 #include <utility>
@@ -32,6 +33,9 @@ class ProtoDatabase {
   using LoadKeysCallback =
       base::OnceCallback<void(bool success,
                               std::unique_ptr<std::vector<std::string>>)>;
+  using LoadKeysAndEntriesCallback =
+      base::OnceCallback<void(bool success,
+                              std::unique_ptr<std::map<std::string, T>>)>;
   using GetCallback =
       base::OnceCallback<void(bool success, std::unique_ptr<T>)>;
   using DestroyCallback = base::OnceCallback<void(bool success)>;
@@ -79,6 +83,17 @@ class ProtoDatabase {
                                      const std::string& target_prefix,
                                      LoadCallback callback) = 0;
 
+  virtual void LoadKeysAndEntries(
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) = 0;
+  virtual void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) = 0;
+  virtual void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      const leveldb::ReadOptions& options,
+      const std::string& target_prefix,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) = 0;
+
   // Asynchronously loads all keys from the database and invokes |callback| with
   // those keys when complete.
   virtual void LoadKeys(LoadKeysCallback callback) = 0;

components/leveldb_proto/proto_database_impl.h

@@ -64,6 +64,16 @@ class ProtoDatabaseImpl : public ProtoDatabase<T> {
       const leveldb::ReadOptions& options,
       const std::string& target_prefix,
       typename ProtoDatabase<T>::LoadCallback callback) override;
+  void LoadKeysAndEntries(
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
+  void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
+  void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      const leveldb::ReadOptions& options,
+      const std::string& target_prefix,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
   void LoadKeys(typename ProtoDatabase<T>::LoadKeysCallback callback) override;
   void GetEntry(const std::string& key,
                 typename ProtoDatabase<T>::GetCallback callback) override;
@@ -109,6 +119,14 @@ void RunLoadCallback(typename ProtoDatabase<T>::LoadCallback callback,
   std::move(callback).Run(*success, std::move(entries));
 }
 
+template <typename T>
+void RunLoadKeysAndEntriesCallback(
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback,
+    bool* success,
+    std::unique_ptr<std::map<std::string, T>> keys_entries) {
+  std::move(callback).Run(*success, std::move(keys_entries));
+}
+
 template <typename T>
 void RunLoadKeysCallback(typename ProtoDatabase<T>::LoadKeysCallback callback,
                          std::unique_ptr<bool> success,
@@ -184,32 +202,48 @@ void UpdateEntriesWithRemoveFilterFromTaskRunner(
 }
 
 template <typename T>
-void LoadEntriesFromTaskRunner(LevelDB* database,
-                               const LevelDB::KeyFilter& filter,
-                               const leveldb::ReadOptions& options,
-                               const std::string& target_prefix,
-                               std::vector<T>* entries,
-                               bool* success) {
+void LoadKeysAndEntriesFromTaskRunner(LevelDB* database,
+                                      const LevelDB::KeyFilter& filter,
+                                      const leveldb::ReadOptions& options,
+                                      const std::string& target_prefix,
+                                      std::map<std::string, T>* keys_entries,
+                                      bool* success) {
   DCHECK(success);
-  DCHECK(entries);
+  DCHECK(keys_entries);
 
-  entries->clear();
+  keys_entries->clear();
 
-  std::vector<std::string> loaded_entries;
-  *success =
-      database->LoadWithFilter(filter, &loaded_entries, options, target_prefix);
+  std::map<std::string, std::string> loaded_entries;
+  *success = database->LoadKeysAndEntriesWithFilter(filter, &loaded_entries,
+                                                    options, target_prefix);
 
-  for (const auto& serialized_entry : loaded_entries) {
+  for (const auto& pair : loaded_entries) {
     T entry;
-    if (!entry.ParseFromString(serialized_entry)) {
+    if (!entry.ParseFromString(pair.second)) {
       DLOG(WARNING) << "Unable to parse leveldb_proto entry";
       // TODO(cjhopman): Decide what to do about un-parseable entries.
     }
 
-    entries->push_back(entry);
+    keys_entries->insert(std::make_pair(pair.first, entry));
   }
 }
 
+template <typename T>
+void LoadEntriesFromTaskRunner(LevelDB* database,
+                               const LevelDB::KeyFilter& filter,
+                               const leveldb::ReadOptions& options,
+                               const std::string& target_prefix,
+                               std::vector<T>* entries,
+                               bool* success) {
+  entries->clear();
+
+  std::map<std::string, T> keys_entries;
+  LoadKeysAndEntriesFromTaskRunner<T>(database, filter, options, target_prefix,
+                                      &keys_entries, success);
+  for (const auto& pair : keys_entries)
+    entries->push_back(pair.second);
+}
+
 inline void LoadKeysFromTaskRunner(LevelDB* database,
                                    std::vector<std::string>* keys,
                                    bool* success) {
@@ -360,7 +394,7 @@ void ProtoDatabaseImpl<T>::LoadEntriesWithFilter(
   DCHECK(thread_checker_.CalledOnValidThread());
   bool* success = new bool(false);
 
-  std::unique_ptr<std::vector<T>> entries(new std::vector<T>());
+  auto entries = std::make_unique<std::vector<T>>();
   // Get this pointer before entries is std::move()'d so we can use it below.
   std::vector<T>* entries_ptr = entries.get();
 
@@ -372,6 +406,42 @@ void ProtoDatabaseImpl<T>::LoadEntriesWithFilter(
                      base::Owned(success), std::move(entries)));
 }
 
+template <typename T>
+void ProtoDatabaseImpl<T>::LoadKeysAndEntries(
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  LoadKeysAndEntriesWithFilter(LevelDB::KeyFilter(), std::move(callback));
+}
+
+template <typename T>
+void ProtoDatabaseImpl<T>::LoadKeysAndEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  LoadKeysAndEntriesWithFilter(key_filter, leveldb::ReadOptions(),
+                               std::string(), std::move(callback));
+}
+
+template <typename T>
+void ProtoDatabaseImpl<T>::LoadKeysAndEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    const leveldb::ReadOptions& options,
+    const std::string& target_prefix,
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  bool* success = new bool(false);
+
+  auto keys_entries = std::make_unique<std::map<std::string, T>>();
+  // Get this pointer before entries is std::move()'d so we can use it below.
+  std::map<std::string, T>* keys_entries_ptr = keys_entries.get();
+
+  task_runner_->PostTaskAndReply(
+      FROM_HERE,
+      base::BindOnce(LoadKeysAndEntriesFromTaskRunner<T>,
+                     base::Unretained(db_.get()), key_filter, options,
+                     target_prefix, keys_entries_ptr, success),
+      base::BindOnce(RunLoadKeysAndEntriesCallback<T>, std::move(callback),
+                     base::Owned(success), std::move(keys_entries)));
+}
+
 template <typename T>
 void ProtoDatabaseImpl<T>::LoadKeys(
     typename ProtoDatabase<T>::LoadKeysCallback callback) {

components/leveldb_proto/proto_database_impl_unittest.cc

@@ -62,6 +62,14 @@ class MockDB : public LevelDB {
                     std::vector<std::string>*,
                     const leveldb::ReadOptions&,
                     const std::string&));
+  MOCK_METHOD1(LoadKeysAndEntries, bool(std::map<std::string, std::string>*));
+  MOCK_METHOD2(LoadKeysAndEntriesWithFilter,
+               bool(const KeyFilter&, std::map<std::string, std::string>*));
+  MOCK_METHOD4(LoadKeysAndEntriesWithFilter,
+               bool(const KeyFilter&,
+                    std::map<std::string, std::string>*,
+                    const leveldb::ReadOptions&,
+                    const std::string&));
   MOCK_METHOD3(Get, bool(const std::string&, bool*, std::string*));
   MOCK_METHOD0(Destroy, bool());
 
@@ -78,6 +86,13 @@ class MockDatabaseCaller {
     LoadCallback1(success, entries.get());
   }
   MOCK_METHOD2(LoadCallback1, void(bool, std::vector<TestProto>*));
+  void LoadKeysAndEntriesCallback(
+      bool success,
+      std::unique_ptr<std::map<std::string, TestProto>> keys_entries) {
+    LoadKeysAndEntriesCallback1(success, keys_entries.get());
+  }
+  MOCK_METHOD2(LoadKeysAndEntriesCallback1,
+               void(bool, std::map<std::string, TestProto>*));
   void GetCallback(bool success, std::unique_ptr<TestProto> entry) {
     GetCallback1(success, entry.get());
   }
@@ -265,11 +280,27 @@ ACTION_P(AppendLoadEntries, model) {
   return true;
 }
 
+ACTION_P(AppendLoadKeysAndEntries, model) {
+  std::map<std::string, std::string>* output = arg1;
+  for (const auto& pair : model)
+    output->insert(std::make_pair(pair.first, pair.second.SerializeAsString()));
+
+  return true;
+}
+
 ACTION_P(VerifyLoadEntries, expected) {
   std::vector<TestProto>* actual = arg1;
   ExpectEntryPointersEquals(expected, *actual);
 }
 
+ACTION_P(VerifyLoadKeysAndEntries, expected) {
+  std::map<std::string, TestProto>* actual_map = arg1;
+  std::vector<TestProto> actual;
+  for (const auto& pair : *actual_map)
+    actual.push_back(pair.second);
+  ExpectEntryPointersEquals(expected, actual);
+}
+
 // Test that ProtoDatabaseImpl calls Load on the underlying database and that
 // the caller's LoadCallback is called with the correct success value. Also
 // confirms that on success, the expected entries are passed to the caller's
@@ -287,8 +318,8 @@ TEST_F(ProtoDatabaseImplTest, TestDBLoadSuccess) {
                         base::BindOnce(&MockDatabaseCaller::InitCallback,
                                        base::Unretained(&caller)));
 
-  EXPECT_CALL(*mock_db, LoadWithFilter(_, _, _, _))
-      .WillOnce(AppendLoadEntries(model));
+  EXPECT_CALL(*mock_db, LoadKeysAndEntriesWithFilter(_, _, _, _))
+      .WillOnce(AppendLoadKeysAndEntries(model));
   EXPECT_CALL(caller, LoadCallback1(true, _))
       .WillOnce(VerifyLoadEntries(testing::ByRef(model)));
   db_->LoadEntries(base::BindOnce(&MockDatabaseCaller::LoadCallback,
@@ -309,7 +340,8 @@ TEST_F(ProtoDatabaseImplTest, TestDBLoadFailure) {
                         base::BindOnce(&MockDatabaseCaller::InitCallback,
                                        base::Unretained(&caller)));
 
-  EXPECT_CALL(*mock_db, LoadWithFilter(_, _, _, _)).WillOnce(Return(false));
+  EXPECT_CALL(*mock_db, LoadKeysAndEntriesWithFilter(_, _, _, _))
+      .WillOnce(Return(false));
   EXPECT_CALL(caller, LoadCallback1(false, _));
   db_->LoadEntries(base::BindOnce(&MockDatabaseCaller::LoadCallback,
                                   base::Unretained(&caller)));
@@ -738,6 +770,35 @@ TEST_F(ProtoDatabaseImplLevelDBTest, TestDBLoadWithFilter) {
   EXPECT_EQ(entry.SerializeAsString(), model["0"].SerializeAsString());
 }
 
+TEST_F(ProtoDatabaseImplLevelDBTest, TestDBLoadKeysAndEntries) {
+  ScopedTempDir temp_dir;
+  ASSERT_TRUE(temp_dir.CreateUniqueTempDir());
+
+  EntryMap model = GetSmallModel();
+
+  KeyValueVector save_entries;
+  std::map<std::string, std::string> load_keys_entries;
+  KeyVector remove_keys;
+
+  for (const auto& pair : model) {
+    save_entries.push_back(
+        std::make_pair(pair.second.id(), pair.second.SerializeAsString()));
+  }
+
+  std::unique_ptr<LevelDB> db(new LevelDB(kTestLevelDBClientName));
+  EXPECT_TRUE(db->Init(temp_dir.GetPath(), CreateSimpleOptions()));
+  EXPECT_TRUE(db->Save(save_entries, remove_keys));
+
+  EXPECT_TRUE(db->LoadKeysAndEntries(&load_keys_entries));
+
+  EXPECT_EQ(load_keys_entries.size(), model.size());
+  for (const auto& pair : load_keys_entries) {
+    TestProto entry;
+    ASSERT_TRUE(entry.ParseFromString(pair.second));
+    EXPECT_EQ(entry.SerializeAsString(), model[pair.first].SerializeAsString());
+  }
+}
+
 TEST_F(ProtoDatabaseImplLevelDBTest, TestDBInitFail) {
   ScopedTempDir temp_dir;
   ASSERT_TRUE(temp_dir.CreateUniqueTempDir());

components/leveldb_proto/testing/fake_db.h

@@ -53,6 +53,16 @@ class FakeDB : public ProtoDatabase<T> {
       const leveldb::ReadOptions& options,
       const std::string& target_prefix,
       typename ProtoDatabase<T>::LoadCallback callback) override;
+  void LoadKeysAndEntries(
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
+  void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
+  void LoadKeysAndEntriesWithFilter(
+      const LevelDB::KeyFilter& filter,
+      const leveldb::ReadOptions& options,
+      const std::string& target_prefix,
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) override;
   void LoadKeys(typename ProtoDatabase<T>::LoadKeysCallback callback) override;
   void GetEntry(const std::string& key,
                 typename ProtoDatabase<T>::GetCallback callback) override;
@@ -78,6 +88,10 @@ class FakeDB : public ProtoDatabase<T> {
   static void RunLoadCallback(typename ProtoDatabase<T>::LoadCallback callback,
                               std::unique_ptr<typename std::vector<T>> entries,
                               bool success);
+  static void RunLoadKeysAndEntriesCallback(
+      typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback,
+      std::unique_ptr<typename std::map<std::string, T>> entries,
+      bool success);
 
   static void RunLoadKeysCallback(
       typename ProtoDatabase<T>::LoadKeysCallback callback,
@@ -179,6 +193,38 @@ void FakeDB<T>::LoadEntriesWithFilter(
       base::BindOnce(RunLoadCallback, std::move(callback), std::move(entries));
 }
 
+template <typename T>
+void FakeDB<T>::LoadKeysAndEntries(
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  LoadKeysAndEntriesWithFilter(LevelDB::KeyFilter(), std::move(callback));
+}
+
+template <typename T>
+void FakeDB<T>::LoadKeysAndEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  LoadKeysAndEntriesWithFilter(key_filter, leveldb::ReadOptions(),
+                               std::string(), std::move(callback));
+}
+
+template <typename T>
+void FakeDB<T>::LoadKeysAndEntriesWithFilter(
+    const LevelDB::KeyFilter& key_filter,
+    const leveldb::ReadOptions& options,
+    const std::string& target_prefix,
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback) {
+  auto keys_entries = std::make_unique<std::map<std::string, T>>();
+  for (const auto& pair : *db_) {
+    if (key_filter.is_null() || key_filter.Run(pair.first)) {
+      if (pair.first.compare(0, target_prefix.length(), target_prefix) == 0)
+        keys_entries->insert(pair);
+    }
+  }
+
+  load_callback_ = base::BindOnce(RunLoadKeysAndEntriesCallback,
+                                  std::move(callback), std::move(keys_entries));
+}
+
 template <typename T>
 void FakeDB<T>::LoadKeys(typename ProtoDatabase<T>::LoadKeysCallback callback) {
   std::unique_ptr<std::vector<std::string>> keys(
@@ -252,6 +298,15 @@ void FakeDB<T>::RunLoadCallback(
   std::move(callback).Run(success, std::move(entries));
 }
 
+// static
+template <typename T>
+void FakeDB<T>::RunLoadKeysAndEntriesCallback(
+    typename ProtoDatabase<T>::LoadKeysAndEntriesCallback callback,
+    std::unique_ptr<typename std::map<std::string, T>> keys_entries,
+    bool success) {
+  std::move(callback).Run(success, std::move(keys_entries));
+}
+
 // static
 template <typename T>
 void FakeDB<T>::RunLoadKeysCallback(