remote_suggestions_database.cc 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372
  1. // Copyright 2016 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/ntp_snippets/remote/remote_suggestions_database.h"
  5. #include <utility>
  6. #include "base/bind.h"
  7. #include "base/files/file_path.h"
  8. #include "base/system/sys_info.h"
  9. #include "base/task/thread_pool.h"
  10. #include "components/leveldb_proto/public/proto_database_provider.h"
  11. #include "components/ntp_snippets/remote/proto/ntp_snippets.pb.h"
  12. using leveldb_proto::ProtoDatabase;
  13. using leveldb_proto::ProtoDatabaseProvider;
  14. namespace {
  15. const char kSnippetDatabaseFolder[] = "snippets";
  16. const char kImageDatabaseFolder[] = "images";
  17. const size_t kDatabaseWriteBufferSizeBytes = 128 << 10;
  18. } // namespace
  19. namespace ntp_snippets {
  20. RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
  21. leveldb_proto::ProtoDatabaseProvider* proto_database_provider,
  22. const base::FilePath& database_dir)
  23. : RemoteSuggestionsDatabase(
  24. proto_database_provider,
  25. database_dir,
  26. base::ThreadPool::CreateSequencedTaskRunner(
  27. {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
  28. base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN})) {}
  29. RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
  30. leveldb_proto::ProtoDatabaseProvider* proto_database_provider,
  31. const base::FilePath& database_dir,
  32. scoped_refptr<base::SequencedTaskRunner> task_runner)
  33. : RemoteSuggestionsDatabase(
  34. proto_database_provider->GetDB<SnippetProto>(
  35. leveldb_proto::ProtoDbType::REMOTE_SUGGESTIONS_DATABASE,
  36. database_dir.AppendASCII(kSnippetDatabaseFolder),
  37. task_runner),
  38. proto_database_provider->GetDB<SnippetImageProto>(
  39. leveldb_proto::ProtoDbType::REMOTE_SUGGESTIONS_IMAGE_DATABASE,
  40. database_dir.AppendASCII(kImageDatabaseFolder),
  41. task_runner)) {}
  42. RemoteSuggestionsDatabase::RemoteSuggestionsDatabase(
  43. std::unique_ptr<ProtoDatabase<SnippetProto>> database,
  44. std::unique_ptr<ProtoDatabase<SnippetImageProto>> image_database)
  45. : database_(std::move(database)),
  46. database_initialized_(false),
  47. image_database_(std::move(image_database)),
  48. image_database_initialized_(false) {
  49. leveldb_env::Options options = leveldb_proto::CreateSimpleOptions();
  50. options.reuse_logs = false; // Consumes less RAM over time.
  51. options.write_buffer_size = kDatabaseWriteBufferSizeBytes;
  52. database_->Init(options,
  53. base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseInited,
  54. weak_ptr_factory_.GetWeakPtr()));
  55. image_database_->Init(
  56. options, base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseInited,
  57. weak_ptr_factory_.GetWeakPtr()));
  58. }
  59. RemoteSuggestionsDatabase::~RemoteSuggestionsDatabase() = default;
  60. bool RemoteSuggestionsDatabase::IsInitialized() const {
  61. return !IsErrorState() && database_initialized_ &&
  62. image_database_initialized_;
  63. }
  64. bool RemoteSuggestionsDatabase::IsErrorState() const {
  65. return !database_ || !image_database_;
  66. }
  67. void RemoteSuggestionsDatabase::SetErrorCallback(
  68. const base::RepeatingClosure& error_callback) {
  69. error_callback_ = error_callback;
  70. }
  71. void RemoteSuggestionsDatabase::LoadSnippets(SnippetsCallback callback) {
  72. if (IsInitialized()) {
  73. LoadSnippetsImpl(std::move(callback));
  74. } else {
  75. pending_snippets_callbacks_.emplace_back(std::move(callback));
  76. }
  77. }
  78. void RemoteSuggestionsDatabase::SaveSnippet(const RemoteSuggestion& snippet) {
  79. if (IsErrorState()) {
  80. DVLOG(0) << "Attempted save snippet but db is in an error state, aborting";
  81. return;
  82. }
  83. std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
  84. // OnDatabaseLoaded relies on the detail that the primary snippet id goes
  85. // first in the protocol representation.
  86. DCHECK_EQ(snippet.ToProto().ids(0), snippet.id());
  87. entries_to_save->emplace_back(snippet.id(), snippet.ToProto());
  88. SaveSnippetsImpl(std::move(entries_to_save));
  89. }
  90. void RemoteSuggestionsDatabase::SaveSnippets(
  91. const RemoteSuggestion::PtrVector& snippets) {
  92. if (IsErrorState()) {
  93. DVLOG(0) << "Attempted save snippets but db is in an error state, aborting";
  94. return;
  95. }
  96. std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
  97. for (const std::unique_ptr<RemoteSuggestion>& snippet : snippets) {
  98. // OnDatabaseLoaded relies on the detail that the primary snippet id goes
  99. // first in the protocol representation.
  100. DCHECK_EQ(snippet->ToProto().ids(0), snippet->id());
  101. entries_to_save->emplace_back(snippet->id(), snippet->ToProto());
  102. }
  103. SaveSnippetsImpl(std::move(entries_to_save));
  104. }
  105. void RemoteSuggestionsDatabase::DeleteSnippet(const std::string& snippet_id) {
  106. DeleteSnippets(std::make_unique<std::vector<std::string>>(1, snippet_id));
  107. }
  108. void RemoteSuggestionsDatabase::DeleteSnippets(
  109. std::unique_ptr<std::vector<std::string>> snippet_ids) {
  110. if (IsErrorState()) {
  111. DVLOG(0)
  112. << "Attempted delete snippets but db is in an error state, aborting";
  113. return;
  114. }
  115. std::unique_ptr<KeyEntryVector> entries_to_save(new KeyEntryVector());
  116. database_->UpdateEntries(
  117. std::move(entries_to_save), std::move(snippet_ids),
  118. base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseSaved,
  119. weak_ptr_factory_.GetWeakPtr()));
  120. }
  121. void RemoteSuggestionsDatabase::LoadImage(const std::string& snippet_id,
  122. SnippetImageCallback callback) {
  123. if (IsInitialized()) {
  124. LoadImageImpl(snippet_id, std::move(callback));
  125. } else {
  126. pending_image_callbacks_.emplace_back(snippet_id, std::move(callback));
  127. }
  128. }
  129. void RemoteSuggestionsDatabase::SaveImage(const std::string& snippet_id,
  130. const std::string& image_data) {
  131. if (IsErrorState()) {
  132. DVLOG(0) << "Attempted save image but db is in an error state, aborting";
  133. return;
  134. }
  135. SnippetImageProto image_proto;
  136. image_proto.set_data(image_data);
  137. std::unique_ptr<ImageKeyEntryVector> entries_to_save(
  138. new ImageKeyEntryVector());
  139. entries_to_save->emplace_back(snippet_id, std::move(image_proto));
  140. image_database_->UpdateEntries(
  141. std::move(entries_to_save), std::make_unique<std::vector<std::string>>(),
  142. base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseSaved,
  143. weak_ptr_factory_.GetWeakPtr()));
  144. }
  145. void RemoteSuggestionsDatabase::DeleteImage(const std::string& snippet_id) {
  146. DeleteImages(std::make_unique<std::vector<std::string>>(1, snippet_id));
  147. }
  148. void RemoteSuggestionsDatabase::DeleteImages(
  149. std::unique_ptr<std::vector<std::string>> snippet_ids) {
  150. if (IsErrorState()) {
  151. DVLOG(0) << "Attempted delete images but db is in an error state, aborting";
  152. return;
  153. }
  154. image_database_->UpdateEntries(
  155. std::make_unique<ImageKeyEntryVector>(), std::move(snippet_ids),
  156. base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseSaved,
  157. weak_ptr_factory_.GetWeakPtr()));
  158. }
  159. void RemoteSuggestionsDatabase::GarbageCollectImages(
  160. std::unique_ptr<std::set<std::string>> alive_snippet_ids) {
  161. if (IsErrorState()) {
  162. DVLOG(0) << "Attempted gc but db is in an error state, aborting";
  163. return;
  164. }
  165. image_database_->LoadKeys(base::BindOnce(
  166. &RemoteSuggestionsDatabase::DeleteUnreferencedImages,
  167. weak_ptr_factory_.GetWeakPtr(), std::move(alive_snippet_ids)));
  168. }
  169. void RemoteSuggestionsDatabase::OnDatabaseInited(
  170. leveldb_proto::Enums::InitStatus status) {
  171. DCHECK(!database_initialized_);
  172. if (status != leveldb_proto::Enums::InitStatus::kOK) {
  173. DVLOG(1) << "RemoteSuggestionsDatabase init failed.";
  174. OnDatabaseError();
  175. return;
  176. }
  177. database_initialized_ = true;
  178. if (IsInitialized()) {
  179. ProcessPendingLoads();
  180. }
  181. }
  182. void RemoteSuggestionsDatabase::OnDatabaseLoaded(
  183. SnippetsCallback callback,
  184. bool success,
  185. std::unique_ptr<std::vector<SnippetProto>> entries) {
  186. if (!success) {
  187. DVLOG(1) << "RemoteSuggestionsDatabase load failed.";
  188. OnDatabaseError();
  189. return;
  190. }
  191. std::unique_ptr<std::vector<std::string>> keys_to_remove(
  192. new std::vector<std::string>());
  193. RemoteSuggestion::PtrVector snippets;
  194. for (const SnippetProto& proto : *entries) {
  195. std::unique_ptr<RemoteSuggestion> snippet =
  196. RemoteSuggestion::CreateFromProto(proto);
  197. if (snippet) {
  198. snippets.emplace_back(std::move(snippet));
  199. } else {
  200. if (proto.ids_size() > 0) {
  201. LOG(WARNING) << "Invalid proto from DB " << proto.ids(0);
  202. keys_to_remove->emplace_back(proto.ids(0));
  203. } else {
  204. LOG(WARNING)
  205. << "Loaded proto without ID from the DB. Cannot clean this up.";
  206. }
  207. }
  208. }
  209. std::move(callback).Run(std::move(snippets));
  210. // If any of the snippet protos couldn't be converted to actual snippets,
  211. // clean them up now.
  212. if (!keys_to_remove->empty()) {
  213. DeleteSnippets(std::move(keys_to_remove));
  214. }
  215. }
  216. void RemoteSuggestionsDatabase::OnDatabaseSaved(bool success) {
  217. if (!success) {
  218. DVLOG(1) << "RemoteSuggestionsDatabase save failed.";
  219. OnDatabaseError();
  220. }
  221. }
  222. void RemoteSuggestionsDatabase::OnImageDatabaseInited(
  223. leveldb_proto::Enums::InitStatus status) {
  224. DCHECK(!image_database_initialized_);
  225. if (status != leveldb_proto::Enums::InitStatus::kOK) {
  226. DVLOG(1) << "RemoteSuggestionsDatabase init failed.";
  227. OnDatabaseError();
  228. return;
  229. }
  230. image_database_initialized_ = true;
  231. if (IsInitialized()) {
  232. ProcessPendingLoads();
  233. }
  234. }
  235. void RemoteSuggestionsDatabase::OnImageDatabaseLoaded(
  236. SnippetImageCallback callback,
  237. bool success,
  238. std::unique_ptr<SnippetImageProto> entry) {
  239. if (!success) {
  240. DVLOG(1) << "RemoteSuggestionsDatabase load failed.";
  241. OnDatabaseError();
  242. return;
  243. }
  244. if (!entry) {
  245. std::move(callback).Run(std::string());
  246. return;
  247. }
  248. std::unique_ptr<std::string> data(entry->release_data());
  249. std::move(callback).Run(std::move(*data));
  250. }
  251. void RemoteSuggestionsDatabase::OnImageDatabaseSaved(bool success) {
  252. if (!success) {
  253. DVLOG(1) << "RemoteSuggestionsDatabase save failed.";
  254. OnDatabaseError();
  255. }
  256. }
  257. void RemoteSuggestionsDatabase::OnDatabaseError() {
  258. database_.reset();
  259. image_database_.reset();
  260. if (!error_callback_.is_null()) {
  261. error_callback_.Run();
  262. }
  263. }
  264. void RemoteSuggestionsDatabase::ProcessPendingLoads() {
  265. DCHECK(IsInitialized());
  266. for (auto& callback : pending_snippets_callbacks_) {
  267. LoadSnippetsImpl(std::move(callback));
  268. }
  269. pending_snippets_callbacks_.clear();
  270. for (auto& id_callback : pending_image_callbacks_) {
  271. LoadImageImpl(id_callback.first, std::move(id_callback.second));
  272. }
  273. pending_image_callbacks_.clear();
  274. }
  275. void RemoteSuggestionsDatabase::LoadSnippetsImpl(SnippetsCallback callback) {
  276. DCHECK(IsInitialized());
  277. database_->LoadEntries(
  278. base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseLoaded,
  279. weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
  280. }
  281. void RemoteSuggestionsDatabase::SaveSnippetsImpl(
  282. std::unique_ptr<KeyEntryVector> entries_to_save) {
  283. DCHECK(IsInitialized());
  284. std::unique_ptr<std::vector<std::string>> keys_to_remove(
  285. new std::vector<std::string>());
  286. database_->UpdateEntries(
  287. std::move(entries_to_save), std::move(keys_to_remove),
  288. base::BindOnce(&RemoteSuggestionsDatabase::OnDatabaseSaved,
  289. weak_ptr_factory_.GetWeakPtr()));
  290. }
  291. void RemoteSuggestionsDatabase::LoadImageImpl(const std::string& snippet_id,
  292. SnippetImageCallback callback) {
  293. DCHECK(IsInitialized());
  294. image_database_->GetEntry(
  295. snippet_id,
  296. base::BindOnce(&RemoteSuggestionsDatabase::OnImageDatabaseLoaded,
  297. weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
  298. }
  299. void RemoteSuggestionsDatabase::DeleteUnreferencedImages(
  300. std::unique_ptr<std::set<std::string>> references,
  301. bool load_keys_success,
  302. std::unique_ptr<std::vector<std::string>> image_keys) {
  303. if (!load_keys_success) {
  304. DVLOG(1) << "RemoteSuggestionsDatabase garbage collection failed.";
  305. OnDatabaseError();
  306. return;
  307. }
  308. auto keys_to_remove = std::make_unique<std::vector<std::string>>();
  309. for (const std::string& key : *image_keys) {
  310. if (references->count(key) == 0) {
  311. keys_to_remove->emplace_back(key);
  312. }
  313. }
  314. if (keys_to_remove->empty())
  315. return;
  316. DeleteImages(std::move(keys_to_remove));
  317. }
  318. } // namespace ntp_snippets