[OT-PW] Implement PaintWorkletImageCache entry lookup

When the raster task pool is running the raster tasks, it would invoke
the PaintWorkletImageCache::PaintImageInTask(). Currently, if an entry
for a certain PaintWorkletInput already exist, we would still call the
Paint() function to produce a PaintRecord and replace that entry, and
then reset the ref-count to 0. This is wrong. This CL fixes it by
looking at the cache, and if an entry already exist for the
PaintWorkletInput, we can just skip the Paint() call and never reset
the ref-count. Unit testing is added.

Bug: 907897
Change-Id: I50f6b6c0ae2a551cdc4a6b6aaf2c2e9bc2f6f30b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1497032
Commit-Queue: Xida Chen <[email protected]>
Reviewed-by: Khushal <[email protected]>
Reviewed-by: vmpstr <[email protected]>
Reviewed-by: Stephen McGruer <[email protected]>
Cr-Commit-Position: refs/heads/master@{#639271}

cc/tiles/image_controller.cc

@@ -163,8 +163,8 @@ void ImageController::ConvertPaintWorkletImagesToTask(
     }
     scoped_refptr<TileTask> result =
         paint_worklet_image_cache_.GetTaskForPaintWorkletImage(*it);
-    DCHECK(result);
-    tasks->push_back(std::move(result));
+    if (result)
+      tasks->push_back(std::move(result));
     // Remove it so that there is no need to check whether an image is
     // PaintWorklet generated or not in TileManager's
     // work_to_schedule->extra_prepaint_images.insert.

cc/tiles/paint_worklet_image_cache.cc

@@ -40,27 +40,59 @@ PaintWorkletImageCache::~PaintWorkletImageCache() {
 
 void PaintWorkletImageCache::SetPaintWorkletLayerPainter(
     std::unique_ptr<PaintWorkletLayerPainter> painter) {
+  DCHECK(!painter_);
   painter_ = std::move(painter);
 }
 
 scoped_refptr<TileTask> PaintWorkletImageCache::GetTaskForPaintWorkletImage(
     const DrawImage& image) {
+  // As described in crbug.com/939192, the |painter_| could be null, and we
+  // should not create any raster task.
+  if (!painter_)
+    return nullptr;
+  DCHECK(image.paint_image().IsPaintWorklet());
   return base::MakeRefCounted<PaintWorkletTaskImpl>(this, image.paint_image());
 }
 
-// TODO(xidachen): dispatch the work to a worklet thread, invoke JS callback.
-// Do check the cache first. If there is already a cache entry for this input,
-// then there is no need to call the Paint() function.
+// TODO(xidachen): we might need to consider the animated property value and the
+// PaintWorkletInput to decide whether we need to call Paint() function or not.
 void PaintWorkletImageCache::PaintImageInTask(const PaintImage& paint_image) {
+  // TODO(crbug.com/939009): When creating a TileTask for a given PaintImage at
+  // GetTaskForPaintWorkletImage, we should not create a new TileTask if there
+  // is already a TileTask for this PaintImage.
+  {
+    base::AutoLock hold(records_lock_);
+    if (records_.find(paint_image.paint_worklet_input()) != records_.end())
+      return;
+  }
+  // Because the compositor could be waiting on the lock in NotifyPrepareTiles,
+  // we unlock here such that the compositor won't be blocked on potentially
+  // slow Paint function.
   // TODO(xidachen): ensure that the canvas operations in the PaintRecord
   // matches the PaintGeneratedImage::Draw.
   sk_sp<PaintRecord> record = painter_->Paint();
-  records_[paint_image.paint_worklet_input()] =
-      PaintWorkletImageCacheValue(std::move(record), 0);
+  {
+    base::AutoLock hold(records_lock_);
+    // It is possible for two or more threads to both pass through the first
+    // lock and arrive here. To avoid ref-count issues caused by potential
+    // racing among threads, we use insert such that if an entry already exists
+    // for a particular key, the value won't be overridden.
+    records_.insert(
+        std::make_pair(paint_image.paint_worklet_input(),
+                       PaintWorkletImageCacheValue(std::move(record), 0)));
+  }
 }
 
 std::pair<PaintRecord*, base::OnceCallback<void()>>
 PaintWorkletImageCache::GetPaintRecordAndRef(PaintWorkletInput* input) {
+  base::AutoLock hold(records_lock_);
+  // If the |painter_| is null, then GetTaskForPaintWorkletImage will return a
+  // null TileTask, and hence there will be no cache entry for this input.
+  if (!painter_) {
+    return std::make_pair(sk_make_sp<PaintOpBuffer>().get(),
+                          base::OnceCallback<void()>());
+  }
+  DCHECK(records_.find(input) != records_.end());
   records_[input].used_ref_count++;
   records_[input].num_of_frames_not_accessed = 0u;
   // The PaintWorkletImageCache object lives as long as the LayerTreeHostImpl,
@@ -78,6 +110,7 @@ void PaintWorkletImageCache::SetNumOfFramesToPurgeCacheEntryForTest(
 }
 
 void PaintWorkletImageCache::DecrementCacheRefCount(PaintWorkletInput* input) {
+  base::AutoLock hold(records_lock_);
   auto it = records_.find(input);
   DCHECK(it != records_.end());
 
@@ -87,12 +120,14 @@ void PaintWorkletImageCache::DecrementCacheRefCount(PaintWorkletInput* input) {
 }
 
 void PaintWorkletImageCache::NotifyDidPrepareTiles() {
+  base::AutoLock hold(records_lock_);
   base::EraseIf(
       records_,
       [this](
           const std::pair<PaintWorkletInput*, PaintWorkletImageCacheValue>& t) {
         return t.second.num_of_frames_not_accessed >=
-               num_of_frames_to_purge_cache_entry_;
+                   num_of_frames_to_purge_cache_entry_ &&
+               t.second.used_ref_count == 0;
       });
   for (auto& pair : records_)
     pair.second.num_of_frames_not_accessed++;

cc/tiles/paint_worklet_image_cache.h

@@ -8,6 +8,7 @@
 #include <utility>
 
 #include "base/containers/flat_map.h"
+#include "base/synchronization/lock.h"
 #include "cc/cc_export.h"
 #include "cc/paint/draw_image.h"
 #include "cc/paint/paint_record.h"
@@ -62,12 +63,17 @@ class CC_EXPORT PaintWorkletImageCache {
 
  private:
   void DecrementCacheRefCount(PaintWorkletInput* input);
+
   // This is a map of paint worklet inputs to a pair of paint record and a
   // reference count. The paint record is the representation of the worklet
   // output based on the input, and the reference count is the number of times
   // that it is used for tile rasterization.
-  // TODO(xidachen): use a struct instead of std::pair.
   base::flat_map<PaintWorkletInput*, PaintWorkletImageCacheValue> records_;
+
+  // The |records_| can be accessed from compositor and raster worker threads at
+  // the same time. To prevent race, we need to lock on it.
+  base::Lock records_lock_;
+
   // The PaintWorkletImageCache is owned by ImageController, which has the same
   // life time as the LayerTreeHostImpl, that guarantees that the painter will
   // live as long as the LayerTreeHostImpl.

cc/tiles/paint_worklet_image_cache_unittest.cc

@@ -2,6 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <memory>
 #include <utility>
 
 #include "cc/tiles/paint_worklet_image_cache.h"
@@ -47,9 +48,6 @@ scoped_refptr<TileTask> GetTaskForPaintWorkletImage(
       paint_image, SkIRect::MakeWH(paint_image.width(), paint_image.height()),
       kNone_SkFilterQuality, CreateMatrix(SkSize::Make(1.f, 1.f), true),
       PaintImage::kDefaultFrameIndex);
-  std::unique_ptr<TestPaintWorkletLayerPainter> painter =
-      std::make_unique<TestPaintWorkletLayerPainter>();
-  cache->SetPaintWorkletLayerPainter(std::move(painter));
   return cache->GetTaskForPaintWorkletImage(draw_image);
 }
 
@@ -64,6 +62,9 @@ void TestPaintRecord(const PaintRecord* record) {
 
 TEST(PaintWorkletImageCacheTest, GetTaskForImage) {
   TestPaintWorkletImageCache cache;
+  std::unique_ptr<TestPaintWorkletLayerPainter> painter =
+      std::make_unique<TestPaintWorkletLayerPainter>();
+  cache.SetPaintWorkletLayerPainter(std::move(painter));
   PaintImage paint_image = CreatePaintImage(100, 100);
   scoped_refptr<TileTask> task =
       GetTaskForPaintWorkletImage(paint_image, &cache);
@@ -116,8 +117,39 @@ TEST(PaintWorkletImageCacheTest, GetTaskForImage) {
   EXPECT_EQ(records[paint_image.paint_worklet_input()].used_ref_count, 0u);
 }
 
+TEST(PaintWorkletImageCacheTest, EntryWithNonZeroRefCountNotPurged) {
+  TestPaintWorkletImageCache cache;
+  std::unique_ptr<TestPaintWorkletLayerPainter> painter =
+      std::make_unique<TestPaintWorkletLayerPainter>();
+  cache.SetPaintWorkletLayerPainter(std::move(painter));
+  PaintImage paint_image = CreatePaintImage(100, 100);
+  scoped_refptr<TileTask> task =
+      GetTaskForPaintWorkletImage(paint_image, &cache);
+  EXPECT_TRUE(task);
+
+  TestTileTaskRunner::ProcessTask(task.get());
+
+  PaintWorkletImageProvider provider(&cache);
+  ImageProvider::ScopedResult result =
+      provider.GetPaintRecordResult(paint_image.paint_worklet_input());
+  base::flat_map<PaintWorkletInput*,
+                 PaintWorkletImageCache::PaintWorkletImageCacheValue>
+      records = cache.GetRecordsForTest();
+  EXPECT_EQ(records[paint_image.paint_worklet_input()].used_ref_count, 1u);
+
+  cache.NotifyDidPrepareTiles();
+  cache.NotifyDidPrepareTiles();
+  cache.NotifyDidPrepareTiles();
+
+  records = cache.GetRecordsForTest();
+  EXPECT_EQ(records.size(), 1u);
+}
+
 TEST(PaintWorkletImageCacheTest, MultipleRecordsInCache) {
   TestPaintWorkletImageCache cache;
+  std::unique_ptr<TestPaintWorkletLayerPainter> painter =
+      std::make_unique<TestPaintWorkletLayerPainter>();
+  cache.SetPaintWorkletLayerPainter(std::move(painter));
   PaintImage paint_image1 = CreatePaintImage(100, 100);
   scoped_refptr<TileTask> task1 =
       GetTaskForPaintWorkletImage(paint_image1, &cache);
@@ -186,5 +218,44 @@ TEST(PaintWorkletImageCacheTest, MultipleRecordsInCache) {
       2u);
 }
 
+// This test ensures that if an entry already exist, then the PaintImageInTask
+// will not replace it with a new entry and reset its ref count.
+TEST(PaintWorkletImageCacheTest, CacheEntryLookup) {
+  TestPaintWorkletImageCache cache;
+  std::unique_ptr<TestPaintWorkletLayerPainter> painter =
+      std::make_unique<TestPaintWorkletLayerPainter>();
+  cache.SetPaintWorkletLayerPainter(std::move(painter));
+  PaintImage paint_image = CreatePaintImage(100, 100);
+  scoped_refptr<TileTask> task =
+      GetTaskForPaintWorkletImage(paint_image, &cache);
+  EXPECT_TRUE(task);
+  PaintWorkletImageProvider provider(&cache);
+
+  TestTileTaskRunner::ProcessTask(task.get());
+
+  {
+    ImageProvider::ScopedResult result =
+        provider.GetPaintRecordResult(paint_image.paint_worklet_input());
+    EXPECT_TRUE(result.paint_record());
+    TestPaintRecord(result.paint_record());
+
+    base::flat_map<PaintWorkletInput*,
+                   PaintWorkletImageCache::PaintWorkletImageCacheValue>
+        records = cache.GetRecordsForTest();
+    // Test the ref count.
+    EXPECT_EQ(records[paint_image.paint_worklet_input()].used_ref_count, 1u);
+
+    // Create a new task with the same PaintWorkletInput as the previous task.
+    // Then ProcessTask will invoke PaintWorkletImageCache::PaintImageInTask,
+    // and it should early exit, without replacing the existing PaintRecord and
+    // resetting the ref count.
+    scoped_refptr<TileTask> task_with_the_same_input =
+        GetTaskForPaintWorkletImage(paint_image, &cache);
+    EXPECT_TRUE(task);
+    TestTileTaskRunner::ProcessTask(task_with_the_same_input.get());
+    EXPECT_EQ(records[paint_image.paint_worklet_input()].used_ref_count, 1u);
+  }
+}
+
 }  // namespace
 }  // namespace cc