Fix trailing frame capture after recording ended

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

@@ -18,8 +18,6 @@ import android.view.PixelCopy
 import android.view.Window
 import com.facebook.proguard.annotations.DoNotStripAny
 import java.io.ByteArrayOutputStream
-import kotlin.coroutines.resume
-import kotlin.coroutines.suspendCoroutine
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -30,20 +28,38 @@ internal class FrameTimingsObserver(
     private val onFrameTimingSequence: (sequence: FrameTimingSequence) -> Unit,
 ) {
   private val isSupported = Build.VERSION.SDK_INT >= Build.VERSION_CODES.N
+  private val mainHandler = Handler(Looper.getMainLooper())
 
-  private val handler = Handler(Looper.getMainLooper())
   private var frameCounter: Int = 0
-  private var bitmapBuffer: Bitmap? = null
-  private var isStarted: Boolean = false
-
+  @Volatile private var isTracing: Boolean = false
   @Volatile private var currentWindow: Window? = null
 
-  private val frameMetricsListener =
-      Window.OnFrameMetricsAvailableListener { _, frameMetrics, _dropCount ->
-        val beginTimestamp = frameMetrics.getMetric(FrameMetrics.VSYNC_TIMESTAMP)
-        val endTimestamp = beginTimestamp + frameMetrics.getMetric(FrameMetrics.TOTAL_DURATION)
-        emitFrameTiming(beginTimestamp, endTimestamp)
-      }
+  fun start() {
+    if (!isSupported) {
+      return
+    }
+
+    frameCounter = 0
+    isTracing = 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)
+
+    currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, mainHandler)
+  }
+
+  fun stop() {
+    if (!isSupported) {
+      return
+    }
+
+    isTracing = false
+
+    currentWindow?.removeOnFrameMetricsAvailableListener(frameMetricsListener)
+    mainHandler.removeCallbacksAndMessages(null)
+  }
 
   fun setCurrentWindow(window: Window?) {
     if (!isSupported || currentWindow === window) {
@@ -52,123 +68,107 @@ internal class FrameTimingsObserver(
 
     currentWindow?.removeOnFrameMetricsAvailableListener(frameMetricsListener)
     currentWindow = window
-    if (isStarted) {
-      currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, handler)
+    if (isTracing) {
+      currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, mainHandler)
+    }
+  }
+
+  private val frameMetricsListener =
+      Window.OnFrameMetricsAvailableListener { _, frameMetrics, _ ->
+        // Guard against calls arriving after stop() has ended tracing. Async work scheduled from
+        // previous frames will still finish.
+        if (!isTracing) {
+          return@OnFrameMetricsAvailableListener
+        }
+        val beginTimestamp = frameMetrics.getMetric(FrameMetrics.VSYNC_TIMESTAMP)
+        val endTimestamp = beginTimestamp + frameMetrics.getMetric(FrameMetrics.TOTAL_DURATION)
+        emitFrameTiming(beginTimestamp, endTimestamp)
+      }
+
+  private fun emitFrameTiming(beginTimestamp: Long, endTimestamp: Long) {
+    val frameId = frameCounter++
+    val threadId = Process.myTid()
+
+    if (screenshotsEnabled) {
+      // Initiate PixelCopy immediately on the main thread, while still in the current frame,
+      // then process and emit asynchronously once the copy is complete.
+      captureScreenshot { screenshot ->
+        CoroutineScope(Dispatchers.Default).launch {
+          onFrameTimingSequence(
+              FrameTimingSequence(frameId, threadId, beginTimestamp, endTimestamp, screenshot)
+          )
+        }
+      }
+    } else {
+      CoroutineScope(Dispatchers.Default).launch {
+        onFrameTimingSequence(
+            FrameTimingSequence(frameId, threadId, beginTimestamp, endTimestamp, null)
+        )
+      }
     }
   }
 
-  private suspend fun captureScreenshot(): String? = suspendCoroutine { continuation ->
+  // Must be called from the main thread so that PixelCopy captures the current frame.
+  private fun captureScreenshot(callback: (String?) -> Unit) {
     if (Build.VERSION.SDK_INT < Build.VERSION_CODES.O) {
-      continuation.resume(null)
-      return@suspendCoroutine
+      callback(null)
+      return
     }
 
     val window = currentWindow
     if (window == null) {
-      continuation.resume(null)
-      return@suspendCoroutine
+      callback(null)
+      return
     }
 
     val decorView = window.decorView
     val width = decorView.width
     val height = decorView.height
-
-    // Reuse bitmap if dimensions haven't changed
-    val bitmap =
-        bitmapBuffer?.let {
-          if (it.width == width && it.height == height) {
-            it
-          } else {
-            it.recycle()
-            null
-          }
-        } ?: Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888).also { bitmapBuffer = it }
+    val bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
 
     PixelCopy.request(
         window,
         bitmap,
         { copyResult ->
           if (copyResult == PixelCopy.SUCCESS) {
             CoroutineScope(Dispatchers.Default).launch {
-              var scaledBitmap: Bitmap? = null
-              try {
-                val scaleFactor = 0.25f
-                val scaledWidth = (width * scaleFactor).toInt()
-                val scaledHeight = (height * scaleFactor).toInt()
-                scaledBitmap = Bitmap.createScaledBitmap(bitmap, scaledWidth, scaledHeight, true)
-
-                val compressFormat =
-                    if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R)
-                        Bitmap.CompressFormat.WEBP_LOSSY
-                    else Bitmap.CompressFormat.WEBP
-
-                val base64 =
-                    ByteArrayOutputStream().use { outputStream ->
-                      scaledBitmap.compress(compressFormat, 0, outputStream)
-                      Base64.encodeToString(outputStream.toByteArray(), Base64.NO_WRAP)
-                    }
-
-                continuation.resume(base64)
-              } catch (e: Exception) {
-                continuation.resume(null)
-              } finally {
-                scaledBitmap?.recycle()
-              }
+              callback(encodeScreenshot(window, bitmap, width, height))
             }
           } else {
-            continuation.resume(null)
+            bitmap.recycle()
+            callback(null)
           }
         },
-        handler,
+        mainHandler,
     )
   }
 
