Count site data size for important sites

chrome/browser/engagement/important_sites_util.cc

 // 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 "chrome/browser/engagement/important_sites_util.h"
 
 #include <algorithm>
 #include <map>
 #include <memory>
 #include <set>
 #include <utility>
 
+#include "base/bind_helpers.h"
 #include "base/containers/hash_tables.h"
 #include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/stl_util.h"
 #include "base/time/time.h"
 #include "base/values.h"
 #include "chrome/browser/banners/app_banner_settings_helper.h"
 #include "chrome/browser/bookmarks/bookmark_model_factory.h"
 #include "chrome/browser/content_settings/host_content_settings_map_factory.h"
 #include "chrome/browser/engagement/site_engagement_score.h"
 #include "chrome/browser/engagement/site_engagement_service.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/common/pref_names.h"
 #include "components/bookmarks/browser/bookmark_model.h"
 #include "components/content_settings/core/browser/host_content_settings_map.h"
 #include "components/content_settings/core/common/content_settings.h"
 #include "components/pref_registry/pref_registry_syncable.h"
 #include "components/prefs/pref_service.h"
 #include "components/prefs/scoped_user_pref_update.h"
+#include "content/public/browser/browser_thread.h"
+#include "content/public/browser/dom_storage_context.h"
+#include "content/public/browser/local_storage_usage_info.h"
 #include "net/base/registry_controlled_domains/registry_controlled_domain.h"
+#include "storage/browser/quota/quota_manager.h"
 #include "third_party/WebKit/public/platform/site_engagement.mojom.h"
 #include "url/gurl.h"
+#include "url/url_util.h"
 
 namespace {
 using bookmarks::BookmarkModel;
+using content::BrowserThread;
 using ImportantDomainInfo = ImportantSitesUtil::ImportantDomainInfo;
 using ImportantReason = ImportantSitesUtil::ImportantReason;
 
 // Note: These values are stored on both the per-site content settings
 // dictionary and the dialog preference dictionary.
 
 static const char kTimeLastIgnored[] = "TimeLastIgnored";
 static const int kBlacklistExpirationTimeDays = 30 * 5;
 
 static const char kNumTimesIgnoredName[] = "NumTimesIgnored";
@@ skipping 86 matching lines @@
 
 std::string GetRegisterableDomainOrIP(const GURL& url) {
   std::string registerable_domain =
       net::registry_controlled_domains::GetDomainAndRegistry(
           url, net::registry_controlled_domains::INCLUDE_PRIVATE_REGISTRIES);
   if (registerable_domain.empty() && url.HostIsIPAddress())
     registerable_domain = url.host();
   return registerable_domain;
 }
 
+std::string GetRegisterableDomainOrIPFromHost(const std::string& host) {
+  std::string registerable_domain =
+      net::registry_controlled_domains::GetDomainAndRegistry(
+          host, net::registry_controlled_domains::INCLUDE_PRIVATE_REGISTRIES);
+  if (registerable_domain.empty() && url::HostIsIPAddress(host))
+    registerable_domain = host;
+  return registerable_domain;
+}
+
 void MaybePopulateImportantInfoForReason(
     const GURL& origin,
     std::set<GURL>* visited_origins,
     ImportantReason reason,
     base::hash_map<std::string, ImportantDomainInfo>* output) {
   if (!origin.is_valid() || !visited_origins->insert(origin).second)
     return;
   std::string registerable_domain = GetRegisterableDomainOrIP(origin);
   ImportantDomainInfo& info = (*output)[registerable_domain];
   info.reason_bitfield |= 1 << reason;
@@ skipping 178 matching lines @@
   base::Time now = base::Time::Now();
   for (const ContentSettingPatternSource& site : content_settings_list) {
     GURL origin(site.primary_pattern.ToString());
     if (!AppBannerSettingsHelper::WasLaunchedRecently(profile, origin, now))
       continue;
     MaybePopulateImportantInfoForReason(origin, &content_origins,
                                         ImportantReason::HOME_SCREEN, output);
   }
 }
 
+class UsageReceiver {
+ public:
+  UsageReceiver(ImportantSitesUtil::UsageCallback done,
+                std::vector<ImportantDomainInfo> sites,
+                storage::QuotaManager* quota_manager,
+                content::DOMStorageContext* dom_storage_context)
+      : done_(done),
+        sites_(sites),
+        quota_manager_(quota_manager),
+        dom_storage_context_(dom_storage_context),
+        tasks_(-1) {
+    for (auto& site : sites_) {
+      site.usage = 0;
+      site_map[site.registerable_domain] = &site;
+    }
+  }
+
+  void RunAndDestroySelf() {
+    BrowserThread::PostTask(
+        BrowserThread::IO, FROM_HERE,
+        base::Bind(&UsageReceiver::RunOnIOThread, base::Unretained(this)));
+  }
+
+ private:
+  void RunOnIOThread() {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+    tasks_ = 1;  // Task for this method
+    tasks_ += 1;
+    quota_manager_->GetUsageInfo(
+        base::Bind(&UsageReceiver::ReceiveQuotaUsage, base::Unretained(this)));
+
+    tasks_ += 1;
+    dom_storage_context_->GetLocalStorageUsage(base::Bind(
+        &UsageReceiver::ReceiveLocalStorageUsage, base::Unretained(this)));
+    Done();
+  }
+
+  void ReceiveQuotaUsage(const std::vector<storage::UsageInfo>& usage_infos) {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+    for (const auto& info : usage_infos) {
+      CountUsage(GetRegisterableDomainOrIPFromHost(info.host), info.usage);
+    }
+    Done();
+  }
+
+  void ReceiveLocalStorageUsage(
+      const std::vector<content::LocalStorageUsageInfo>& storage_infos) {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+    for (const auto& info : storage_infos) {
+      CountUsage(GetRegisterableDomainOrIP(info.origin), info.data_size);
+    }
+    Done();
+  }
+
+  // Look up the corresponding ImportantDomainInfo for |url| and increase its
+  // usage by |size|.
+  void CountUsage(const std::string& domain, int64_t size) {
+    auto it = site_map.find(domain);
+    if (it != site_map.end()) {
+      it->second->usage += size;
+    }
+  }
+
+  void Done() {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+    DCHECK(tasks_ > 0);
+    if (--tasks_ == 0) {
+      BrowserThread::PostTask(
+          BrowserThread::UI, FROM_HERE,
+          base::Bind(&UsageReceiver::Finish, base::Unretained(this)));
+    }
+  }
+
+  void Finish() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    done_.Run(std::move(sites_));
+    base::ThreadTaskRunnerHandle::Get()->DeleteSoon(FROM_HERE, this);
+  }
+
+ private:
+  ImportantSitesUtil::UsageCallback done_;
+  std::vector<ImportantDomainInfo> sites_;
+  std::map<std::string, ImportantDomainInfo*> site_map;

[dmurph, 2017/03/27 22:22:45]

unordered_map should be more suitable here -- or even a base::SmallMap<
std::unordered_map<string, ImportantDomainInfo*> >.

[dullweber, 2017/03/28 14:06:45]

I changed it to unordered_map.

This map can have up to kMaxImportantSites = 10 entries but is accessed very
often (for every localstorage and quota entry). 
The default size, up to which SmallMap handles the map differently, is
kArraySize = 4.
I'm not sure if SmallMap could have an advantage.

[dmurph, 2017/03/28 21:51:13]

Then SmallMap will definitely be faster. Iterating a vector with size < 200 is
faster than using map, as iterating is super fast. Also we aren't making
modifications to the map, so there is no vector resizing.

+  storage::QuotaManager* quota_manager_;
+  content::DOMStorageContext* dom_storage_context_;
+  int tasks_;
+};
+
 }  // namespace
 
 bool ImportantSitesUtil::IsDialogDisabled(Profile* profile) {
   PrefService* service = profile->GetPrefs();
   DictionaryPrefUpdate update(service, prefs::kImportantSitesDialogHistory);
 
   return ShouldSuppressItem(update.Get());
 }
 
 void ImportantSitesUtil::RegisterProfilePrefs(
@@ skipping 41 matching lines @@
     final_list.push_back(domain_info.second);
     RECORD_UMA_FOR_IMPORTANT_REASON(
         "Storage.ImportantSites.GeneratedReason",
         "Storage.ImportantSites.GeneratedReasonCount",
         domain_info.second.reason_bitfield);
   }
 
   return final_list;
 }
 
