Split initial load of cookies by domains

chrome/browser/net/sqlite_persistent_cookie_store.cc

Index: chrome/browser/net/sqlite_persistent_cookie_store.cc
===================================================================
--- chrome/browser/net/sqlite_persistent_cookie_store.cc	(revision 100935)
+++ chrome/browser/net/sqlite_persistent_cookie_store.cc	(working copy)
@@ -4,8 +4,6 @@
 
 #include "chrome/browser/net/sqlite_persistent_cookie_store.h"
 

[erikwright (departed), 2011/09/17 02:22:10]

Add "#include <list>" back in.
Also include <map>, <set>, <pair>, "base/synchronization/lock.h", ("include what
you use").
Move the include of basictypes.h to sqlite_persistent_cookie_store.h (for
DISALLOW_COPY_AND_ASSIGN).

[guohui, 2011/09/20 18:34:46]

Done.

-#include <list>
-
 #include "base/basictypes.h"
 #include "base/bind.h"
 #include "base/file_path.h"
@@ -20,6 +18,7 @@
 #include "chrome/browser/diagnostics/sqlite_diagnostics.h"
 #include "content/browser/browser_thread.h"
 #include "googleurl/src/gurl.h"
+#include "net/base/registry_controlled_domain.h"
 #include "sql/meta_table.h"
 #include "sql/statement.h"
 #include "sql/transaction.h"
@@ -45,8 +44,12 @@
   }
 
   // Creates or load the SQLite database.
-  bool Load(const LoadedCallback& loaded_callback);
+  void Load(const LoadedCallback& loaded_callback);
 
+  // Load cookies for the domain key (eTLD+1).
+  void LoadCookiesForKey(const std::string& domain,
+      const LoadedCallback& loaded_callback);
+
   // Batch a cookie addition.
   void AddCookie(const net::CookieMonster::CanonicalCookie& cc);
 
@@ -100,16 +103,27 @@
  private:
   // Creates or load the SQLite database on DB thread.
   void LoadAndNotifyOnDBThread(const LoadedCallback& loaded_callback);
-  // Notify the CookieMonster when loading complete.
+
+  // Load cookies for the domain key (eTLD+1) on DB thread.
+  void LoadKeyAndNotifyOnDBThread(const std::string& domains,
+      const LoadedCallback& loaded_callback);
+
+  // Notify the CookieMonster when loading complete for a specific domain key
+  // or for all domain keys. It triggers the callback and passes it all cookies
+  // that have been loaded from DB since last IO notification.
   void NotifyOnIOThread(
       const LoadedCallback& loaded_callback,
-      bool load_success,
-      const std::vector<net::CookieMonster::CanonicalCookie*>& cookies);
+      bool load_success);
   // Initialize the data base.
+
   bool InitializeDatabase();
-  // Load cookies to the data base, and read cookies.
-  bool LoadInternal(std::vector<net::CookieMonster::CanonicalCookie*>* cookies);
 
+  // Chain load cookies from DB by domain key.
+  void ChainLoadCookies(const LoadedCallback& loaded_callback);
+
+  // Load all cookies for a set of domains/hosts
+  bool LoadCookiesForDomains(const std::set<std::string>& key);
+
   // Batch a cookie operation (add or delete)
   void BatchOperation(PendingOperation::OperationType op,
                       const net::CookieMonster::CanonicalCookie& cc);
@@ -130,6 +144,15 @@
   // Guard |pending_|, |num_pending_| and |clear_local_state_on_exit_|.
   base::Lock lock_;
 
+  // Temporary buffer for cookies loaded from DB.

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

I'd like a bit more detail on this comment, something like "accumulated here to
reduce the number of messages to the IO thread; sent back in response to
requests or when all loading complete".

[guohui, 2011/09/20 18:34:46]

Done.

+  std::vector<net::CookieMonster::CanonicalCookie*> cookies_;
+
+  // Map of domain keys(eTLD+1) to domains/hosts that are to be loaded from DB.
+  std::map<std::string, std::set<std::string>> keys_to_load_;
+
+  // Indicates if DB has been initialized.
+  bool initialized_;
+
   DISALLOW_COPY_AND_ASSIGN(Backend);
 };
 
