Fix memory leak and bug in the RunsOnCreationTaskRunner check

fml/memory/task_runner_checker.cc

@@ -7,14 +7,18 @@
 namespace fml {
 
 TaskRunnerChecker::TaskRunnerChecker()
-    : initialized_queue_id_(InitTaskQueueId()){};
+    : initialized_queue_id_(InitTaskQueueId()),
+      subsumed_queue_id_(
+          MessageLoopTaskQueues::GetInstance()->GetSubsumedTaskQueueId(
+              initialized_queue_id_)){};
 
 TaskRunnerChecker::~TaskRunnerChecker() = default;
 
 bool TaskRunnerChecker::RunsOnCreationTaskRunner() const {
   FML_CHECK(fml::MessageLoop::IsInitializedForCurrentThread());
   const auto current_queue_id = MessageLoop::GetCurrentTaskQueueId();
-  return RunsOnTheSameThread(current_queue_id, initialized_queue_id_);
+  return RunsOnTheSameThread(current_queue_id, initialized_queue_id_) ||
+         RunsOnTheSameThread(current_queue_id, subsumed_queue_id_);
 };
 
 bool TaskRunnerChecker::RunsOnTheSameThread(TaskQueueId queue_a,

fml/memory/task_runner_checker.h

@@ -22,6 +22,7 @@ class TaskRunnerChecker final {
 
  private:
   TaskQueueId initialized_queue_id_;
+  TaskQueueId subsumed_queue_id_;
 
   TaskQueueId InitTaskQueueId();
 };

fml/memory/task_runner_checker_unittest.cc

@@ -115,5 +115,57 @@ TEST(TaskRunnerCheckerTests, MergedTaskRunnersRunsOnTheSameThread) {
   thread2.join();
 }
 
+TEST(TaskRunnerCheckerTests,
+     PassesRunsOnCreationTaskRunnerIfOnDifferentTaskRunner) {
+  fml::MessageLoop* loop1 = nullptr;
+  fml::AutoResetWaitableEvent latch1;
+  std::thread thread1([&]() {
+    fml::MessageLoop::EnsureInitializedForCurrentThread();
+    loop1 = &fml::MessageLoop::GetCurrent();
+    latch1.Signal();
+    loop1->Run();
+  });
+
+  fml::MessageLoop* loop2 = nullptr;
+  fml::AutoResetWaitableEvent latch2;
+  std::thread thread2([&]() {
+    fml::MessageLoop::EnsureInitializedForCurrentThread();
+    loop2 = &fml::MessageLoop::GetCurrent();
+    latch2.Signal();
+    loop2->Run();
+  });
+
+  latch1.Wait();
+  latch2.Wait();
+
+  fml::TaskQueueId qid1 = loop1->GetTaskRunner()->GetTaskQueueId();
+  fml::TaskQueueId qid2 = loop2->GetTaskRunner()->GetTaskQueueId();
+  fml::MessageLoopTaskQueues::GetInstance()->Merge(qid1, qid2);
+
+  std::unique_ptr<TaskRunnerChecker> checker;
+
+  fml::AutoResetWaitableEvent latch3;
+  loop2->GetTaskRunner()->PostTask([&]() {
+    checker = std::make_unique<TaskRunnerChecker>();
+    EXPECT_EQ(checker->RunsOnCreationTaskRunner(), true);
+    latch3.Signal();
+  });
+  latch3.Wait();
+
+  fml::MessageLoopTaskQueues::GetInstance()->Unmerge(qid1);
+
+  fml::AutoResetWaitableEvent latch4;
+  loop2->GetTaskRunner()->PostTask([&]() {
+    EXPECT_EQ(checker->RunsOnCreationTaskRunner(), true);
+    latch4.Signal();
+  });
+  latch4.Wait();
+
+  loop1->Terminate();
+  loop2->Terminate();
+  thread1.join();
+  thread2.join();
+}
+
 }  // namespace testing
 }  // namespace fml

fml/message_loop_task_queues.cc

