Count site data size for important sites

chrome/browser/engagement/important_sites_util.cc

Index: chrome/browser/engagement/important_sites_util.cc
diff --git a/chrome/browser/engagement/important_sites_util.cc b/chrome/browser/engagement/important_sites_util.cc
index e31a809d8242b91e6264d4cd2497911401ca6cff..ea319a2bfbf430f9a243e08fffc6fe9609f57e19 100644
--- a/chrome/browser/engagement/important_sites_util.cc
+++ b/chrome/browser/engagement/important_sites_util.cc
@@ -10,6 +10,7 @@
 #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"
@@ -29,12 +30,18 @@
 #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;
 
@@ -141,6 +148,15 @@ std::string GetRegisterableDomainOrIP(const GURL& url) {
   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,
@@ -339,6 +355,94 @@ void PopulateInfoMapWithHomeScreen(
   }
 }
 
+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) {
@@ -400,6 +504,15 @@ ImportantSitesUtil::GetImportantRegisterableDomains(Profile* profile,
   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,