+void ImportantSitesUtil::PopulateUsage(storage::QuotaManager* quota_manager,
+                                       content::DOMStorageContext* dom_storage,
+                                       std::vector<ImportantDomainInfo> sites,
+                                       UsageCallback callback) {
+  UsageReceiver* usage_receiver =
+      new UsageReceiver(callback, std::move(sites), quota_manager, dom_storage);
+  usage_receiver->RunAndDestroySelf();
+}
+
 void ImportantSitesUtil::RecordBlacklistedAndIgnoredImportantSites(
     Profile* profile,
     const std::vector<std::string>& blacklisted_sites,
     const std::vector<int32_t>& blacklisted_sites_reason_bitfield,
     const std::vector<std::string>& ignored_sites,
     const std::vector<int32_t>& ignored_sites_reason_bitfield) {
   // First, record the metrics for blacklisted and ignored sites.
   for (int32_t reason_bitfield : blacklisted_sites_reason_bitfield) {
     RECORD_UMA_FOR_IMPORTANT_REASON(
         "Storage.ImportantSites.CBDChosenReason",
@@ skipping 59 matching lines @@
                                                          const GURL& origin) {
   SiteEngagementScore::SetParamValuesForTesting();
   // First get data from site engagement.
   SiteEngagementService* site_engagement_service =
       SiteEngagementService::Get(profile);
   site_engagement_service->ResetBaseScoreForURL(
       origin, SiteEngagementScore::GetMediumEngagementBoundary());
   DCHECK(site_engagement_service->IsEngagementAtLeast(
       origin, blink::mojom::EngagementLevel::MEDIUM));
 }