Consistent content placement method for dom-distiller

components/dom_distiller/content/dom_distiller_viewer_source.cc

Index: components/dom_distiller/content/dom_distiller_viewer_source.cc
diff --git a/components/dom_distiller/content/dom_distiller_viewer_source.cc b/components/dom_distiller/content/dom_distiller_viewer_source.cc
index a3ae496db86455a07c886dfe45e3b366ece85493..c42b0d6fdda005f24737c44ceebe12067adfa476 100644
--- a/components/dom_distiller/content/dom_distiller_viewer_source.cc
+++ b/components/dom_distiller/content/dom_distiller_viewer_source.cc
@@ -8,6 +8,7 @@
 #include <string>
 #include <vector>
 
+#include "base/logging.h"
 #include "base/memory/ref_counted_memory.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_loop.h"
@@ -25,9 +26,50 @@
 #include "content/public/browser/user_metrics.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/browser/web_contents_observer.h"
+#include "content/public/browser/web_contents_user_data.h"
 #include "net/base/url_util.h"
 #include "net/url_request/url_request.h"
 
+namespace {
+
+// Handles sending the first page of content after the template has been
+// loaded successfully.
+class WebContentsTemplateObserver
+    : public content::WebContentsObserver,
+      public content::WebContentsUserData<WebContentsTemplateObserver> {
+ public:
+  explicit WebContentsTemplateObserver(content::WebContents* web_contents) {
+    content::WebContentsObserver::Observe(web_contents);
+  }
+
+  // Once the page has finished loading, send the page(s).
+  void DocumentOnLoadCompletedInMainFrame() override {
+    web_contents()->GetMainFrame()->ExecuteJavaScript(
+        base::UTF8ToUTF16(js_page_content_));
+    // First page has been sent, no need to keep observing.
+    content::WebContentsObserver::Observe(NULL);
+  }
+
+  // Set the content that should be loaded immediately after the template.
+  void setInitialContent(const std::string &content) {
+    js_page_content_ = content;
+  }
+
+ private:
+  friend class content::WebContentsUserData<WebContentsTemplateObserver>;
+
+  ~WebContentsTemplateObserver() override {
+    content::WebContentsObserver::Observe(NULL);
+  }
+
+  // The JavaScript that puts the initial content in the page template.
+  std::string js_page_content_;
+
+  DISALLOW_COPY_AND_ASSIGN(WebContentsTemplateObserver);
+};
+
+}  // namespace
+
 namespace dom_distiller {
 
 // Handles receiving data asynchronously for a specific entry, and passing
@@ -104,6 +146,9 @@ class DomDistillerViewerSource::RequestViewerHandle
   // Temporary store of pending JavaScript if the page isn't ready to receive
   // data from distillation.
   std::string buffer_;
+
+  // WebContentsObserver for detecting the load completion of the page template.
+  WebContentsTemplateObserver* template_observer_;
 };
 
 DomDistillerViewerSource::RequestViewerHandle::RequestViewerHandle(
@@ -120,6 +165,7 @@ DomDistillerViewerSource::RequestViewerHandle::RequestViewerHandle(
       waiting_for_page_ready_(true) {
   content::WebContentsObserver::Observe(web_contents);
   distilled_page_prefs_->AddObserver(this);
+  template_observer_ = new WebContentsTemplateObserver(web_contents);
 }
 
 DomDistillerViewerSource::RequestViewerHandle::~RequestViewerHandle() {
@@ -150,7 +196,6 @@ void DomDistillerViewerSource::RequestViewerHandle::DidNavigateMainFrame(
   }
 
   Cancel();
-
 }
 
 void DomDistillerViewerSource::RequestViewerHandle::RenderProcessGone(
@@ -188,13 +233,16 @@ void DomDistillerViewerSource::RequestViewerHandle::DidFinishLoad(
 void DomDistillerViewerSource::RequestViewerHandle::OnArticleReady(
     const DistilledArticleProto* article_proto) {
   if (page_count_ == 0) {
-    // This is a single-page article.
-    std::string unsafe_page_html =
-        viewer::GetUnsafeArticleHtml(
-            article_proto,
-            distilled_page_prefs_->GetTheme(),
-            distilled_page_prefs_->GetFontFamily());
+    template_observer_->setInitialContent(
+        viewer::GetUnsafeArticleContentJs(article_proto));
+
+    std::string unsafe_page_html = viewer::GetUnsafeArticleTemplateHtml(
+        &article_proto->pages(0),
+        distilled_page_prefs_->GetTheme(),
+        distilled_page_prefs_->GetFontFamily());
     callback_.Run(base::RefCountedString::TakeString(&unsafe_page_html));
+
+    // Send the actual article content to the page.
   } else if (page_count_ == article_proto->pages_size()) {
     // We may still be showing the "Loading" indicator.
     SendJavaScript(viewer::GetToggleLoadingIndicatorJs(true));
@@ -220,15 +268,18 @@ void DomDistillerViewerSource::RequestViewerHandle::OnArticleUpdated(
     const DistilledPageProto& page =
         article_update.GetDistilledPage(page_count_);
     if (page_count_ == 0) {
+      template_observer_->setInitialContent(
+          viewer::GetUnsafeIncrementalDistilledPageJs(&page, false));
+
       // This is the first page, so send Viewer page scaffolding too.
-      std::string unsafe_page_html = viewer::GetUnsafePartialArticleHtml(
+      std::string unsafe_page_html = viewer::GetUnsafeArticleTemplateHtml(
           &page,
           distilled_page_prefs_->GetTheme(),
           distilled_page_prefs_->GetFontFamily());
       callback_.Run(base::RefCountedString::TakeString(&unsafe_page_html));
     } else {
-      SendJavaScript(
-          viewer::GetUnsafeIncrementalDistilledPageJs(&page, false));
+      // Send the page content to the client.
+      SendJavaScript(viewer::GetUnsafeIncrementalDistilledPageJs(&page, false));
     }
   }
 }
@@ -314,6 +365,10 @@ void DomDistillerViewerSource::StartDataRequest(
     // |RequestViewerHandle| will never be called, so clean up now.
     delete request_viewer_handle;
 
+    WebContentsTemplateObserver* template_observer =

[cjhopman, 2015/03/25 02:30:50]

Who owns this? When is it deleted?

[mdjones, 2015/03/25 21:29:20]

The observer now deletes itself once the first page is sent, the web contents
are destroyed or the renderer is gone.

+        new WebContentsTemplateObserver(web_contents);
+    template_observer->setInitialContent(viewer::GetErrorPageJs());
+
     std::string error_page_html = viewer::GetErrorPageHtml(
         dom_distiller_service_->GetDistilledPagePrefs()->GetTheme(),
         dom_distiller_service_->GetDistilledPagePrefs()->GetFontFamily());