Support partitioning of storage contexts based on render_id.

content/browser/browser_context.cc

 // Copyright (c) 2012 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 "content/public/browser/browser_context.h"
 
+#include "base/file_path.h"
+#include "base/stl_util.h"
 #include "content/browser/appcache/chrome_appcache_service.h"
 #include "content/browser/dom_storage/dom_storage_context_impl.h"
 #include "content/browser/download/download_file_manager.h"
 #include "content/browser/download/download_manager_impl.h"
 #include "content/browser/fileapi/browser_file_system_helper.h"
 #include "content/browser/in_process_webkit/indexed_db_context_impl.h"
 #include "content/browser/renderer_host/resource_dispatcher_host_impl.h"
 #include "content/browser/resource_context_impl.h"
+#include "content/common/child_process_host_impl.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/content_browser_client.h"
 #include "content/public/common/content_constants.h"
 #include "net/base/server_bound_cert_service.h"
 #include "net/base/server_bound_cert_store.h"
 #include "net/cookies/cookie_monster.h"
 #include "net/cookies/cookie_store.h"
 #include "net/url_request/url_request_context.h"
 #include "webkit/database/database_tracker.h"
 #include "webkit/quota/quota_manager.h"
 
 using appcache::AppCacheService;
 using base::UserDataAdapter;
 using content::BrowserThread;
 using fileapi::FileSystemContext;
 using quota::QuotaManager;
 using webkit_database::DatabaseTracker;
 
 // Key names on BrowserContext.
-static const char* kAppCacheServicKeyName = "content_appcache_service_tracker";
-static const char* kDatabaseTrackerKeyName = "content_database_tracker";
-static const char* kDOMStorageContextKeyName = "content_dom_storage_context";
 static const char* kDownloadManagerKeyName = "download_manager";
-static const char* kFileSystemContextKeyName = "content_file_system_context";
-static const char* kIndexedDBContextKeyName = "content_indexed_db_context";
-static const char* kQuotaManagerKeyName = "content_quota_manager";
+static const char* kStorageParitionMapKeyName = "content_storage_partition_map";
+
+// Dirname for storing persistent data for renderers with partitioned storage.

[Charlie Reis, 2012/07/02 23:21:23]

partitioned -> isolated

I think this would be more consistent with the other descriptions, while still
letting us talk about a "partition" as the place where isolated storage resides.

[awong, 2012/07/09 20:37:43]

sgtm.

