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..613f763f00ff78e577ce6ca1c2283e4e1c825d88 100644
--- a/chrome/browser/engagement/important_sites_util.cc
+++ b/chrome/browser/engagement/important_sites_util.cc
@@ -8,9 +8,12 @@
 #include <map>
 #include <memory>
 #include <set>
+#include <unordered_map>
 #include <utility>
 
+#include "base/bind_helpers.h"
 #include "base/containers/hash_tables.h"
+#include "base/containers/small_map.h"
 #include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #include "base/stl_util.h"
@@ -29,14 +32,23 @@
 #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;

[dominickn, 2017/04/10 05:06:12]

Nit: the other files in engagement/ explicitly spell out content::BrowserThread
rather than call using, so perhaps do that for consistency.

[dullweber, 2017/04/10 13:29:38]

Done.

 using ImportantDomainInfo = ImportantSitesUtil::ImportantDomainInfo;
 using ImportantReason = ImportantSitesUtil::ImportantReason;
+using ImportantSiteMap =
+    base::SmallMap<std::unordered_map<std::string, ImportantDomainInfo*>,
+                   ImportantSitesUtil::kMaxImportantSites>;
 
 // Note: These values are stored on both the per-site content settings
 // dictionary and the dialog preference dictionary.
@@ -141,6 +153,15 @@ std::string GetRegisterableDomainOrIP(const GURL& url) {
   return registerable_domain;
 }
 
+std::string GetRegisterableDomainOrIPFromHost(const std::string& host) {

[dominickn, 2017/04/10 05:06:12]

Would it be cheaper / more clear to just call
GetRegisterableDomainOrIP(GURL(host)); instead of having this method? It's
almost identical to the former, with the need to call the url::HostIsIPAddress
method the only difference.

[dullweber, 2017/04/10 13:29:38]

The issue is that I can't create a GURL from hosts because hosts are missing the
scheme. I will go the other way round and let the GetRegisterableDomainOrIP
function call GetRegisterableDomainOrIPFromHost with gurl.host_piece(). This
removes the duplicated code. (And that's what GetDomainAndRegistry() and
url.HostIsIPAdress() do anyway.)

+  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 +360,96 @@ void PopulateInfoMapWithHomeScreen(
   }
 }
 
+// A helper class that retrieves the localstorage and quota usage for each
+// domain in |ImportantDomainInfo| and populates |ImportantDomainInfo::usage|.

[dominickn, 2017/04/10 05:06:12]

This is a pretty big class (100 lines). It might be worth splitting it into a
new file, and unit testing it separately.

[dullweber, 2017/04/10 13:29:38]

I moved it a separate file. That cleaned up this file and the unittest but I had
to make the  GetRegisterableDomainOrIPFrom* public static methods to share them
between ImportantSitesUtil and the new class.

+class UsageReceiver {
+ public:

[dominickn, 2017/04/10 05:06:12]

You should probably document that the receiver will take ownership of |sites|
and modify it, so the caller can't use it again.

[dullweber, 2017/04/10 13:29:39]

The |sites| parameter is a std::vector<ImportantDomainInfo> and not a pointer or
reference. I use it by moving the vector into this class, which automatically
prevents me from using it again. If the caller wouldn't move the vector, then it
also shouldn't see the modifications? I think this is a good way to have cheap
transfer of the vector but making it clear which place currently owns it. 

But maybe I'm wrong with my interpretation of how std::move will behave or is
what I am doing an uncommon usage of it?

[dominickn, 2017/04/11 00:18:20]

D'oh, I misread and thought that param was an rvalue reference. You're right on
this one.

+  UsageReceiver(ImportantSitesUtil::UsageCallback done,

[dominickn, 2017/04/10 05:06:12]

Nit: call this "callback" or "finish_callback". "done_" looks like a boolean
flag.

[dullweber, 2017/04/10 13:29:38]

Done.

+                std::vector<ImportantDomainInfo> sites,
+                storage::QuotaManager* quota_manager,
+                content::DOMStorageContext* dom_storage_context)
+      : done_(done),
+        sites_(std::move(sites)),
+        quota_manager_(quota_manager),
+        dom_storage_context_(dom_storage_context),
+        tasks_(-1) {
+    for (auto& site : sites_) {

[dominickn, 2017/04/10 05:06:12]

It's not clear what the type is here, so I'd prefer ImportantDomainInfo&
domain_info : sites_

[dullweber, 2017/04/10 13:29:38]

Done.

+      site.usage = 0;
+      site_map_[site.registerable_domain] = &site;
+    }
+  }
+
+  void RunAndDestroySelf() {
+    DCHECK_CURRENTLY_ON(BrowserThread::UI);
+    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(

[dominickn, 2017/04/10 05:06:12]

I can't find any other place in the codebase which calls this method on the IO
thread. I'm fairly sure that this runs in a non-blocking background task, but it
would be good to confirm that. :)

[dullweber, 2017/04/10 13:29:39]

Yes it looks like localstorage should be queried from the UI thread. Thanks!

+        &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) {

[dominickn, 2017/04/10 05:06:12]

Nit: no braces for 1 line conditional.

[dullweber, 2017/04/10 13:29:38]

Done.

+      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) {

[dominickn, 2017/04/10 05:06:12]

Nit: no braces for 1 line conditional.

[dullweber, 2017/04/10 13:29:38]

Done.

+      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) {

[dominickn, 2017/04/10 05:06:12]

Nit: This is more accurately named "IncrementUsage"

[dullweber, 2017/04/10 13:29:39]

Done.

+    auto it = site_map_.find(domain);

[dominickn, 2017/04/10 05:06:11]

It looks to me like the only reason for site_map_ is to provide a map interface
to sites_. But you've declared site_map_ as a SmallMap, which is backed by an
array anyway. And there will only be max 10 entries too right?

Why not just do this using std::find_if and a lambda directly on sites_, and
remove site_map_ all together? Saves the extra overhead of the additional map
and the weird ownership semantics where the map is holding pointers to things
that are in a vector in this object anyway.

[dullweber, 2017/04/10 13:29:38]

That sounds like a good idea, thanks!

+    if (it != site_map_.end()) {

[dominickn, 2017/04/10 05:06:12]

Nit: no braces for a 1 line conditional.

[dullweber, 2017/04/10 13:29:39]

Done.

+      it->second->usage += size;
+    }
+  }
+
+  void Done() {
+    DCHECK_CURRENTLY_ON(BrowserThread::IO);
+    DCHECK(tasks_ > 0);

[dominickn, 2017/04/10 05:06:12]

DCHECK_GE(tasks_, 0);

[dullweber, 2017/04/10 13:29:39]

Done.

+    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);
+  }
+
+  ImportantSitesUtil::UsageCallback done_;
+  std::vector<ImportantDomainInfo> sites_;
+  ImportantSiteMap site_map_;
+  storage::QuotaManager* quota_manager_;
+  content::DOMStorageContext* dom_storage_context_;
+  int tasks_;
+};
+
 }  // namespace
 
 bool ImportantSitesUtil::IsDialogDisabled(Profile* profile) {
@@ -400,6 +511,16 @@ ImportantSitesUtil::GetImportantRegisterableDomains(Profile* profile,
   return final_list;
 }
 
+void ImportantSitesUtil::PopulateUsage(storage::QuotaManager* quota_manager,
+                                       content::DOMStorageContext* dom_storage,
+                                       std::vector<ImportantDomainInfo> sites,
+                                       UsageCallback callback) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+  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,