Introduce `TaskRunnerAffineWeakPtrFactory` to generate `TaskRunnerAffineWeakPtr`s (#18346)

fml/memory/weak_ptr.h

@@ -30,17 +30,18 @@ struct DebugTaskRunnerChecker {
 template <typename T>
 class WeakPtrFactory;
 
-// Class for "weak pointers" that can be invalidated. Valid weak pointers can
-// only originate from a |WeakPtrFactory| (see below), though weak pointers are
-// copyable and movable.
+// Class for "weak pointers" that can be invalidated. Valid weak pointers
+// can only originate from a |WeakPtrFactory| (see below), though weak
+// pointers are copyable and movable.
 //
-// Weak pointers are not in general thread-safe. They may only be *used* on a
-// single thread, namely the same thread as the "originating" |WeakPtrFactory|
-// (which can invalidate the weak pointers that it generates).
+// Weak pointers are not in general thread-safe. They may only be *used* on
+// a single thread, namely the same thread as the "originating"
+// |WeakPtrFactory| (which can invalidate the weak pointers that it
+// generates).
 //
 // However, weak pointers may be passed to other threads, reset on other
-// threads, or destroyed on other threads. They may also be reassigned on other
-// threads (in which case they should then only be used on the thread
+// threads, or destroyed on other threads. They may also be reassigned on
+// other threads (in which case they should then only be used on the thread
 // corresponding to the new "originating" |WeakPtrFactory|).
 template <typename T>
 class WeakPtr {
@@ -117,9 +118,6 @@ class WeakPtr {
   explicit WeakPtr(T* ptr, fml::RefPtr<fml::internal::WeakPtrFlag>&& flag)
       : ptr_(ptr), flag_(std::move(flag)) {}
 
-  T* ptr_;
-  fml::RefPtr<fml::internal::WeakPtrFlag> flag_;
-
   virtual void CheckThreadSafety() const {
     FML_DCHECK_CREATION_THREAD_IS_CURRENT(checker_.checker);
   }
@@ -134,12 +132,17 @@ class WeakPtr {
                    fml::RefPtr<fml::internal::WeakPtrFlag>&& flag,
                    DebugThreadChecker checker)
       : ptr_(ptr), flag_(std::move(flag)), checker_(checker) {}
-
+  T* ptr_;
+  fml::RefPtr<fml::internal::WeakPtrFlag> flag_;
   DebugThreadChecker checker_;
 
   // Copy/move construction/assignment supported.
 };
 
+// Forward declaration, so |TaskRunnerAffineWeakPtr<T>| can friend it.
+template <typename T>
+class TaskRunnerAffineWeakPtrFactory;
+
 // A weak pointer that can be used in different threads as long as
 // the threads are belong to the same |TaskRunner|.
 //
@@ -153,14 +156,13 @@ class TaskRunnerAffineWeakPtr : public WeakPtr<T> {
 
   template <typename U>
   TaskRunnerAffineWeakPtr(const TaskRunnerAffineWeakPtr<U>& r)
-      : WeakPtr<T>(static_cast<T*>(r.ptr_), r.flag_), checker_(r.checker_) {}
+      : WeakPtr<T>(r), checker_(r.checker_) {}
 
   TaskRunnerAffineWeakPtr(TaskRunnerAffineWeakPtr<T>&& r) = default;
 
   template <typename U>
   TaskRunnerAffineWeakPtr(TaskRunnerAffineWeakPtr<U>&& r)
-      : WeakPtr<T>(static_cast<T*>(r.ptr_), std::move(r.flag_)),
-        checker_(r.checker_) {}
+      : WeakPtr<T>(r), checker_(r.checker_) {}
 
   ~TaskRunnerAffineWeakPtr() = default;
 
@@ -176,7 +178,9 @@ class TaskRunnerAffineWeakPtr : public WeakPtr<T> {
   }
 
  private:
-  friend class WeakPtrFactory<T>;
+  template <typename U>
+  friend class TaskRunnerAffineWeakPtr;
+  friend class TaskRunnerAffineWeakPtrFactory<T>;
 
   explicit TaskRunnerAffineWeakPtr(
       T* ptr,
@@ -248,28 +252,55 @@ class WeakPtrFactory {
     return WeakPtr<T>(ptr_, flag_.Clone(), checker_);
   }
 
+ private:
+  // Note: See weak_ptr_internal.h for an explanation of why we store the
+  // pointer here, instead of in the "flag".
+  T* const ptr_;
+  fml::RefPtr<fml::internal::WeakPtrFlag> flag_;
+
+  void CheckThreadSafety() const {
+    FML_DCHECK_CREATION_THREAD_IS_CURRENT(checker_.checker);
+  }
+
+  DebugThreadChecker checker_;
+
+  FML_DISALLOW_COPY_AND_ASSIGN(WeakPtrFactory);
+};
+
+// A type of |WeakPtrFactory| that produces |TaskRunnerAffineWeakPtr| instead of
+// |WeakPtr|.
+template <typename T>
+class TaskRunnerAffineWeakPtrFactory {
+ public:
+  explicit TaskRunnerAffineWeakPtrFactory(T* ptr)
+      : ptr_(ptr), flag_(fml::MakeRefCounted<fml::internal::WeakPtrFlag>()) {
+    FML_DCHECK(ptr_);
+  }
+
+  ~TaskRunnerAffineWeakPtrFactory() {
+    CheckThreadSafety();
+    flag_->Invalidate();
+  }
+
   // Gets a new weak pointer, which will be valid until either
   // |InvalidateWeakPtrs()| is called or this object is destroyed.
-  TaskRunnerAffineWeakPtr<T> GetTaskRunnerAffineWeakPtr() const {
-    return TaskRunnerAffineWeakPtr<T>(ptr_, flag_.Clone(),
-                                      task_runner_checker_);
+  TaskRunnerAffineWeakPtr<T> GetWeakPtr() const {
+    return TaskRunnerAffineWeakPtr<T>(ptr_, flag_.Clone(), checker_);
   }
 
- protected:
+ private:
   // Note: See weak_ptr_internal.h for an explanation of why we store the
   // pointer here, instead of in the "flag".
   T* const ptr_;
   fml::RefPtr<fml::internal::WeakPtrFlag> flag_;
 
   void CheckThreadSafety() const {
-    FML_DCHECK_CREATION_THREAD_IS_CURRENT(checker_.checker);
+    FML_DCHECK_TASK_RUNNER_IS_CURRENT(checker_.checker);
   }
 
- private:
-  DebugThreadChecker checker_;
-  DebugTaskRunnerChecker task_runner_checker_;
+  DebugTaskRunnerChecker checker_;
 
-  FML_DISALLOW_COPY_AND_ASSIGN(WeakPtrFactory);
+  FML_DISALLOW_COPY_AND_ASSIGN(TaskRunnerAffineWeakPtrFactory);
 };
 
 }  // namespace fml

fml/memory/weak_ptr_unittest.cc

@@ -188,13 +188,13 @@ TEST(TaskRunnerAffineWeakPtrTest, ShouldNotCrashIfRunningOnTheSameTaskRunner) {
                        &loop2_task_start_latch]() {
     fml::MessageLoop::EnsureInitializedForCurrentThread();
     int data = 0;
-    WeakPtrFactory<int> factory(&data);
+    TaskRunnerAffineWeakPtrFactory<int> factory(&data);
     loop2 = &fml::MessageLoop::GetCurrent();
 
     loop2->GetTaskRunner()->PostTask([&]() {
       latch2.Signal();
       loop2_task_start_latch.Wait();
-      TaskRunnerAffineWeakPtr<int> ptr = factory.GetTaskRunnerAffineWeakPtr();
+      TaskRunnerAffineWeakPtr<int> ptr = factory.GetWeakPtr();
       EXPECT_EQ(*ptr, data);
       loop2_task_finish_latch.Signal();
     });

shell/common/rasterizer.cc

@@ -54,10 +54,11 @@ Rasterizer::Rasterizer(
 Rasterizer::~Rasterizer() = default;
 
 fml::TaskRunnerAffineWeakPtr<Rasterizer> Rasterizer::GetWeakPtr() const {
-  return weak_factory_.GetTaskRunnerAffineWeakPtr();
+  return weak_factory_.GetWeakPtr();
 }
 
-fml::WeakPtr<SnapshotDelegate> Rasterizer::GetSnapshotDelegate() const {
+fml::TaskRunnerAffineWeakPtr<SnapshotDelegate> Rasterizer::GetSnapshotDelegate()
+    const {
   return weak_factory_.GetWeakPtr();
 }
 
@@ -155,7 +156,7 @@ void Rasterizer::Draw(fml::RefPtr<Pipeline<flutter::LayerTree>> pipeline) {
   switch (consume_result) {
     case PipelineConsumeResult::MoreAvailable: {
       task_runners_.GetRasterTaskRunner()->PostTask(
-          [weak_this = weak_factory_.GetTaskRunnerAffineWeakPtr(), pipeline]() {
+          [weak_this = weak_factory_.GetWeakPtr(), pipeline]() {
             if (weak_this) {
               weak_this->Draw(pipeline);
             }

shell/common/rasterizer.h

@@ -186,7 +186,7 @@ class Rasterizer final : public SnapshotDelegate {
   ///
   fml::TaskRunnerAffineWeakPtr<Rasterizer> GetWeakPtr() const;
 
-  fml::WeakPtr<SnapshotDelegate> GetSnapshotDelegate() const;
+  fml::TaskRunnerAffineWeakPtr<SnapshotDelegate> GetSnapshotDelegate() const;
 
   //----------------------------------------------------------------------------
   /// @brief      Sometimes, it may be necessary to render the same frame again
@@ -430,7 +430,7 @@ class Rasterizer final : public SnapshotDelegate {
   fml::closure next_frame_callback_;
   bool user_override_resource_cache_bytes_;
   std::optional<size_t> max_cache_bytes_;
-  fml::WeakPtrFactory<Rasterizer> weak_factory_;
+  fml::TaskRunnerAffineWeakPtrFactory<Rasterizer> weak_factory_;
   fml::RefPtr<fml::RasterThreadMerger> raster_thread_merger_;
 
   // |SnapshotDelegate|

shell/common/shell.cc

@@ -338,7 +338,7 @@ Shell::Shell(DartVMRef vm, TaskRunners task_runners, Settings settings)
   fml::TaskRunner::RunNowOrPostTask(
       task_runners_.GetRasterTaskRunner(), fml::MakeCopyable([this]() mutable {
         this->weak_factory_gpu_ =
-            std::make_unique<fml::WeakPtrFactory<Shell>>(this);
+            std::make_unique<fml::TaskRunnerAffineWeakPtrFactory<Shell>>(this);
       }));
 
   // Install service protocol handlers.
@@ -1156,7 +1156,7 @@ void Shell::OnFrameRasterized(const FrameTiming& timing) {
     // never be reported until the next animation starts.
     frame_timings_report_scheduled_ = true;
     task_runners_.GetRasterTaskRunner()->PostDelayedTask(
-        [self = weak_factory_gpu_->GetTaskRunnerAffineWeakPtr()]() {
+        [self = weak_factory_gpu_->GetWeakPtr()]() {
           if (!self.get()) {
             return;
           }

shell/common/shell.h

@@ -580,7 +580,7 @@ class Shell final : public PlatformView::Delegate,
 
   // For accessing the Shell via the raster thread, necessary for various
   // rasterizer callbacks.
-  std::unique_ptr<fml::WeakPtrFactory<Shell>> weak_factory_gpu_;
+  std::unique_ptr<fml::TaskRunnerAffineWeakPtrFactory<Shell>> weak_factory_gpu_;
 
   friend class testing::ShellTest;
 

shell/gpu/gpu_surface_gl.cc

@@ -257,8 +257,8 @@ std::unique_ptr<SurfaceFrame> GPUSurfaceGL::AcquireFrame(const SkISize& size) {
   surface->getCanvas()->setMatrix(root_surface_transformation);
 
   SurfaceFrame::SubmitCallback submit_callback =
-      [weak = weak_factory_.GetTaskRunnerAffineWeakPtr()](
-          const SurfaceFrame& surface_frame, SkCanvas* canvas) {
+      [weak = weak_factory_.GetWeakPtr()](const SurfaceFrame& surface_frame,
+                                          SkCanvas* canvas) {
         return weak ? weak->PresentSurface(canvas) : false;
       };
 

shell/gpu/gpu_surface_gl.h

@@ -58,7 +58,7 @@ class GPUSurfaceGL : public Surface {
   // external view embedder is present.
   const bool render_to_surface_;
   bool valid_ = false;
-  fml::WeakPtrFactory<GPUSurfaceGL> weak_factory_;
+  fml::TaskRunnerAffineWeakPtrFactory<GPUSurfaceGL> weak_factory_;
 
   bool CreateOrUpdateSurfaces(const SkISize& size);
 

shell/gpu/gpu_surface_software.cc

@@ -57,8 +57,8 @@ std::unique_ptr<SurfaceFrame> GPUSurfaceSoftware::AcquireFrame(
   canvas->resetMatrix();
 
   SurfaceFrame::SubmitCallback on_submit =
-      [self = weak_factory_.GetTaskRunnerAffineWeakPtr()](
-          const SurfaceFrame& surface_frame, SkCanvas* canvas) -> bool {
+      [self = weak_factory_.GetWeakPtr()](const SurfaceFrame& surface_frame,
+                                          SkCanvas* canvas) -> bool {
     // If the surface itself went away, there is nothing more to do.
     if (!self || !self->IsValid() || canvas == nullptr) {
       return false;

shell/gpu/gpu_surface_software.h

@@ -41,7 +41,7 @@ class GPUSurfaceSoftware : public Surface {
   // hack to make avoid allocating resources for the root surface when an
   // external view embedder is present.
   const bool render_to_surface_;
-  fml::WeakPtrFactory<GPUSurfaceSoftware> weak_factory_;
+  fml::TaskRunnerAffineWeakPtrFactory<GPUSurfaceSoftware> weak_factory_;
 
   FML_DISALLOW_COPY_AND_ASSIGN(GPUSurfaceSoftware);
 };