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..853911cbddcbabdfdc4a1223f1bb065fe874860b 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,77 @@
 #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 conent to pages that are in an error state (i.e. the URL
+// was bad or could not be loaded).
+class WebContentsErrorObserver
+    : public content::WebContentsObserver,
+      public content::WebContentsUserData<WebContentsErrorObserver> {
+ public:
+  // Once the page has finished loading, send an error message.
+  void DocumentOnLoadCompletedInMainFrame() override {
+    web_contents()->GetMainFrame()->ExecuteJavaScript(
+        base::UTF8ToUTF16(dom_distiller::viewer::GetErrorPageJs()));
+  }
+
+ private:
+  friend class content::WebContentsUserData<WebContentsErrorObserver>;
+  explicit WebContentsErrorObserver(content::WebContents* web_contents)
+      : WebContentsObserver(web_contents) {}
+
+  ~WebContentsErrorObserver() override {
+    content::WebContentsObserver::Observe(NULL);
+  }
+
+  DISALLOW_COPY_AND_ASSIGN(WebContentsErrorObserver);
+};
+
+// 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<WebContentsErrorObserver>;
+
+  ~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
+
+DEFINE_WEB_CONTENTS_USER_DATA_KEY(WebContentsErrorObserver);
+
 namespace dom_distiller {
 
 // Handles receiving data asynchronously for a specific entry, and passing
@@ -104,6 +173,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 +192,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 +223,6 @@ void DomDistillerViewerSource::RequestViewerHandle::DidNavigateMainFrame(
   }
 
   Cancel();
-
 }
 
 void DomDistillerViewerSource::RequestViewerHandle::RenderProcessGone(
@@ -188,13 +260,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 +295,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));

[cjhopman, 2015/03/24 01:45:36]

Isn't it possible now that the second page has it's javascript called before the
first page?

I feel like we should be able to do this with just one observer and delaying all
SendJavaScript() calls (and then just using SendJavaScript as normal for both
the error case and the first page).

[mdjones, 2015/03/24 18:42:52]

Done/See general comment.

     }
   }
 }
@@ -314,6 +392,8 @@ void DomDistillerViewerSource::StartDataRequest(
     // |RequestViewerHandle| will never be called, so clean up now.
     delete request_viewer_handle;
 
+    WebContentsErrorObserver::CreateForWebContents(web_contents);
+
     std::string error_page_html = viewer::GetErrorPageHtml(
         dom_distiller_service_->GetDistilledPagePrefs()->GetTheme(),
         dom_distiller_service_->GetDistilledPagePrefs()->GetFontFamily());