-  fun start() {
-    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)
-
-    currentWindow?.addOnFrameMetricsAvailableListener(frameMetricsListener, handler)
-  }
-
-  private fun emitFrameTiming(beginTimestamp: Long, endTimestamp: Long) {
-    val frameId = frameCounter++
-    val threadId = Process.myTid()
-
-    CoroutineScope(Dispatchers.Default).launch {
-      val screenshot = if (screenshotsEnabled) captureScreenshot() else null
-
-      onFrameTimingSequence(
-          FrameTimingSequence(
-              frameId,
-              threadId,
-              beginTimestamp,
-              endTimestamp,
-              screenshot,
-          )
-      )
+  private fun encodeScreenshot(window: Window, bitmap: Bitmap, width: Int, height: Int): String? {
+    var scaledBitmap: Bitmap? = null
+    return try {
+      val density = window.context.resources.displayMetrics.density
+      val scaledWidth = (width / density * SCREENSHOT_SCALE_FACTOR).toInt()
+      val scaledHeight = (height / density * SCREENSHOT_SCALE_FACTOR).toInt()
+      scaledBitmap = Bitmap.createScaledBitmap(bitmap, scaledWidth, scaledHeight, true)
+
+      val compressFormat =
+          if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) Bitmap.CompressFormat.WEBP_LOSSY
+          else Bitmap.CompressFormat.JPEG
+
+      ByteArrayOutputStream().use { outputStream ->
+        scaledBitmap.compress(compressFormat, SCREENSHOT_QUALITY, outputStream)
+        Base64.encodeToString(outputStream.toByteArray(), Base64.NO_WRAP)
+      }
+    } catch (e: Exception) {
+      null
+    } finally {
+      scaledBitmap?.recycle()
+      bitmap.recycle()
     }
   }
 
-  fun stop() {
-    if (!isSupported) {
-      return
-    }
-
-    isStarted = false
-
-    currentWindow?.removeOnFrameMetricsAvailableListener(frameMetricsListener)
-    handler.removeCallbacksAndMessages(null)
-
-    bitmapBuffer?.recycle()
-    bitmapBuffer = null
+  companion object {
+    private const val SCREENSHOT_SCALE_FACTOR = 0.75f
+    private const val SCREENSHOT_QUALITY = 80
   }
 }