Track batch execution via inspector. (#361)

instacart · Mar 18, 2024 · d106b81 · d106b81
1 parent 3089282
commit d106b81
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 18 deletions.
diff --git a/formula/src/main/java/com/instacart/formula/FormulaRuntime.kt b/formula/src/main/java/com/instacart/formula/FormulaRuntime.kt
@@ -178,6 +178,9 @@ class FormulaRuntime<Input : Any, Output : Any>(
             synchronizedUpdateQueue.postUpdate {
                 // We disable run until all batch updates are processed
                 isRunEnabled = false
+
+                inspector?.onBatchStarted(updates.size)
+
                 for (update in updates) {
                     update()
                 }
@@ -188,6 +191,8 @@ class FormulaRuntime<Input : Any, Output : Any>(
                 isRunEnabled = true
 
                 runIfNeeded()
+
+                inspector?.onBatchFinished()
             }
         }
     }

diff --git a/formula/src/main/java/com/instacart/formula/internal/ListInspector.kt b/formula/src/main/java/com/instacart/formula/internal/ListInspector.kt
@@ -47,6 +47,14 @@ internal class ListInspector(
         forEachInspector { onRunFinished() }
     }
 
+    override fun onBatchStarted(updateCount: Int) {
+        forEachInspector { onBatchStarted(updateCount) }
+    }
+
+    override fun onBatchFinished() {
+        forEachInspector { onBatchFinished() }
+    }
+
     private inline fun forEachInspector(callback: Inspector.() -> Unit) {
         for (inspector in inspectors) {
             inspector.callback()

diff --git a/formula/src/main/java/com/instacart/formula/plugin/Inspector.kt b/formula/src/main/java/com/instacart/formula/plugin/Inspector.kt
@@ -83,4 +83,14 @@ interface Inspector {
      * executing actions.
      */
     fun onRunFinished() = Unit
+
+    /**
+     * Called when batch execution started.
+     */
+    fun onBatchStarted(updateCount: Int) = Unit
+
+    /**
+     * Called when batch execution finished
+     */
+    fun onBatchFinished() = Unit
 }
diff --git a/formula/src/test/java/com/instacart/formula/FormulaRuntimeTest.kt b/formula/src/test/java/com/instacart/formula/FormulaRuntimeTest.kt
@@ -1732,6 +1732,8 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
 
 
     @Test fun `batched formulas are executed as part of a single evaluation`() {
+
+
         val childFormulaCount = 100
 
         val batchScheduler = StateBatchScheduler()
@@ -1864,5 +1866,64 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
 
         assertThat(batchScheduler.batchesOutsideOfScope).hasSize(0)
     }
+
+    @Test fun `batched events notify the inspector of start and stop`() {
+        val globalInspector = TestInspector()
+        FormulaPlugins.setPlugin(object : Plugin {
+            override fun inspector(type: KClass<*>): Inspector {
+                return globalInspector
+            }
+        })
+        val localInspector = TestInspector()
+
+        val childFormulaCount = 3
+
+        val batchScheduler = StateBatchScheduler()
+        val relay = runtime.newRelay()
+
+        val childFormula = IncrementActionFormula(
+            incrementRelay = relay,
+            executionType = Transition.Batched(batchScheduler)
+        )
+
+        val rootFormula = object : StatelessFormula<Unit, Int>() {
+            override fun Snapshot<Unit, Unit>.evaluate(): Evaluation<Int> {
+                val sum = (0 until childFormulaCount).sumOf { id ->
+                    context.key(id) {
+                        context.child(childFormula, input)
+                    }
+                }
+                return Evaluation(
+                    output = sum,
+                )
+            }
+        }
+
+        val subject = runtime.test(rootFormula, Unit, localInspector)
+        batchScheduler.performUpdate { relay.triggerEvent() }
+
+        // Inspect!
+        for (inspector in listOf(localInspector, globalInspector)) {
+            // Filtering out logs before "batch-started"
+            val events = inspector.events.dropWhile { !it.contains("batch-started") }
+            assertThat(events).containsExactly(
+                "batch-started: 3 updates",
+                "state-changed: com.instacart.formula.types.IncrementActionFormula",
+                "state-changed: com.instacart.formula.types.IncrementActionFormula",
+                "state-changed: com.instacart.formula.types.IncrementActionFormula",
+                "formula-run-started",
+                "evaluate-started: null",
+                "evaluate-started: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-finished: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-started: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-finished: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-started: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-finished: com.instacart.formula.types.IncrementActionFormula",
+                "evaluate-finished: null",
+                "formula-run-finished",
+                "batch-finished",
+            ).inOrder()
+        }
+    }
 }
 
diff --git a/formula/src/test/java/com/instacart/formula/internal/TestInspector.kt b/formula/src/test/java/com/instacart/formula/internal/TestInspector.kt
@@ -56,4 +56,14 @@ class TestInspector : Inspector {
         super.onFormulaFinished(formulaType)
         events.add("formula-finished: ${formulaType.qualifiedName}")
     }
+
+    override fun onBatchStarted(updateCount: Int) {
+        super.onBatchStarted(updateCount)
+        events.add("batch-started: $updateCount updates")
+    }
+
+    override fun onBatchFinished() {
+        super.onBatchFinished()
+        events.add("batch-finished")
+    }
 }
diff --git a/formula/src/test/java/com/instacart/formula/types/IncrementActionFormula.kt b/formula/src/test/java/com/instacart/formula/types/IncrementActionFormula.kt
@@ -1,36 +1,27 @@
 package com.instacart.formula.types
 
 import com.instacart.formula.Evaluation
+import com.instacart.formula.Formula
 import com.instacart.formula.Snapshot
-import com.instacart.formula.StatelessFormula
 import com.instacart.formula.Transition
 import com.instacart.formula.test.Relay
 
 class IncrementActionFormula(
     private val incrementRelay: Relay,
     private val executionType: Transition.ExecutionType? = null,
-) : StatelessFormula<Unit, Int>() {
+) : Formula<Unit, Int, Int>() {
+    override fun initialState(input: Unit): Int {
+        return 0
+    }
 
-    private val actionInput = ActionDelegateFormula.Input(
-        delegateAction = {
-            if (executionType == null) {
-                incrementRelay.action().onEvent {
-                    transition(state + 1)
-                }
-            } else {
+    override fun Snapshot<Unit, Int>.evaluate(): Evaluation<Int> {
+        return Evaluation(
+            output = state,
+            actions = context.actions {
                 incrementRelay.action().onEventWithExecutionType(executionType) {
                     transition(state + 1)
                 }
             }
-        },
-        onAction = {}
-    )
-
-    private val actionFormula = ActionDelegateFormula()
-
-    override fun Snapshot<Unit, Unit>.evaluate(): Evaluation<Int> {
-        return Evaluation(
-            output = context.child(actionFormula, actionInput)
         )
     }
 }