Make formula internals thread safe.

formula-coroutines/src/main/java/com/instacart/formula/coroutines/FlowRuntime.kt

@@ -3,7 +3,6 @@ package com.instacart.formula.coroutines
 import com.instacart.formula.FormulaRuntime
 import com.instacart.formula.IFormula
 import com.instacart.formula.Inspector
-import com.instacart.formula.internal.ThreadChecker
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.channels.trySendBlocking
@@ -22,20 +21,16 @@ object FlowRuntime {
         inspector: Inspector? = null,
         isValidationEnabled: Boolean = false,
     ): Flow<Output> {
-        val threadChecker = ThreadChecker(formula)
         return callbackFlow<Output> {
-            threadChecker.check("Need to subscribe on main thread.")
-
             val runtime = FormulaRuntime(
-                threadChecker = threadChecker,
                 formula = formula,
                 onOutput = this::trySendBlocking,
                 onError = this::close,
                 inspector = inspector,
                 isValidationEnabled = isValidationEnabled,
             )
 
-            input.onEach { input -> runtime.onInput(input) }.launchIn(this)
+            input.onEach(runtime::onInput).launchIn(this)
 
             awaitClose {
                 runtime.terminate()

formula-rxjava3/src/main/java/com/instacart/formula/rxjava3/RxJavaRuntime.kt

@@ -1,10 +1,8 @@
 package com.instacart.formula.rxjava3
 
-import com.instacart.formula.FormulaPlugins
 import com.instacart.formula.FormulaRuntime
 import com.instacart.formula.IFormula
 import com.instacart.formula.Inspector
-import com.instacart.formula.internal.ThreadChecker
 import io.reactivex.rxjava3.core.Observable
 import io.reactivex.rxjava3.disposables.CompositeDisposable
 import io.reactivex.rxjava3.disposables.FormulaDisposableHelper
@@ -16,12 +14,8 @@ object RxJavaRuntime {
         inspector: Inspector? = null,
         isValidationEnabled: Boolean = false,
     ): Observable<Output> {
-        val threadChecker = ThreadChecker(formula)
-        return Observable.create { emitter ->
-            threadChecker.check("Need to subscribe on main thread.")
-
+        return Observable.create<Output> { emitter ->
             val runtime = FormulaRuntime(
-                threadChecker = threadChecker,
                 formula = formula,
                 onOutput = emitter::onNext,
                 onError = emitter::onError,
@@ -30,11 +24,9 @@ object RxJavaRuntime {
             )
 
             val disposables = CompositeDisposable()
-            disposables.add(input.subscribe({ input ->
-                runtime.onInput(input)
-            }, emitter::onError))
-
+            disposables.add(input.subscribe(runtime::onInput, emitter::onError))
             disposables.add(FormulaDisposableHelper.fromRunnable(runtime::terminate))
+
             emitter.setDisposable(disposables)
         }.distinctUntilChanged()
     }

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

@@ -3,30 +3,27 @@ package com.instacart.formula
 import com.instacart.formula.internal.FormulaManager
 import com.instacart.formula.internal.FormulaManagerImpl
 import com.instacart.formula.internal.ManagerDelegate
-import com.instacart.formula.internal.ThreadChecker
+import com.instacart.formula.internal.SynchronizedUpdateQueue
 import java.util.LinkedList
 
 /**
  * Takes a [Formula] and creates an Observable<Output> from it.
  */
 class FormulaRuntime<Input : Any, Output : Any>(
-    private val threadChecker: ThreadChecker,
     private val formula: IFormula<Input, Output>,
     private val onOutput: (Output) -> Unit,
     private val onError: (Throwable) -> Unit,
     private val isValidationEnabled: Boolean = false,
     inspector: Inspector? = null,
 ) : ManagerDelegate {
+    private val synchronizedUpdateQueue = SynchronizedUpdateQueue()
+    private val inspector = FormulaPlugins.inspector(type = formula.type(), local = inspector)
     private val implementation = formula.implementation()
+
     private var manager: FormulaManagerImpl<Input, *, Output>? = null
-    private val inspector = FormulaPlugins.inspector(
-        type = formula.type(),
-        local = inspector,
-    )
 
     private var emitOutput = false
     private var lastOutput: Output? = null
-
     private var input: Input? = null
     private var key: Any? = null
 
@@ -43,8 +40,7 @@ class FormulaRuntime<Input : Any, Output : Any>(
     private var inputId: Int = 0
 
     /**
-     * Global transition effect queue which executes side-effects
-     * after all formulas are idle.
+     * Global transition effect queue which executes side-effects after all formulas are idle.
      */
     private var globalEffectQueue = LinkedList<Effects>()
 
@@ -66,8 +62,10 @@ class FormulaRuntime<Input : Any, Output : Any>(
     }
 
     fun onInput(input: Input) {
-        threadChecker.check("Input arrived on a wrong thread.")
+        synchronizedUpdateQueue.postUpdate { onInputInternal(input) }
+    }
 
+    private fun onInputInternal(input: Input) {
         if (isRuntimeTerminated) return
 
         val isKeyValid = isKeyValid(input)
@@ -105,8 +103,10 @@ class FormulaRuntime<Input : Any, Output : Any>(
     }
 
     fun terminate() {
-        threadChecker.check("Need to unsubscribe on the main thread.")
+        synchronizedUpdateQueue.postUpdate(this::terminateInternal)
+    }
 
+    private fun terminateInternal() {
         if (isRuntimeTerminated) return
         isRuntimeTerminated = true
 
@@ -127,8 +127,6 @@ class FormulaRuntime<Input : Any, Output : Any>(
     }
 
     override fun onPostTransition(effects: Effects?, evaluate: Boolean) {
-        threadChecker.check("Only thread that created it can post transition result")
-
         effects?.let {
             globalEffectQueue.addLast(effects)
         }
@@ -271,6 +269,7 @@ class FormulaRuntime<Input : Any, Output : Any>(
 
     private fun initManager(initialInput: Input): FormulaManagerImpl<Input, *, Output> {
         return FormulaManagerImpl(
+            queue = synchronizedUpdateQueue,
             delegate = this,
             formula = implementation,
             initialInput = initialInput,

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

@@ -110,6 +110,7 @@ internal class ChildrenManager(
         val childFormulaHolder = children.findOrInit(key) {
             val implementation = formula.implementation()
             FormulaManagerImpl(
+                queue = delegate.queue,
                 delegate = delegate,
                 formula = implementation,
                 initialInput = input,

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

@@ -18,6 +18,7 @@ import kotlin.reflect.KClass
  * a state change, it will rerun [Formula.evaluate].
  */
 internal class FormulaManagerImpl<Input, State, Output>(
+    val queue: SynchronizedUpdateQueue,
     private val delegate: ManagerDelegate,
     private val formula: Formula<Input, State, Output>,
     initialInput: Input,

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

@@ -18,8 +18,10 @@ internal class ListenerImpl<Input, State, EventT>(internal var key: Any) : Liste
         // TODO: log if null listener (it might be due to formula removal or due to callback removal)
         val manager = manager ?: return
 
-        val deferredTransition = DeferredTransition(this, transition, event)
-        manager.onPendingTransition(deferredTransition)
+        manager.queue.postUpdate {
+            val deferredTransition = DeferredTransition(this, transition, event)
+            manager.onPendingTransition(deferredTransition)
+        }
     }
 
     fun disable() {

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

@@ -0,0 +1,116 @@
+package com.instacart.formula.internal
+
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * We can only process one formula update at a time. To enable thread-safety we use a
+ * non-blocking event queue with a serial confinement strategy for queue processing. All external
+ * formula events use [postUpdate] function which adds the update to [updateQueue], and then,
+ * tries to start processing the queue via atomic [threadRunning] variable. If another thread
+ * was first to take over [threadRunning], we let that thread continue and we exit out. Given
+ * all formula state access is gated via [threadRunning] atomic reference, we are able to ensure
+ * that there is happens-before relationship between each thread and memory changes are visible
+ * between them.
+ */
+class SynchronizedUpdateQueue {
+    /**
+     * Defines a thread currently executing formula update. Null value indicates idle queue.
+     *
+     * To ensure that memory changes within formula internals are synchronized between threads,
+     * we piggyback on the internal synchronization of this variable. Modification to this
+     * variable wraps around every formula update:
+     * - threadRunning = MyThread
+     * - formulaUpdate()
+     * - threadRunning = null
+     *
+     * This creates happens-before relationship between multiple threads and makes sure that
+     * all modifications within formulaUpdate() block are visible to the next thread.
+     */
+    private val threadRunning = AtomicReference<Thread>()
+
+    /**
+     * A non-blocking thread-safe FIFO queue that tracks pending updates.
+     */
+    private val updateQueue = ConcurrentLinkedQueue<() -> Unit>()
+
+    /**
+     * To ensure that we execute one update at a time, all external formula events use this
+     * function to post updates. We add the update to a queue and then try to start processing.
+     * Failure to start processing indicates that another thread was first and we allow that
+     * thread to continue.
+     */
+    fun postUpdate(update: () -> Unit) {
+        val currentThread = Thread.currentThread()
+        val owner = threadRunning.get()
+        if (owner == currentThread) {
+            // This indicates a nested update where an update triggers another update. Given we
+            // are already thread gated, we can execute this update immediately without a need
+            // for any extra synchronization.
+            update()
+            return
+        }
+
+        val success = if (updateQueue.peek() == null) {
+            // No pending update, let's try to run our update immediately
+            takeOver(currentThread) { update() }
+        } else {
+            false
+        }
+
+        if (!success) {
+            updateQueue.add(update)
+        }
+        tryToDrainQueue(currentThread)
+    }
+
+    /**
+     * Tries to drain the update queue. It will process one update at a time until
+     * queue is empty or another thread takes over processing.
+     */
+    private fun tryToDrainQueue(currentThread: Thread) {
+        while (true) {
+            // First, we peek to see if there is a value to process.
+            val peekUpdate = updateQueue.peek()
+            if (peekUpdate != null) {
+                // Since there is a pending update, we try to process it.
+                val success = takeOver(currentThread) {
+                    // We successfully set ourselves as the running thread
+                    // We poll the queue to get the latest value (it could have changed). It
+                    // also removes the value from the queue.
+                    val actualUpdate = updateQueue.poll()
+                    actualUpdate?.invoke()
+                }
+
+                if (!success) {
+                    return
+                }
+            } else {
+                return
+            }
+        }
+    }
+
+    /**
+     * Tries to take over the processing and execute an [update].
+     *
+     * Returns true if it was able to successfully claim the ownership and execute the
+     * update. Otherwise, returns false (this indicates another thread claimed the right first).
+     */
+    private inline fun takeOver(currentThread: Thread, crossinline update: () -> Unit): Boolean {
+        return if (threadRunning.compareAndSet(null, currentThread)) {
+            // We took over the processing, let's execute the [update]
+            try {
+                update()
+            } finally {
+                // We reset the running thread. To ensure happens-before relationship, this must
+                // always happen after the [update].
+                threadRunning.set(null)
+            }
+            true
+        } else {
+            // Another thread is running, so we return false.
+            false
+        }
+    }
+}