Restore sessionStorage when chrome restarts.

chrome/browser/sessions/session_service.cc

Index: chrome/browser/sessions/session_service.cc
diff --git a/chrome/browser/sessions/session_service.cc b/chrome/browser/sessions/session_service.cc
index 7ec80a4ce19f92227e97578a6438b2a3fa96190a..7d6acc2d3a690eea2b248fdf071188b7a10ce00f 100644
--- a/chrome/browser/sessions/session_service.cc
+++ b/chrome/browser/sessions/session_service.cc
@@ -2,6 +2,8 @@
 // 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_thread.h"
+
 #include "chrome/browser/sessions/session_service.h"
 
 #include <algorithm>
@@ -32,6 +34,10 @@
 #include "chrome/browser/ui/tab_contents/tab_contents_wrapper.h"
 #include "chrome/common/chrome_notification_types.h"
 #include "chrome/common/extensions/extension.h"
+#include "content/browser/in_process_webkit/dom_storage_context.h"
+#include "content/browser/in_process_webkit/dom_storage_namespace.h"
+#include "content/browser/in_process_webkit/session_storage_namespace.h"
+#include "content/browser/in_process_webkit/webkit_context.h"
 #include "content/public/browser/navigation_details.h"
 #include "content/public/browser/navigation_entry.h"
 #include "content/public/browser/notification_service.h"
@@ -65,6 +71,7 @@ static const SessionCommand::id_type
 static const SessionCommand::id_type kCommandSetPinnedState = 12;
 static const SessionCommand::id_type kCommandSetExtensionAppID = 13;
 static const SessionCommand::id_type kCommandSetWindowBounds3 = 14;
+static const SessionCommand::id_type kCommandSessionStorageCreated = 15;
 
 // Every kWritesPerReset commands triggers recreating the file.
 static const int kWritesPerReset = 250;
@@ -142,6 +149,56 @@ struct PinnedStatePayload {
 
 }  // namespace
 
+// SessionService needs to read data from DOMStorageContext and it can be only
+// done on the WEBKIT_DEPRECATED thread. This class is for reading the data. It
+// also takes into account that SessionService might get deleted before we have
+// a change to call it back.
+class SessionService::SessionStorageAssociationReader
+    : public base::RefCountedThreadSafe<
+        SessionService::SessionStorageAssociationReader> {
+ public:
+  SessionStorageAssociationReader(SessionService* session_service,
+                                  int64 namespace_id,
+                                  WebKitContext* webkit_context)
+      : session_service_(session_service),
+        namespace_id_(namespace_id),
+        webkit_context_(webkit_context) {
+    content::BrowserThread::PostTask(
+        content::BrowserThread::WEBKIT_DEPRECATED, FROM_HERE,
+        base::Bind(&SessionStorageAssociationReader::ReadData, this));
+  }
+  ~SessionStorageAssociationReader() { }
+
+ private:
+  friend class SessionService;
+
+  void ReadData() {
+    DOMStorageNamespace* dom_storage_namespace = webkit_context_->
+        dom_storage_context()->GetStorageNamespace(namespace_id_, false);
+    if (dom_storage_namespace != NULL) {
+      session_storage_directory_ =
+          dom_storage_namespace->session_storage_directory();
+    }
+    content::BrowserThread::PostTask(
+        content::BrowserThread::UI, FROM_HERE,
+        base::Bind(&SessionStorageAssociationReader::ForwardData, this));
+  }
+
+  void ForwardData() {
+    // If SessionService gets destroyed before this object, it nullifies this
+    // pointer.
+    if (session_service_) {
+      session_service_->AssociateSessionStorage(this, namespace_id_,
+                                                session_storage_directory_);
+    }
+  }
+
+  SessionService* session_service_;
+  int64 namespace_id_;
+  WebKitContext* webkit_context_;
+  FilePath session_storage_directory_;
+};
+
 // SessionService -------------------------------------------------------------
 
 SessionService::SessionService(Profile* profile)
