Make session restore understand that tabs have multiple SessionStorageNamespaces.

chrome/browser/sessions/session_service.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 "chrome/browser/sessions/session_service.h"
 
 #include <algorithm>
 #include <set>
 #include <utility>
 #include <vector>
@@ skipping 600 matching lines @@
                    tab->restore_tab_helper()->session_id());
       if (tab->extension_tab_helper()->extension_app()) {
         SetTabExtensionAppID(
             tab->restore_tab_helper()->window_id(),
             tab->restore_tab_helper()->session_id(),
             tab->extension_tab_helper()->extension_app()->id());
       }
 
       // Record the association between the SessionStorageNamespace and the
       // tab.
-      //
-      // TODO(ajwong): This should be processing the whole map rather than
-      // just the default. This in particular will not work for tabs with only
-      // isolated apps which won't have a default partition.
-      content::SessionStorageNamespace* session_storage_namespace =
-          tab->web_contents()->GetController().
-              GetDefaultSessionStorageNamespace();
-      ScheduleCommand(CreateSessionStorageAssociatedCommand(
-          tab->restore_tab_helper()->session_id(),
-          session_storage_namespace->persistent_id()));
-      session_storage_namespace->SetShouldPersist(true);
+      // Record the association between all the SessionStorageNamespaces and
+      // their containing tab.
+      const content::SessionStorageNamespaceMap& session_storage_namespace_map =
+          tab->web_contents()->GetController().GetSessionStorageNamespaceMap();
+      LOG(ERROR) << "Associating map of size " << session_storage_namespace_map.size();
+      for (content::SessionStorageNamespaceMap::const_iterator it =
+               session_storage_namespace_map.begin();
+           it != session_storage_namespace_map.end();
+           ++it) {
+        LOG(ERROR) << "Associating --" << it->first << "++ with --" << it->second->persistent_id();
+        ScheduleCommand(CreateSessionStorageAssociatedCommand(
+            tab->restore_tab_helper()->session_id(),
+            it->second->persistent_id(),
+            it->first));
+        it->second->SetShouldPersist(true);
+      }
       break;
     }
 
     case chrome::NOTIFICATION_TAB_CONTENTS_DESTROYED: {
       TabContents* tab = content::Source<TabContents>(source).ptr();
       if (!tab || tab->profile() != profile())
         return;
       // Allow the associated sessionStorage to get deleted; it won't be needed
       // in the session restore.
-      content::SessionStorageNamespace* session_storage_namespace =
-          tab->web_contents()->GetController().
-              GetDefaultSessionStorageNamespace();
-      session_storage_namespace->SetShouldPersist(false);
+     const content::SessionStorageNamespaceMap& session_storage_namespace_map =
+         tab->web_contents()->GetController().GetSessionStorageNamespaceMap();
+      for (content::SessionStorageNamespaceMap::const_iterator it =
+               session_storage_namespace_map.begin();
+           it != session_storage_namespace_map.end();
+           ++it) {
+        it->second->SetShouldPersist(false);
+      }
       TabClosed(tab->restore_tab_helper()->window_id(),
                 tab->restore_tab_helper()->session_id(),
                 tab->web_contents()->GetClosedByUserGesture());
       RecordSessionUpdateHistogramData(type, &last_updated_tab_closed_time_);
       break;
     }
 
     case content::NOTIFICATION_NAV_LIST_PRUNED: {
       TabContents* tab = TabContents::FromWebContents(
           content::Source<content::NavigationController>(source).ptr()->
@@ skipping 201 matching lines @@
   payload.tab_id = tab_id.id();
   payload.pinned_state = is_pinned;
   SessionCommand* command =
       new SessionCommand(kCommandSetPinnedState, sizeof(payload));
   memcpy(command->contents(), &payload, sizeof(payload));
   return command;
 }
 
 SessionCommand* SessionService::CreateSessionStorageAssociatedCommand(
     const SessionID& tab_id,
-    const std::string& session_storage_persistent_id) {
+    const std::string& session_storage_persistent_id,
+    const std::string& partiton_id) {
   Pickle pickle;
   pickle.WriteInt(tab_id.id());
   pickle.WriteString(session_storage_persistent_id);
+  pickle.WriteString(partiton_id);
   return new SessionCommand(kCommandSessionStorageAssociated, pickle);
 }
 
 void SessionService::OnGotSessionCommands(
     Handle handle,
     scoped_refptr<InternalGetCommandsRequest> request) {
   if (request->canceled())
     return;
 
   ScopedVector<SessionWindow> valid_windows;
@@ skipping 383 matching lines @@
         }
 
         GetTab(tab_id, tabs)->user_agent_override.swap(user_agent_override);
         break;
       }
 
       case kCommandSessionStorageAssociated: {
         scoped_ptr<Pickle> command_pickle(command->PayloadAsPickle());
         SessionID::id_type command_tab_id;
         std::string session_storage_persistent_id;
+        std::string partition_id;
         PickleIterator iter(*command_pickle.get());
         if (!command_pickle->ReadInt(&iter, &command_tab_id) ||
-            !command_pickle->ReadString(&iter, &session_storage_persistent_id))
-          return true;
+            !command_pickle->ReadString(&iter, &session_storage_persistent_id) ||
+            !command_pickle->ReadString(&iter, &partition_id))
+          continue;  // We can safely just lose SessionStorage on a restore.
         // Associate the session storage back.
