Refactor FrameTimingsObserver for multi-activity support

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/devsupport/inspector/FrameTimingsObserver.kt

@@ -26,13 +26,17 @@ import kotlinx.coroutines.launch
 
 @DoNotStripAny
 internal class FrameTimingsObserver(
-    private val window: Window,
     private val screenshotsEnabled: Boolean,
     private val onFrameTimingSequence: (sequence: FrameTimingSequence) -> Unit,
 ) {
+  private val isSupported = Build.VERSION.SDK_INT >= Build.VERSION_CODES.N
+
   private val handler = Handler(Looper.getMainLooper())
   private var frameCounter: Int = 0
   private var bitmapBuffer: Bitmap? = null
+  private var isStarted: Boolean = false
+
+  @Volatile private var currentWindow: Window? = null
 
   private val frameMetricsListener =
       Window.OnFrameMetricsAvailableListener { _, frameMetrics, _dropCount ->
@@ -41,12 +45,30 @@ internal class FrameTimingsObserver(
         emitFrameTiming(beginTimestamp, endTimestamp)
       }
 
+  fun setCurrentWindow(window: Window?) {
+    if (!isSupported || currentWindow === window) {
+      return
+    }
+
+    currentWindow?.removeOnFrameMetricsAvailableListener(frameMetricsListener)
+    currentWindow = window
+    if (isStarted) {
+      currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, handler)
+    }
+  }
+
   private suspend fun captureScreenshot(): String? = suspendCoroutine { continuation ->
     if (Build.VERSION.SDK_INT < Build.VERSION_CODES.O) {
       continuation.resume(null)
       return@suspendCoroutine
     }
 
+    val window = currentWindow
+    if (window == null) {
+      continuation.resume(null)
+      return@suspendCoroutine
+    }
+
     val decorView = window.decorView
     val width = decorView.width
     val height = decorView.height
@@ -102,17 +124,19 @@ internal class FrameTimingsObserver(
   }
 
   fun start() {
-    frameCounter = 0
-    if (Build.VERSION.SDK_INT < Build.VERSION_CODES.N) {
+    if (!isSupported) {
       return
     }
 
+    frameCounter = 0
+    isStarted = true
+
     // Capture initial screenshot to ensure there's always at least one frame
     // recorded at the start of tracing, even if no UI changes occur
     val timestamp = System.nanoTime()
     emitFrameTiming(timestamp, timestamp)
 
-    window.addOnFrameMetricsAvailableListener(frameMetricsListener, handler)
+    currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, handler)
   }
 
   private fun emitFrameTiming(beginTimestamp: Long, endTimestamp: Long) {
@@ -135,11 +159,13 @@ internal class FrameTimingsObserver(
   }
 
   fun stop() {
-    if (Build.VERSION.SDK_INT < Build.VERSION_CODES.N) {
+    if (!isSupported) {
       return
     }
 
-    window.removeOnFrameMetricsAvailableListener(frameMetricsListener)
+    isStarted = false
+
+    currentWindow?.removeOnFrameMetricsAvailableListener(frameMetricsListener)
     handler.removeCallbacksAndMessages(null)
 
     bitmapBuffer?.recycle()

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/runtime/ReactHostImpl.kt

@@ -250,6 +250,7 @@ public class ReactHostImpl(
     stateTracker.enterState("onHostResume(activity)")
 
     currentActivity = activity
+    frameTimingsObserver?.setCurrentWindow(activity?.window)
 
     maybeEnableDevSupport(true)
     reactLifecycleStateManager.moveToOnHostResume(currentReactContext, activity)
@@ -851,6 +852,7 @@ public class ReactHostImpl(
   private fun moveToHostDestroy(currentContext: ReactContext?) {
     reactLifecycleStateManager.moveToOnHostDestroy(currentContext)
     currentActivity = null
+    frameTimingsObserver?.setCurrentWindow(null)
   }
 
   private fun raiseSoftException(
@@ -1563,18 +1565,16 @@ public class ReactHostImpl(
           when (state) {
             TracingState.ENABLED_IN_BACKGROUND_MODE,
             TracingState.ENABLED_IN_CDP_MODE -> {
-              currentActivity?.window?.let { window ->
-                val observer =
-                    FrameTimingsObserver(
-                        window,
-                        _screenshotsEnabled,
-                        { frameTimingsSequence ->
-                          inspectorTarget.recordFrameTimings(frameTimingsSequence)
-                        },
-                    )
-                observer.start()
-                frameTimingsObserver = observer
-              }
+              val observer =
+                  FrameTimingsObserver(
+                      _screenshotsEnabled,
+                      { frameTimingsSequence ->
+                        inspectorTarget.recordFrameTimings(frameTimingsSequence)
+                      },
+                  )
+              observer.setCurrentWindow(currentActivity?.window)
+              observer.start()
+              frameTimingsObserver = observer
             }
             TracingState.DISABLED -> {
               frameTimingsObserver?.stop()