123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372 |
- // Copyright 2016 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.
- #include "components/ntp_snippets/remote/remote_suggestions_database.h"
- #include <utility>
- #include "base/bind.h"
- #include "base/files/file_path.h"
- #include "base/system/sys_info.h"
- #include "base/task/thread_pool.h"
- #include "components/leveldb_proto/public/proto_database_provider.h"
- #include "components/ntp_snippets/remote/proto/ntp_snippets.pb.h"
- using leveldb_proto::ProtoDatabase;
- using leveldb_proto::ProtoDatabaseProvider;
- namespace {
- const char kSnippetDatabaseFolder[] = "snippets";
- const char kImageDatabaseFolder[] = "images";
- const size_t kDatabaseWriteBufferSizeBytes = 128 << 10;
- } // namespace
- namespace ntp_snippets {
- RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
- leveldb_proto::ProtoDatabaseProvider* proto_database_provider,
- const base::FilePath& database_dir)
- : RemoteSuggestionsDatabase(
- proto_database_provider,
- database_dir,
- base::ThreadPool::CreateSequencedTaskRunner(
- {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
- base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN})) {}
- RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
- leveldb_proto::ProtoDatabaseProvider* proto_database_provider,
- const base::FilePath& database_dir,
- scoped_refptr<base::SequencedTaskRunner> task_runner)
- : RemoteSuggestionsDatabase(
- proto_database_provider->GetDB<SnippetProto>(
- leveldb_proto::ProtoDbType::REMOTE_SUGGESTIONS_DATABASE,
- database_dir.AppendASCII(kSnippetDatabaseFolder),
- task_runner),
- proto_database_provider->GetDB<SnippetImageProto>(
- leveldb_proto::ProtoDbType::REMOTE_SUGGESTIONS_IMAGE_DATABASE,
- database_dir.AppendASCII(kImageDatabaseFolder),
- task_runner)) {}
- RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
- std::unique_ptr<ProtoDatabase<SnippetProto>> database,
- std::unique_ptr<ProtoDatabase<SnippetImageProto>> image_database)
- : database_(std::move(database)),
- database_initialized_(false),
- image_database_(std::move(image_database)),
- image_database_initialized_(false) {
- leveldb_env::Options options = leveldb_proto::CreateSimpleOptions();
- options.reuse_logs = false; // Consumes less RAM over time.
- options.write_buffer_size = kDatabaseWriteBufferSizeBytes;
- database_->Init(options,
- base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseInited,
- weak_ptr_factory_.GetWeakPtr()));
- image_database_->Init(
- options, base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseInited,
- weak_ptr_factory_.GetWeakPtr()));
- }
- RemoteSuggestionsDatabase::~RemoteSuggestionsDatabase() = default;
- bool RemoteSuggestionsDatabase::IsInitialized() const {
- return !IsErrorState() && database_initialized_ &&
- image_database_initialized_;
- }
- bool RemoteSuggestionsDatabase::IsErrorState() const {
- return !database_ || !image_database_;
- }
- void RemoteSuggestionsDatabase::SetErrorCallback(
- const base::RepeatingClosure& error_callback) {
- error_callback_ = error_callback;
- }
- void RemoteSuggestionsDatabase::LoadSnippets(SnippetsCallback callback) {
- if (IsInitialized()) {
- LoadSnippetsImpl(std::move(callback));
- } else {
- pending_snippets_callbacks_.emplace_back(std::move(callback));
- }
- }
- void RemoteSuggestionsDatabase::SaveSnippet(const RemoteSuggestion& snippet) {
- if (IsErrorState()) {
- DVLOG(0) << "Attempted save snippet but db is in an error state, aborting";
- return;
- }
- std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
- // OnDatabaseLoaded relies on the detail that the primary snippet id goes
- // first in the protocol representation.
- DCHECK_EQ(snippet.ToProto().ids(0), snippet.id());
- entries_to_save->emplace_back(snippet.id(), snippet.ToProto());
- SaveSnippetsImpl(std::move(entries_to_save));
- }
- void RemoteSuggestionsDatabase::SaveSnippets(
- const RemoteSuggestion::PtrVector& snippets) {
- if (IsErrorState()) {
- DVLOG(0) << "Attempted save snippets but db is in an error state, aborting";
- return;
- }
- std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
- for (const std::unique_ptr<RemoteSuggestion>& snippet : snippets) {
- // OnDatabaseLoaded relies on the detail that the primary snippet id goes
- // first in the protocol representation.
- DCHECK_EQ(snippet->ToProto().ids(0), snippet->id());
- entries_to_save->emplace_back(snippet->id(), snippet->ToProto());
- }
- SaveSnippetsImpl(std::move(entries_to_save));
- }
- void RemoteSuggestionsDatabase::DeleteSnippet(const std::string& snippet_id) {
- DeleteSnippets(std::make_unique<std::vector<std::string>>(1, snippet_id));
- }
- void RemoteSuggestionsDatabase::DeleteSnippets(
- std::unique_ptr<std::vector<std::string>> snippet_ids) {
- if (IsErrorState()) {
- DVLOG(0)
- << "Attempted delete snippets but db is in an error state, aborting";
- return;
- }
- std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
- database_->UpdateEntries(
- std::move(entries_to_save), std::move(snippet_ids),
- base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseSaved,
- weak_ptr_factory_.GetWeakPtr()));
- }
- void RemoteSuggestionsDatabase::LoadImage(const std::string& snippet_id,
- SnippetImageCallback callback) {
- if (IsInitialized()) {
- LoadImageImpl(snippet_id, std::move(callback));
- } else {
- pending_image_callbacks_.emplace_back(snippet_id, std::move(callback));
- }
- }
- void RemoteSuggestionsDatabase::SaveImage(const std::string& snippet_id,
- const std::string& image_data) {
- if (IsErrorState()) {
- DVLOG(0) << "Attempted save image but db is in an error state, aborting";
- return;
- }
- SnippetImageProto image_proto;
- image_proto.set_data(image_data);
- std::unique_ptr<ImageKeyEntryVector> entries_to_save(
- new ImageKeyEntryVector());
- entries_to_save->emplace_back(snippet_id, std::move(image_proto));
- image_database_->UpdateEntries(
- std::move(entries_to_save), std::make_unique<std::vector<std::string>>(),
- base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseSaved,
- weak_ptr_factory_.GetWeakPtr()));
- }
- void RemoteSuggestionsDatabase::DeleteImage(const std::string& snippet_id) {
- DeleteImages(std::make_unique<std::vector<std::string>>(1, snippet_id));
- }
- void RemoteSuggestionsDatabase::DeleteImages(
- std::unique_ptr<std::vector<std::string>> snippet_ids) {
- if (IsErrorState()) {
- DVLOG(0) << "Attempted delete images but db is in an error state, aborting";
- return;
- }
- image_database_->UpdateEntries(
- std::make_unique<ImageKeyEntryVector>(), std::move(snippet_ids),
- base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseSaved,
- weak_ptr_factory_.GetWeakPtr()));
- }
- void RemoteSuggestionsDatabase::GarbageCollectImages(
- std::unique_ptr<std::set<std::string>> alive_snippet_ids) {
- if (IsErrorState()) {
- DVLOG(0) << "Attempted gc but db is in an error state, aborting";
- return;
- }
- image_database_->LoadKeys(base::BindOnce(
- &RemoteSuggestionsDatabase::DeleteUnreferencedImages,
- weak_ptr_factory_.GetWeakPtr(), std::move(alive_snippet_ids)));
- }
- void RemoteSuggestionsDatabase::OnDatabaseInited(
- leveldb_proto::Enums::InitStatus status) {
- DCHECK(!database_initialized_);
- if (status != leveldb_proto::Enums::InitStatus::kOK) {
- DVLOG(1) << "RemoteSuggestionsDatabase init failed.";
- OnDatabaseError();
- return;
- }
- database_initialized_ = true;
- if (IsInitialized()) {
- ProcessPendingLoads();
- }
- }
- void RemoteSuggestionsDatabase::OnDatabaseLoaded(
- SnippetsCallback callback,
- bool success,
- std::unique_ptr<std::vector<SnippetProto>> entries) {
- if (!success) {
- DVLOG(1) << "RemoteSuggestionsDatabase load failed.";
- OnDatabaseError();
- return;
- }
- std::unique_ptr<std::vector<std::string>> keys_to_remove(
- new std::vector<std::string>());
- RemoteSuggestion::PtrVector snippets;
- for (const SnippetProto& proto : *entries) {
- std::unique_ptr<RemoteSuggestion> snippet =
- RemoteSuggestion::CreateFromProto(proto);
- if (snippet) {
- snippets.emplace_back(std::move(snippet));
- } else {
- if (proto.ids_size() > 0) {
- LOG(WARNING) << "Invalid proto from DB " << proto.ids(0);
- keys_to_remove->emplace_back(proto.ids(0));
- } else {
- LOG(WARNING)
- << "Loaded proto without ID from the DB. Cannot clean this up.";
- }
- }
- }
- std::move(callback).Run(std::move(snippets));
- // If any of the snippet protos couldn't be converted to actual snippets,
- // clean them up now.
- if (!keys_to_remove->empty()) {
- DeleteSnippets(std::move(keys_to_remove));
- }
- }
- void RemoteSuggestionsDatabase::OnDatabaseSaved(bool success) {
- if (!success) {
- DVLOG(1) << "RemoteSuggestionsDatabase save failed.";
- OnDatabaseError();
- }
- }
- void RemoteSuggestionsDatabase::OnImageDatabaseInited(
- leveldb_proto::Enums::InitStatus status) {
- DCHECK(!image_database_initialized_);
- if (status != leveldb_proto::Enums::InitStatus::kOK) {
- DVLOG(1) << "RemoteSuggestionsDatabase init failed.";
- OnDatabaseError();
- return;
- }
- image_database_initialized_ = true;
- if (IsInitialized()) {
- ProcessPendingLoads();
- }
- }
- void RemoteSuggestionsDatabase::OnImageDatabaseLoaded(
- SnippetImageCallback callback,
- bool success,
- std::unique_ptr<SnippetImageProto> entry) {
- if (!success) {
- DVLOG(1) << "RemoteSuggestionsDatabase load failed.";
- OnDatabaseError();
- return;
- }
- if (!entry) {
- std::move(callback).Run(std::string());
- return;
- }
- std::unique_ptr<std::string> data(entry->release_data());
- std::move(callback).Run(std::move(*data));
- }
- void RemoteSuggestionsDatabase::OnImageDatabaseSaved(bool success) {
- if (!success) {
- DVLOG(1) << "RemoteSuggestionsDatabase save failed.";
- OnDatabaseError();
- }
- }
- void RemoteSuggestionsDatabase::OnDatabaseError() {
- database_.reset();
- image_database_.reset();
- if (!error_callback_.is_null()) {
- error_callback_.Run();
- }
- }
- void RemoteSuggestionsDatabase::ProcessPendingLoads() {
- DCHECK(IsInitialized());
- for (auto& callback : pending_snippets_callbacks_) {
- LoadSnippetsImpl(std::move(callback));
- }
- pending_snippets_callbacks_.clear();
- for (auto& id_callback : pending_image_callbacks_) {
- LoadImageImpl(id_callback.first, std::move(id_callback.second));
- }
- pending_image_callbacks_.clear();
- }
- void RemoteSuggestionsDatabase::LoadSnippetsImpl(SnippetsCallback callback) {
- DCHECK(IsInitialized());
- database_->LoadEntries(
- base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseLoaded,
- weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
- }
- void RemoteSuggestionsDatabase::SaveSnippetsImpl(
- std::unique_ptr<KeyEntryVector> entries_to_save) {
- DCHECK(IsInitialized());
- std::unique_ptr<std::vector<std::string>> keys_to_remove(
- new std::vector<std::string>());
- database_->UpdateEntries(
- std::move(entries_to_save), std::move(keys_to_remove),
- base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseSaved,
- weak_ptr_factory_.GetWeakPtr()));
- }
- void RemoteSuggestionsDatabase::LoadImageImpl(const std::string& snippet_id,
- SnippetImageCallback callback) {
- DCHECK(IsInitialized());
- image_database_->GetEntry(
- snippet_id,
- base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseLoaded,
- weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
- }
- void RemoteSuggestionsDatabase::DeleteUnreferencedImages(
- std::unique_ptr<std::set<std::string>> references,
- bool load_keys_success,
- std::unique_ptr<std::vector<std::string>> image_keys) {
- if (!load_keys_success) {
- DVLOG(1) << "RemoteSuggestionsDatabase garbage collection failed.";
- OnDatabaseError();
- return;
- }
- auto keys_to_remove = std::make_unique<std::vector<std::string>>();
- for (const std::string& key : *image_keys) {
- if (references->count(key) == 0) {
- keys_to_remove->emplace_back(key);
- }
- }
- if (keys_to_remove->empty())
- return;
- DeleteImages(std::move(keys_to_remove));
- }
- } // namespace ntp_snippets
|