Handling errors from child formulas (#364)

* Handling errors from child formulas

* Add termination side effects

* Poll transition queue on termination

* Try finally block for isRunning

* Add some tests

* Add another test

* Implement termination without executing local transition queue and docs update

* Fix typo

formula/src/main/java/com/instacart/formula/FormulaContext.kt

@@ -93,6 +93,20 @@ abstract class FormulaContext<out Input, State> internal constructor(
         return child(formula, Unit)
     }
 
+    /**
+     * A convenience method to run a formula that takes no input. Returns the latest output of
+     * the [child] formula or null if the child formula has encountered an exception. Child formulas
+     * that encounter exceptions will be be terminated and will not be run again. Formula runtime
+     * ensures the [formula] is running, manages its internal state and will trigger `evaluate`
+     * if needed.
+     */
+    fun <ChildOutput> child(
+        formula: IFormula<Unit, ChildOutput>,
+        onError: (Throwable) -> Unit,
+    ): ChildOutput? {
+        return child(formula, Unit, onError)
+    }
+
     /**
      * Returns the latest output of the [child] formula. Formula runtime ensures the [child]
      * is running, manages its internal state and will trigger `evaluate` if needed.
@@ -102,6 +116,18 @@ abstract class FormulaContext<out Input, State> internal constructor(
         input: ChildInput
     ): ChildOutput
 
+    /**
+     * Returns the latest output of the [child] formula or null if the child formula has encountered
+     * an exception. Child formulas that encounter exceptions will be be terminated and will not
+     * be run again. Formula runtime ensures the [child] is running, manages its internal state
+     * and will trigger `evaluate` if needed.
+     */
+    abstract fun <ChildInput, ChildOutput> child(
+        formula: IFormula<ChildInput, ChildOutput>,
+        input: ChildInput,
+        onError: (Throwable) -> Unit,
+    ): ChildOutput?
+
     /**
      * Builds a list of deferred actions that will be executed by Formula runtime.
      */

formula/src/main/java/com/instacart/formula/internal/ChildrenManager.kt

@@ -45,8 +45,8 @@ internal class ChildrenManager(
         children?.forEachValue { it.markAsTerminated() }
     }
 
-    fun performTerminationSideEffects() {
-        children?.forEachValue { it.performTerminationSideEffects() }
+    fun performTerminationSideEffects(executeTransitionQueue: Boolean) {
+        children?.forEachValue { it.performTerminationSideEffects(executeTransitionQueue) }
     }
 
     fun <ChildInput, ChildOutput> findOrInitChild(

formula/src/main/java/com/instacart/formula/internal/FormulaManager.kt

@@ -19,6 +19,8 @@ interface FormulaManager<Input, Output> {
      */
     fun run(input: Input): Evaluation<Output>
 
+    fun isTerminated(): Boolean
+
     /**
      * Called when [Formula] is removed. This is should not trigger any external side-effects,
      * only mark itself and its children as terminated.
@@ -27,6 +29,11 @@ interface FormulaManager<Input, Output> {
 
     /**
      * Called when we are ready to perform termination side-effects.
+     * @param executeTransitionQueue whether the formula should execute remaining transitions
+     *      that accumulated while it was running. False is passed in the case where the
+     *      formula threw an exception and we want to avoid executing transitions that have
+     *      an elevated chance of also throwing an exception. Other termination side
+     *      effects such as listener and action termination still happen.
      */
-    fun performTerminationSideEffects()
+    fun performTerminationSideEffects(executeTransitionQueue: Boolean = true)
 }

formula/src/main/java/com/instacart/formula/internal/FormulaManagerImpl.kt

@@ -80,7 +80,7 @@ internal class FormulaManagerImpl<Input, State, Output>(
         return globalEvaluationId != evaluationId
     }
 
-    fun isTerminated(): Boolean {
+    override fun isTerminated(): Boolean {
         return terminated
     }
 
@@ -130,28 +130,30 @@ internal class FormulaManagerImpl<Input, State, Output>(
 
         var result: Evaluation<Output>? = null
         isRunning = true
-        var firstRun = true
-        while (result == null) {
-            val lastFrame = frame
-            val evaluationId = globalEvaluationId
-            val evaluation = if (firstRun) {
-                firstRun = false
-                evaluation(input, evaluationId)
-            } else if (lastFrame == null || isEvaluationNeeded(lastFrame.associatedEvaluationId)) {
-                evaluation(input, evaluationId)
-            } else {
-                lastFrame.evaluation
-            }
+        try {
+            var firstRun = true
+            while (result == null) {
+                val lastFrame = frame
+                val evaluationId = globalEvaluationId
+                val evaluation = if (firstRun) {
+                    firstRun = false
+                    evaluation(input, evaluationId)
+                } else if (lastFrame == null || isEvaluationNeeded(lastFrame.associatedEvaluationId)) {
+                    evaluation(input, evaluationId)
+                } else {
+                    lastFrame.evaluation
+                }
 
-            if (postEvaluation(evaluationId)) {
-                continue
-            }
+                if (postEvaluation(evaluationId)) {
+                    continue
+                }
 
-            result = evaluation
+                result = evaluation
+            }
+            return result
+        } finally {
+            isRunning = false
         }
-        isRunning = false
-
-        return result
     }
 
     /**
@@ -245,18 +247,58 @@ internal class FormulaManagerImpl<Input, State, Output>(
         return manager.run(input).output
     }
 
+    fun <ChildInput, ChildOutput> child(
+        key: Any,
+        formula: IFormula<ChildInput, ChildOutput>,
+        input: ChildInput,
+        onError: (Throwable) -> Unit,
+    ): ChildOutput? {
+        val childrenManager = getOrInitChildrenManager()
+        val manager = childrenManager.findOrInitChild(key, formula, input)
+
+        // If termination happens while running, we might still be initializing child formulas. To
+        // ensure correct behavior, we mark each requested child manager as terminated to avoid
+        // starting new actions.
+        if (isTerminated()) {
+            manager.markAsTerminated()
+        }
+        manager.setValidationRun(isValidationEnabled)
+
+        return try {
+            // If manager.run(input) previously threw an exception it would be marked as
+            // terminated in the catch block. We avoid running it again because there is a decent
+            // chance it will continue to throw the same exception and cause performance issues
+            if (!manager.isTerminated()) {
+                manager.run(input).output
+            } else {
+                null
+            }
+        } catch (e: ValidationException) {
+            throw e
+        } catch (e: Throwable) {
+            manager.markAsTerminated()
+            onError(e)
+            manager.performTerminationSideEffects()
+            null
+        }
+    }
+
     override fun markAsTerminated() {
         terminated = true
         childrenManager?.markAsTerminated()
     }
 
-    override fun performTerminationSideEffects() {
-        childrenManager?.performTerminationSideEffects()
+    override fun performTerminationSideEffects(executeTransitionQueue: Boolean) {
+        childrenManager?.performTerminationSideEffects(executeTransitionQueue)
         actionManager.terminate()
 
-        // Execute deferred transitions
-        for (transition in transitionQueue) {
-            transition.execute()
+        if (executeTransitionQueue) {
+            // Execute deferred transitions
+            while (transitionQueue.isNotEmpty()) {
+                transitionQueue.pollFirst().execute()
+            }
+        } else {
+            transitionQueue.clear()
         }
 
         listeners.disableAll()

formula/src/main/java/com/instacart/formula/internal/SnapshotImpl.kt

@@ -43,6 +43,20 @@ internal class SnapshotImpl<out Input, State> internal constructor(
         return delegate.child(key, formula, input)
     }
 
+    override fun <ChildInput, ChildOutput> child(
+        formula: IFormula<ChildInput, ChildOutput>,
+        input: ChildInput,
+        onError: (Throwable) -> Unit,
+    ): ChildOutput? {
+        ensureNotRunning()
+
+        val key = createScopedKey(
+            type = formula.type(),
+            key = formula.key(input)
+        )
+        return delegate.child(key, formula, input, onError)
+    }
+
     override fun <Event> eventListener(
         key: Any,
         useIndex: Boolean,

formula/src/test/java/com/instacart/formula/FormulaRuntimeTest.kt

@@ -13,6 +13,7 @@ import com.instacart.formula.plugin.Inspector
 import com.instacart.formula.plugin.Plugin
 import com.instacart.formula.rxjava3.RxAction
 import com.instacart.formula.subjects.ChildActionFiresParentEventOnStart
+import com.instacart.formula.subjects.ChildErrorAfterToggleFormula
 import com.instacart.formula.subjects.ChildMessageNoParentStateChange
 import com.instacart.formula.subjects.ChildMessageTriggersEventTransitionInParent
 import com.instacart.formula.subjects.ChildMessageWithParentStateChange
@@ -32,6 +33,7 @@ import com.instacart.formula.subjects.EventFormula
 import com.instacart.formula.subjects.ExtremelyNestedFormula
 import com.instacart.formula.subjects.FromObservableWithInputFormula
 import com.instacart.formula.subjects.HasChildFormula
+import com.instacart.formula.subjects.HasChildrenFormula
 import com.instacart.formula.subjects.IncrementingDispatcher
 import com.instacart.formula.subjects.InputChangeWhileFormulaRunningRobot
 import com.instacart.formula.subjects.KeyFormula
@@ -1375,8 +1377,78 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
         assertThat(error?.message).isEqualTo("crashed")
     }
 
+    @Test
+    fun `errored child formula can fail in isolation when evaluation throws`() {
+        val childFormula = object : StatelessFormula<HasChildrenFormula.ChildParamsInput<*>, Int>() {
+            override fun Snapshot<HasChildrenFormula.ChildParamsInput<*>, Unit>.evaluate(): Evaluation<Int> {
+                if (input.index == 1) throw RuntimeException()
+                return Evaluation(output = input.index)
+            }
+        }
+        val formula = HasChildrenFormula(childCount = 3, childFormula)
+        runtime.test(formula, 0)
+            .output {
+                assertThat(childOutputs).isEqualTo(listOf(0, 2))
+                assertThat(errors).hasSize(1)
+            }
+    }
+
+    @Test
+    fun `errored child formula can fail in isolation when action throws`() {
+        val childFormula = object : StatelessFormula<HasChildrenFormula.ChildParamsInput<*>, Int>() {
+            override fun Snapshot<HasChildrenFormula.ChildParamsInput<*>, Unit>.evaluate(): Evaluation<Int> {
+                return Evaluation(
+                    output = input.index,
+                    actions = context.actions {
+                        Action.onInit().onEvent {
+                            if (input.index == 1) throw RuntimeException()
+                            transition(Unit) {}
+                        }
+                    }
+                )
+            }
+        }
+        val formula = HasChildrenFormula(childCount = 3, childFormula)
+        runtime.test(formula, 0)
+            .output {
+                assertThat(childOutputs).isEqualTo(listOf(0, 2))
+                assertThat(errors).hasSize(1)
+            }
+    }
+    @Test
+    fun `errored child event listener disabled`() {
+        val indexToExecutionCount = mutableMapOf<Int, Int>()
+        val listener = { index: Int ->
+            indexToExecutionCount[index] = indexToExecutionCount.getOrDefault(index, 0) + 1
+        }
+        val formula = HasChildrenFormula(
+            childCount = 3,
+            child = ChildErrorAfterToggleFormula(),
+            createChildInput = { params ->
+                HasChildrenFormula.ChildParamsInput(
+                    index = params.index,
+                    run = params.index,
+                    value = callback { transition { listener(params.index) }}
+                )
+            },
+        )
+        runtime.test(formula, 0)
+            .output {
+                childOutputs.forEach { it.listener() }
+                childOutputs[1].errorToggle()
+                childOutputs.forEach { it.listener() }
+            }
+        val expected = mapOf(
+            0 to 2,
+            1 to 1,
+            2 to 2,
+        )
+        assertThat(indexToExecutionCount).containsExactlyEntriesIn(expected)
+    }
+
     @Test
     fun `initialize 100 levels nested formula`() {
+
         val inspector = CountingInspector()
         val formula = ExtremelyNestedFormula.nested(100)
         runtime.test(formula, Unit, inspector).output {

formula/src/test/java/com/instacart/formula/subjects/ChildErrorAfterToggleFormula.kt

@@ -0,0 +1,32 @@
+package com.instacart.formula.subjects
+
+import com.instacart.formula.Evaluation
+import com.instacart.formula.Formula
+import com.instacart.formula.Snapshot
+import com.instacart.formula.subjects.ChildErrorAfterToggleFormula.Output
+import com.instacart.formula.subjects.ChildErrorAfterToggleFormula.State
+
+class ChildErrorAfterToggleFormula: Formula<HasChildrenFormula.ChildParamsInput<()-> Unit>, State, Output>()  {
+    data class Output(
+        val errorToggle: () -> Unit,
+        val listener: () -> Unit
+    )
+
+    data class State(val throwError: Boolean = false)
+
+    override fun initialState(input: HasChildrenFormula.ChildParamsInput<()-> Unit>) = State()
+
+    override fun Snapshot<HasChildrenFormula.ChildParamsInput<()-> Unit>, State>.evaluate(): Evaluation<Output> {
+        if (state.throwError) throw RuntimeException()
+        return Evaluation(
+            output = Output(
+                errorToggle = context.callback() {
+                    transition(state.copy(throwError = !state.throwError))
+                },
+                listener = context.callback {
+                    transition { input.value() }
+                }
+            ),
+        )
+    }
+}