[Code Health] Clarify PaintWorkletPaintDispatcher behaviour

This CL adds some documentation for PaintWorkletPaintDispatcher, and
also changes up the |painter_map_| to allow direct lookup using the
worklet id.

Bug: None
Change-Id: I4f60ddce561b9fbc415f12f7f0bfcdd7a2c173d6
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1625462Reviewed-by: Xida Chen <xidachen@chromium.org>
Reviewed-by: Robert Flack <flackr@chromium.org>
Commit-Queue: Stephen McGruer <smcgruer@chromium.org>
Cr-Commit-Position: refs/heads/master@{#662634}

third_party/blink/renderer/modules/csspaint/paint_worklet_proxy_client.cc

@@ -22,23 +22,28 @@ namespace blink {
 const char PaintWorkletProxyClient::kSupplementName[] =
     "PaintWorkletProxyClient";
 
+// static
+PaintWorkletProxyClient* PaintWorkletProxyClient::From(WorkerClients* clients) {
+  return Supplement<WorkerClients>::From<PaintWorkletProxyClient>(clients);
+}
+
 // static
 PaintWorkletProxyClient* PaintWorkletProxyClient::Create(Document* document,
                                                          int worklet_id) {
   WebLocalFrameImpl* local_frame =
       WebLocalFrameImpl::FromFrame(document->GetFrame());
   PaintWorklet* paint_worklet = PaintWorklet::From(*document->domWindow());
-  scoped_refptr<PaintWorkletPaintDispatcher> compositor_painter_dispatcher =
+  scoped_refptr<PaintWorkletPaintDispatcher> compositor_paint_dispatcher =
       local_frame->LocalRootFrameWidget()->EnsureCompositorPaintDispatcher();
   return MakeGarbageCollected<PaintWorkletProxyClient>(
-      worklet_id, paint_worklet, std::move(compositor_painter_dispatcher));
+      worklet_id, paint_worklet, std::move(compositor_paint_dispatcher));
 }
 
 PaintWorkletProxyClient::PaintWorkletProxyClient(
     int worklet_id,
     PaintWorklet* paint_worklet,
-    scoped_refptr<PaintWorkletPaintDispatcher> compositor_paintee)
-    : compositor_paintee_(std::move(compositor_paintee)),
+    scoped_refptr<PaintWorkletPaintDispatcher> paint_dispatcher)
+    : paint_dispatcher_(std::move(paint_dispatcher)),
       worklet_id_(worklet_id),
       state_(RunState::kUninitialized),
       main_thread_runner_(Thread::MainThread()->GetTaskRunner()),
@@ -46,12 +51,6 @@ PaintWorkletProxyClient::PaintWorkletProxyClient(
   DCHECK(IsMainThread());
 }
 