@@ -167,43 +190,90 @@
   // so we want those people to get it. Ignore errors, since it may exist.
   db->Execute("CREATE INDEX IF NOT EXISTS cookie_times ON cookies"
               " (creation_utc)");
+
+  db->Execute("CREATE INDEX IF NOT EXISTS domain ON cookies(host_key)");
+
   return true;
 }
 
 }  // namespace
 
-bool SQLitePersistentCookieStore::Backend::Load(
+void SQLitePersistentCookieStore::Backend::Load(
     const LoadedCallback& loaded_callback) {
   // This function should be called only once per instance.
   DCHECK(!db_.get());
   BrowserThread::PostTask(
       BrowserThread::DB, FROM_HERE,
-      base::Bind(&Backend::LoadAndNotifyOnDBThread, base::Unretained(this),
-                 loaded_callback));
-  return true;
+      base::Bind(&Backend::LoadAndNotifyOnDBThread, this, loaded_callback));
 }
 
+void SQLitePersistentCookieStore::Backend::LoadCookiesForKey(
+    const std::string& key,
+    const LoadedCallback& loaded_callback) {
+  BrowserThread::PostTask(
+    BrowserThread::DB, FROM_HERE,
+    base::Bind(&Backend::LoadKeyAndNotifyOnDBThread, this,
+    key,
+    loaded_callback));
+}
+
 void SQLitePersistentCookieStore::Backend::LoadAndNotifyOnDBThread(
     const LoadedCallback& loaded_callback) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::DB));
-  std::vector<net::CookieMonster::CanonicalCookie*> cookies;
 
-  bool load_success = LoadInternal(&cookies);
+  if (!InitializeDatabase()) {
+    BrowserThread::PostTask(
+      BrowserThread::IO, FROM_HERE,
+      base::Bind(&SQLitePersistentCookieStore::Backend::NotifyOnIOThread,
+                 this, loaded_callback, false));
+  } else {
+    ChainLoadCookies(loaded_callback);
+  }
+}
 
