1 files changed, 33 insertions, 16 deletions

diff --git a/components/dom_distiller/core/task_tracker.cc b/components/dom_distiller/core/task_tracker.cc
index faf6fc1..ccd5656 100644
--- a/components/dom_distiller/core/task_tracker.cc
+++ b/components/dom_distiller/core/task_tracker.cc
@@ -48,17 +48,13 @@ void TaskTracker::StartBlobFetcher() {
   // |entry_| and asynchronously notify |this| when it is done.
 }
 
-void TaskTracker::SetSaveCallback(SaveCallback callback) {
-  save_callback_ = callback;
-  if (save_callback_.is_null()) {
-    MaybeCancel();
-  } else if (distilled_page_) {
+void TaskTracker::AddSaveCallback(const SaveCallback& callback) {
+  DCHECK(!callback.is_null());
+  save_callbacks_.push_back(callback);
+  if (distilled_page_) {
     // Distillation for this task has already completed, and so it can be
     // immediately saved.
-    base::MessageLoop::current()->PostTask(
-        FROM_HERE,
-        base::Bind(&TaskTracker::DoSaveCallback,
-                   weak_ptr_factory_.GetWeakPtr()));
+    ScheduleSaveCallbacks(true);
   }
 }
 
@@ -77,6 +73,8 @@ scoped_ptr<ViewerHandle> TaskTracker::AddViewer(ViewRequestDelegate* delegate) {
       &TaskTracker::RemoveViewer, weak_ptr_factory_.GetWeakPtr(), delegate)));
 }
 
+const std::string& TaskTracker::GetEntryId() const { return entry_.entry_id(); }
+
 bool TaskTracker::HasEntryId(const std::string& entry_id) const {
   return entry_.entry_id() == entry_id;
 }
@@ -98,7 +96,7 @@ void TaskTracker::RemoveViewer(ViewRequestDelegate* delegate) {
 }
 
 void TaskTracker::MaybeCancel() {
-  if (!save_callback_.is_null() || !viewers_.empty()) {
+  if (!save_callbacks_.empty() || !viewers_.empty()) {
     // There's still work to be done.
     return;
   }
@@ -110,11 +108,27 @@ void TaskTracker::MaybeCancel() {
   DCHECK(self);
 }
 
-void TaskTracker::DoSaveCallback() {
-  DCHECK(distilled_page_);
-  if (!save_callback_.is_null()) {
-    save_callback_.Run(entry_, distilled_page_.get());
-    save_callback_.Reset();
+void TaskTracker::CancelSaveCallbacks() { ScheduleSaveCallbacks(false); }
+
+void TaskTracker::ScheduleSaveCallbacks(bool distillation_succeeded) {
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&TaskTracker::DoSaveCallbacks,
+                 weak_ptr_factory_.GetWeakPtr(),
+                 distillation_succeeded));
+}
+
+void TaskTracker::DoSaveCallbacks(bool distillation_succeeded) {
+  if (!save_callbacks_.empty()) {
+    DistilledPageProto* distilled_proto =
+        distillation_succeeded ? distilled_page_.get() : NULL;
+
+    for (size_t i = 0; i < save_callbacks_.size(); ++i) {
+      DCHECK(!save_callbacks_[i].is_null());
+      save_callbacks_[i].Run(entry_, distilled_proto, distillation_succeeded);
+    }
+
+    save_callbacks_.clear();
     MaybeCancel();
   }
 }
@@ -126,12 +140,15 @@ void TaskTracker::NotifyViewer(ViewRequestDelegate* delegate) {
 
 void TaskTracker::OnDistilledDataReady(scoped_ptr<DistilledPageProto> proto) {
   distilled_page_ = proto.Pass();
+  DCHECK(distilled_page_);
 
   entry_.set_title(distilled_page_->title());
   for (size_t i = 0; i < viewers_.size(); ++i) {
     NotifyViewer(viewers_[i]);
   }
-  DoSaveCallback();
+
+  // Already inside a callback run SaveCallbacks directly.
+  DoSaveCallbacks(true);
 }
 
 }  // namespace dom_distiller