Take snapshot on subtree update

Summary:
DecorView owns a BitmapPool to take snapshots of the view. These snapshots are later on serialised by the manager.

There's a couple of unrelated changes in this diff but that were already in place making it hard to split.

(1) Renamed 'traverseAndSend' to 'processUpdate'. Why?

The observers as a whole shouldn't necessary know that their 'observation' is being sent to any place. Future changes should move the send logic altogether from the observer too. But that can be made within the scope of a different diff.

(2) There was a bug for nodes that were being observed but then unsubscribed from. If the nodes were being marked for observation and observer was already into place, these were not being told to subscribe again for changes.

Reviewed By: LukeDefeo

Differential Revision: D39812943

fbshipit-source-id: af98c5caf54e1c69f97043bae95049395a2e4545

android/plugins/litho/src/main/java/com/facebook/flipper/plugins/uidebugger/litho/LithoObserver.kt

@@ -19,7 +19,7 @@ class LithoViewTreeObserver(val context: Context) : TreeObserver<LithoView>() {
 
   override val type = "Litho"
 
-  var nodeRef: LithoView? = null
+  private var nodeRef: LithoView? = null
 
   override fun subscribe(node: Any) {
 
@@ -27,7 +27,7 @@ class LithoViewTreeObserver(val context: Context) : TreeObserver<LithoView>() {
 
     nodeRef = node as LithoView
 
-    val listener: (view: LithoView) -> Unit = { traverseAndSend(context, node) }
+    val listener: (view: LithoView) -> Unit = { processUpdate(context, node) }
     node.setOnDirtyMountListener(listener)
 
     listener(node)

android/src/main/java/com/facebook/flipper/plugins/uidebugger/model/Events.kt

@@ -17,7 +17,13 @@ data class InitEvent(val rootId: Id) {
 }
 
 @kotlinx.serialization.Serializable
-data class SubtreeUpdateEvent(val txId: Long, val observerType: String, val nodes: List<Node>) {
+data class SubtreeUpdateEvent(
+    val txId: Long,
+    val observerType: String,
+    val rootId: Id,
+    val nodes: List<Node>,
+    val snapshot: String? = null
+) {
   companion object {
     const val name = "subtreeUpdate"
   }

android/src/main/java/com/facebook/flipper/plugins/uidebugger/observers/ApplicationTreeObserver.kt

@@ -35,18 +35,13 @@ class ApplicationTreeObserver(val context: Context) : TreeObserver<ApplicationRe
 
           override fun onRootViewsChanged(rootViews: List<View>) {
             Log.i(LogTag, "Root views updated, num ${rootViews.size}")
-            traverseAndSend(context, applicationRef)
+            processUpdate(context, applicationRef)
           }
         }
     context.applicationRef.rootsResolver.attachListener(rootViewListener)
     // On subscribe, trigger a traversal on whatever roots we have
     rootViewListener.onRootViewsChanged(applicationRef.rootViews)
 
-    // TODO: Subscribing to root view changes but not to activity changes.
-    // Obviously changes in activities have an effect on root views, but
-    // then it may make sense to unsubscribe the root views listener instead
-    // of activities.
-
     Log.i(LogTag, "${context.applicationRef.rootViews.size} root views")
     Log.i(LogTag, "${context.applicationRef.activitiesStack.size} activities")
   }

android/src/main/java/com/facebook/flipper/plugins/uidebugger/observers/DecorViewTreeObserver.kt

@@ -11,6 +11,7 @@ import android.util.Log
 import android.view.View
 import android.view.ViewTreeObserver
 import com.facebook.flipper.plugins.uidebugger.LogTag
+import com.facebook.flipper.plugins.uidebugger.common.BitmapPool
 import com.facebook.flipper.plugins.uidebugger.core.Context
 import com.facebook.flipper.plugins.uidebugger.scheduler.throttleLatest
 import java.lang.ref.WeakReference
@@ -37,14 +38,17 @@ class DecorViewObserver(val context: Context) : TreeObserver<DecorView>() {
 
     Log.i(LogTag, "Subscribing to decor view changes")
 
-    val throttleSend =
+    val throttledUpdate =
         throttleLatest<WeakReference<View>?>(throttleTimeMs, waitScope, mainScope) { weakView ->
-          weakView?.get()?.let { view -> traverseAndSend(context, view) }
+          if (node.width > 0 && node.height > 0) {
+            bitmapPool = BitmapPool(node.width, node.height)
+          }
+          weakView?.get()?.let { view -> processUpdateWithSnapshot(context, view) }
         }
 
     listener =
         ViewTreeObserver.OnPreDrawListener {
-          throttleSend(nodeRef)
+          throttledUpdate(nodeRef)
           true
         }
 
@@ -52,7 +56,7 @@ class DecorViewObserver(val context: Context) : TreeObserver<DecorView>() {
 
     // It can be the case that the DecorView the current observer owns has already
     // drawn. In this case, manually trigger an update.
-    throttleSend(nodeRef)
+    throttledUpdate(nodeRef)
   }
 
   override fun unsubscribe() {
@@ -61,8 +65,12 @@ class DecorViewObserver(val context: Context) : TreeObserver<DecorView>() {
     listener.let {
       nodeRef?.get()?.viewTreeObserver?.removeOnPreDrawListener(it)
       listener = null
-      nodeRef = null
     }
+
+    nodeRef = null
+
+    bitmapPool?.recycle()
+    bitmapPool = null
   }
 }
 

android/src/main/java/com/facebook/flipper/plugins/uidebugger/observers/TreeObserver.kt

@@ -9,7 +9,10 @@ package com.facebook.flipper.plugins.uidebugger.observers
 
 import android.util.Log
 import com.facebook.flipper.plugins.uidebugger.LogTag
+import com.facebook.flipper.plugins.uidebugger.common.BitmapPool
 import com.facebook.flipper.plugins.uidebugger.core.Context
+import com.facebook.flipper.plugins.uidebugger.descriptors.Id
+import com.facebook.flipper.plugins.uidebugger.descriptors.NodeDescriptor
 import com.facebook.flipper.plugins.uidebugger.descriptors.nodeId
 
 /*
@@ -23,11 +26,12 @@ import com.facebook.flipper.plugins.uidebugger.descriptors.nodeId
  * If while traversing it encounters a node type which has its own TreeObserver, it
  * does not traverse that, instead it sets up a Tree observer responsible for that subtree
  *
- * The parent is responsible for detecting when a child observer needs to be cleaned up
+ * The parent is responsible for detecting when a child observer needs to be cleaned up.
  */
 abstract class TreeObserver<T> {
 
   protected val children: MutableMap<Int, TreeObserver<*>> = mutableMapOf()
+  protected var bitmapPool: BitmapPool? = null
 
   abstract val type: String
 
@@ -35,44 +39,63 @@ abstract class TreeObserver<T> {
 
   abstract fun unsubscribe()
 
-  /**
+  fun processUpdate(context: Context, root: Any) = this.processUpdate(context, root, false)
-   * Optional helper method that traverses the layout hierarchy while managing any encountered child
+  fun processUpdateWithSnapshot(context: Context, root: Any) =
-   * observers correctly
+      this.processUpdate(context, root, true)
-   */
+
-  fun traverseAndSend(context: Context, root: Any) {
+  /** Traverses the layout hierarchy while managing any encountered child observers. */
-    val start = System.currentTimeMillis()
+  fun processUpdate(context: Context, root: Any, takeSnapshot: Boolean) {
-    val (visitedNodes, observerRootsNodes) = context.layoutTraversal.traverse(root)
+    val startTimestamp = System.currentTimeMillis()
+    val (visitedNodes, observableRoots) = context.layoutTraversal.traverse(root)
 
     // Add any new observers
-    for (observerRoot in observerRootsNodes) {
+    observableRoots.forEach { observable ->
-      if (!children.containsKey(observerRoot.nodeId())) {
+      if (!children.containsKey(observable.nodeId())) {
-        context.observerFactory.createObserver(observerRoot, context)?.let { childObserver ->
+        context.observerFactory.createObserver(observable, context)?.let { observer ->
           Log.d(
               LogTag,
-              "Observer ${this.type} discovered new child of type ${childObserver.type} Node ID ${observerRoot.nodeId()}")
+              "Observer ${this.type} discovered new child of type ${observer.type} Node ID ${observable.nodeId()}")
-          childObserver.subscribe(observerRoot)
+          observer.subscribe(observable)
-          children[observerRoot.nodeId()] = childObserver
+          children[observable.nodeId()] = observer
         }
+      } else {
+        children[observable.nodeId()]?.subscribe(observable)
       }
     }
 
     // Remove any old observers
-    val observerRootIds = observerRootsNodes.map { it.nodeId() }
+    val observableRootsIdentifiers = observableRoots.map { it.nodeId() }
-    for (childKey in children.keys) {
+    val removables = mutableListOf<Id>()
-      if (!observerRootIds.contains(childKey)) {
+    children.keys.forEach { key ->
-        children[childKey]?.let { childObserver ->
+      if (!observableRootsIdentifiers.contains(key)) {
+        children[key]?.let { observer ->
           Log.d(
               LogTag,
-              "Observer ${this.type} cleaning up child of type ${childObserver.type} Node ID $childKey")
+              "Observer ${this.type} cleaning up child of type ${observer.type} Node ID $key")
 
-          childObserver.cleanUpRecursive()
+          observer.cleanUpRecursive()
         }
+        removables.add(key)
       }
     }
+    removables.forEach { key -> children.remove(key) }
 
     Log.d(LogTag, "For Observer ${this.type} Sending ${visitedNodes.size}")
-    context.treeObserverManager.send(
+
-        SubtreeUpdate(type, visitedNodes, start, System.currentTimeMillis()))
+    var recyclableBitmap: BitmapPool.RecyclableBitmap? = null
+    if (takeSnapshot && bitmapPool != null) {
+      @Suppress("unchecked_cast")
+      val descriptor =
+          context.descriptorRegister.descriptorForClassUnsafe(root::class.java)
+              as NodeDescriptor<Any>
+      recyclableBitmap = bitmapPool?.getBitmap()
+      descriptor.getSnapshot(root, recyclableBitmap?.bitmap)
+    }
+
+    val endTimestamp = System.currentTimeMillis()
+    context.treeObserverManager.enqueueUpdate(
+        SubtreeUpdate(
+            type, root.nodeId(), visitedNodes, startTimestamp, endTimestamp, recyclableBitmap))
   }
 
   fun cleanUpRecursive() {

android/src/main/java/com/facebook/flipper/plugins/uidebugger/observers/TreeObserverManager.kt

@@ -7,12 +7,19 @@
 
 package com.facebook.flipper.plugins.uidebugger.observers
 
+import android.annotation.SuppressLint
+import android.graphics.Bitmap
+import android.util.Base64
+import android.util.Base64OutputStream
 import android.util.Log
 import com.facebook.flipper.plugins.uidebugger.LogTag
+import com.facebook.flipper.plugins.uidebugger.common.BitmapPool
 import com.facebook.flipper.plugins.uidebugger.core.Context
+import com.facebook.flipper.plugins.uidebugger.descriptors.Id
 import com.facebook.flipper.plugins.uidebugger.model.Node
 import com.facebook.flipper.plugins.uidebugger.model.PerfStatsEvent
 import com.facebook.flipper.plugins.uidebugger.model.SubtreeUpdateEvent
+import java.io.ByteArrayOutputStream
 import java.util.concurrent.atomic.AtomicInteger
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.Channel
@@ -20,9 +27,11 @@ import kotlinx.serialization.json.Json
 
 data class SubtreeUpdate(
     val observerType: String,
+    val rootId: Id,
     val nodes: List<Node>,
     val startTime: Long,
-    val traversalCompleteTime: Long
+    val traversalCompleteTime: Long,
+    val snapshot: BitmapPool.RecyclableBitmap?
 )
 
 /** Holds the root observer and manages sending updates to desktop */
@@ -34,7 +43,7 @@ class TreeObserverManager(val context: Context) {
   private val workerScope = CoroutineScope(Dispatchers.IO)
   private val txId = AtomicInteger()
 
-  fun send(update: SubtreeUpdate) {
+  fun enqueueUpdate(update: SubtreeUpdate) {
     treeUpdates.trySend(update)
   }
 
@@ -42,6 +51,7 @@ class TreeObserverManager(val context: Context) {
    * 1. Sets up the root observer
    * 2. Starts worker to listen to channel, which serializers and sends data over connection
    */
+  @SuppressLint("NewApi")
   fun start() {
 
     treeUpdates = Channel(Channel.UNLIMITED)
@@ -51,16 +61,34 @@ class TreeObserverManager(val context: Context) {
         workerScope.launch {
           while (isActive) {
             try {
-
               val treeUpdate = treeUpdates.receive()
-
               val onWorkerThread = System.currentTimeMillis()
-
               val txId = txId.getAndIncrement().toLong()
-              val serialized =
+
-                  Json.encodeToString(
+              var serialized: String?
-                      SubtreeUpdateEvent.serializer(),
+              if (treeUpdate.snapshot == null) {
-                      SubtreeUpdateEvent(txId, treeUpdate.observerType, treeUpdate.nodes))
+                serialized =
+                    Json.encodeToString(
+                        SubtreeUpdateEvent.serializer(),
+                        SubtreeUpdateEvent(
+                            txId, treeUpdate.observerType, treeUpdate.rootId, treeUpdate.nodes))
+              } else {
+                val stream = ByteArrayOutputStream()
+                val base64Stream = Base64OutputStream(stream, Base64.DEFAULT)
+                treeUpdate.snapshot.bitmap?.compress(Bitmap.CompressFormat.JPEG, 100, base64Stream)
+                val snapshot = stream.toString()
+                serialized =
+                    Json.encodeToString(
+                        SubtreeUpdateEvent.serializer(),
+                        SubtreeUpdateEvent(
+                            txId,
+                            treeUpdate.observerType,
+                            treeUpdate.rootId,
+                            treeUpdate.nodes,
+                            snapshot))
+
+                treeUpdate.snapshot.recycle()
+              }
 
               val serializationEnd = System.currentTimeMillis()
 
@@ -80,13 +108,13 @@ class TreeObserverManager(val context: Context) {
                       serializationComplete = serializationEnd,
                       socketComplete = socketEnd,
                       nodesCount = treeUpdate.nodes.size)
+
               context.connectionRef.connection?.send(
                   PerfStatsEvent.name, Json.encodeToString(PerfStatsEvent.serializer(), perfStats))
             } catch (e: CancellationException) {} catch (e: java.lang.Exception) {
               Log.e(LogTag, "Unexpected Error in channel ", e)
             }
           }
-
           Log.i(LogTag, "Shutting down worker")
         }
   }