-        GetTab(command_tab_id, tabs)->session_storage_persistent_id =
-            session_storage_persistent_id;
+        GetTab(command_tab_id, tabs)->
+            session_storage_persistent_id_map[partition_id] =
+                session_storage_persistent_id;
         break;
       }
 
       default:
         VLOG(1) << "Failed reading an unknown command " << command->id();
         return true;
     }
   }
   return true;
 }
@@ skipping 53 matching lines @@
     }
   }
   commands->push_back(
       CreateSetSelectedNavigationIndexCommand(session_id, current_index));
 
   if (index_in_window != -1) {
     commands->push_back(
         CreateSetTabIndexInWindowCommand(session_id, index_in_window));
   }
 
-  // Record the association between the sessionStorage namespace and the tab.
-  content::SessionStorageNamespace* session_storage_namespace =
-      tab->web_contents()->GetController().GetDefaultSessionStorageNamespace();
-  ScheduleCommand(CreateSessionStorageAssociatedCommand(
-      tab->restore_tab_helper()->session_id(),
-      session_storage_namespace->persistent_id()));
+  // Record the association between the SessionStorageNamespaces and the tab.
+  const content::SessionStorageNamespaceMap& session_storage_namespace_map =
+      tab->web_contents()->GetController().GetSessionStorageNamespaceMap();
+  for (content::SessionStorageNamespaceMap::const_iterator it =
+       session_storage_namespace_map.begin();
+       it != session_storage_namespace_map.end();
+       ++it) {
+    // TODO(ajwong): Again I wonder if we should do this?
+    it->second->SetShouldPersist(true);
+    ScheduleCommand(CreateSessionStorageAssociatedCommand(
+        tab->restore_tab_helper()->session_id(),
+        it->second->persistent_id(), it->first));
+  }
 }
 
 void SessionService::BuildCommandsForBrowser(
     Browser* browser,
     std::vector<SessionCommand*>* commands,
     IdToRange* tab_to_available_range,
     std::set<SessionID::id_type>* windows_to_track) {
   DCHECK(browser && commands);
   DCHECK(browser->session_id().id());
 
@@ skipping 362 matching lines @@
       50);
   if (use_long_period) {
     std::string long_name_("SessionRestore.SaveLongPeriod");
     UMA_HISTOGRAM_CUSTOM_TIMES(long_name_,
         delta,
         save_delay_in_mins_,
         save_delay_in_hrs_,
         50);
   }
 }