Support server categories in NTPSnippetsService.

components/ntp_snippets/ntp_snippets_service.cc

Index: components/ntp_snippets/ntp_snippets_service.cc
diff --git a/components/ntp_snippets/ntp_snippets_service.cc b/components/ntp_snippets/ntp_snippets_service.cc
index 282c5d9765c0bd9ab31306e02dc9bb735ae968bb..7f6f9420fe85047c0e4937cb98b49f166cff7b86 100644
--- a/components/ntp_snippets/ntp_snippets_service.cc
+++ b/components/ntp_snippets/ntp_snippets_service.cc
@@ -200,9 +200,10 @@ NTPSnippetsService::NTPSnippetsService(
     std::unique_ptr<NTPSnippetsStatusService> status_service)
     : ContentSuggestionsProvider(observer, category_factory),
       state_(State::NOT_INITED),
-      category_status_(CategoryStatus::INITIALIZING),
       pref_service_(pref_service),
       suggestions_service_(suggestions_service),
+      articles_category_(
+          category_factory->FromKnownCategory(KnownCategories::ARTICLES)),
       application_language_code_(application_language_code),
       scheduler_(scheduler),
       snippets_fetcher_(std::move(snippets_fetcher)),
@@ -211,14 +212,16 @@ NTPSnippetsService::NTPSnippetsService(
       database_(std::move(database)),
       snippets_status_service_(std::move(status_service)),
       fetch_after_load_(false),
-      provided_category_(
-          category_factory->FromKnownCategory(KnownCategories::ARTICLES)),
       thumbnail_requests_throttler_(
           pref_service,
           RequestThrottler::RequestType::CONTENT_SUGGESTION_THUMBNAIL) {
-  observer->OnCategoryStatusChanged(this, provided_category_, category_status_);
+  // Articles category always exists; others will be added as needed.
+  categories_.emplace(articles_category_, CategoryContent());

[Marc Treib, 2016/08/22 15:06:46]

map::emplace isn't available everywhere yet, see
https://chromium-cpp.appspot.com/

[sfiera, 2016/08/24 14:35:56]

Done.

+  observer->OnCategoryStatusChanged(this, articles_category_,
+                                    CategoryStatus::INITIALIZING);

[Marc Treib, 2016/08/22 15:06:47]

Use the actual value from categories_, so the default state isn't hardcoded in
two places

[sfiera, 2016/08/24 14:35:56]

Done.

   if (database_->IsErrorState()) {
-    EnterState(State::ERROR_OCCURRED, CategoryStatus::LOADING_ERROR);
+    EnterState(State::ERROR_OCCURRED);
+    UpdateAllCategoryStatus(CategoryStatus::LOADING_ERROR);
     return;
   }
 
@@ -254,8 +257,11 @@ void NTPSnippetsService::FetchSnippetsFromHosts(
   if (!ready())
     return;
 
-  if (snippets_.empty())
-    UpdateCategoryStatus(CategoryStatus::AVAILABLE_LOADING);
+  // Empty categories are marked as loading; others are unchanged.
+  for (const auto& category : categories_) {

[Marc Treib, 2016/08/22 15:06:47]

Use the actual type? auto isn't clear here, since this isn't actually a
Category.

[sfiera, 2016/08/24 14:35:56]

I've switched this to what I do in most places, which is to break the pair out
immediately into locals.

+    if (category.second.snippets.empty())
+      UpdateCategoryStatus(category.first, CategoryStatus::AVAILABLE_LOADING);
+  }
 
   snippets_fetcher_->FetchSnippetsFromHosts(
       hosts, application_language_code_, kMaxSnippetCount, interactive_request);
@@ -277,11 +283,13 @@ void NTPSnippetsService::RescheduleFetching() {
 }
 
 CategoryStatus NTPSnippetsService::GetCategoryStatus(Category category) {
-  DCHECK(category.IsKnownCategory(KnownCategories::ARTICLES));
-  return category_status_;
+  DCHECK(categories_.find(category) != categories_.end());
+  return categories_[category].status;
 }
 
 CategoryInfo NTPSnippetsService::GetCategoryInfo(Category category) {
+  DCHECK(categories_.find(category) != categories_.end());
+  // TODO(sfiera): pass back localized category titles.

[Marc Treib, 2016/08/22 15:06:47]

nit: mention that this refers to non-articles categories, otherwise it's a bit
confusing: we *do* return a localized string!

[sfiera, 2016/08/24 14:35:56]

Changed.

   return CategoryInfo(
       l10n_util::GetStringUTF16(IDS_NTP_ARTICLE_SUGGESTIONS_SECTION_HEADER),
       ContentSuggestionsCardLayout::FULL_CARD,
@@ -293,14 +301,18 @@ void NTPSnippetsService::DismissSuggestion(const std::string& suggestion_id) {
   if (!ready())
     return;
 
+  Category category_id = GetCategoryFromUniqueID(suggestion_id);

[Marc Treib, 2016/08/22 15:06:47]

Just "category" please (that's what it's called everywhere else)

[sfiera, 2016/08/24 14:35:56]

Done.

   std::string snippet_id = GetWithinCategoryIDFromUniqueID(suggestion_id);
 
+  DCHECK(categories_.find(category_id) != categories_.end());
+
+  CategoryContent* category = &categories_[category_id];

[Marc Treib, 2016/08/22 15:06:46]

call this something else, like "content"

[sfiera, 2016/08/24 14:35:56]

Done.

   auto it =
-      std::find_if(snippets_.begin(), snippets_.end(),
+      std::find_if(category->snippets.begin(), category->snippets.end(),
                    [&snippet_id](const std::unique_ptr<NTPSnippet>& snippet) {
                      return snippet->id() == snippet_id;
                    });
-  if (it == snippets_.end())
+  if (it == category->snippets.end())
     return;
 
   (*it)->set_dismissed(true);
@@ -308,44 +320,48 @@ void NTPSnippetsService::DismissSuggestion(const std::string& suggestion_id) {
   database_->SaveSnippet(**it);
   database_->DeleteImage((*it)->id());
 
-  dismissed_snippets_.push_back(std::move(*it));
-  snippets_.erase(it);
+  category->dismissed.push_back(std::move(*it));
+  category->snippets.erase(it);
 }
 
 void NTPSnippetsService::FetchSuggestionImage(
     const std::string& suggestion_id,
     const ImageFetchedCallback& callback) {
-  std::string snippet_id = GetWithinCategoryIDFromUniqueID(suggestion_id);
   database_->LoadImage(
-      snippet_id,
+      suggestion_id,
       base::Bind(&NTPSnippetsService::OnSnippetImageFetchedFromDatabase,
-                 base::Unretained(this), callback, snippet_id));
+                 base::Unretained(this), callback, suggestion_id));
 }
 
-void NTPSnippetsService::ClearCachedSuggestionsForDebugging(Category category) {
-  DCHECK_EQ(category, provided_category_);
+void NTPSnippetsService::ClearCachedSuggestionsForDebugging(
+    Category category_id) {
   if (!initialized())
     return;
 
-  if (snippets_.empty())
+  if (categories_.find(category_id) == categories_.end())
+    return;
+  CategoryContent* category = &categories_[category_id];
+  if (category->snippets.empty())
     return;
 
-  database_->DeleteSnippets(snippets_);
-  snippets_.clear();
+  if (category_id == articles_category_)
+    database_->DeleteSnippets(category->snippets);
+  category->snippets.clear();
 
   NotifyNewSuggestions();
 }
 
 std::vector<ContentSuggestion>
-NTPSnippetsService::GetDismissedSuggestionsForDebugging(Category category) {
-  DCHECK_EQ(category, provided_category_);
+NTPSnippetsService::GetDismissedSuggestionsForDebugging(Category category_id) {
+  DCHECK(categories_.find(category_id) != categories_.end());
+
   std::vector<ContentSuggestion> result;
-  for (const std::unique_ptr<NTPSnippet>& snippet : dismissed_snippets_) {
+  const CategoryContent& category = categories_[category_id];
+  for (const std::unique_ptr<NTPSnippet>& snippet : category.dismissed) {
     if (!snippet->is_complete())
       continue;
-    ContentSuggestion suggestion(
-        MakeUniqueID(provided_category_, snippet->id()),
-        snippet->best_source().url);
+    ContentSuggestion suggestion(MakeUniqueID(category_id, snippet->id()),
+                                 snippet->best_source().url);
     suggestion.set_amp_url(snippet->best_source().amp_url);
     suggestion.set_title(base::UTF8ToUTF16(snippet->title()));
     suggestion.set_snippet_text(base::UTF8ToUTF16(snippet->snippet()));
@@ -359,16 +375,19 @@ NTPSnippetsService::GetDismissedSuggestionsForDebugging(Category category) {
 }
 
 void NTPSnippetsService::ClearDismissedSuggestionsForDebugging(
-    Category category) {
-  DCHECK_EQ(category, provided_category_);
+    Category category_id) {
+  DCHECK(categories_.find(category_id) != categories_.end());
+
   if (!initialized())
     return;
 
-  if (dismissed_snippets_.empty())
+  CategoryContent* category = &categories_[category_id];
+  if (category->dismissed.empty())
     return;
 
-  database_->DeleteSnippets(dismissed_snippets_);
-  dismissed_snippets_.clear();
+  if (category_id == articles_category_)
+    database_->DeleteSnippets(category->dismissed);
+  category->dismissed.clear();
 }
 
 std::set<std::string> NTPSnippetsService::GetSuggestionsHosts() const {
@@ -390,47 +409,61 @@ int NTPSnippetsService::GetMaxSnippetCountForTesting() {
 // Private methods
 
 // image_fetcher::ImageFetcherDelegate implementation.
-void NTPSnippetsService::OnImageDataFetched(const std::string& snippet_id,
+void NTPSnippetsService::OnImageDataFetched(const std::string& suggestion_id,
                                             const std::string& image_data) {
   if (image_data.empty())
     return;
 
+  Category category_id = GetCategoryFromUniqueID(suggestion_id);
+  std::string snippet_id = GetWithinCategoryIDFromUniqueID(suggestion_id);
+
+  auto category_it = categories_.find(category_id);
+  if (category_it == categories_.end())
+    return;
+
+  const CategoryContent& category = category_it->second;
+
   // Only save the image if the corresponding snippet still exists.
   auto it =
-      std::find_if(snippets_.begin(), snippets_.end(),
+      std::find_if(category.snippets.begin(), category.snippets.end(),
                    [&snippet_id](const std::unique_ptr<NTPSnippet>& snippet) {
                      return snippet->id() == snippet_id;
                    });
-  if (it == snippets_.end())
+  if (it == category.snippets.end())
     return;
 
-  database_->SaveImage(snippet_id, image_data);
+  database_->SaveImage(suggestion_id, image_data);
 }
 
 void NTPSnippetsService::OnDatabaseLoaded(NTPSnippet::PtrVector snippets) {
   if (state_ == State::ERROR_OCCURRED)
     return;
   DCHECK(state_ == State::NOT_INITED);
-  DCHECK(snippets_.empty());
-  DCHECK(dismissed_snippets_.empty());
+  DCHECK(categories_.size() == 1);  // Only articles category, so far.
+  DCHECK(categories_.find(articles_category_) != categories_.end());
+
+  // TODO(sfiera): support non-article categories in database.
+  CategoryContent* category = &categories_[articles_category_];
   for (std::unique_ptr<NTPSnippet>& snippet : snippets) {
     if (snippet->is_dismissed())
-      dismissed_snippets_.emplace_back(std::move(snippet));
+      category->dismissed.emplace_back(std::move(snippet));
     else
-      snippets_.emplace_back(std::move(snippet));
+      category->snippets.emplace_back(std::move(snippet));
   }
-  std::sort(snippets_.begin(), snippets_.end(),
+
+  std::sort(category->snippets.begin(), category->snippets.end(),
             [](const std::unique_ptr<NTPSnippet>& lhs,
                const std::unique_ptr<NTPSnippet>& rhs) {
               return lhs->score() > rhs->score();
             });
 
-  ClearExpiredSnippets();
+  ClearExpiredSnippets(articles_category_);
   FinishInitialization();
 }
 
 void NTPSnippetsService::OnDatabaseError() {
-  EnterState(State::ERROR_OCCURRED, CategoryStatus::LOADING_ERROR);
+  EnterState(State::ERROR_OCCURRED);
+  UpdateAllCategoryStatus(CategoryStatus::LOADING_ERROR);
 }
 
 // TODO(dgn): name clash between content suggestions and suggestions hosts.
@@ -444,15 +477,20 @@ void NTPSnippetsService::OnSuggestionsChanged(
     return;
 
   // Remove existing snippets that aren't in the suggestions anymore.
+  //
   // TODO(treib,maybelle): If there is another source with an allowed host,
   // then we should fall back to that.
   // First, move them over into |to_delete|.

[Marc Treib, 2016/08/22 15:06:46]

This line isn't part of the TODO; move it to the end?

[sfiera, 2016/08/24 14:35:56]

Done.

+  //
+  // TODO(sfiera): determine when non-article categories should restrict hosts,
+  // and apply the same logic to them here. Maybe never?

[Marc Treib, 2016/08/22 15:06:47]

Eh, since host restricts are likely going away anyway, probably not worth
bothering with. (It's fine to leave the TODO there though)

[sfiera, 2016/08/24 14:35:56]

Acknowledged.

+  CategoryContent* category = &categories_[articles_category_];
   NTPSnippet::PtrVector to_delete;
-  for (std::unique_ptr<NTPSnippet>& snippet : snippets_) {
+  for (std::unique_ptr<NTPSnippet>& snippet : category->snippets) {
     if (!hosts.count(snippet->best_source().url.host()))
       to_delete.emplace_back(std::move(snippet));
   }
-  Compact(&snippets_);
+  Compact(&category->snippets);
   // Then delete the removed snippets from the database.
   database_->DeleteSnippets(to_delete);
 
@@ -471,48 +509,84 @@ void NTPSnippetsService::OnFetchFinished(
   if (!ready())
     return;
 
-  DCHECK(category_status_ == CategoryStatus::AVAILABLE ||
-         category_status_ == CategoryStatus::AVAILABLE_LOADING);
+  // If snippets were fetched successfully, update our |categories_| from each
+  // source and report the category as AVAILABLE.
+  if (snippets) {
+    for (std::pair<const Category, NTPSnippet::PtrVector>& item : *snippets) {
+      Category category_id = item.first;
+      NTPSnippet::PtrVector& new_snippets = item.second;
+      // Sparse histogram used because the number of snippets is small (bound by
+      // kMaxSnippetCount).
+      DCHECK_LE(snippets->size(), static_cast<size_t>(kMaxSnippetCount));
+      // TODO(sfiera): per-category histograms
+      if (category_id == articles_category_)

[Marc Treib, 2016/08/22 15:06:47]

Braces please

[sfiera, 2016/08/24 14:35:56]

Done.

+        UMA_HISTOGRAM_SPARSE_SLOWLY("NewTabPage.Snippets.NumArticlesFetched",
+                                    new_snippets.size());
 
-  // TODO(sfiera): support more than just the provided_category_ ARTICLES.
-  if (snippets && (snippets->find(provided_category_) != snippets->end())) {
-    // Sparse histogram used because the number of snippets is small (bound by
-    // kMaxSnippetCount).
-    DCHECK_LE(snippets->size(), static_cast<size_t>(kMaxSnippetCount));
-    UMA_HISTOGRAM_SPARSE_SLOWLY("NewTabPage.Snippets.NumArticlesFetched",
-                                (*snippets)[provided_category_].size());
-    MergeSnippets(std::move((*snippets)[provided_category_]));
+      MergeSnippets(category_id, std::move(new_snippets));
+      ClearExpiredSnippets(category_id);
+      // If there are more snippets than we want to show, delete the extra ones.
+      CategoryContent* category = &categories_[category_id];
+      if (category->snippets.size() > kMaxSnippetCount) {
+        NTPSnippet::PtrVector to_delete(
+            std::make_move_iterator(category->snippets.begin() +
+                                    kMaxSnippetCount),
+            std::make_move_iterator(category->snippets.end()));
+        category->snippets.resize(kMaxSnippetCount);
+        database_->DeleteSnippets(to_delete);
+      }
+
+      // Report category as available, even if empty. It is provided by the
+      // service, and is no longer loading.
+      UpdateCategoryStatus(category_id, CategoryStatus::AVAILABLE);
+    }
   }
 
-  ClearExpiredSnippets();
+  // Go back and expire articles out of categories that weren't updated. If
+  // |snippets| was not provided due to an error, that's all categories.
+  std::vector<Category> categories_to_remove_;
+  for (auto& item : categories_) {
+    Category category_id = item.first;
+    CategoryContent& category = item.second;
+    if (!snippets || (snippets->find(category_id) == snippets->end()))
+      continue;  // skip anything we processed in the first loop.

[Marc Treib, 2016/08/22 15:06:46]

Why?
I think the main purpose of calling ClearExpiredSnippets here was to schedule
the next expiry run, so we want this particularly for the newly-updated stuff.
For the categories that weren't updated, it should already be scheduled.

[sfiera, 2016/08/24 14:35:56]

I'd sort of misunderstood ClearExpiredSnippets(), and had broken it. In
particular, my understanding of OneShotTimer now is that only one callback can
be scheduled, so the callbacks shouldn't be category-specific.

The logic for dropping categories has moved there too.

+
+    ClearExpiredSnippets(category_id);
+    if (category.snippets.empty() && (category_id != articles_category_)) {
+      categories_to_remove_.push_back(category_id);
+    }
+  }
 
-  // If there are more snippets than we want to show, delete the extra ones.
-  if (snippets_.size() > kMaxSnippetCount) {
-    NTPSnippet::PtrVector to_delete(
-        std::make_move_iterator(snippets_.begin() + kMaxSnippetCount),
-        std::make_move_iterator(snippets_.end()));
-    snippets_.resize(kMaxSnippetCount);
-    database_->DeleteSnippets(to_delete);
+  // If a category (other than ARTICLES) is now empty, and it was not present in
+  // |snippets|, then report that the snippets service no longer provides that
+  // category. If the category is now |empty| but it was in |snippets|, then we

[Marc Treib, 2016/08/22 15:06:46]

nit: no pipes around empty

[sfiera, 2016/08/24 14:35:57]

Done.

+  // will continue to report it as AVAILABLE.
+  for (Category category : categories_to_remove_) {
+    UpdateCategoryStatus(category, CategoryStatus::NOT_PROVIDED);
+    categories_.erase(category);
   }
 
+  // TODO(sfiera): equivalent metrics for non-articles.
+  const CategoryContent& category = categories_[articles_category_];
   UMA_HISTOGRAM_SPARSE_SLOWLY("NewTabPage.Snippets.NumArticles",
-                              snippets_.size());
-  if (snippets_.empty() && !dismissed_snippets_.empty()) {
+                              category.snippets.size());
+  if (category.snippets.empty() && !category.dismissed.empty()) {
     UMA_HISTOGRAM_COUNTS("NewTabPage.Snippets.NumArticlesZeroDueToDiscarded",
-                         dismissed_snippets_.size());
+                         category.dismissed.size());
   }
 
-  UpdateCategoryStatus(CategoryStatus::AVAILABLE);
   NotifyNewSuggestions();
 }
 
-void NTPSnippetsService::MergeSnippets(NTPSnippet::PtrVector new_snippets) {
+void NTPSnippetsService::MergeSnippets(Category category_id,
+                                       NTPSnippet::PtrVector new_snippets) {
   DCHECK(ready());
+  CategoryContent* category = &categories_[category_id];
 
   // Remove new snippets that we already have, or that have been dismissed.
   std::set<std::string> old_snippet_ids;
-  InsertAllIDs(dismissed_snippets_, &old_snippet_ids);
-  InsertAllIDs(snippets_, &old_snippet_ids);
+  InsertAllIDs(category->dismissed, &old_snippet_ids);
+  InsertAllIDs(category->snippets, &old_snippet_ids);
   new_snippets.erase(
       std::remove_if(
           new_snippets.begin(), new_snippets.end(),
@@ -560,13 +634,15 @@ void NTPSnippetsService::MergeSnippets(NTPSnippet::PtrVector new_snippets) {
     }
   }
 
-  // Save the new snippets to the DB.
-  database_->SaveSnippets(new_snippets);
+  // Save new articles to the DB.
+  // TODO(sfiera): save non-articles to DB too.
+  if (category_id == articles_category_)
+    database_->SaveSnippets(new_snippets);
 
   // Insert the new snippets at the front.
-  snippets_.insert(snippets_.begin(),
-                   std::make_move_iterator(new_snippets.begin()),
-                   std::make_move_iterator(new_snippets.end()));
+  category->snippets.insert(category->snippets.begin(),
+                            std::make_move_iterator(new_snippets.begin()),
+                            std::make_move_iterator(new_snippets.end()));
 }
 
 std::set<std::string> NTPSnippetsService::GetSnippetHostsFromPrefs() const {
@@ -589,94 +665,109 @@ void NTPSnippetsService::StoreSnippetHostsToPrefs(
   pref_service_->Set(prefs::kSnippetHosts, list);
 }
 
-void NTPSnippetsService::ClearExpiredSnippets() {
+void NTPSnippetsService::ClearExpiredSnippets(Category category_id) {
+  DCHECK(categories_.find(category_id) != categories_.end());
+  CategoryContent* category = &categories_[category_id];
+
   base::Time expiry = base::Time::Now();
 
   // Move expired snippets over into |to_delete|.
   NTPSnippet::PtrVector to_delete;
-  for (std::unique_ptr<NTPSnippet>& snippet : snippets_) {
+  for (std::unique_ptr<NTPSnippet>& snippet : category->snippets) {
     if (snippet->expiry_date() <= expiry)
       to_delete.emplace_back(std::move(snippet));
   }
-  Compact(&snippets_);
+  Compact(&category->snippets);
 
   // Move expired dismissed snippets over into |to_delete| as well.
-  for (std::unique_ptr<NTPSnippet>& snippet : dismissed_snippets_) {
+  for (std::unique_ptr<NTPSnippet>& snippet : category->dismissed) {
     if (snippet->expiry_date() <= expiry)
       to_delete.emplace_back(std::move(snippet));
   }
-  Compact(&dismissed_snippets_);
+  Compact(&category->dismissed);
 
   // Finally, actually delete the removed snippets from the DB.
-  database_->DeleteSnippets(to_delete);
+  if (category_id == articles_category_)
+    database_->DeleteSnippets(to_delete);
 
-  // If there are any snippets left, schedule a timer for the next expiry.
-  if (snippets_.empty() && dismissed_snippets_.empty())
+  // Unless there are no snippets left, schedule a timer for the next expiry.
+  if (category->snippets.empty() && category->dismissed.empty())
     return;
 
   base::Time next_expiry = base::Time::Max();
-  for (const auto& snippet : snippets_) {
+  for (const auto& snippet : category->snippets) {
     if (snippet->expiry_date() < next_expiry)
       next_expiry = snippet->expiry_date();
   }
-  for (const auto& snippet : dismissed_snippets_) {
+  for (const auto& snippet : category->dismissed) {
     if (snippet->expiry_date() < next_expiry)
       next_expiry = snippet->expiry_date();
   }
+
   DCHECK_GT(next_expiry, expiry);
   expiry_timer_.Start(FROM_HERE, next_expiry - expiry,
                       base::Bind(&NTPSnippetsService::ClearExpiredSnippets,
-                                 base::Unretained(this)));
+                                 base::Unretained(this), category_id));
 }
 
 void NTPSnippetsService::OnSnippetImageFetchedFromDatabase(
     const ImageFetchedCallback& callback,
-    const std::string& snippet_id,
+    const std::string& suggestion_id,
     std::string data) {
   // |image_decoder_| is null in tests.
   if (image_decoder_ && !data.empty()) {
     image_decoder_->DecodeImage(
         std::move(data),
         base::Bind(&NTPSnippetsService::OnSnippetImageDecodedFromDatabase,
-                   base::Unretained(this), callback, snippet_id));
+                   base::Unretained(this), callback, suggestion_id));
     return;
   }
 
   // Fetching from the DB failed; start a network fetch.
-  FetchSnippetImageFromNetwork(snippet_id, callback);
+  FetchSnippetImageFromNetwork(suggestion_id, callback);
 }
 
 void NTPSnippetsService::OnSnippetImageDecodedFromDatabase(
     const ImageFetchedCallback& callback,
-    const std::string& snippet_id,
+    const std::string& suggestion_id,
     const gfx::Image& image) {
   if (!image.IsEmpty()) {
-    callback.Run(MakeUniqueID(provided_category_, snippet_id), image);
+    callback.Run(suggestion_id, image);
     return;
   }
 
   // If decoding the image failed, delete the DB entry.
-  database_->DeleteImage(snippet_id);
+  database_->DeleteImage(suggestion_id);
 
-  FetchSnippetImageFromNetwork(snippet_id, callback);
+  FetchSnippetImageFromNetwork(suggestion_id, callback);
 }
 
 void NTPSnippetsService::FetchSnippetImageFromNetwork(
-    const std::string& snippet_id,
+    const std::string& suggestion_id,
     const ImageFetchedCallback& callback) {
+  Category category_id = GetCategoryFromUniqueID(suggestion_id);
+  std::string snippet_id = GetWithinCategoryIDFromUniqueID(suggestion_id);
+
+  auto category_it = categories_.find(category_id);
+  if (category_it == categories_.end()) {
+    callback.Run(suggestion_id, gfx::Image());
+    return;
+  }
+
+  const CategoryContent& category = category_it->second;
   auto it =
-      std::find_if(snippets_.begin(), snippets_.end(),
+      std::find_if(category.snippets.begin(), category.snippets.end(),
                    [&snippet_id](const std::unique_ptr<NTPSnippet>& snippet) {
                      return snippet->id() == snippet_id;
                    });
 
-  if (it == snippets_.end() ||
+  if (it == category.snippets.end() ||
       !thumbnail_requests_throttler_.DemandQuotaForRequest(
           /*interactive_request=*/true)) {
     // Return an empty image. Directly, this is never synchronous with the
     // original FetchSuggestionImage() call - an asynchronous database query has
     // happened in the meantime.
-    OnSnippetImageDecodedFromNetwork(callback, snippet_id, gfx::Image());
+    OnSnippetImageDecodedFromNetwork(callback, suggestion_id, gfx::Image());
     return;
   }
 
@@ -689,16 +780,16 @@ void NTPSnippetsService::FetchSnippetImageFromNetwork(
   // is an ImageFetcherDelegate) and then also
   // OnSnippetImageDecodedFromNetwork() after the raw data gets decoded.
   image_fetcher_->StartOrQueueNetworkRequest(
-      snippet.id(), snippet.salient_image_url(),
+      suggestion_id, snippet.salient_image_url(),
       base::Bind(&NTPSnippetsService::OnSnippetImageDecodedFromNetwork,
                  base::Unretained(this), callback));
 }
 
 void NTPSnippetsService::OnSnippetImageDecodedFromNetwork(
     const ImageFetchedCallback& callback,
-    const std::string& snippet_id,
+    const std::string& suggestion_id,
     const gfx::Image& image) {
-  callback.Run(MakeUniqueID(provided_category_, snippet_id), image);
+  callback.Run(suggestion_id, image);
 }
 
 void NTPSnippetsService::EnterStateEnabled(bool fetch_snippets) {
@@ -707,8 +798,9 @@ void NTPSnippetsService::EnterStateEnabled(bool fetch_snippets) {
 
   // FetchSnippets should set the status to |AVAILABLE_LOADING| if relevant,
   // otherwise we transition to |AVAILABLE| here.
-  if (category_status_ != CategoryStatus::AVAILABLE_LOADING)
-    UpdateCategoryStatus(CategoryStatus::AVAILABLE);
+  if (categories_[articles_category_].status !=
+      CategoryStatus::AVAILABLE_LOADING)

[Marc Treib, 2016/08/22 15:06:46]

Braces please

[sfiera, 2016/08/24 14:35:57]

Done.

+    UpdateCategoryStatus(articles_category_, CategoryStatus::AVAILABLE);
 
   // If host restrictions are enabled, register for host list updates.
   // |suggestions_service_| can be null in tests.
@@ -722,8 +814,22 @@ void NTPSnippetsService::EnterStateEnabled(bool fetch_snippets) {
 }
 
 void NTPSnippetsService::EnterStateDisabled() {
-  ClearCachedSuggestionsForDebugging(provided_category_);
-  ClearDismissedSuggestionsForDebugging(provided_category_);
+  std::vector<Category> category_list;
+  for (const auto& item : categories_) {
+    Category category_id = item.first;
+    category_list.push_back(category_id);
+  }
+
+  // Empty the ARTICLES category and remove all others, since they may or may
+  // not be personalized.
+  for (Category category_id : category_list) {
+    ClearCachedSuggestionsForDebugging(category_id);
+    ClearDismissedSuggestionsForDebugging(category_id);
+    if (category_id != articles_category_) {
+      UpdateCategoryStatus(category_id, CategoryStatus::NOT_PROVIDED);
+      categories_.erase(category_id);
+    }
+  }
 
   expiry_timer_.Stop();
   suggestions_service_subscription_.reset();
@@ -763,22 +869,22 @@ void NTPSnippetsService::OnDisabledReasonChanged(
   switch (disabled_reason) {
     case DisabledReason::NONE:
       // Do not change the status. That will be done in EnterStateEnabled()
-      EnterState(State::READY, category_status_);
+      EnterState(State::READY);
       break;
 
     case DisabledReason::EXPLICITLY_DISABLED:
-      EnterState(State::DISABLED, CategoryStatus::CATEGORY_EXPLICITLY_DISABLED);
+      EnterState(State::DISABLED);
+      UpdateAllCategoryStatus(CategoryStatus::CATEGORY_EXPLICITLY_DISABLED);
       break;
 
     case DisabledReason::SIGNED_OUT:
-      EnterState(State::DISABLED, CategoryStatus::SIGNED_OUT);
+      EnterState(State::DISABLED);
+      UpdateAllCategoryStatus(CategoryStatus::SIGNED_OUT);
       break;
   }
 }
 
-void NTPSnippetsService::EnterState(State state, CategoryStatus status) {
-  UpdateCategoryStatus(status);
-
+void NTPSnippetsService::EnterState(State state) {
   if (state == state_)
     return;
 
@@ -791,7 +897,10 @@ void NTPSnippetsService::EnterState(State state, CategoryStatus status) {
     case State::READY: {
       DCHECK(state_ == State::NOT_INITED || state_ == State::DISABLED);
 
-      bool fetch_snippets = snippets_.empty() || fetch_after_load_;
+      // TODO(sfiera): always fetch, because there might be server-side
+      // categories we're interested in?
+      bool fetch_snippets =
+          categories_[articles_category_].snippets.empty() || fetch_after_load_;
       DVLOG(1) << "Entering state: READY";
       state_ = State::READY;
       fetch_after_load_ = false;
@@ -816,32 +925,50 @@ void NTPSnippetsService::EnterState(State state, CategoryStatus status) {
 }
 
 void NTPSnippetsService::NotifyNewSuggestions() {
-  std::vector<ContentSuggestion> result;
-  for (const std::unique_ptr<NTPSnippet>& snippet : snippets_) {
-    if (!snippet->is_complete())
-      continue;
-    ContentSuggestion suggestion(
-        MakeUniqueID(provided_category_, snippet->id()),
-        snippet->best_source().url);
-    suggestion.set_amp_url(snippet->best_source().amp_url);
-    suggestion.set_title(base::UTF8ToUTF16(snippet->title()));
-    suggestion.set_snippet_text(base::UTF8ToUTF16(snippet->snippet()));
-    suggestion.set_publish_date(snippet->publish_date());
-    suggestion.set_publisher_name(
-        base::UTF8ToUTF16(snippet->best_source().publisher_name));
-    suggestion.set_score(snippet->score());
-    result.emplace_back(std::move(suggestion));
+  for (auto& item : categories_) {
+    Category category_id = item.first;
+    CategoryContent& category = item.second;
+
+    std::vector<ContentSuggestion> result;
+    for (const std::unique_ptr<NTPSnippet>& snippet : category.snippets) {
+      // TODO(sfiera): if a snippet is not going to be displayed, move it
+      // directly to category.dismissed on fetch. Otherwise, we might prune
+      // other snippets to get down to kMaxSnippetCount, only to hide one of the
+      // incomplete ones we kept.
+      if (!snippet->is_complete())
+        continue;
+      ContentSuggestion suggestion(MakeUniqueID(category_id, snippet->id()),
+                                   snippet->best_source().url);
+      suggestion.set_amp_url(snippet->best_source().amp_url);
+      suggestion.set_title(base::UTF8ToUTF16(snippet->title()));
+      suggestion.set_snippet_text(base::UTF8ToUTF16(snippet->snippet()));
+      suggestion.set_publish_date(snippet->publish_date());
+      suggestion.set_publisher_name(
+          base::UTF8ToUTF16(snippet->best_source().publisher_name));
+      suggestion.set_score(snippet->score());
+      result.emplace_back(std::move(suggestion));
+    }
+    // TODO(sfiera): detect when a category is unchanged and don't notify the
+    // observer in that case.
+    observer()->OnNewSuggestions(this, category_id, std::move(result));
   }
-  observer()->OnNewSuggestions(this, provided_category_, std::move(result));
 }
 
-void NTPSnippetsService::UpdateCategoryStatus(CategoryStatus status) {
-  if (status == category_status_)
+void NTPSnippetsService::UpdateCategoryStatus(Category category_id,
+                                              CategoryStatus status) {
+  DCHECK(categories_.find(category_id) != categories_.end());
+  CategoryContent& category = categories_[category_id];
+  if (status == category.status)
     return;
 
-  category_status_ = status;
-  observer()->OnCategoryStatusChanged(this, provided_category_,
-                                      category_status_);
+  category.status = status;
+  observer()->OnCategoryStatusChanged(this, category_id, category.status);
+}
+
+void NTPSnippetsService::UpdateAllCategoryStatus(CategoryStatus status) {
+  for (const auto& category : categories_) {
+    UpdateCategoryStatus(category.first, status);
+  }
 }
 
 }  // namespace ntp_snippets