-void PaintWorkletProxyClient::Trace(blink::Visitor* visitor) {
-  visitor->Trace(document_definition_map_);
-  Supplement<WorkerClients>::Trace(visitor);
-  PaintWorkletPainter::Trace(visitor);
-}
-
 void PaintWorkletProxyClient::AddGlobalScope(WorkletGlobalScope* global_scope) {
   DCHECK(global_scope);
   DCHECK(global_scope->IsContextThread());
@@ -66,14 +65,14 @@ void PaintWorkletProxyClient::AddGlobalScope(WorkletGlobalScope* global_scope) {
     return;
   }
 
-  // All the global scopes that share a single PaintWorkletProxyClient are
-  // running on the same thread, and so we can just grab the task runner from
-  // the last one to call this function and use that.
+  // All the global scopes that share a single PaintWorkletProxyClient run on
+  // the same thread with the same scheduler. As such we can just grab a task
+  // runner from the last one to register.
   scoped_refptr<base::SingleThreadTaskRunner> global_scope_runner =
       global_scope->GetThread()->GetTaskRunner(TaskType::kMiscPlatformAPI);
   state_ = RunState::kWorking;
 
-  compositor_paintee_->RegisterPaintWorkletPainter(this, global_scope_runner);
+  paint_dispatcher_->RegisterPaintWorkletPainter(this, global_scope_runner);
 }
 
 void PaintWorkletProxyClient::RegisterCSSPaintDefinition(
@@ -122,9 +121,9 @@ void PaintWorkletProxyClient::RegisterCSSPaintDefinition(
 
 void PaintWorkletProxyClient::Dispose() {
   if (state_ == RunState::kWorking) {
-    compositor_paintee_->UnregisterPaintWorkletPainter(this);
+    paint_dispatcher_->UnregisterPaintWorkletPainter(worklet_id_);
   }
-  compositor_paintee_ = nullptr;
+  paint_dispatcher_ = nullptr;
 
   state_ = RunState::kDisposed;
 
@@ -133,6 +132,12 @@ void PaintWorkletProxyClient::Dispose() {
   global_scopes_.clear();
 }
 
+void PaintWorkletProxyClient::Trace(blink::Visitor* visitor) {
+  visitor->Trace(document_definition_map_);
+  Supplement<WorkerClients>::Trace(visitor);
+  PaintWorkletPainter::Trace(visitor);
+}
+
 sk_sp<PaintRecord> PaintWorkletProxyClient::Paint(
     CompositorPaintWorkletInput* compositor_input) {
   // TODO: Can this happen? We don't register till all are here.
@@ -162,11 +167,6 @@ sk_sp<PaintRecord> PaintWorkletProxyClient::Paint(
                            style_map, nullptr);
 }
 
-// static
-PaintWorkletProxyClient* PaintWorkletProxyClient::From(WorkerClients* clients) {
-  return Supplement<WorkerClients>::From<PaintWorkletProxyClient>(clients);
-}
-
 void ProvidePaintWorkletProxyClientTo(WorkerClients* clients,
                                       PaintWorkletProxyClient* client) {
   clients->ProvideSupplement(client);

third_party/blink/renderer/modules/csspaint/paint_worklet_proxy_client.h

@@ -18,12 +18,14 @@ namespace blink {
 
 class WorkletGlobalScope;
 
-// Mediates between a worklet-thread bound PaintWorkletGlobalScope and its
-// associated dispatchers. A PaintWorkletProxyClient is associated with a single
-// global scope and one dispatcher to the compositor thread.
+// Mediates between the (multiple) PaintWorkletGlobalScopes on the worklet
+// thread and the (single) PaintWorkletPaintDispatcher on the non-worklet
+// threads. PaintWorkletProxyClient is responsible both for informing the
+// dispatcher about its existence once all global scopes are registered, as well
+// as choosing the global scope to use for any given paint request.
 //
-// This is constructed on the main thread but it is used in the worklet backing
-// thread.
+// This class is constructed on the main thread but it is used in the worklet
+// backing thread.
 class MODULES_EXPORT PaintWorkletProxyClient
     : public GarbageCollectedFinalized<PaintWorkletProxyClient>,
       public Supplement<WorkerClients>,
@@ -32,8 +34,13 @@ class MODULES_EXPORT PaintWorkletProxyClient
   DISALLOW_COPY_AND_ASSIGN(PaintWorkletProxyClient);
 
  public:
+  // blink::Supplement hook to retrieve the PaintWorkletProxyClient for a given
+  // WorkerClients.
   static const char kSupplementName[];
+  static PaintWorkletProxyClient* From(WorkerClients*);
 
+  // Create the PaintWorkletProxyClient for a given PaintWorklet, represented by
+  // its unique |worklet_id|.
   static PaintWorkletProxyClient* Create(Document*, int worklet_id);
 
   PaintWorkletProxyClient(
@@ -42,26 +49,31 @@ class MODULES_EXPORT PaintWorkletProxyClient
       scoped_refptr<PaintWorkletPaintDispatcher> compositor_paintee);
   ~PaintWorkletProxyClient() override = default;
 
-  void Trace(blink::Visitor*) override;
-
-  // PaintWorkletPainter implementation
+  // PaintWorkletPainter implementation.
   int GetWorkletId() const override { return worklet_id_; }
   sk_sp<PaintRecord> Paint(CompositorPaintWorkletInput*) override;
 
+  // Add a global scope to the PaintWorkletProxyClient.
   virtual void AddGlobalScope(WorkletGlobalScope*);
-  const Vector<CrossThreadPersistent<PaintWorkletGlobalScope>>&
-  GetGlobalScopesForTesting() const {
-    return global_scopes_;
-  }
 
+  // Register a paint definition for this PaintWorklet.
+  // See https://drafts.css-houdini.org/css-paint-api-1/#paint-definition
   void RegisterCSSPaintDefinition(const String& name,
                                   CSSPaintDefinition*,
                                   ExceptionState&);
 
+  // Dispose of the PaintWorkletProxyClient. Called when the worklet global
+  // scopes are being torn down. May be called once per global scope - calls
+  // after the first have no effect.
   void Dispose();
 
-  static PaintWorkletProxyClient* From(WorkerClients*);
+  void Trace(blink::Visitor*) override;
 
+  // Hooks for testing.
+  const Vector<CrossThreadPersistent<PaintWorkletGlobalScope>>&
+  GetGlobalScopesForTesting() const {
+    return global_scopes_;
+  }
   const HeapHashMap<String, Member<DocumentPaintDefinition>>&
   DocumentDefinitionMapForTesting() const {
     return document_definition_map_;
@@ -81,16 +93,38 @@ class MODULES_EXPORT PaintWorkletProxyClient
   FRIEND_TEST_ALL_PREFIXES(PaintWorkletProxyClientTest,
                            PaintWorkletProxyClientConstruction);
 
-  scoped_refptr<PaintWorkletPaintDispatcher> compositor_paintee_;
+  // The |paint_dispatcher_| is shared between all PaintWorklets on the same
+  // Renderer process, and is responsible for dispatching paint calls from the
+  // non-worklet threads to the correct PaintWorkletProxyClient on its worklet
+  // thread. PaintWorkletProxyClient requires a reference to the dispatcher in
+  // order to register and unregister itself.
+  scoped_refptr<PaintWorkletPaintDispatcher> paint_dispatcher_;
+
+  // The unique id for the PaintWorklet that this class is a proxy client for.
   const int worklet_id_;
+
+  // The set of global scopes registered for this PaintWorklet. Multiple global
+  // scopes are used to enforce statelessness - paint instances may have their
+  // global scope changed at random which means they cannot easily store state.
   Vector<CrossThreadPersistent<PaintWorkletGlobalScope>> global_scopes_;
+
+  // The current state of the proxy client. PaintWorkletProxyClient is initially
+  // uninitialized. Once all global scopes are registered, it is considered
+  // working - unless it is disposed of before this happens in which case it
+  // stays in the disposed state.
   enum RunState { kUninitialized, kWorking, kDisposed } state_;
 
-  // Stores the definitions for each paint that is registered.
+  // Stores the paint definitions as they are registered from the global scopes.
+  // For a given named paint definition, all global scopes must report the same
+  // DocumentPaintDefinition or the definition is invalid. Additionally we
+  // cannot tell the main thread about a paint definition until all global
+  // scopes have registered it.
   HeapHashMap<String, Member<DocumentPaintDefinition>> document_definition_map_;
 
-  // Used for OffMainThreadPaintWorklet; we post to the main-thread PaintWorklet
-  // instance to store the definitions for each paint that is registered.
+  // The main thread needs to know about registered paint definitions so that it
+  // can invalidate any associated paint objects and correctly create the paint
+  // instance input state for the object, etc. We communicate with it via a
+  // handle to the PaintWorklet called via a stored task runner.
   scoped_refptr<base::SingleThreadTaskRunner> main_thread_runner_;
   CrossThreadWeakPersistent<PaintWorklet> paint_worklet_;
 };

third_party/blink/renderer/modules/csspaint/paint_worklet_proxy_client_test.cc

@@ -108,7 +108,7 @@ TEST_F(PaintWorkletProxyClientTest, PaintWorkletProxyClientConstruction) {
   PaintWorkletProxyClient* proxy_client =
       MakeGarbageCollected<PaintWorkletProxyClient>(1, nullptr, nullptr);
   EXPECT_EQ(proxy_client->worklet_id_, 1);
-  EXPECT_EQ(proxy_client->compositor_paintee_, nullptr);
+  EXPECT_EQ(proxy_client->paint_dispatcher_, nullptr);
 
   scoped_refptr<PaintWorkletPaintDispatcher> dispatcher =
       base::MakeRefCounted<PaintWorkletPaintDispatcher>();
@@ -116,7 +116,7 @@ TEST_F(PaintWorkletProxyClientTest, PaintWorkletProxyClientConstruction) {
   proxy_client = MakeGarbageCollected<PaintWorkletProxyClient>(
       1, nullptr, std::move(dispatcher));
   EXPECT_EQ(proxy_client->worklet_id_, 1);
-  EXPECT_NE(proxy_client->compositor_paintee_, nullptr);
+  EXPECT_NE(proxy_client->paint_dispatcher_, nullptr);
 }
 
 void RunAddGlobalScopesTestOnWorklet(

third_party/blink/renderer/platform/graphics/paint_worklet_paint_dispatcher.cc

@@ -49,21 +49,20 @@ void PaintWorkletPaintDispatcher::RegisterPaintWorkletPainter(
   TRACE_EVENT0("cc",
                "PaintWorkletPaintDispatcher::RegisterPaintWorkletPainter");
 
-  DCHECK(painter);
+  int worklet_id = painter->GetWorkletId();
   MutexLocker lock(painter_map_mutex_);
-  DCHECK(painter_map_.find(painter) == painter_map_.end());
-  painter_map_.insert(painter, painter_runner);
+  DCHECK(painter_map_.find(worklet_id) == painter_map_.end());
+  painter_map_.insert(worklet_id, std::make_pair(painter, painter_runner));
 }
 
 void PaintWorkletPaintDispatcher::UnregisterPaintWorkletPainter(
-    PaintWorkletPainter* painter) {
+    int worklet_id) {
   TRACE_EVENT0("cc",
                "PaintWorkletPaintDispatcher::"
                "UnregisterPaintWorkletPainter");
-  DCHECK(painter);
   MutexLocker lock(painter_map_mutex_);
-  DCHECK(painter_map_.find(painter) != painter_map_.end());
-  painter_map_.erase(painter);
+  DCHECK(painter_map_.find(worklet_id) != painter_map_.end());
+  painter_map_.erase(worklet_id);
 }
 
 // TODO(xidachen): we should bundle all PaintWorkletInputs and send them to the
@@ -86,15 +85,14 @@ sk_sp<cc::PaintRecord> PaintWorkletPaintDispatcher::Paint(
 
   base::WaitableEvent done_event;
 
-  for (auto& pair : copied_painter_map) {
-    if (pair.key->GetWorkletId() != input->WorkletId())
-      continue;
-    PaintWorkletPainter* painter = pair.key;
-    scoped_refptr<base::SingleThreadTaskRunner> task_runner = pair.value;
+  auto it = copied_painter_map.find(input->WorkletId());
+  if (it == copied_painter_map.end())
+    return output;
 
+  PaintWorkletPainter* painter = it->value.first;
+  scoped_refptr<base::SingleThreadTaskRunner> task_runner = it->value.second;
   DCHECK(!task_runner->BelongsToCurrentThread());
-    std::unique_ptr<AutoSignal> done =
-        std::make_unique<AutoSignal>(&done_event);
+  std::unique_ptr<AutoSignal> done = std::make_unique<AutoSignal>(&done_event);
 
   PostCrossThreadTask(
       *task_runner, FROM_HERE,
@@ -106,7 +104,6 @@ sk_sp<cc::PaintRecord> PaintWorkletPaintDispatcher::Paint(
           },
           WrapCrossThreadPersistent(painter), CrossThreadUnretained(input),
           WTF::Passed(std::move(done)), CrossThreadUnretained(&output)));
-  }
 
   done_event.Wait();
 

third_party/blink/renderer/platform/graphics/paint_worklet_paint_dispatcher.h

@@ -20,9 +20,17 @@
 
 namespace blink {
 
-// The dispatcher receives JS paint callback from the compositor, and dispatch
-// the callback to the painter (on the paint worklet thread) that is associated
-// with the given paint image.
+// PaintWorkletPaintDispatcher is responsible for mediating between the raster
+// threads and the PaintWorklet thread(s). It receives requests from raster
+// threads to paint a paint class instance represented by a PaintWorkletInput,
+// dispatches the input to the appropriate PaintWorklet, synchronously receives
+// the result, and passes it back to the raster thread.
+//
+// Each PaintWorklet (there is one per frame, either same-origin or
+// same-process-cross-origin) has a backing thread, which may be shared between
+// worklets, and a scheduler, which is not shared. All PaintWorklets for a
+// single renderer process share one PaintWorkletPaintDispatcher on the
+// compositor side.
 class PLATFORM_EXPORT PaintWorkletPaintDispatcher
     : public ThreadSafeRefCounted<PaintWorkletPaintDispatcher> {
  public:
@@ -32,30 +40,39 @@ class PLATFORM_EXPORT PaintWorkletPaintDispatcher
 
   PaintWorkletPaintDispatcher() = default;
 
-  // Interface for use by the PaintWorklet thread(s) to request calls.
-  // (To the given Painter on the given TaskRunner.)
-  void RegisterPaintWorkletPainter(
-      PaintWorkletPainter*,
-      scoped_refptr<base::SingleThreadTaskRunner> mutator_runner);
-
-  void UnregisterPaintWorkletPainter(PaintWorkletPainter*);
-
+  // Dispatches a single paint class instance - represented by a
+  // PaintWorkletInput - to the appropriate PaintWorklet thread, and blocks
+  // until it receives the result.
   sk_sp<cc::PaintRecord> Paint(cc::PaintWorkletInput*);
 
+  // Register and unregister a PaintWorklet (represented in this context by a
+  // PaintWorkletPainter). A given PaintWorklet is registered once all its
+  // global scopes have been created, and is usually only unregistered when the
+  // associated PaintWorklet thread is being torn down.
+  //
+  // The passed in PaintWorkletPainter* should only be used on the given
+  // base::SingleThreadTaskRunner.
+  using PaintWorkletId = int;
+  void RegisterPaintWorkletPainter(PaintWorkletPainter*,
+                                   scoped_refptr<base::SingleThreadTaskRunner>);
+  void UnregisterPaintWorkletPainter(PaintWorkletId);
+
   using PaintWorkletPainterToTaskRunnerMap =
-      HashMap<CrossThreadPersistent<PaintWorkletPainter>,
-              scoped_refptr<base::SingleThreadTaskRunner>>;
+      HashMap<PaintWorkletId,
+              std::pair<CrossThreadPersistent<PaintWorkletPainter>,
+                        scoped_refptr<base::SingleThreadTaskRunner>>>;
   const PaintWorkletPainterToTaskRunnerMap& PainterMapForTesting() const {
     return painter_map_;
   }
 
  private:
-  // We can have more than one task-runner because using a worklet inside a
-  // frame with a different origin causes a new global scope => new thread.
+  // This class handles paint class instances for multiple PaintWorklets. These
+  // are disambiguated via the PaintWorklets unique id; this map exists to do
+  // that disambiguation.
   PaintWorkletPainterToTaskRunnerMap painter_map_;
 
   // The (Un)registerPaintWorkletPainter comes from the worklet thread, and the
-  // Paint call is initiated from the raster threads, this mutex ensures that
+  // Paint call is initiated from the raster threads - this mutex ensures that
   // accessing / updating the |painter_map_| is thread safe.
   Mutex painter_map_mutex_;