@@ -166,6 +223,17 @@ SessionService::SessionService(const FilePath& save_path)
 
 SessionService::~SessionService() {
   Save();
+  DCHECK(content::BrowserThread::CurrentlyOn(content::BrowserThread::UI));
+  for (SessionStorageAssociationReaders::const_iterator it =
+           session_storage_association_readers_.begin();
+       it != session_storage_association_readers_.end();
+       ++it) {
+    // This is done in the UI thread. SessionStorageAssociationReader only
+    // accesses session_service_ in the UI thread, so this is safe.
+    it->first->session_service_ = NULL;
+  }
+  // References from SessionService to SessionStorageAssociationReader will be
+  // automatically dropped as session_storage_association_readers_ is destroyed.
 }
 
 bool SessionService::RestoreIfNecessary(const std::vector<GURL>& urls_to_open) {
@@ -457,6 +525,8 @@ void SessionService::Init() {
   registrar_.Add(
       this, chrome::NOTIFICATION_TAB_CONTENTS_APPLICATION_EXTENSION_CHANGED,
       content::NotificationService::AllSources());
+  registrar_.Add(this, content::NOTIFICATION_SESSION_STORAGE_NAMESPACE_CREATED,
+                 content::NotificationService::AllSources());
 }
 
 bool SessionService::ShouldNewWindowStartSession() {
@@ -519,6 +589,15 @@ void SessionService::Observe(int type,
           content::Source<TabContentsWrapper>(source).ptr();
       if (tab->profile() != profile())
         return;
+
+      // Record the association between the sessionStorage namespace and the
+      // tab.
+      SessionStorageNamespace* session_storage_namespace =
+          tab->web_contents()->GetController().GetSessionStorageNamespace();
+      session_storage_namespace_map_[session_storage_namespace->id()] =
+          std::make_pair(tab->restore_tab_helper()->window_id(),
+                         tab->restore_tab_helper()->session_id());
+
       SetTabWindow(tab->restore_tab_helper()->window_id(),
                    tab->restore_tab_helper()->session_id());
       if (tab->extension_tab_helper()->extension_app()) {
@@ -627,6 +706,15 @@ void SessionService::Observe(int type,
       break;
     }
 
+    case content::NOTIFICATION_SESSION_STORAGE_NAMESPACE_CREATED: {
+      content::Details<SessionStorageCreatedDetails>
+          namespace_created_details(details);
+      CreateAndScheduleSessionStorageCreatedCommand(
+          namespace_created_details->id,
+          namespace_created_details->session_storage_directory);
+      break;
+    }
+
     default:
       NOTREACHED();
   }
@@ -760,6 +848,21 @@ SessionCommand* SessionService::CreatePinnedStateCommand(
   return command;
 }
 
+void SessionService::CreateAndScheduleSessionStorageCreatedCommand(
+    int64 namespace_id,
+    const FilePath& directory) {
+  SessionStorageNamespaceMap::const_iterator it =
+      session_storage_namespace_map_.find(namespace_id);
+  if (it != session_storage_namespace_map_.end()) {
+    Pickle session_storage_created_pickle;
+    session_storage_created_pickle.WriteInt(it->second.first.id());
+    session_storage_created_pickle.WriteInt(it->second.second.id());
+    session_storage_created_pickle.WriteString(directory.value());
+    ScheduleCommand(new SessionCommand(kCommandSessionStorageCreated,
+                                       session_storage_created_pickle));
+  }
+}
+
 void SessionService::OnGotSessionCommands(
     Handle handle,
     scoped_refptr<InternalGetCommandsRequest> request) {
@@ -779,11 +882,20 @@ void SessionService::RestoreSessionFromCommands(
   std::map<int, SessionTab*> tabs;
   std::map<int, SessionWindow*> windows;
 
-  if (CreateTabsAndWindows(commands, &tabs, &windows)) {
+  // Keeps track of which sessionStorages are in use after session restore, so
+  // that the rest can be deleted.
+  std::set<std::string> needed_session_storages;
+
+  if (CreateTabsAndWindows(commands, &tabs, &windows,
+                           &needed_session_storages)) {
     AddTabsToWindows(&tabs, &windows);
     SortTabsBasedOnVisualOrderAndPrune(&windows, valid_windows);
     UpdateSelectedTabIndex(valid_windows);
   }
+
+  profile()->GetWebKitContext()->DeleteUnneededSessionStorages(
+      needed_session_storages);
+
   STLDeleteValues(&tabs);
   // Don't delete conents of windows, that is done by the caller as all
   // valid windows are added to valid_windows.
@@ -921,7 +1033,8 @@ void SessionService::AddTabsToWindows(std::map<int, SessionTab*>* tabs,
 bool SessionService::CreateTabsAndWindows(
     const std::vector<SessionCommand*>& data,
     std::map<int, SessionTab*>* tabs,
-    std::map<int, SessionWindow*>* windows) {
+    std::map<int, SessionWindow*>* windows,
+    std::set<std::string>* needed_session_storages) {
   // If the file is corrupt (command with wrong size, or unknown command), we
   // still return true and attempt to restore what we we can.
 
@@ -1097,6 +1210,20 @@ bool SessionService::CreateTabsAndWindows(
         GetTab(tab_id, tabs)->extension_app_id.swap(extension_app_id);
         break;
       }
+      case kCommandSessionStorageCreated: {
+        scoped_ptr<Pickle> pickle(command->PayloadAsPickle());
+        SessionID::id_type window_id;
+        SessionID::id_type tab_id;
+        std::string session_storage_directory;
+        void* iter = NULL;
+        pickle->ReadInt(&iter, &window_id);
+        pickle->ReadInt(&iter, &tab_id);
+        pickle->ReadString(&iter, &session_storage_directory);
+        GetTab(tab_id, tabs)->session_storage_directory =
+            FilePath(session_storage_directory);
+        needed_session_storages->insert(session_storage_directory);
+        break;
+      }
 
       default:
         return true;
@@ -1115,6 +1242,16 @@ void SessionService::BuildCommandsForTab(
   DCHECK(tab && commands && window_id.id());
   const SessionID& session_id(tab->restore_tab_helper()->session_id());
   commands->push_back(CreateSetTabWindowCommand(window_id, session_id));
+
+  // Schedule reading the session storage directory on the WEBKIT_DEPRECATED
+  // thread.
+  int64 namespace_id =
+      tab->web_contents()->GetController().GetSessionStorageNamespace()->id();
+  WebKitContext* webkit_context = profile()->GetWebKitContext();
+  scoped_refptr<SessionStorageAssociationReader> reader =
+      new SessionStorageAssociationReader(this, namespace_id, webkit_context);
+  session_storage_association_readers_[reader.get()] = reader;
+
   const int current_index =
       tab->web_contents()->GetController().GetCurrentEntryIndex();
   const int min_index = std::max(0,
@@ -1519,3 +1656,14 @@ void SessionService::RecordUpdatedSaveTime(base::TimeDelta delta,
         50);
   }
 }
+
+void SessionService::AssociateSessionStorage(
+    SessionStorageAssociationReader* reader,
+    int64 namespace_id,
+    const FilePath& directory) {
+  if (!directory.empty()) {
+    CreateAndScheduleSessionStorageCreatedCommand(
+        namespace_id, directory);
+  }
+  session_storage_association_readers_.erase(reader);
+}