+const FilePath::CharType kStoragePartitionDirName[] =
+    FILE_PATH_LITERAL("Storage Partitions");
 
 namespace content {
 
 namespace {
 
-void CreateQuotaManagerAndClients(BrowserContext* context) {
-  if (context->GetUserData(kQuotaManagerKeyName)) {
-    DCHECK(context->GetUserData(kDatabaseTrackerKeyName));
-    DCHECK(context->GetUserData(kDOMStorageContextKeyName));
-    DCHECK(context->GetUserData(kFileSystemContextKeyName));
-    DCHECK(context->GetUserData(kIndexedDBContextKeyName));
+class StoragePartition {

[Charlie Reis, 2012/07/02 23:21:23]

Please give this a class-level comment.

[awong, 2012/07/09 20:37:43]

Done.

+ public:
+  ~StoragePartition() {
+    // These message loop checks are just to avoid leaks in unittests.
+    if (database_tracker() &&
+        BrowserThread::IsMessageLoopValid(BrowserThread::FILE)) {
+      BrowserThread::PostTask(
+          BrowserThread::FILE, FROM_HERE,
+          base::Bind(&webkit_database::DatabaseTracker::Shutdown,
+                     database_tracker()));
+    }
+
+    if (dom_storage_context())
+      dom_storage_context()->Shutdown();
+  }
+
+  static StoragePartition* Create(BrowserContext* context,
+                                  FilePath partition_path) {
+    // All of the clients have to be created and registered with the
+    // QuotaManager prior to the QuotaManger being used. So we do them
+    // all together here prior to handing out a reference to anything
+    // that utlizes the QuotaManager.

[Charlie Reis, 2012/07/02 23:21:23]

nits: Drop the "So", and utlizes -> utilizes

I don't understand if "clients," "consumers," and "contexts" are all the same
thing or not.  What are examples?

(I realize these are from pre-existing code, but it would be nice to clean up if
possible.)

[awong, 2012/07/09 20:37:43]

Done.

+    scoped_refptr<QuotaManager> quota_manager = new quota::QuotaManager(
+        context->IsOffTheRecord(), partition_path,
+        BrowserThread::GetMessageLoopProxyForThread(BrowserThread::IO),
+        BrowserThread::GetMessageLoopProxyForThread(BrowserThread::DB),
+        context->GetSpecialStoragePolicy());
+
+    // Each consumer is responsible for registering its QuotaClient during
+    // its construction.
+    scoped_refptr<FileSystemContext> filesystem_context =
+        CreateFileSystemContext(partition_path, context->IsOffTheRecord(),
+                                context->GetSpecialStoragePolicy(),
+                                quota_manager->proxy());
+
+    scoped_refptr<DatabaseTracker> database_tracker = new DatabaseTracker(
+        partition_path, context->IsOffTheRecord(),
+        context->GetSpecialStoragePolicy(), quota_manager->proxy(),
+        BrowserThread::GetMessageLoopProxyForThread(BrowserThread::FILE));
+
+    FilePath path = context->IsOffTheRecord() ? FilePath() : partition_path;
+    scoped_refptr<DOMStorageContextImpl> dom_storage_context =
+        new DOMStorageContextImpl(path, context->GetSpecialStoragePolicy());
+
+    scoped_refptr<IndexedDBContextImpl> indexed_db_context =
+        new IndexedDBContextImpl(path, context->GetSpecialStoragePolicy(),
+                                 quota_manager->proxy(),
+                                 BrowserThread::GetMessageLoopProxyForThread(
+                                     BrowserThread::WEBKIT_DEPRECATED));
+
+    scoped_refptr<ChromeAppCacheService> appcache_service =
+        new ChromeAppCacheService(quota_manager->proxy());
+
+    return new StoragePartition(context,
+                                partition_path,
+                                quota_manager,
+                                appcache_service,
+                                filesystem_context,
+                                database_tracker,
+                                dom_storage_context,
+                                indexed_db_context);
+  }
+
+  quota::QuotaManager* quota_manager() {
+    return quota_manager_;
+  }
+  ChromeAppCacheService* appcache_service() {
+    return appcache_service_;
+  }
+  fileapi::FileSystemContext* filesystem_context() {
+    return filesystem_context_;
+  }
+  webkit_database::DatabaseTracker* database_tracker() {
+    return database_tracker_;
+  }
+  DOMStorageContextImpl* dom_storage_context() {
+    return dom_storage_context_;
+  }
+  IndexedDBContext* indexed_db_context() {
+    return indexed_db_context_;
+  }
+
+ private:
+  // TODO(ajwong): Break the direct dependency on browser_context. We only
+  // need 2 pieces of info from it for use in the Initialize() function.
+  StoragePartition(BrowserContext* browser_context,
+                   const FilePath& partition_path,
+                   quota::QuotaManager* quota_manager,
+                   ChromeAppCacheService* appcache_service,
+                   fileapi::FileSystemContext* filesystem_context,
+                   webkit_database::DatabaseTracker* database_tracker,
+                   DOMStorageContextImpl* dom_storage_context,
+                   IndexedDBContext* indexed_db_context)
+      : browser_context_(browser_context),
+        partition_path_(partition_path),
+        quota_manager_(quota_manager),
+        appcache_service_(appcache_service),
+        filesystem_context_(filesystem_context),
+        database_tracker_(database_tracker),
+        dom_storage_context_(dom_storage_context),
+        indexed_db_context_(indexed_db_context) {
+  }
+
+  BrowserContext* browser_context_;
+  FilePath partition_path_;
+  scoped_refptr<quota::QuotaManager> quota_manager_;
+  scoped_refptr<ChromeAppCacheService> appcache_service_;
+  scoped_refptr<fileapi::FileSystemContext> filesystem_context_;
+  scoped_refptr<webkit_database::DatabaseTracker> database_tracker_;
+  scoped_refptr<DOMStorageContextImpl> dom_storage_context_;
+  scoped_refptr<IndexedDBContext> indexed_db_context_;
+};
+
+class StoragePartitionMap : public base::SupportsUserData::Data {

[Charlie Reis, 2012/07/02 23:21:23]

Class level comment.

[awong, 2012/07/09 20:37:43]

Done.

+ public:
+  explicit StoragePartitionMap(BrowserContext* browser_context)
+      : browser_context_(browser_context) {
+  }
+
+  virtual ~StoragePartitionMap() {
+    STLDeleteContainerPairSecondPointers(partitions_.begin(),
+                                         partitions_.end());
+  }
+
+  StoragePartition* Get(const std::string& partition_id) {
+    // Find the previously created partition if it's available.
+    std::map<std::string, StoragePartition*>::const_iterator it =
+        partitions_.find(partition_id);
+    if (it != partitions_.end()) {

[Charlie Reis, 2012/07/02 23:21:23]

nit: No braces.

[awong, 2012/07/09 20:37:43]

Done.

+      return it->second;
+    }
+
+    // There was no previous partition, so let's make a new one.
+    FilePath partition_path = browser_context_->GetPath();
+    if (!partition_id.empty()) {
+      CHECK(IsStringASCII(partition_id));

[Charlie Reis, 2012/07/02 23:21:23]

Are we ever taking this in from an untrusted context?  I wonder if a DCHECK is
sufficient.

[awong, 2012/07/09 20:37:43]

We aren't taking it from an untrusted context now, but the consequence of future
modifications feeding this bad data scares me.  

If this invaraint is broken, AppendASCII will just produce a corrupt path
otherwise which may do all sorts of unpredictable things. This seems like a
place where always dying hard is a good thing.

Biasing towards CHECK is always safer anyways :)

+      partition_path = partition_path.Append(kStoragePartitionDirName)
+          .AppendASCII(partition_id);
+    }
+
+    StoragePartition* storage_partition =
+        StoragePartition::Create(browser_context_, partition_path);
+    partitions_[partition_id] = storage_partition;
+
+    PostAddInitialization(storage_partition, partition_path);
+
+    // TODO(ajwong): We need to remove this conditional by making
+    // InitializeResourceContext() understand having different parition data

[Charlie Reis, 2012/07/02 23:21:23]

nit: partition

[awong, 2012/07/09 20:37:43]

Done.

+    // based on the renderer_id.
+    if (partition_id.empty()) {
+      InitializeResourceContext(browser_context_);
+    }
+
+    return storage_partition;
+  }
+
+  void ForEach(const base::Callback<void(StoragePartition*)>& callback) {
+    for (std::map<std::string, StoragePartition*>::const_iterator it =
+             partitions_.begin();
+         it != partitions_.end();
+         ++it) {
+      callback.Run(it->second);
+    }
+  }
+
+ private:
+  // TODO(ajwong): This must always be called *after* it's been added to the
+  // partition map.  This feels dangerous.  Should this not be in this class?
+  void PostAddInitialization(StoragePartition* partition,

[Charlie Reis, 2012/07/02 23:21:23]

PostCreateInitialization?  Seems more consistent than Add.

[awong, 2012/07/09 20:37:43]

Good point. Done.

I hate this function BTW. It's logic is really confusing, but I haven't figured
out how to clean it up yet. :(

+                             const FilePath& partition_path) {
+    // Check first to avoid memory leak in unittests.
+    if (BrowserThread::IsMessageLoopValid(BrowserThread::IO)) {
+      BrowserThread::PostTask(
+          BrowserThread::IO, FROM_HERE,
+          base::Bind(&ChromeAppCacheService::InitializeOnIOThread,
+                     partition->appcache_service(),
+                     browser_context_->IsOffTheRecord() ? FilePath() :
+                         partition_path.Append(content::kAppCacheDirname),
+                         browser_context_->GetResourceContext(),
+                     make_scoped_refptr(
+                         browser_context_->GetSpecialStoragePolicy())));
+    }
+  }
+
+  BrowserContext* browser_context_;
+  std::map<std::string, StoragePartition*> partitions_;
+};
+
+StoragePartition* GetStoragePartition(BrowserContext* browser_context,
+                                      int renderer_child_id) {
+  StoragePartitionMap* partition_map = static_cast<StoragePartitionMap*>(
+      browser_context->GetUserData(kStorageParitionMapKeyName));
+  if (!partition_map) {
+    partition_map = new StoragePartitionMap(browser_context);
+    browser_context->SetUserData(kStorageParitionMapKeyName, partition_map);
+  }
+
+  const std::string& partition_id =
+      GetContentClient()->browser()->GetStoragePartitionForChildProcess(
+          browser_context,
+          renderer_child_id);
+
+  return partition_map->Get(partition_id);
+}
+
+void ForEachStoragePartition(
+    BrowserContext* browser_context,
+    const base::Callback<void(StoragePartition*)>& callback) {
+  StoragePartitionMap* partition_map = static_cast<StoragePartitionMap*>(
+      browser_context->GetUserData(kStorageParitionMapKeyName));
+  if (!partition_map) {
     return;
   }
 
-  // All of the clients have to be created and registered with the
-  // QuotaManager prior to the QuotaManger being used. So we do them
-  // all together here prior to handing out a reference to anything
-  // that utlizes the QuotaManager.
-  scoped_refptr<QuotaManager> quota_manager = new quota::QuotaManager(
-    context->IsOffTheRecord(), context->GetPath(),
-    BrowserThread::GetMessageLoopProxyForThread(BrowserThread::IO),
-    BrowserThread::GetMessageLoopProxyForThread(BrowserThread::DB),
-    context->GetSpecialStoragePolicy());
-  context->SetUserData(kQuotaManagerKeyName,
-                       new UserDataAdapter<QuotaManager>(quota_manager));
-
-  // Each consumer is responsible for registering its QuotaClient during
-  // its construction.
-  scoped_refptr<FileSystemContext> filesystem_context = CreateFileSystemContext(
-      context->GetPath(), context->IsOffTheRecord(),
-      context->GetSpecialStoragePolicy(), quota_manager->proxy());
-  context->SetUserData(
-      kFileSystemContextKeyName,
-      new UserDataAdapter<FileSystemContext>(filesystem_context));
-
-  scoped_refptr<DatabaseTracker> db_tracker = new DatabaseTracker(
-      context->GetPath(), context->IsOffTheRecord(),
-      context->GetSpecialStoragePolicy(), quota_manager->proxy(),
-      BrowserThread::GetMessageLoopProxyForThread(BrowserThread::FILE));
-  context->SetUserData(kDatabaseTrackerKeyName,
-                       new UserDataAdapter<DatabaseTracker>(db_tracker));
-
-  FilePath path = context->IsOffTheRecord() ? FilePath() : context->GetPath();
-  scoped_refptr<DOMStorageContextImpl> dom_storage_context =
-      new DOMStorageContextImpl(path, context->GetSpecialStoragePolicy());
-  context->SetUserData(
-      kDOMStorageContextKeyName,
-      new UserDataAdapter<DOMStorageContextImpl>(dom_storage_context));
-
-  scoped_refptr<IndexedDBContext> indexed_db_context = new IndexedDBContextImpl(
-      path, context->GetSpecialStoragePolicy(), quota_manager->proxy(),
-      BrowserThread::GetMessageLoopProxyForThread(
-          BrowserThread::WEBKIT_DEPRECATED));
-  context->SetUserData(
-      kIndexedDBContextKeyName,
-      new UserDataAdapter<IndexedDBContext>(indexed_db_context));
-
-  scoped_refptr<ChromeAppCacheService> appcache_service =
-      new ChromeAppCacheService(quota_manager->proxy());
-  context->SetUserData(
-      kAppCacheServicKeyName,
-      new UserDataAdapter<ChromeAppCacheService>(appcache_service));
-
-  InitializeResourceContext(context);
-
-  // Check first to avoid memory leak in unittests.
-  if (BrowserThread::IsMessageLoopValid(BrowserThread::IO)) {
-    BrowserThread::PostTask(
-        BrowserThread::IO, FROM_HERE,
-        base::Bind(&ChromeAppCacheService::InitializeOnIOThread,
-                   appcache_service,
-                   context->IsOffTheRecord() ? FilePath() :
-                       context->GetPath().Append(content::kAppCacheDirname),
-                   context->GetResourceContext(),
-                   make_scoped_refptr(context->GetSpecialStoragePolicy())));
-  }
+  partition_map->ForEach(callback);
+}
+
+void DomStorageAdapter(

[Charlie Reis, 2012/07/02 23:21:23]

These three names aren't making sense to me.  If they're functions, can we make
them verbs?  For example, DOMStorageAdapter sounds like a class.

[awong, 2012/07/09 20:37:43]

Called this one ProcessDomStorageContext().

As for replacing the ForEach...I don't know what else I would call it.
RunOnStoragePartition? ExecuteOnEachPartition? ProcessEachPartition?  It's
really running the callback on each partition like a for loop...ForEeach

[Charlie Reis, 2012/07/10 21:37:48]

Thanks.  ForEach* is ok.

+    const base::Callback<void(DOMStorageContextImpl*)>& callback,
+    StoragePartition* partition) {
+  callback.Run(partition->dom_storage_context());
+}
+
+void DOMStorageContextForEach(

[nasko, 2012/06/29 23:08:56]

nit: In ForEachStoragePartition the "ForEach" part is at the beginning, here it
is at the end. Any reason? It is nice to be consistent.

[awong, 2012/07/09 20:37:43]

Done.

+      BrowserContext* browser_context,
+      const base::Callback<void(DOMStorageContextImpl*)>& callback) {
+  ForEachStoragePartition(browser_context,
+                          base::Bind(&DomStorageAdapter, callback));
 }
 
 void SaveSessionStateOnIOThread(ResourceContext* resource_context) {
   resource_context->GetRequestContext()->cookie_store()->GetCookieMonster()->
       SetForceKeepSessionState();
   resource_context->GetRequestContext()->server_bound_cert_service()->
       GetCertStore()->SetForceKeepSessionState();
   ResourceContext::GetAppCacheService(resource_context)->
       set_force_keep_session_state();
 }
 
 void SaveSessionStateOnWebkitThread(
     scoped_refptr<IndexedDBContextImpl> indexed_db_context) {
   indexed_db_context->SetForceKeepSessionState();
 }
 
 void PurgeMemoryOnIOThread(ResourceContext* resource_context) {
   ResourceContext::GetAppCacheService(resource_context)->PurgeMemory();
 }
 
-DOMStorageContextImpl* GetDOMStorageContextImpl(BrowserContext* context) {
+DOMStorageContextImpl* GetDefaultDOMStorageContextImpl(
+    BrowserContext* context) {
   return static_cast<DOMStorageContextImpl*>(
-      BrowserContext::GetDOMStorageContext(context));
+      BrowserContext::GetDefaultDOMStorageContext(context));
 }
 
 }  // namespace
 
 DownloadManager* BrowserContext::GetDownloadManager(
     BrowserContext* context) {
   if (!context->GetUserData(kDownloadManagerKeyName)) {
     ResourceDispatcherHostImpl* rdh = ResourceDispatcherHostImpl::Get();
     DCHECK(rdh);
     DownloadFileManager* file_manager = rdh->download_file_manager();
     DCHECK(file_manager);
     scoped_refptr<DownloadManager> download_manager =
         new DownloadManagerImpl(
             file_manager,
             scoped_ptr<DownloadItemFactory>(),
             GetContentClient()->browser()->GetNetLog());
 
     context->SetUserData(
         kDownloadManagerKeyName,
         new UserDataAdapter<DownloadManager>(download_manager));
     download_manager->SetDelegate(context->GetDownloadManagerDelegate());
     download_manager->Init(context);
   }
 
   return UserDataAdapter<DownloadManager>::Get(
       context, kDownloadManagerKeyName);
 }
 
-QuotaManager* BrowserContext::GetQuotaManager(BrowserContext* context) {
-  CreateQuotaManagerAndClients(context);
-  return UserDataAdapter<QuotaManager>::Get(context, kQuotaManagerKeyName);
+QuotaManager* BrowserContext::GetQuotaManager(BrowserContext* browser_context) {

[Charlie Reis, 2012/07/02 23:21:23]

Should this be GetDefaultQuotaManager?  Or do we not need a process ID specific
one?

[awong, 2012/07/09 20:37:43]

There should be no "Default" anythings, but the transition has to be staged
cause it's just too complicated to remove the default accessor at once.

Right now all the "unmigrated" storage APIs just return a "default" under the
covers.  Later CLs are going to migrate each of these APIs one by one so they
always require a process ID.

+  StoragePartition* partition =
+      GetStoragePartition(browser_context,
+                          ChildProcessHostImpl::kInvalidChildProcessId);

[Charlie Reis, 2012/07/02 23:21:23]

This looks odd to me, but I suppose it makes sense.  Can we document this
behavior (i.e., that an invalid ID returns the default context)?

Should we be special casing this in GetStoragePartition so that we don't have to
rely on the ContentBrowserClient to map the invalid ID to an empty string?

[awong, 2012/07/09 20:37:43]

I think we actually want ContentBrowserClient to do the mapping.  This code as
is is a temporary thing. I added a TODO to each place that uses
kInvalidChildProcess.

+  return partition->quota_manager();
+}
+
+DOMStorageContext* BrowserContext::GetDefaultDOMStorageContext(
+    BrowserContext* browser_context) {
+  return GetDOMStorageContext(browser_context,
+                              ChildProcessHostImpl::kInvalidChildProcessId);
 }
 
 DOMStorageContext* BrowserContext::GetDOMStorageContext(
-    BrowserContext* context) {
-  CreateQuotaManagerAndClients(context);
-  return UserDataAdapter<DOMStorageContextImpl>::Get(
-      context, kDOMStorageContextKeyName);
+    BrowserContext* browser_context,
+    int render_child_id) {
+  StoragePartition* partition =
+      GetStoragePartition(browser_context, render_child_id);
+  return partition->dom_storage_context();
 }
 
-IndexedDBContext* BrowserContext::GetIndexedDBContext(BrowserContext* context) {
-  CreateQuotaManagerAndClients(context);
-  return UserDataAdapter<IndexedDBContext>::Get(
-      context, kIndexedDBContextKeyName);
+IndexedDBContext* BrowserContext::GetIndexedDBContext(
+    BrowserContext* browser_context) {
+  StoragePartition* partition =
+      GetStoragePartition(browser_context,
+                          ChildProcessHostImpl::kInvalidChildProcessId);
+  return partition->indexed_db_context();
 }
 
-DatabaseTracker* BrowserContext::GetDatabaseTracker(BrowserContext* context) {
-  CreateQuotaManagerAndClients(context);
-  return UserDataAdapter<DatabaseTracker>::Get(
-      context, kDatabaseTrackerKeyName);
+DatabaseTracker* BrowserContext::GetDatabaseTracker(
+    BrowserContext* browser_context) {
+  StoragePartition* partition =
+      GetStoragePartition(browser_context,
+                          ChildProcessHostImpl::kInvalidChildProcessId);
+  return partition->database_tracker();
 }
 
 AppCacheService* BrowserContext::GetAppCacheService(
     BrowserContext* browser_context) {
-  CreateQuotaManagerAndClients(browser_context);
-  return UserDataAdapter<ChromeAppCacheService>::Get(
-      browser_context, kAppCacheServicKeyName);
+  StoragePartition* partition =
+      GetStoragePartition(browser_context,
+                          ChildProcessHostImpl::kInvalidChildProcessId);
+  return partition->appcache_service();
 }
 
 FileSystemContext* BrowserContext::GetFileSystemContext(
     BrowserContext* browser_context) {
-  CreateQuotaManagerAndClients(browser_context);
-  return UserDataAdapter<FileSystemContext>::Get(
-      browser_context, kFileSystemContextKeyName);
+  StoragePartition* partition =
+      GetStoragePartition(browser_context,
+                          ChildProcessHostImpl::kInvalidChildProcessId);
+  return partition->filesystem_context();
 }
 
 void BrowserContext::EnsureResourceContextInitialized(BrowserContext* context) {
   // This will be enough to tickle initialization of BrowserContext if
   // necessary, which initializes ResourceContext. The reason we don't call
   // ResourceContext::InitializeResourceContext directly here is that if
   // ResourceContext ends up initializing it will call back into BrowserContext
-  // and when that call return it'll end rewriting its UserData map (with the
+  // and when that call returns it'll end rewriting its UserData map (with the
   // same value) but this causes a race condition. See http://crbug.com/115678.
-  CreateQuotaManagerAndClients(context);
+  GetStoragePartition(context, ChildProcessHostImpl::kInvalidChildProcessId);
 }
 
 void BrowserContext::SaveSessionState(BrowserContext* browser_context) {
   GetDatabaseTracker(browser_context)->SetForceKeepSessionState();
 
   if (BrowserThread::IsMessageLoopValid(BrowserThread::IO)) {
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
         base::Bind(&SaveSessionStateOnIOThread,
                    browser_context->GetResourceContext()));
   }
 
-  GetDOMStorageContextImpl(browser_context)->SetForceKeepSessionState();
+  // TODO(ajwong): This is the only usage of GetDefaultDOMStorageContextImpl().
+  // After we migrate this to support multiple DOMStorageContexts, don't forget
+  // to remove the GetDefaultDOMStorageContextImpl() function as well.
+  GetDefaultDOMStorageContextImpl(browser_context)->SetForceKeepSessionState();
 
   if (BrowserThread::IsMessageLoopValid(BrowserThread::WEBKIT_DEPRECATED)) {
     IndexedDBContextImpl* indexed_db = static_cast<IndexedDBContextImpl*>(
         GetIndexedDBContext(browser_context));
     BrowserThread::PostTask(
         BrowserThread::WEBKIT_DEPRECATED, FROM_HERE,
         base::Bind(&SaveSessionStateOnWebkitThread,
                    make_scoped_refptr(indexed_db)));
   }
 }
 
 void BrowserContext::PurgeMemory(BrowserContext* browser_context) {
   if (BrowserThread::IsMessageLoopValid(BrowserThread::IO)) {
     BrowserThread::PostTask(
         BrowserThread::IO, FROM_HERE,
         base::Bind(&PurgeMemoryOnIOThread,
                    browser_context->GetResourceContext()));
   }
 
-  GetDOMStorageContextImpl(browser_context)->PurgeMemory();
+  DOMStorageContextForEach(browser_context,
+                           base::Bind(&DOMStorageContextImpl::PurgeMemory));
 }
 
 BrowserContext::~BrowserContext() {
-  // These message loop checks are just to avoid leaks in unittests.
-  if (GetUserData(kDatabaseTrackerKeyName) &&
-      BrowserThread::IsMessageLoopValid(BrowserThread::FILE)) {
-    BrowserThread::PostTask(
-        BrowserThread::FILE, FROM_HERE,
-        base::Bind(&webkit_database::DatabaseTracker::Shutdown,
-                   GetDatabaseTracker(this)));
-  }
-
-  if (GetUserData(kDOMStorageContextKeyName))
-    GetDOMStorageContextImpl(this)->Shutdown();
-
   if (GetUserData(kDownloadManagerKeyName))
     GetDownloadManager(this)->Shutdown();
 }
 
 }  // namespace content