-  BrowserThread::PostTask(BrowserThread::IO, FROM_HERE, base::Bind(
-      &SQLitePersistentCookieStore::Backend::NotifyOnIOThread,
-      base::Unretained(this), loaded_callback, load_success, cookies));
+void SQLitePersistentCookieStore::Backend::LoadKeyAndNotifyOnDBThread(
+    const std::string& key,
+    const LoadedCallback& loaded_callback) {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::DB));
+
+  bool success = false;
+  if (InitializeDatabase()) {
+    std::map<std::string, std::set<std::string> >::iterator
+      it = keys_to_load_.find(key);
+    if (it != keys_to_load_.end()) {
+      success = LoadCookiesForDomains(it->second);
+      keys_to_load_.erase(it);
+    } else {
+      success = true;
+    }
+  }
+
+  BrowserThread::PostTask(
+    BrowserThread::IO, FROM_HERE,
+    base::Bind(&SQLitePersistentCookieStore::Backend::NotifyOnIOThread,
+    this, loaded_callback, success));
 }
 
 void SQLitePersistentCookieStore::Backend::NotifyOnIOThread(
     const LoadedCallback& loaded_callback,
-    bool load_success,
-    const std::vector<net::CookieMonster::CanonicalCookie*>& cookies) {
+    bool load_success) {

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

We appear to be completely ignoring failure.  That wasn't introduced by this CL,
but it isn't good--maybe a crbug filing + comment?

[guohui, 2011/09/20 18:34:46]

A crbug 97224 is filed for this issue.

On 2011/09/19 00:58:57, rdsmith wrote:

   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
+
+  std::vector<net::CookieMonster::CanonicalCookie*> cookies;
+  {
+    base::AutoLock locked(lock_);

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

Expand the comment on lock_ in the header file to make clear it protects
cookies_ as well as its other targets.

[guohui, 2011/09/20 18:34:46]

Done.

+    cookies = cookies_;
+    cookies_.clear();
+  }
+
   loaded_callback.Run(cookies);
 }
 
 bool SQLitePersistentCookieStore::Backend::InitializeDatabase() {
+  if (initialized_) {
+    return true;
+  }
+
   const FilePath dir = path_.DirName();
   if (!file_util::PathExists(dir) && !file_util::CreateDirectory(dir)) {
     return false;
@@ -225,45 +295,104 @@
   }
 
   db_->Preload();
+
+  // Retrieve all the domains
+  sql::Statement smt(db_->GetUniqueStatement(
+    "SELECT DISTINCT host_key FROM cookies"));
+
+  if (!smt) {
+    NOTREACHED() << "select statement prep failed";
+    db_.reset();
+    return false;
+  }
+
+  // Build a map of domain keys (eTLD+1) to domains.

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

I believe the code is correct, but I kept getting tripped up a bit reading this
because the CookieMonster code is written so that the map key can be either
eTLD+1 or just straight domain (originally so that we could switch back easily
if we needed to).   It's worth being aware in writing comments and choosing
names that a reader of this code may be confused about whether your keys are
similar to the CookieMonsters or are always eTLD+1s (my reading of the code is
the latter).

[guohui, 2011/09/20 18:34:46]

Done.

+  while (smt.Step()) {
+    std::string domain = smt.ColumnString(0);
+    std::string key =
+      net::RegistryControlledDomainService::GetDomainAndRegistry(domain);
+
+    std::map<std::string, std::set<std::string> >::iterator it =
+      keys_to_load_.find(key);
+    if (it == keys_to_load_.end()) {
+      std::set<std::string> domains;
+      domains.insert(domain);
+      keys_to_load_.insert(
+        std::pair<std::string, std::set<std::string> >(key, domains));
+    } else {
+      it->second.insert(domain);
+    }

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

nit, suggestion: I'd find this code slightly easier to read if it was
(pseudo-code):

If key not in map:
   Insert key in map with null set
map[key].insert(domain)

But it's just a preference and you're welcome to do whatever you prefer.

[grt (UTC plus 2), 2011/09/19 02:20:40]

owing to the awesomeness of STL containers, you can replace lines 315 through
324 with:

keys_to_load_.insert(std::make_pair(key,
std::set<std::string>())).first->second.insert(domain);

the downside to this is that it will needlessly construct and destroy a std::set
when |key| is already in |keys_to_load_|.  that should be pretty inexpensive,
though, especially if the find+insert pair codepath above is common.

[guohui, 2011/09/20 18:34:46]

Thanks for the nice STL tip. Your code is indeed much neater. However I think in
this scenario find+insert path is rare, since one domain key(eTLD+1) usually has
multiple associated cookies. The performance penalty is small, but still
unnecessary. So I fixed as suggested by randy.

On 2011/09/19 02:20:40, grt wrote:

+  }
+
+  initialized_ = true;
   return true;
 }
 
