[Chromeshine] Switch to multiple ProtoDatabases.

Now that the binary size is less significantly impacted by introducing
additional ProtoDatabase instances (crrev.com/c/1419199), switch to
using a unique ProtoDatabase for each proto message type.

Bug: 921133
Change-Id: Ieae1fc5ba98c574c8e4f9a2587bffa4333fa80df
Reviewed-on: https://chromium-review.googlesource.com/c/1471266
Commit-Queue: Natalie Chouinard <chouinard@chromium.org>
Reviewed-by: Patrick Noland <pnoland@chromium.org>
Cr-Commit-Position: refs/heads/master@{#631968}

chrome/browser/android/usage_stats/usage_stats_database.cc

@@ -16,8 +16,10 @@
 
 namespace usage_stats {
 
+using leveldb_proto::ProtoDatabaseProvider;
+using leveldb_proto::ProtoDatabaseProviderFactory;
+
 const char kNamespace[] = "usage_stats";
-const char kTypePrefix[] = "usage_stats";
 
 const char kKeySeparator[] = "_";
 
@@ -30,27 +32,37 @@ const int kUnixTimeDigits = 11;
 // zero-filled (example: 1548353315 => 01548353315).
 const char kUnixTimeFormat[] = "%011d";
 
-const int kFqdnPosition = base::size(kWebsiteEventPrefix) + kUnixTimeDigits + 1;
-
 UsageStatsDatabase::UsageStatsDatabase(Profile* profile)
     : weak_ptr_factory_(this) {
-  leveldb_proto::ProtoDatabaseProvider* db_provider =
-      leveldb_proto::ProtoDatabaseProviderFactory::GetInstance()
-          ->GetForBrowserContext(profile);
+  ProtoDatabaseProvider* db_provider =
+      ProtoDatabaseProviderFactory::GetInstance()->GetForBrowserContext(
+          profile);
 
   base::FilePath usage_stats_dir = profile->GetPath().Append(kNamespace);
 
-  // TODO(crbug/921133): Switch back to separate dbs for each message type when
-  // possible.
-  proto_db_ = db_provider->GetDB<UsageStat>(
-      kNamespace, kTypePrefix, usage_stats_dir,
+  scoped_refptr<base::SequencedTaskRunner> db_task_runner =
       base::CreateSequencedTaskRunnerWithTraits(
-          {base::MayBlock(), base::TaskPriority::BEST_EFFORT}));
+          {base::MayBlock(), base::TaskPriority::BEST_EFFORT});
+
+  website_event_db_ = db_provider->GetDB<WebsiteEvent>(
+      kNamespace, kWebsiteEventPrefix, usage_stats_dir, db_task_runner);
+
+  suspension_db_ = db_provider->GetDB<Suspension>(
+      kNamespace, kSuspensionPrefix, usage_stats_dir, db_task_runner);
+
+  token_mapping_db_ = db_provider->GetDB<TokenMapping>(
+      kNamespace, kTokenMappingPrefix, usage_stats_dir, db_task_runner);
 }
 
+// Used for testing.
 UsageStatsDatabase::UsageStatsDatabase(
-    std::unique_ptr<leveldb_proto::ProtoDatabase<UsageStat>> proto_db)
-    : proto_db_(std::move(proto_db)), weak_ptr_factory_(this) {}
+    std::unique_ptr<ProtoDatabase<WebsiteEvent>> website_event_db,
+    std::unique_ptr<ProtoDatabase<Suspension>> suspension_db,
+    std::unique_ptr<ProtoDatabase<TokenMapping>> token_mapping_db)
+    : website_event_db_(std::move(website_event_db)),
+      suspension_db_(std::move(suspension_db)),
+      token_mapping_db_(std::move(token_mapping_db)),
+      weak_ptr_factory_(this) {}
 
 UsageStatsDatabase::~UsageStatsDatabase() = default;
 
