Adds Trace Listener API to DebugEvents

Summary: Adds Trace Listener API to DebugEvents

Reviewed By: Katalune

Differential Revision: D55142166

fbshipit-source-id: 01bd3bcca32fded0eb71bd5f8505cfa6cd31bf09

litho-rendercore/src/main/java/com/facebook/rendercore/debug/DebugEvents.kt

@@ -18,6 +18,7 @@ package com.facebook.rendercore.debug
 
 import com.facebook.kotlin.compilerplugins.dataclassgenerate.annotation.DataClassGenerate
 import com.facebook.rendercore.LogLevel
+import com.facebook.rendercore.debug.DebugEvent.Companion.All
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.atomic.AtomicReference
@@ -67,12 +68,14 @@ sealed class DebugEvent(
       |  type = '$type',
       |  renderStateId = '$renderStateId',
       |  thread = '$threadName',
-      |  attributes = ${attributes.entries.joinToString(
-        prefix = "{\n",
-        separator = ",\n",
-        postfix = "\n|  }",
-        transform = { e -> "|    ${e.key} = ${e.value}" }
-    )}
+      |  attributes = ${
+      attributes.entries.joinToString(
+          prefix = "{\n",
+          separator = ",\n",
+          postfix = "\n|  }",
+          transform = { e -> "|    ${e.key} = ${e.value}" }
+      )
+    }
     """
         .trimMargin()
   }
@@ -152,6 +155,20 @@ abstract class DebugEventSubscriber(vararg val events: String) {
   abstract fun onEvent(event: DebugEvent)
 }
 
+/**
+ * Any [DebugEventSubscriber] that implements this interface will receive callbacks just before and
+ * right after the a trace event. The subscriber can return a token from [onTraceStart] and the
+ * dispatcher will return the token in [onTraceEnd].
+ */
+interface TraceListener<T : Any?> {
+
+  // Called right before trace event starts
+  fun onTraceStart(type: String): T
+
+  // Called after the trace event ends
+  fun onTraceEnd(token: T, event: DebugProcessEvent)
+}
+
 /** Object to dispatch debug events */
 object DebugEventDispatcher {
 
@@ -341,9 +358,14 @@ object DebugEventDispatcher {
     }
 
     // find the subscribers listening for this event
+    val traceListeners = mutableListOf<TraceListener<Any?>>()
     val subscribersToNotify =
         subscribers.filter { subscriber ->
-          subscriber.events.contains(type) || subscriber.events.contains(DebugEvent.All)
+          val matches = subscriber.events.contains(type) || subscriber.events.contains(All)
+          if (subscriber is TraceListener<*>) {
+            traceListeners.add(subscriber as TraceListener<Any?>)
+          }
+          return@filter matches
         }
 
     // run the block if there are no subscribers
@@ -354,6 +376,7 @@ object DebugEventDispatcher {
     val attributes = LinkedHashMap<String, Any?>()
     attributesAccumulator(attributes)
 
+    val tokens: List<Any?> = traceListeners.map { it.onTraceStart(type) }
     val timestamp = System.currentTimeMillis()
     val startTime = System.nanoTime()
     val res = block(TraceScope(attributes = attributes))
@@ -368,6 +391,7 @@ object DebugEventDispatcher {
             attributes = attributes,
         )
 
+    traceListeners.forEachIndexed { i, it -> it.onTraceEnd(token = tokens[i], event = event) }
     subscribersToNotify.forEach { it.onEvent(event) }
 
     return res

litho-rendercore/src/test/java/com/facebook/rendercore/debug/DebugEventsTest.kt

@@ -420,6 +420,24 @@ class DebugEventsTest {
     assertThat(events).hasSize(1)
   }
 
+  @Test
+  fun `trace event should be dispatched to trace listeners`() {
+    val traceListener1 = TestTraceListener()
+    DebugEventBus.subscribe(traceListener1)
+
+    DebugEventDispatcher.trace(type = TestEvent, renderStateId = { TestRenderStateId }) {}
+
+    assertThat(traceListener1.records.size).isEqualTo(1)
+
+    val traceListener2 = TestTraceListener()
+    DebugEventBus.subscribe(traceListener2)
+
+    DebugEventDispatcher.trace(type = TestEvent, renderStateId = { TestRenderStateId }) {}
+
+    assertThat(traceListener1.records.size).isEqualTo(2)
+    assertThat(traceListener2.records.size).isEqualTo(1)
+  }
+
   class TestEventSubscriber(val listener: (DebugEvent) -> Unit) : DebugEventSubscriber(TestEvent) {
     override fun onEvent(event: DebugEvent) {
       listener(event)
@@ -434,4 +452,24 @@ class DebugEventsTest {
       listener(event)
     }
   }
+
+  class TestTraceListener : TraceListener<Int>, DebugEventSubscriber(TestEvent) {
+
+    val records = mutableListOf<Any>()
+
+    private var counter = 0
+
+    override fun onEvent(event: DebugEvent): Unit = Unit
+
+    override fun onTraceStart(type: String): Int {
+      records.add(type)
+      val token = counter
+      counter++
+      return token
+    }
+
+    override fun onTraceEnd(token: Int, event: DebugProcessEvent) {
+      assertThat(records[token]).isEqualTo(event.type)
+    }
+  }
 }