More coverage improvements.

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

@@ -26,11 +26,7 @@ internal class ChildrenManager(
     fun prepareForPostEvaluation() {
         indexes?.clear()
 
-        children?.clearUnrequested {
-            pendingRemoval = pendingRemoval ?: mutableListOf()
-            it.markAsTerminated()
-            pendingRemoval?.add(it)
-        }
+        children?.clearUnrequested(this::prepareForTermination)
     }
 
     fun terminateChildren(evaluationId: Long): Boolean {
@@ -95,6 +91,12 @@ internal class ChildrenManager(
         }
     }
 
+    private fun prepareForTermination(it: FormulaManager<*, *>) {
+        pendingRemoval = pendingRemoval ?: mutableListOf()
+        it.markAsTerminated()
+        pendingRemoval?.add(it)
+    }
+
     private fun <ChildInput, ChildOutput>  childFormulaHolder(
         key: Any,
         formula: IFormula<ChildInput, ChildOutput>,

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

@@ -1,7 +1,5 @@
 package com.instacart.formula.internal
 
-import com.instacart.formula.Transition
-
 /**
  * A deferred transition contains an [event] and a related [transition] which can
  * be executed using the [execute] function. If the formula is ready, it will be
@@ -11,11 +9,10 @@ import com.instacart.formula.Transition
  */
 class DeferredTransition<Input, State, EventT> internal constructor(
     private val listener: ListenerImpl<Input, State, EventT>,
-    private val transition: Transition<Input, State, EventT>,
     private val event: EventT,
 ) {
 
     fun execute() {
-        listener.snapshotImpl?.dispatch(transition, event)
+        listener.applyInternal(event)
     }
 }

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

@@ -57,7 +57,7 @@ internal class FormulaManagerImpl<Input, State, Output>(
      * each formula output with an identifier value and compare it for validity with
      * the global value.
      */
-    var globalEvaluationId: Long = 0
+    private var globalEvaluationId: Long = 0
 
     /**
      * Determines if we are executing within [run] block. Enables optimizations
@@ -194,7 +194,6 @@ internal class FormulaManagerImpl<Input, State, Output>(
         val snapshot = SnapshotImpl(
             input = input,
             state = state,
-            associatedEvaluationId = evaluationId,
             listeners = listeners,
             delegate = this,
         )
@@ -220,7 +219,7 @@ internal class FormulaManagerImpl<Input, State, Output>(
         listeners.prepareForPostEvaluation()
         childrenManager?.prepareForPostEvaluation()
 
-        snapshot.running = true
+        snapshot.markRunning()
         if (!isValidationEnabled) {
             inspector?.onEvaluateFinished(loggingType, newFrame.evaluation.output, evaluated = true)
         }

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

@@ -8,13 +8,13 @@ import com.instacart.formula.plugin.Dispatcher
  * Note: this class is not a data class because equality is based on instance and not [key].
  */
 @PublishedApi
-internal class ListenerImpl<Input, State, EventT>(internal var key: Any) : Listener<EventT> {
+internal class ListenerImpl<Input, State, EventT>(val key: Any) : Listener<EventT> {
 
-    @Volatile internal var manager: FormulaManagerImpl<Input, State, *>? = null
-    @Volatile internal var snapshotImpl: SnapshotImpl<Input, State>? = null
-    @Volatile internal var executionType: Transition.ExecutionType? = null
+    @Volatile private var manager: FormulaManagerImpl<Input, State, *>? = null
+    @Volatile private var snapshotImpl: SnapshotImpl<Input, State>? = null
+    @Volatile private var executionType: Transition.ExecutionType? = null
 
-    internal lateinit var transition: Transition<Input, State, EventT>
+    private lateinit var transition: Transition<Input, State, EventT>
 
     override fun invoke(event: EventT) {
         // TODO: log if null listener (it might be due to formula removal or due to callback removal)
@@ -29,16 +29,32 @@ internal class ListenerImpl<Input, State, EventT>(internal var key: Any) : Liste
         }
     }
 
+    fun setDependencies(
+        manager: FormulaManagerImpl<Input, State, *>?,
+        snapshot: SnapshotImpl<Input, State>?,
+        executionType: Transition.ExecutionType?,
+        transition: Transition<Input, State, EventT>,
+    ) {
+        this.manager = manager
+        this.snapshotImpl = snapshot
+        this.executionType = executionType
+        this.transition = transition
+    }
+
     fun disable() {
         manager = null
         snapshotImpl = null
     }
 
+    fun applyInternal(event: EventT) {
+        snapshotImpl?.dispatch(transition, event)
+    }
+
     private fun handleBatched(type: Transition.Batched, event: EventT) {
         val manager = manager ?: return
 
         manager.batchManager.add(type, event) {
-            val deferredTransition = DeferredTransition(this, transition, event)
+            val deferredTransition = DeferredTransition(this, event)
             manager.onPendingTransition(deferredTransition)
         }
     }
@@ -48,7 +64,7 @@ internal class ListenerImpl<Input, State, EventT>(internal var key: Any) : Liste
 
         dispatcher.dispatch {
             manager.queue.postUpdate {
-                val deferredTransition = DeferredTransition(this, transition, event)
+                val deferredTransition = DeferredTransition(this, event)
                 manager.onPendingTransition(deferredTransition)
             }
         }

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

@@ -35,11 +35,7 @@ internal class Listeners {
      */
     fun prepareForPostEvaluation() {
         indexes?.clear()
-
-        listeners?.clearUnrequested {
-            // TODO log that disabled listener was invoked.
-            it.disable()
-        }
+        listeners?.clearUnrequested(this::disableListener)
     }
 
     fun disableAll() {
@@ -50,6 +46,10 @@ internal class Listeners {
         listeners?.clear()
     }
 
+    private fun disableListener(listener: ListenerImpl<*, *, *>) {
+        listener.disable()
+    }
+
     /**
      * Function which returns next index for a given key. It will
      * mutate the [indexes] map.

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

@@ -23,7 +23,7 @@ internal class SingleRequestHolder<T>(val value: T) {
 
 internal typealias SingleRequestMap<Key, Value> = MutableMap<Key, SingleRequestHolder<Value>>
 
-internal inline fun <Value> SingleRequestMap<*, Value>.clearUnrequested(onUnrequested: (Value) -> Unit) {
+internal fun <Value> SingleRequestMap<*, Value>.clearUnrequested(onUnrequested: (Value) -> Unit) {
     val callbackIterator = this.iterator()
     while (callbackIterator.hasNext()) {
         val callback = callbackIterator.next()

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

@@ -8,20 +8,18 @@ import com.instacart.formula.Listener
 import com.instacart.formula.Snapshot
 import com.instacart.formula.Transition
 import com.instacart.formula.TransitionContext
-import com.instacart.formula.plugin.Dispatcher
 import java.lang.IllegalStateException
 import kotlin.reflect.KClass
 
 internal class SnapshotImpl<out Input, State> internal constructor(
     override val input: Input,
     override val state: State,
-    private val associatedEvaluationId: Long,
     listeners: Listeners,
     private val delegate: FormulaManagerImpl<Input, State, *>,
 ) : FormulaContext<Input, State>(listeners), Snapshot<Input, State>, TransitionContext<Input, State> {
 
     private var scopeKey: Any? = null
-    var running = false
+    private var running = false
 
     override val context: FormulaContext<Input, State> = this
 
@@ -53,10 +51,7 @@ internal class SnapshotImpl<out Input, State> internal constructor(
     ): Listener<Event> {
         ensureNotRunning()
         val listener = listeners.initOrFindListener<Input, State, Event>(key, useIndex)
-        listener.manager = delegate
-        listener.snapshotImpl = this
-        listener.executionType = executionType
-        listener.transition = transition
+        listener.setDependencies(delegate, this, executionType, transition)
         return listener
     }
 
@@ -91,15 +86,6 @@ internal class SnapshotImpl<out Input, State> internal constructor(
     }
 
     fun <Event> dispatch(transition: Transition<Input, State, Event>, event: Event) {
-        if (!running) {
-            throw IllegalStateException("Transitions are not allowed during evaluation")
-        }
-
-        if (!delegate.isTerminated() && delegate.isEvaluationNeeded(associatedEvaluationId)) {
-            // We have already transitioned, this should not happen.
-            throw IllegalStateException("Transition already happened. This is using old event listener: $transition & $event. Transition: $associatedEvaluationId != ${delegate.globalEvaluationId}")
-        }
-
         val result = transition.toResult(this, event)
         if (TransitionUtils.isEmpty(result)) {
             return
@@ -108,6 +94,10 @@ internal class SnapshotImpl<out Input, State> internal constructor(
         delegate.handleTransitionResult(event, result)
     }
 
+    fun markRunning() {
+        running = true
+    }
+
     private fun ensureNotRunning() {
         if (running) {
             throw IllegalStateException("Cannot call this transition after evaluation finished. See https://instacart.github.io/formula/faq/#after-evaluation-finished")

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

@@ -8,6 +8,7 @@ import com.instacart.formula.internal.ClearPluginsRule
 import com.instacart.formula.internal.FormulaKey
 import com.instacart.formula.internal.TestInspector
 import com.instacart.formula.internal.Try
+import com.instacart.formula.plugin.Dispatcher
 import com.instacart.formula.plugin.Inspector
 import com.instacart.formula.plugin.Plugin
 import com.instacart.formula.rxjava3.RxAction
@@ -653,6 +654,28 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
             }
     }
 
+    @Test fun `listener removed while dispatching an event will drop the event`() {
+        var observer: TestFormulaObserver<Unit, OptionalCallbackFormula.Output, OptionalCallbackFormula>? = null
+        FormulaPlugins.setPlugin(object : Plugin {
+            override fun backgroundThreadDispatcher(): Dispatcher {
+                return object : Dispatcher {
+                    override fun dispatch(executable: () -> Unit) {
+                        // We disable callback before executing increment
+                        observer?.output { toggleCallback() }
+                        executable()
+                    }
+                }
+            }
+        })
+
+        val root = OptionalCallbackFormula(
+            incrementExecutionType = Transition.Background
+        )
+        observer = runtime.test(root, Unit)
+        observer.output { listener?.invoke() }
+        observer.output { assertThat(state).isEqualTo(0) }
+    }
+
     @Test
     fun `listeners are not the same after removing then adding it again`() {
         runtime.test(OptionalCallbackFormula(), Unit)
@@ -1503,6 +1526,12 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
 
     @Test fun `input changed inspector event`() {
         val localInspector = TestInspector()
+        val globalInspector = TestInspector()
+        FormulaPlugins.setPlugin(object : Plugin {
+            override fun inspector(type: KClass<*>): Inspector {
+                return globalInspector
+            }
+        })
 
         val formula = object : StatelessFormula<Int, Int>() {
             val delegate = InputIdentityFormula<Int>()
@@ -1515,25 +1544,27 @@ class FormulaRuntimeTest(val runtime: TestableRuntime, val name: String) {
         val subject = runtime.test(formula, 0, localInspector)
         subject.input(1)
 
-        assertThat(localInspector.events).containsExactly(
-            "formula-run-started",
-            "formula-started: null",
-            "evaluate-started: null",
-            "formula-started: com.instacart.formula.types.InputIdentityFormula",
-            "evaluate-started: com.instacart.formula.types.InputIdentityFormula",
-            "evaluate-finished: com.instacart.formula.types.InputIdentityFormula",
-            "evaluate-finished: null",
-            "formula-run-finished",
-
-            "formula-run-started",
-            "evaluate-started: null",
-            "input-changed: null",
-            "evaluate-started: com.instacart.formula.types.InputIdentityFormula",
-            "input-changed: com.instacart.formula.types.InputIdentityFormula",
-            "evaluate-finished: com.instacart.formula.types.InputIdentityFormula",
-            "evaluate-finished: null",
-            "formula-run-finished"
-        ).inOrder()
+        for (inspector in listOf(globalInspector, localInspector)) {
+            assertThat(inspector.events).containsExactly(
+                "formula-run-started",
+                "formula-started: null",
+                "evaluate-started: null",
+                "formula-started: com.instacart.formula.types.InputIdentityFormula",
+                "evaluate-started: com.instacart.formula.types.InputIdentityFormula",
+                "evaluate-finished: com.instacart.formula.types.InputIdentityFormula",
+                "evaluate-finished: null",
+                "formula-run-finished",
+
+                "formula-run-started",
+                "evaluate-started: null",
+                "input-changed: null",
+                "evaluate-started: com.instacart.formula.types.InputIdentityFormula",
+                "input-changed: com.instacart.formula.types.InputIdentityFormula",
+                "evaluate-finished: com.instacart.formula.types.InputIdentityFormula",
+                "evaluate-finished: null",
+                "formula-run-finished"
+            ).inOrder()
+        }
     }
 
     @Test

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

@@ -4,9 +4,12 @@ import com.instacart.formula.Evaluation
 import com.instacart.formula.Formula
 import com.instacart.formula.Listener
 import com.instacart.formula.Snapshot
+import com.instacart.formula.Transition
 
-class OptionalCallbackFormula :
-    Formula<Unit, OptionalCallbackFormula.State, OptionalCallbackFormula.Output>() {
+class OptionalCallbackFormula(
+    private val toggleExecutionType: Transition.ExecutionType? = null,
+    private val incrementExecutionType: Transition.ExecutionType? = null,
+) : Formula<Unit, OptionalCallbackFormula.State, OptionalCallbackFormula.Output>() {
     data class State(
         val callbackEnabled: Boolean = true,
         val state: Int = 0
@@ -22,7 +25,9 @@ class OptionalCallbackFormula :
 
     override fun Snapshot<Unit, State>.evaluate(): Evaluation<Output> {
         val callback = if (state.callbackEnabled) {
-            context.onEvent<Unit> { transition(state.copy(state = state.state + 1)) }
+            context.onEventWithExecutionType<Unit>(incrementExecutionType) {
+                transition(state.copy(state = state.state + 1))
+            }
         } else {
             null
         }
@@ -31,7 +36,7 @@ class OptionalCallbackFormula :
             output = Output(
                 state = state.state,
                 listener = callback,
-                toggleCallback = context.onEvent {
+                toggleCallback = context.onEventWithExecutionType(toggleExecutionType) {
                     transition(state.copy(callbackEnabled = !state.callbackEnabled))
                 }
             )