@@ -63,7 +75,7 @@ bool DoesNotContainFilter(const base::flat_set<std::string>& set,
 
 bool KeyContainsDomainFilter(const base::flat_set<std::string>& domains,
                              const std::string& key) {
-  return domains.contains(key.substr(kFqdnPosition));
+  return domains.contains(key.substr(kUnixTimeDigits + 1));
 }
 
 UsageStatsDatabase::Error ToError(bool isSuccess) {
@@ -79,18 +91,15 @@ std::string CreateWebsiteEventKey(int64_t seconds, const std::string& fqdn) {
       base::strings::SafeSPrintf(unixTime, kUnixTimeFormat, seconds);
   DCHECK(printed == kUnixTimeDigits);
 
-  // Create the key from the prefix, time, and fqdn (example:
-  // website_event_01548276551_foo.com).
-  return base::StrCat(
-      {kWebsiteEventPrefix, kKeySeparator, unixTime, kKeySeparator, fqdn});
+  // Create the key from the time and fqdn (example: 01548276551_foo.com).
+  return base::StrCat({unixTime, kKeySeparator, fqdn});
 }
 
 }  // namespace
 
 void UsageStatsDatabase::GetAllEvents(EventsCallback callback) {
-  // Load all UsageStats with the website events prefix.
-  proto_db_->LoadEntriesWithFilter(
-      leveldb_proto::KeyFilter(), leveldb::ReadOptions(), kWebsiteEventPrefix,
+  // Load all WebsiteEvents.
+  website_event_db_->LoadEntries(
       base::BindOnce(&UsageStatsDatabase::OnLoadEntriesForGetAllEvents,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
@@ -98,12 +107,12 @@ void UsageStatsDatabase::GetAllEvents(EventsCallback callback) {
 void UsageStatsDatabase::QueryEventsInRange(int64_t start,
                                             int64_t end,
                                             EventsCallback callback) {
-  // Load all UsageStats with the website events prefix, where the timestamp is
+  // Load all WebsiteEvents where the timestamp is
   // in the specified range. Function accepts a half-open range [start, end) as
   // input, but the database operates on fully-closed ranges. Because the
   // timestamps are represented by integers, [start, end) is equivalent to
   // [start, end - 1].
-  proto_db_->LoadKeysAndEntriesInRange(
+  website_event_db_->LoadKeysAndEntriesInRange(
       CreateWebsiteEventKey(start, ""), CreateWebsiteEventKey(end - 1, ""),
       base::BindOnce(&UsageStatsDatabase::OnLoadEntriesForQueryEventsInRange,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
@@ -111,35 +120,29 @@ void UsageStatsDatabase::QueryEventsInRange(int64_t start,
 
 void UsageStatsDatabase::AddEvents(std::vector<WebsiteEvent> events,
                                    StatusCallback callback) {
-  auto entries = std::make_unique<
-      leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>();
+  auto entries =
+      std::make_unique<ProtoDatabase<WebsiteEvent>::KeyEntryVector>();
   entries->reserve(events.size());
 
   for (WebsiteEvent event : events) {
     std::string key =
         CreateWebsiteEventKey(event.timestamp().seconds(), event.fqdn());
 
-    UsageStat value;
-    WebsiteEvent* website_event = value.mutable_website_event();
-    *website_event = event;
-
-    entries->emplace_back(key, value);
+    entries->emplace_back(key, event);
   }
 
   // Add all entries created from input vector.
-  proto_db_->UpdateEntries(
+  website_event_db_->UpdateEntries(
       std::move(entries), std::make_unique<std::vector<std::string>>(),
       base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
 
 void UsageStatsDatabase::DeleteAllEvents(StatusCallback callback) {
-  // Remove all entries with the website event prefix.
-  proto_db_->UpdateEntriesWithRemoveFilter(
-      std::make_unique<
-          leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>(),
+  // Remove all WebsiteEvent entries.
+  website_event_db_->UpdateEntriesWithRemoveFilter(
+      std::make_unique<ProtoDatabase<WebsiteEvent>::KeyEntryVector>(),
       base::BindRepeating([](const std::string& key) { return true; }),
-      kWebsiteEventPrefix,
       base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
@@ -150,12 +153,12 @@ void UsageStatsDatabase::DeleteEventsInRange(int64_t start,
   // function, we should consolidate these two proto_db_ calls into a single
   // call.
 
-  // Load all keys with the website events prefix, where the timestamp is
-  // in the specified range, passing a callback to delete those entries.
-  // Function accepts a half-open range [start, end) as input, but the database
-  // operates on fully-closed ranges. Because the timestamps are represented by
-  // integers, [start, end) is equivalent to [start, end - 1].
-  proto_db_->LoadKeysAndEntriesInRange(
+  // Load all keys where the timestamp is in the specified range, passing a
+  // callback to delete those entries. Function accepts a half-open range
+  // [start, end) as input, but the database operates on fully-closed ranges.
+  // Because the timestamps are represented by integers, [start, end) is
+  // equivalent to [start, end - 1].
+  website_event_db_->LoadKeysAndEntriesInRange(
       CreateWebsiteEventKey(start, ""), CreateWebsiteEventKey(end - 1, ""),
       base::BindOnce(&UsageStatsDatabase::OnLoadEntriesForDeleteEventsInRange,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
@@ -165,57 +168,42 @@ void UsageStatsDatabase::DeleteEventsWithMatchingDomains(
     base::flat_set<std::string> domains,
     StatusCallback callback) {
   // Remove all events with domains in the given set.
-  proto_db_->UpdateEntriesWithRemoveFilter(
-      std::make_unique<
-          leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>(),
+  website_event_db_->UpdateEntriesWithRemoveFilter(
+      std::make_unique<ProtoDatabase<WebsiteEvent>::KeyEntryVector>(),
       base::BindRepeating(&KeyContainsDomainFilter, std::move(domains)),
       base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
 
 void UsageStatsDatabase::GetAllSuspensions(SuspensionsCallback callback) {
-  // Load all UsageStats with the suspension prefix.
-  proto_db_->LoadEntriesWithFilter(
-      leveldb_proto::KeyFilter(), leveldb::ReadOptions(), kSuspensionPrefix,
+  // Load all Suspensions.
+  suspension_db_->LoadEntries(
       base::BindOnce(&UsageStatsDatabase::OnLoadEntriesForGetAllSuspensions,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
 
 void UsageStatsDatabase::SetSuspensions(base::flat_set<std::string> domains,
                                         StatusCallback callback) {
-  std::vector<std::string> keys;
-  keys.reserve(domains.size());
-
-  auto entries = std::make_unique<
-      leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>();
+  auto entries = std::make_unique<ProtoDatabase<Suspension>::KeyEntryVector>();
 
   for (std::string domain : domains) {
-    // Prepend prefix to form key.
-    std::string key = base::StrCat({kSuspensionPrefix, kKeySeparator, domain});
-
-    keys.emplace_back(key);
-
-    UsageStat value;
-    Suspension* suspension = value.mutable_suspension();
-    suspension->set_fqdn(domain);
+    Suspension suspension;
+    suspension.set_fqdn(domain);
 
-    entries->emplace_back(key, value);
+    entries->emplace_back(domain, suspension);
   }
 
-  auto key_set = base::flat_set<std::string>(keys);
-
   // Add all entries created from domain set, remove all entries not in the set.
-  proto_db_->UpdateEntriesWithRemoveFilter(
+  suspension_db_->UpdateEntriesWithRemoveFilter(
       std::move(entries),
-      base::BindRepeating(&DoesNotContainFilter, std::move(key_set)),
+      base::BindRepeating(&DoesNotContainFilter, std::move(domains)),
       base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
 
 void UsageStatsDatabase::GetAllTokenMappings(TokenMappingsCallback callback) {
-  // Load all UsageStats with the token mapping prefix.
-  proto_db_->LoadEntriesWithFilter(
-      leveldb_proto::KeyFilter(), leveldb::ReadOptions(), kTokenMappingPrefix,
+  // Load all TokenMappings.
+  token_mapping_db_->LoadEntries(
       base::BindOnce(&UsageStatsDatabase::OnLoadEntriesForGetAllTokenMappings,
                      weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
 }
@@ -225,28 +213,26 @@ void UsageStatsDatabase::SetTokenMappings(TokenMap mappings,
   std::vector<std::string> keys;
   keys.reserve(mappings.size());
 
-  auto entries = std::make_unique<
-      leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>();
+  auto entries =
+      std::make_unique<ProtoDatabase<TokenMapping>::KeyEntryVector>();
 
   for (const auto& mapping : mappings) {
-    // Prepend prefix to token to form key.
-    std::string key =
-        base::StrCat({kTokenMappingPrefix, kKeySeparator, mapping.first});
+    std::string token = mapping.first;
+    std::string fqdn = mapping.second;
 
-    keys.emplace_back(key);
+    keys.emplace_back(token);
 
-    UsageStat value;
-    TokenMapping* token_mapping = value.mutable_token_mapping();
-    token_mapping->set_token(mapping.first);
-    token_mapping->set_fqdn(mapping.second);
+    TokenMapping token_mapping;
+    token_mapping.set_token(token);
+    token_mapping.set_fqdn(fqdn);
 
-    entries->emplace_back(key, value);
+    entries->emplace_back(token, token_mapping);
   }
 
   auto key_set = base::flat_set<std::string>(keys);
 
   // Add all entries created from map, remove all entries not in the map.
-  proto_db_->UpdateEntriesWithRemoveFilter(
+  token_mapping_db_->UpdateEntriesWithRemoveFilter(
       std::move(entries),
       base::BindRepeating(&DoesNotContainFilter, std::move(key_set)),
       base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
@@ -261,29 +247,24 @@ void UsageStatsDatabase::OnUpdateEntries(StatusCallback callback,
 void UsageStatsDatabase::OnLoadEntriesForGetAllEvents(
     EventsCallback callback,
     bool isSuccess,
-    std::unique_ptr<std::vector<UsageStat>> stats) {
-  std::vector<WebsiteEvent> results;
-
-  if (stats) {
-    results.reserve(stats->size());
-    for (UsageStat stat : *stats) {
-      results.emplace_back(stat.website_event());
-    }
+    std::unique_ptr<std::vector<WebsiteEvent>> events) {
+  if (isSuccess && events) {
+    std::move(callback).Run(ToError(isSuccess), *events);
+  } else {
+    std::move(callback).Run(ToError(isSuccess), std::vector<WebsiteEvent>());
   }
-
-  std::move(callback).Run(ToError(isSuccess), std::move(results));
 }
 
 void UsageStatsDatabase::OnLoadEntriesForQueryEventsInRange(
     EventsCallback callback,
     bool isSuccess,
-    std::unique_ptr<UsageStatMap> stat_map) {
+    std::unique_ptr<std::map<std::string, WebsiteEvent>> event_map) {
   std::vector<WebsiteEvent> results;
 
-  if (stat_map) {
-    results.reserve(stat_map->size());
-    for (const auto& key_stat : *stat_map) {
-      results.emplace_back(key_stat.second.website_event());
+  if (event_map) {
+    results.reserve(event_map->size());
+    for (const auto& entry : *event_map) {
+      results.emplace_back(entry.second);
     }
   }
 
@@ -293,20 +274,19 @@ void UsageStatsDatabase::OnLoadEntriesForQueryEventsInRange(
 void UsageStatsDatabase::OnLoadEntriesForDeleteEventsInRange(
     StatusCallback callback,
     bool isSuccess,
-    std::unique_ptr<UsageStatMap> stat_map) {
-  if (isSuccess && stat_map) {
+    std::unique_ptr<std::map<std::string, WebsiteEvent>> event_map) {
+  if (isSuccess && event_map) {
     // Collect keys found in range to be deleted.
     auto keys_to_delete = std::make_unique<std::vector<std::string>>();
-    keys_to_delete->reserve(stat_map->size());
+    keys_to_delete->reserve(event_map->size());
 
-    for (const auto& key_stat : *stat_map) {
-      keys_to_delete->emplace_back(key_stat.first);
+    for (const auto& entry : *event_map) {
+      keys_to_delete->emplace_back(entry.first);
     }
 
     // Remove all entries found in range.
-    proto_db_->UpdateEntries(
-        std::make_unique<
-            leveldb_proto::ProtoDatabase<UsageStat>::KeyEntryVector>(),
+    website_event_db_->UpdateEntries(
+        std::make_unique<ProtoDatabase<WebsiteEvent>::KeyEntryVector>(),
         std::move(keys_to_delete),
         base::BindOnce(&UsageStatsDatabase::OnUpdateEntries,
                        weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
@@ -318,13 +298,13 @@ void UsageStatsDatabase::OnLoadEntriesForDeleteEventsInRange(
 void UsageStatsDatabase::OnLoadEntriesForGetAllSuspensions(
     SuspensionsCallback callback,
     bool isSuccess,
-    std::unique_ptr<std::vector<UsageStat>> stats) {
+    std::unique_ptr<std::vector<Suspension>> suspensions) {
   std::vector<std::string> results;
 
-  if (stats) {
-    results.reserve(stats->size());
-    for (UsageStat stat : *stats) {
-      results.emplace_back(stat.suspension().fqdn());
+  if (suspensions) {
+    results.reserve(suspensions->size());
+    for (Suspension suspension : *suspensions) {
+      results.emplace_back(suspension.fqdn());
     }
   }
 
@@ -334,12 +314,11 @@ void UsageStatsDatabase::OnLoadEntriesForGetAllSuspensions(
 void UsageStatsDatabase::OnLoadEntriesForGetAllTokenMappings(
     TokenMappingsCallback callback,
     bool isSuccess,
-    std::unique_ptr<std::vector<UsageStat>> stats) {
+    std::unique_ptr<std::vector<TokenMapping>> mappings) {
   TokenMap results;
 
-  if (stats) {
-    for (UsageStat stat : *stats) {
-      TokenMapping mapping = stat.token_mapping();
+  if (mappings) {
+    for (TokenMapping mapping : *mappings) {
       results.emplace(mapping.token(), mapping.fqdn());
     }
   }

chrome/browser/android/usage_stats/usage_stats_database.h

@@ -19,6 +19,8 @@
 
 namespace usage_stats {
 
+using leveldb_proto::ProtoDatabase;
+
 // Stores website events, suspensions and token to fully-qualified domain name
 // (FQDN) mappings in LevelDB.
 class UsageStatsDatabase {
@@ -36,14 +38,14 @@ class UsageStatsDatabase {
 
   using StatusCallback = base::OnceCallback<void(Error)>;
 
-  using UsageStatMap = std::map<std::string, UsageStat>;
-
   // Initializes the database with user |profile|.
   explicit UsageStatsDatabase(Profile* profile);
 
   // Initializes the database with a |ProtoDatabase|. Useful for testing.
   explicit UsageStatsDatabase(
-      std::unique_ptr<leveldb_proto::ProtoDatabase<UsageStat>> proto_db);
+      std::unique_ptr<ProtoDatabase<WebsiteEvent>> website_event_db,
+      std::unique_ptr<ProtoDatabase<Suspension>> suspension_db,
+      std::unique_ptr<ProtoDatabase<TokenMapping>> token_mapping_db);
 
   ~UsageStatsDatabase();
 
@@ -84,29 +86,31 @@ class UsageStatsDatabase {
   void OnLoadEntriesForGetAllEvents(
       EventsCallback callback,
       bool isSuccess,
-      std::unique_ptr<std::vector<UsageStat>> stats);
+      std::unique_ptr<std::vector<WebsiteEvent>> stats);
 
   void OnLoadEntriesForQueryEventsInRange(
       EventsCallback callback,
       bool isSuccess,
-      std::unique_ptr<UsageStatMap> stat_map);
+      std::unique_ptr<std::map<std::string, WebsiteEvent>> event_map);
 
   void OnLoadEntriesForDeleteEventsInRange(
       StatusCallback callback,
       bool isSuccess,
-      std::unique_ptr<UsageStatMap> stat_map);
+      std::unique_ptr<std::map<std::string, WebsiteEvent>> event_map);
 
   void OnLoadEntriesForGetAllSuspensions(
       SuspensionsCallback callback,
       bool isSuccess,
-      std::unique_ptr<std::vector<UsageStat>> stats);
+      std::unique_ptr<std::vector<Suspension>> suspensions);
 
   void OnLoadEntriesForGetAllTokenMappings(
       TokenMappingsCallback callback,
       bool isSuccess,
-      std::unique_ptr<std::vector<UsageStat>> stats);
+      std::unique_ptr<std::vector<TokenMapping>> mappings);
 
-  std::unique_ptr<leveldb_proto::ProtoDatabase<UsageStat>> proto_db_;
+  std::unique_ptr<ProtoDatabase<WebsiteEvent>> website_event_db_;
+  std::unique_ptr<ProtoDatabase<Suspension>> suspension_db_;
+  std::unique_ptr<ProtoDatabase<TokenMapping>> token_mapping_db_;
 
   base::WeakPtrFactory<UsageStatsDatabase> weak_ptr_factory_;
 

chrome/browser/android/usage_stats/usage_stats_database_unittest.cc

@@ -47,20 +47,36 @@ MATCHER_P(EqualsWebsiteEvent, other, "") {
 class UsageStatsDatabaseTest : public testing::Test {
  public:
   UsageStatsDatabaseTest() {
-    auto fake_db = std::make_unique<FakeDB<UsageStat>>(&fake_db_unowned_store_);
-
-    // Maintain a pointer to the FakeDB for test callback execution.
-    fake_db_unowned_ = fake_db.get();
-
-    usage_stats_database_ =
-        std::make_unique<UsageStatsDatabase>(std::move(fake_db));
+    auto fake_website_event_db =
+        std::make_unique<FakeDB<WebsiteEvent>>(&website_event_store_);
+    auto fake_suspension_db =
+        std::make_unique<FakeDB<Suspension>>(&suspension_store_);
+    auto fake_token_mapping_db =
+        std::make_unique<FakeDB<TokenMapping>>(&token_mapping_store_);
+
+    // Maintain pointers to the FakeDBs for test callback execution.
+    website_event_db_unowned_ = fake_website_event_db.get();
+    suspension_db_unowned_ = fake_suspension_db.get();
+    token_mapping_db_unowned_ = fake_token_mapping_db.get();
+
+    usage_stats_database_ = std::make_unique<UsageStatsDatabase>(
+        std::move(fake_website_event_db), std::move(fake_suspension_db),
+        std::move(fake_token_mapping_db));
   }
 
   UsageStatsDatabase* usage_stats_database() {
     return usage_stats_database_.get();
   }
 
-  FakeDB<UsageStat>* fake_db() { return fake_db_unowned_; }
+  FakeDB<WebsiteEvent>* fake_website_event_db() {
+    return website_event_db_unowned_;
+  }
+
+  FakeDB<Suspension>* fake_suspension_db() { return suspension_db_unowned_; }
+
+  FakeDB<TokenMapping>* fake_token_mapping_db() {
+    return token_mapping_db_unowned_;
+  }
 
   MOCK_METHOD1(OnUpdateDone, void(UsageStatsDatabase::Error));
   MOCK_METHOD2(OnGetEventsDone,
@@ -71,8 +87,14 @@ class UsageStatsDatabaseTest : public testing::Test {
                void(UsageStatsDatabase::Error, UsageStatsDatabase::TokenMap));
 
  private:
-  std::map<std::string, UsageStat> fake_db_unowned_store_;
-  FakeDB<UsageStat>* fake_db_unowned_;
+  std::map<std::string, WebsiteEvent> website_event_store_;
+  std::map<std::string, Suspension> suspension_store_;
+  std::map<std::string, TokenMapping> token_mapping_store_;
+
+  FakeDB<WebsiteEvent>* website_event_db_unowned_;
+  FakeDB<Suspension>* suspension_db_unowned_;
+  FakeDB<TokenMapping>* token_mapping_db_unowned_;
+
   std::unique_ptr<UsageStatsDatabase> usage_stats_database_;
 
   DISALLOW_COPY_AND_ASSIGN(UsageStatsDatabaseTest);
@@ -80,7 +102,9 @@ class UsageStatsDatabaseTest : public testing::Test {
 
 TEST_F(UsageStatsDatabaseTest, Initialization) {
   ASSERT_NE(nullptr, usage_stats_database());
-  ASSERT_NE(nullptr, fake_db());
+  ASSERT_NE(nullptr, fake_website_event_db());
+  ASSERT_NE(nullptr, fake_suspension_db());
+  ASSERT_NE(nullptr, fake_token_mapping_db());
 }
 
 // Website Event Tests
@@ -91,7 +115,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllEventsSuccess) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, GetAllEventsFailure) {
@@ -101,7 +125,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllEventsFailure) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(false);
+  fake_website_event_db()->LoadCallback(false);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddEventsEmpty) {
@@ -113,7 +137,7 @@ TEST_F(UsageStatsDatabaseTest, AddEventsEmpty) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddAndGetOneEvent) {
@@ -128,7 +152,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndGetOneEvent) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Get 1 event.
   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
@@ -137,7 +161,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndGetOneEvent) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddAndQueryEventsInRange) {
@@ -154,7 +178,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndQueryEventsInRange) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Get events between time 0 (inclusive) and 9 (exclusive).
   // This test validates the correct lexicographic ordering of timestamps such
@@ -167,7 +191,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndQueryEventsInRange) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetEventsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddAndDeleteAllEvents) {
@@ -182,7 +206,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteAllEvents) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Delete all events.
   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
@@ -190,7 +214,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteAllEvents) {
   usage_stats_database()->DeleteAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnUpdateDone, base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Get all events (expecting none).
   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
@@ -199,7 +223,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteAllEvents) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsInRange) {
@@ -218,7 +242,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsInRange) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Delete events between time 1 (inclusive) and 10 (exclusive).
   EXPECT_CALL(*this, OnUpdateDone(UsageStatsDatabase::Error::kNoError));
@@ -228,8 +252,8 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsInRange) {
       base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->LoadCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Get 1 remaining event outside range (at time 10).
   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
@@ -238,7 +262,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsInRange) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsMatchingDomain) {
@@ -257,7 +281,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsMatchingDomain) {
       events, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                              base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Delete 2 events by FQDN.
   base::flat_set<std::string> domains({kFqdn1});
@@ -267,7 +291,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsMatchingDomain) {
       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                               base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_website_event_db()->UpdateCallback(true);
 
   // Get 1 remaining event with non-matching FQDN.
   EXPECT_CALL(*this, OnGetEventsDone(UsageStatsDatabase::Error::kNoError,
@@ -276,7 +300,7 @@ TEST_F(UsageStatsDatabaseTest, AddAndDeleteEventsMatchingDomain) {
   usage_stats_database()->GetAllEvents(base::BindOnce(
       &UsageStatsDatabaseTest::OnGetEventsDone, base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_website_event_db()->LoadCallback(true);
 }
 
 // Suspension Tests
@@ -289,7 +313,7 @@ TEST_F(UsageStatsDatabaseTest, SetSuspensionsSuccess) {
       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                               base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_suspension_db()->UpdateCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetSuspensionsFailure) {
@@ -301,7 +325,7 @@ TEST_F(UsageStatsDatabaseTest, SetSuspensionsFailure) {
       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                               base::Unretained(this)));
 
-  fake_db()->UpdateCallback(false);
+  fake_suspension_db()->UpdateCallback(false);
 }
 
 TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsSuccess) {
@@ -314,7 +338,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsSuccess) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllSuspensionsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_suspension_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsFailure) {
@@ -327,7 +351,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllSuspensionsFailure) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllSuspensionsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(false);
+  fake_suspension_db()->LoadCallback(false);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetAndGetSuspension) {
@@ -340,7 +364,7 @@ TEST_F(UsageStatsDatabaseTest, SetAndGetSuspension) {
       domains, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                               base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_suspension_db()->UpdateCallback(true);
 
   // Get 1 suspension.
   std::vector<std::string> expected({kFqdn1});
@@ -352,7 +376,7 @@ TEST_F(UsageStatsDatabaseTest, SetAndGetSuspension) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllSuspensionsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_suspension_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetSuspension) {
@@ -365,7 +389,7 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetSuspension) {
       domains1, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_suspension_db()->UpdateCallback(true);
 
   // Insert 1 suspension, and remove the other.
   base::flat_set<std::string> domains2({kFqdn1});
@@ -376,7 +400,7 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetSuspension) {
       domains2, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_suspension_db()->UpdateCallback(true);
 
   // Get 1 suspension.
   std::vector<std::string> expected({kFqdn1});
@@ -388,7 +412,7 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetSuspension) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllSuspensionsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_suspension_db()->LoadCallback(true);
 }
 
 // Token Mapping Tests
@@ -402,7 +426,7 @@ TEST_F(UsageStatsDatabaseTest, SetTokenMappingsSuccess) {
       mappings, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_token_mapping_db()->UpdateCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetTokenMappingsFailure) {
@@ -414,7 +438,7 @@ TEST_F(UsageStatsDatabaseTest, SetTokenMappingsFailure) {
       mappings, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                base::Unretained(this)));
 
-  fake_db()->UpdateCallback(false);
+  fake_token_mapping_db()->UpdateCallback(false);
 }
 
 TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsSuccess) {
@@ -427,7 +451,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsSuccess) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllTokenMappingsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_token_mapping_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsFailure) {
@@ -440,7 +464,7 @@ TEST_F(UsageStatsDatabaseTest, GetAllTokenMappingsFailure) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllTokenMappingsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(false);
+  fake_token_mapping_db()->LoadCallback(false);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetAndGetTokenMapping) {
@@ -453,7 +477,7 @@ TEST_F(UsageStatsDatabaseTest, SetAndGetTokenMapping) {
       mapping, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                               base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_token_mapping_db()->UpdateCallback(true);
 
   // Get 1 token mapping.
   EXPECT_CALL(*this, OnGetAllTokenMappingsDone(
@@ -463,7 +487,7 @@ TEST_F(UsageStatsDatabaseTest, SetAndGetTokenMapping) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllTokenMappingsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_token_mapping_db()->LoadCallback(true);
 }
 
 TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetTokenMapping) {
@@ -477,7 +501,7 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetTokenMapping) {
       mappings1, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                 base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_token_mapping_db()->UpdateCallback(true);
 
   // Re-insert 1 token mapping, and remove the other.apping) {
   UsageStatsDatabase::TokenMap mappings2({{kToken1, kFqdn1}});
@@ -488,7 +512,7 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetTokenMapping) {
       mappings2, base::BindOnce(&UsageStatsDatabaseTest::OnUpdateDone,
                                 base::Unretained(this)));
 
-  fake_db()->UpdateCallback(true);
+  fake_token_mapping_db()->UpdateCallback(true);
 
   // Get 1 remaining token mapping.
   EXPECT_CALL(*this, OnGetAllTokenMappingsDone(
@@ -498,6 +522,6 @@ TEST_F(UsageStatsDatabaseTest, SetRemoveAndGetTokenMapping) {
       base::BindOnce(&UsageStatsDatabaseTest::OnGetAllTokenMappingsDone,
                      base::Unretained(this)));
 
-  fake_db()->LoadCallback(true);
+  fake_token_mapping_db()->LoadCallback(true);
 }
 }  // namespace usage_stats

chrome/browser/android/usage_stats/website_event.proto

@@ -42,13 +42,3 @@ message TokenMapping {
 message Suspension {
   optional string fqdn = 1;
 }
-
-// Allows for any one of {WebsiteEvent, Suspension, TokenMapping} to be stored
-// in the same proto database.
-message UsageStat {
-  oneof value {
-    WebsiteEvent website_event = 1;
-    Suspension suspension = 2;
-    TokenMapping token_mapping = 3;
-  }
-}