[wip] Introduce SiteInstanceKeepAlive.

content/browser/frame_host/navigation_entry_screenshot_manager.cc

Index: content/browser/frame_host/navigation_entry_screenshot_manager.cc
diff --git a/content/browser/frame_host/navigation_entry_screenshot_manager.cc b/content/browser/frame_host/navigation_entry_screenshot_manager.cc
index b2bf429dc0a6c66a34266a5267ed62267a8aacb6..034220abf9a9f9e92d17e289a66a2210b1e1a84d 100644
--- a/content/browser/frame_host/navigation_entry_screenshot_manager.cc
+++ b/content/browser/frame_host/navigation_entry_screenshot_manager.cc
@@ -8,9 +8,11 @@
 #include "base/threading/worker_pool.h"
 #include "content/browser/frame_host/navigation_controller_impl.h"
 #include "content/browser/frame_host/navigation_entry_impl.h"
+#include "content/browser/frame_host/render_frame_host_manager.h"
 #include "content/browser/renderer_host/render_view_host_impl.h"
 #include "content/public/browser/render_widget_host.h"
 #include "content/public/browser/render_widget_host_view.h"
+#include "content/public/browser/site_instance.h"
 #include "content/public/common/content_switches.h"
 #include "third_party/skia/include/core/SkCanvas.h"
 #include "third_party/skia/include/core/SkPaint.h"
@@ -75,6 +77,33 @@ class ScreenshotData : public base::RefCountedThreadSafe<ScreenshotData> {
   DISALLOW_COPY_AND_ASSIGN(ScreenshotData);
 };
 
+// A temporary solution to keep a RenderViewHost and SiteInstanceImpl live
+// (after cross-process/cross-site navigation) for long enough to get the
+// screenshot.
+// TODO(sad): This can be removed when RenderFrameProxy replaces swappedout://
+class SiteInstanceKeepAlive {
+ public:
+  explicit SiteInstanceKeepAlive(scoped_refptr<SiteInstanceImpl> site_instance)
+      : site_instance_(site_instance) {
+    site_instance_->increment_active_view_count();
+    site_instance_->GetProcess()->AddPendingView();
+  }
+
+  ~SiteInstanceKeepAlive() {
+    site_instance_->GetProcess()->RemovePendingView();
+    site_instance_->decrement_active_view_count();
+    if (!site_instance_->active_view_count()) {
+      RenderFrameHostManager::ShutdownRenderFrameHostsInSiteInstance(
+        site_instance_->GetId());
+    }
+  }
+
+ private:
+  scoped_refptr<SiteInstanceImpl> site_instance_;
+
+  DISALLOW_COPY_AND_ASSIGN(SiteInstanceKeepAlive);
+};
+
 NavigationEntryScreenshotManager::NavigationEntryScreenshotManager(
     NavigationControllerImpl* owner)
     : owner_(owner),
@@ -134,11 +163,18 @@ void NavigationEntryScreenshotManager::TakeScreenshotImpl(
     NavigationEntryImpl* entry) {
   DCHECK(host && host->GetView());
   DCHECK(entry);
+  scoped_refptr<SiteInstanceImpl> site_instance(
+      static_cast<SiteInstanceImpl*>(host->GetSiteInstance()));
+  SiteInstanceKeepAlive* keep_alive = new SiteInstanceKeepAlive(site_instance);
+  keep_alives_.push_back(keep_alive);
+  if (!web_contents())
+    Observe(owner_->GetWebContents());
   host->CopyFromBackingStore(gfx::Rect(),
       host->GetView()->GetViewBounds().size(),
       base::Bind(&NavigationEntryScreenshotManager::OnScreenshotTaken,
                  screenshot_factory_.GetWeakPtr(),
-                 entry->GetUniqueID()));
+                 entry->GetUniqueID(),
+                 keep_alive));
 }
 
 void NavigationEntryScreenshotManager::SetMinScreenshotIntervalMS(
@@ -147,9 +183,20 @@ void NavigationEntryScreenshotManager::SetMinScreenshotIntervalMS(
   min_screenshot_interval_ms_ = interval_ms;
 }
 
-void NavigationEntryScreenshotManager::OnScreenshotTaken(int unique_id,
-                                                     bool success,
-                                                     const SkBitmap& bitmap) {
+void NavigationEntryScreenshotManager::OnScreenshotTaken(
+    int unique_id,
+    SiteInstanceKeepAlive* keep_alive,
+    bool success,
+    const SkBitmap& bitmap) {
+  if (keep_alive) {
+    // |keep_alive| can be NULL in tests.
+    keep_alives_.erase(std::find(keep_alives_.begin(), keep_alives_.end(),
+                                 keep_alive));
+    keep_alive = NULL;
+    if (keep_alives_.empty())
+      Observe(NULL);
+  }
+
   NavigationEntryImpl* entry = NULL;
   int entry_count = owner_->GetEntryCount();
   for (int i = 0; i < entry_count; ++i) {
@@ -166,8 +213,7 @@ void NavigationEntryScreenshotManager::OnScreenshotTaken(int unique_id,
   }
 
   if (!success || bitmap.empty() || bitmap.isNull()) {
-    if (!ClearScreenshot(entry))
-      OnScreenshotSet(entry);
+    ClearScreenshot(entry);
     return;
   }
 
@@ -212,8 +258,8 @@ void NavigationEntryScreenshotManager::OnScreenshotEncodeComplete(
 
 void NavigationEntryScreenshotManager::OnScreenshotSet(
     NavigationEntryImpl* entry) {
-  if (entry->screenshot().get())
-    PurgeScreenshotsIfNecessary();
+  CHECK(entry->screenshot().get());
+  PurgeScreenshotsIfNecessary();
 }
 
 bool NavigationEntryScreenshotManager::ClearScreenshot(
@@ -290,4 +336,12 @@ void NavigationEntryScreenshotManager::PurgeScreenshotsIfNecessary() {
   CHECK_LE(screenshot_count, kMaxScreenshots);
 }
 
+void NavigationEntryScreenshotManager::WebContentsDestroyed(
+    WebContents* contents) {
+  // The WebContents is going to be destroyed. So discard all pending screenshot
+  // requests and destroy all the SiteInstanceKeepAlives.
+  screenshot_factory_.InvalidateWeakPtrs();
+  keep_alives_.clear();
+}
+
 }  // namespace content