-bool SQLitePersistentCookieStore::Backend::LoadInternal(
-    std::vector<net::CookieMonster::CanonicalCookie*>* cookies) {
-  if (!InitializeDatabase()) {
-    return false;
+void SQLitePersistentCookieStore::Backend::ChainLoadCookies(
+    const LoadedCallback& loaded_callback) {
+  bool load_success = true;
+
+  if (keys_to_load_.size() > 0) {
+    // Load cookies for the first domain key.
+    std::map<std::string, std::set<std::string> >::iterator
+      it = keys_to_load_.begin();
+    load_success = LoadCookiesForDomains(it->second);
+    keys_to_load_.erase(it);
   }
 
-  // Slurp all the cookies into the out-vector.
-  sql::Statement smt(db_->GetUniqueStatement(
-      "SELECT creation_utc, host_key, name, value, path, expires_utc, secure, "
-      "httponly, last_access_utc FROM cookies"));
+  // If load is successful and there are more domain keys to be loaded,
+  // then post a DB task to continue chain-load;
+  // Otherwise notify on IO thread.
+  if (load_success && keys_to_load_.size() > 0) {
+    BrowserThread::PostTask(
+      BrowserThread::DB, FROM_HERE,
+      base::Bind(&Backend::ChainLoadCookies, this, loaded_callback));
+  } else {
+    BrowserThread::PostTask(
+      BrowserThread::IO, FROM_HERE,
+      base::Bind(&SQLitePersistentCookieStore::Backend::NotifyOnIOThread,
+                 this, loaded_callback, load_success));
+  }
+}
+
+bool SQLitePersistentCookieStore::Backend::LoadCookiesForDomains(
+  const std::set<std::string>& domains) {
+  base::AutoLock locked(lock_);

[Randy Smith (Not in Mondays), 2011/09/19 00:58:57]

What is this protecting?  I twitch at holding a lock across a file system
operation.  Can you hold the lock only across the data structures you are
protecting with the lock?  (I *think* that's just cookies_).

[guohui, 2011/09/20 18:34:46]

Done.

+
+  sql::Statement smt(db_->GetCachedStatement(SQL_FROM_HERE,
+    "SELECT creation_utc, host_key, name, value, path, expires_utc, secure, "
+    "httponly, last_access_utc FROM cookies WHERE host_key = ?"));
   if (!smt) {
     NOTREACHED() << "select statement prep failed";
     db_.reset();
     return false;
   }
 
-  while (smt.Step()) {
-    scoped_ptr<net::CookieMonster::CanonicalCookie> cc(
-        new net::CookieMonster::CanonicalCookie(
-            // The "source" URL is not used with persisted cookies.
-            GURL(),                                         // Source
-            smt.ColumnString(2),                            // name
-            smt.ColumnString(3),                            // value
-            smt.ColumnString(1),                            // domain
-            smt.ColumnString(4),                            // path
-            std::string(),  // TODO(abarth): Persist mac_key
-            std::string(),  // TODO(abarth): Persist mac_algorithm
-            Time::FromInternalValue(smt.ColumnInt64(0)),    // creation_utc
-            Time::FromInternalValue(smt.ColumnInt64(5)),    // expires_utc
-            Time::FromInternalValue(smt.ColumnInt64(8)),    // last_access_utc
-            smt.ColumnInt(6) != 0,                          // secure
-            smt.ColumnInt(7) != 0,                          // httponly
-            true));                                         // has_expires
-    DLOG_IF(WARNING,
-            cc->CreationDate() > Time::Now()) << L"CreationDate too recent";
-    cookies->push_back(cc.release());
+  std::set<std::string>::const_iterator it = domains.begin();
+  for (; it != domains.end(); ++it) {
+    smt.BindString(0, *it);
+    while (smt.Step()) {
+      scoped_ptr<net::CookieMonster::CanonicalCookie> cc(
+          new net::CookieMonster::CanonicalCookie(
+              // The "source" URL is not used with persisted cookies.
+              GURL(),                                         // Source
+              smt.ColumnString(2),                            // name
+              smt.ColumnString(3),                            // value
+              smt.ColumnString(1),                            // domain
+              smt.ColumnString(4),                            // path
+              std::string(),  // TODO(abarth): Persist mac_key
+              std::string(),  // TODO(abarth): Persist mac_algorithm
+              Time::FromInternalValue(smt.ColumnInt64(0)),    // creation_utc
+              Time::FromInternalValue(smt.ColumnInt64(5)),    // expires_utc
+              Time::FromInternalValue(smt.ColumnInt64(8)),    // last_access_utc
+              smt.ColumnInt(6) != 0,                          // secure
+              smt.ColumnInt(7) != 0,                          // httponly
+              true));                                         // has_expires
+      DLOG_IF(WARNING,
+              cc->CreationDate() > Time::Now()) << L"CreationDate too recent";
+      cookies_.push_back(cc.release());
+    }
+    smt.Reset();
   }
 
   return true;
@@ -534,10 +663,16 @@
   }
 }
 
-bool SQLitePersistentCookieStore::Load(const LoadedCallback& loaded_callback) {
-  return backend_->Load(loaded_callback);
+void SQLitePersistentCookieStore::Load(const LoadedCallback& loaded_callback) {
+  backend_->Load(loaded_callback);
 }
 
+void SQLitePersistentCookieStore::LoadCookiesForKey(
+    const std::string& key,
+    const LoadedCallback& loaded_callback) {
+  backend_->LoadCookiesForKey(key, loaded_callback);
+}
+
 void SQLitePersistentCookieStore::AddCookie(
     const net::CookieMonster::CanonicalCookie& cc) {
   if (backend_.get())