[Paint Preview] Change PaintPreviewClient callback thread

The callbacks passed to PaintPreviewClient were getting run on an
arbitrary thread via PostTask. This is a problem as they should be run
on the calling thread such that WeakPtrs bound to them are valid.

Fix
- Make the callbacks run on the same thread that called capture.
- Explicitly run CapturePaintPreviewInternal on the UI thread.

Follow-up
- Make a dedicated SequencedTaskRunner for executing captures.

Bug: 1058079
Change-Id: I1dfa969de613112b4bd1650051e394094dffda5e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2083567
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Cr-Commit-Position: refs/heads/master@{#746517}

components/paint_preview/browser/paint_preview_client.cc

@@ -135,8 +135,13 @@ void PaintPreviewClient::CapturePaintPreview(
   document_data.root_url = render_frame_host->GetLastCommittedURL();
   document_data.source_id =
       ukm::GetSourceIdForWebContentsDocument(web_contents());
+  document_data.called_on_thread = base::SequencedTaskRunnerHandle::Get();
   all_document_data_.insert({params.document_guid, std::move(document_data)});
-  CapturePaintPreviewInternal(params, render_frame_host);
+  base::PostTask(
+      FROM_HERE, {content::BrowserThread::UI},
+      base::BindOnce(&PaintPreviewClient::CapturePaintPreviewInternal,
+                     weak_ptr_factory_.GetWeakPtr(), params,
+                     render_frame_host));
 }
 
 void PaintPreviewClient::CaptureSubframePaintPreview(
@@ -147,7 +152,11 @@ void PaintPreviewClient::CaptureSubframePaintPreview(
   params.document_guid = guid;
   params.clip_rect = rect;
   params.is_main_frame = false;
-  CapturePaintPreviewInternal(params, render_subframe_host);
+  base::PostTask(
+      FROM_HERE, {content::BrowserThread::UI},
+      base::BindOnce(&PaintPreviewClient::CapturePaintPreviewInternal,
+                     weak_ptr_factory_.GetWeakPtr(), params,
+                     render_subframe_host));
 }
 
 void PaintPreviewClient::RenderFrameDeleted(
@@ -215,6 +224,7 @@ mojom::PaintPreviewCaptureParamsPtr PaintPreviewClient::CreateMojoParams(
 void PaintPreviewClient::CapturePaintPreviewInternal(
     const PaintPreviewParams& params,
     content::RenderFrameHost* render_frame_host) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   // Use a frame's embedding token as its GUID. Note that we create a GUID for
   // the main frame so that we can treat it the same as other frames.
   auto token = render_frame_host->GetEmbeddingToken();
@@ -260,9 +270,7 @@ void PaintPreviewClient::RequestCaptureOnUIThread(
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   auto* document_data = &all_document_data_[params.document_guid];
   if (result.error != base::File::FILE_OK) {
-    // Don't block up the UI thread and answer the callback on a different
-    // thread.
-    base::PostTask(
+    document_data->called_on_thread->PostTask(
         FROM_HERE,
         base::BindOnce(std::move(document_data->callback), params.document_guid,
                        mojom::PaintPreviewStatus::kFileCreationError, nullptr));
@@ -385,7 +393,7 @@ void PaintPreviewClient::OnFinished(base::UnguessableToken guid,
     // At a minimum one frame was captured successfully, it is up to the
     // caller to decide if a partial success is acceptable based on what is
     // contained in the proto.
-    base::PostTask(
+    document_data->called_on_thread->PostTask(
         FROM_HERE,
         base::BindOnce(std::move(document_data->callback), guid,
                        document_data->had_error
@@ -394,8 +402,8 @@ void PaintPreviewClient::OnFinished(base::UnguessableToken guid,
                        std::move(document_data->proto)));
   } else {
     // A proto could not be created indicating all frames failed to capture.
-    base::PostTask(FROM_HERE,
-                   base::BindOnce(std::move(document_data->callback), guid,
+    document_data->called_on_thread->PostTask(
+        FROM_HERE, base::BindOnce(std::move(document_data->callback), guid,
                                   mojom::PaintPreviewStatus::kFailed, nullptr));
   }
   all_document_data_.erase(guid);

components/paint_preview/browser/paint_preview_client.h

@@ -57,6 +57,9 @@ class PaintPreviewClient
 
   ~PaintPreviewClient() override;
 
+  // These methods do not need to be run on the UI thread, but do need to be run
+  // from a SequencedTaskRunner.
+
   // Captures a paint preview corresponding to the content of
   // |render_frame_host|. This will work for capturing entire documents if
   // passed the main frame or for just a specific subframe depending on
@@ -104,6 +107,7 @@ class PaintPreviewClient
 
     // Callback that is invoked on completion of data.
     PaintPreviewCallback callback;
+    scoped_refptr<base::SequencedTaskRunner> called_on_thread;
 
     // All the render frames that are still required.
     base::flat_set<base::UnguessableToken> awaiting_subframes;
@@ -148,7 +152,7 @@ class PaintPreviewClient
       base::File file);
 
   // Sets up for a capture of a frame on |render_frame_host| according to
-  // |params|.
+  // |params|. This should be called from the UI thread.
   void CapturePaintPreviewInternal(const PaintPreviewParams& params,
                                    content::RenderFrameHost* render_frame_host);