@@ -238,7 +238,14 @@ bool MessageLoopTaskQueues::Unmerge(TaskQueueId owner) {
 bool MessageLoopTaskQueues::Owns(TaskQueueId owner,
                                  TaskQueueId subsumed) const {
   std::lock_guard guard(queue_mutex_);
-  return subsumed == queue_entries_.at(owner)->owner_of;
+  return owner != _kUnmerged && subsumed != _kUnmerged &&
+         subsumed == queue_entries_.at(owner)->owner_of;
+}
+
+TaskQueueId MessageLoopTaskQueues::GetSubsumedTaskQueueId(
+    TaskQueueId owner) const {
+  std::lock_guard guard(queue_mutex_);
+  return queue_entries_.at(owner)->owner_of;
 }
 
 // Subsumed queues will never have pending tasks.

fml/message_loop_task_queues.h

@@ -122,6 +122,10 @@ class MessageLoopTaskQueues
   // Returns true if owner owns the subsumed task queue.
   bool Owns(TaskQueueId owner, TaskQueueId subsumed) const;
 
+  // Returns the subsumed task queue if any or |TaskQueueId::kUnmerged|
+  // otherwise.
+  TaskQueueId GetSubsumedTaskQueueId(TaskQueueId owner) const;
+
  private:
   class MergedQueuesRunner;
 

fml/message_loop_task_queues_unittests.cc

@@ -185,6 +185,13 @@ TEST(MessageLoopTaskQueue, QueueDoNotOwnItself) {
   ASSERT_FALSE(task_queue->Owns(queue_id, queue_id));
 }
 
+TEST(MessageLoopTaskQueue, QueueDoNotOwnUnmergedTaskQueueId) {
+  auto task_queue = fml::MessageLoopTaskQueues::GetInstance();
+  ASSERT_FALSE(task_queue->Owns(task_queue->CreateTaskQueue(), _kUnmerged));
+  ASSERT_FALSE(task_queue->Owns(_kUnmerged, task_queue->CreateTaskQueue()));
+  ASSERT_FALSE(task_queue->Owns(_kUnmerged, _kUnmerged));
+}
+
 // TODO(chunhtai): This unit-test is flaky and sometimes fails asynchronizely
 // after the test has finished.
 // https://github.com/flutter/flutter/issues/43858

shell/platform/android/io/flutter/plugin/platform/PlatformViewsController.java

@@ -458,6 +458,7 @@ public void detach() {
     if (platformViewsChannel != null) {
       platformViewsChannel.setPlatformViewsHandler(null);
     }
+    destroyOverlaySurfaces();
     platformViewsChannel = null;
     context = null;
     textureRegistry = null;
@@ -761,6 +762,9 @@ public void onDisplayPlatformView(
   }
 
   public void onDisplayOverlaySurface(int id, int x, int y, int width, int height) {
+    if (overlayLayerViews.get(id) == null) {
+      throw new IllegalStateException("The overlay surface (id:" + id + ") doesn't exist");
+    }
     initializeRootImageViewIfNeeded();
 
     final FlutterImageView overlayView = overlayLayerViews.get(id);

shell/platform/android/test/io/flutter/plugin/platform/PlatformViewsControllerTest.java

@@ -562,6 +562,61 @@ public void onEndFrame__revertsFlutterSurface() {
     assertEquals(flutterView.getRenderSurface().getClass().getSimpleName(), "FlutterSurfaceView");
   }
 
+  @Test
+  @Config(shadows = {ShadowFlutterSurfaceView.class, ShadowFlutterJNI.class})
+  public void detach__destroysOverlaySurfaces() {
+    final PlatformViewsController platformViewsController = new PlatformViewsController();
+
+    final int platformViewId = 0;
+    assertNull(platformViewsController.getPlatformViewById(platformViewId));
+
+    final PlatformViewFactory viewFactory = mock(PlatformViewFactory.class);
+    final PlatformView platformView = mock(PlatformView.class);
+    when(platformView.getView()).thenReturn(mock(View.class));
+    when(viewFactory.create(any(), eq(platformViewId), any())).thenReturn(platformView);
+
+    platformViewsController.getRegistry().registerViewFactory("testType", viewFactory);
+
+    final FlutterJNI jni = new FlutterJNI();
+    jni.attachToNative(false);
+    attach(jni, platformViewsController);
+
+    jni.onFirstFrame();
+
+    // Simulate create call from the framework.
+    createPlatformView(jni, platformViewsController, platformViewId, "testType");
+
+    // Produce a frame that displays a platform view and an overlay surface.
+    platformViewsController.onBeginFrame();
+    platformViewsController.onDisplayPlatformView(
+        platformViewId,
+        /* x=*/ 0,
+        /* y=*/ 0,
+        /* width=*/ 10,
+        /* height=*/ 10,
+        /* viewWidth=*/ 10,
+        /* viewHeight=*/ 10,
+        /* mutatorsStack=*/ new FlutterMutatorsStack());
+
+    final FlutterImageView overlayImageView = mock(FlutterImageView.class);
+    when(overlayImageView.acquireLatestImage()).thenReturn(true);
+
+    final FlutterOverlaySurface overlaySurface =
+        platformViewsController.createOverlaySurface(overlayImageView);
+    // This is OK.
+    platformViewsController.onDisplayOverlaySurface(
+        overlaySurface.getId(), /* x=*/ 0, /* y=*/ 0, /* width=*/ 10, /* height=*/ 10);
+
+    platformViewsController.detach();
+
+    assertThrows(
+        IllegalStateException.class,
+        () -> {
+          platformViewsController.onDisplayOverlaySurface(
+              overlaySurface.getId(), /* x=*/ 0, /* y=*/ 0, /* width=*/ 10, /* height=*/ 10);
+        });
+  }
+
   @Test
   public void checkInputConnectionProxy__falseIfViewIsNull() {
     final PlatformViewsController platformViewsController = new PlatformViewsController();

testing/scenario_app/android/app/build.gradle

@@ -1,6 +1,8 @@
 apply plugin: 'com.android.application'
 apply plugin: 'com.facebook.testing.screenshot'
 
+def leakcanary_version = '2.6'
+
 screenshots {
     failureDir = "${rootProject.buildDir}/reports/diff_failures"
     recordDir = "${rootProject.projectDir}/reports/screenshots"
@@ -13,12 +15,14 @@ android {
         targetCompatibility JavaVersion.VERSION_1_8
     }
     defaultConfig {
-        applicationId "dev.flutter.scenarios"
+        applicationId 'dev.flutter.scenarios'
         minSdkVersion 18
         targetSdkVersion 30
         versionCode 1
-        versionName "1.0"
-        testInstrumentationRunner "dev.flutter.TestRunner"
+        versionName '1.0'
+        testInstrumentationRunner 'dev.flutter.TestRunner'
+        testInstrumentationRunnerArgument 'listener', 'leakcanary.FailTestOnLeakRunListener'
+        testInstrumentationRunnerArguments clearPackageData: 'true'
     }
     buildTypes {
         release {
@@ -38,8 +42,10 @@ dependencies {
     implementation 'androidx.constraintlayout:constraintlayout:1.1.3'
     implementation 'com.google.android.material:material:1.0.0'
     implementation 'androidx.lifecycle:lifecycle-common-java8:2.2.0-alpha01'
+    implementation "com.squareup.leakcanary:leakcanary-android:$leakcanary_version"
     testImplementation 'junit:junit:4.12'
     androidTestImplementation 'androidx.test:runner:1.2.0'
     androidTestImplementation 'androidx.test:rules:1.2.0'
     androidTestImplementation 'androidx.test.espresso:espresso-core:3.2.0'
+    androidTestImplementation "com.squareup.leakcanary:leakcanary-android-instrumentation:$leakcanary_version"
 }

testing/scenario_app/android/app/src/androidTest/java/dev/flutter/scenariosui/MemoryLeakTests.java

@@ -0,0 +1,34 @@
+// Copyright 2013 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package dev.flutter.scenariosui;
+
+import android.content.Intent;
+import androidx.test.filters.LargeTest;
+import androidx.test.rule.ActivityTestRule;
+import androidx.test.runner.AndroidJUnit4;
+import dev.flutter.scenarios.TextPlatformViewActivity;
+import leakcanary.FailTestOnLeak;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@LargeTest
+public class MemoryLeakTests {
+  @Rule
+  public ActivityTestRule<TextPlatformViewActivity> activityRule =
+      new ActivityTestRule<>(
+          TextPlatformViewActivity.class, /*initialTouchMode=*/ false, /*launchActivity=*/ false);
+
+  @Test
+  @FailTestOnLeak
+  public void platformViewHybridComposition_launchActivityFinishAndLaunchAgain() throws Exception {
+    Intent intent = new Intent(Intent.ACTION_MAIN);
+    intent.putExtra("scenario", "platform_view");
+    intent.putExtra("use_android_view", true);
+
+    activityRule.launchActivity(intent);
+  }
+}

testing/scenario_app/android/app/src/main/res/values/strings.xml

@@ -1,4 +1,5 @@
 <resources>
     <string name="app_name">Scenarios</string>
     <string name="title_activity_main">MainActivity</string>
+    <string name="leak_canary_test_class_name">assertk.Assert</string>
 </resources>