Fix workflow freezing thread safety.

Author: zach-klippenstein

workflow-runtime/src/appleMain/kotlin/com/squareup/workflow1/internal/Synchronization.apple.kt

@@ -1,6 +1,12 @@
 package com.squareup.workflow1.internal
 
+import kotlinx.cinterop.CPointer
+import kotlinx.cinterop.ExperimentalForeignApi
+import platform.Foundation.NSCopyingProtocol
 import platform.Foundation.NSLock
+import platform.Foundation.NSThread
+import platform.Foundation.NSZone
+import platform.darwin.NSObject
 
 internal actual typealias Lock = NSLock
 
@@ -12,3 +18,35 @@ internal actual inline fun <R> Lock.withLock(block: () -> R): R {
     unlock()
   }
 }
+
+/**
+ * Implementation of [ThreadLocal] that works in a similar way to Java's, based on a thread-specific
+ * map/dictionary.
+ */
+internal actual class ThreadLocal<T>(
+  private val initialValue: () -> T
+) : NSObject(), NSCopyingProtocol {
+
+  private val threadDictionary
+    get() = NSThread.currentThread().threadDictionary
+
+  actual fun get(): T {
+    @Suppress("UNCHECKED_CAST")
+    return (threadDictionary.objectForKey(aKey = this) as T?)
+      ?: initialValue().also(::set)
+  }
+
+  actual fun set(value: T) {
+    threadDictionary.setObject(value, forKey = this)
+  }
+
+  /**
+   * [Docs](https://developer.apple.com/documentation/foundation/nscopying/copy(with:)) say [zone]
+   * is unused.
+   */
+  @OptIn(ExperimentalForeignApi::class)
+  override fun copyWithZone(zone: CPointer<NSZone>?): Any = this
+}
+
+internal actual fun <T> threadLocalOf(initialValue: () -> T): ThreadLocal<T> =
+  ThreadLocal(initialValue)

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/RealRenderContext.kt

@@ -51,8 +51,12 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    * Used to:
    *  - prevent modifications to this object after [freeze] is called.
    *  - prevent sending to sinks before render returns.
+   *
+   * This is a [ThreadLocal] since we only care about preventing calls during rendering from the
+   * thread that is actually doing the rendering. If a background thread happens to send something
+   * into the sink, for example, while the main thread is rendering, it's not a violation.
    */
-  private var frozen = false
+  private var frozen by threadLocalOf { true }
 
   override val actionSink: Sink<WorkflowAction<PropsT, StateT, OutputT>> get() = this
 
@@ -100,7 +104,6 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    * Freezes this context so that any further calls to this context will throw.
    */
   fun freeze() {
-    checkNotFrozen("freeze") { "freeze" }
     frozen = true
   }
 
@@ -117,8 +120,10 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    *
    * @see checkWithKey
    */
-  private inline fun checkNotFrozen(stackTraceKey: Any, lazyMessage: () -> Any) =
-    checkWithKey(!frozen, stackTraceKey) {
-      "RenderContext cannot be used after render method returns: ${lazyMessage()}"
-    }
+  private inline fun checkNotFrozen(
+    stackTraceKey: Any,
+    lazyMessage: () -> Any
+  ) = checkWithKey(!frozen, stackTraceKey) {
+    "RenderContext cannot be used after render method returns: ${lazyMessage()}"
+  }
 }

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/Synchronization.kt

@@ -1,5 +1,29 @@
 package com.squareup.workflow1.internal
 
+import kotlin.reflect.KProperty
+
 internal expect class Lock()
 
 internal expect inline fun <R> Lock.withLock(block: () -> R): R
+
+internal expect class ThreadLocal<T> {
+  fun get(): T
+  fun set(value: T)
+}
+
+internal expect fun <T> threadLocalOf(initialValue: () -> T): ThreadLocal<T>
+
+@Suppress("NOTHING_TO_INLINE")
+internal inline operator fun <T> ThreadLocal<T>.getValue(
+  receiver: Any?,
+  property: KProperty<*>
+): T = get()
+
+@Suppress("NOTHING_TO_INLINE")
+internal inline operator fun <T> ThreadLocal<T>.setValue(
+  receiver: Any?,
+  property: KProperty<*>,
+  value: T
+) {
+  set(value)
+}

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/RealRenderContextTest.kt

@@ -220,8 +220,10 @@ internal class RealRenderContextTest {
 
     val child = Workflow.stateless<Unit, Nothing, Unit> { fail() }
     assertFailsWith<IllegalStateException> { context.renderChild(child) }
-    assertFailsWith<IllegalStateException> { context.freeze() }
     assertFailsWith<IllegalStateException> { context.remember("key", typeOf<String>()) {} }
+
+    // Freeze is the exception, it's idempotent and can be called again.
+    context.freeze()
   }
 
   private fun createdPoisonedContext(): RealRenderContext<String, String, String> {
@@ -234,7 +236,9 @@ internal class RealRenderContextTest {
       eventActionsChannel,
       workflowTracer = null,
       runtimeConfig = emptySet(),
-    )
+    ).apply {
+      unfreeze()
+    }
   }
 
   private fun createTestContext(): RealRenderContext<String, String, String> {
@@ -247,6 +251,8 @@ internal class RealRenderContextTest {
       eventActionsChannel,
       workflowTracer = null,
       runtimeConfig = emptySet(),
-    )
+    ).apply {
+      unfreeze()
+    }
   }
 }

workflow-runtime/src/iosTest/kotlin/com/squareup/workflow1/internal/ThreadLocalTest.kt

@@ -0,0 +1,82 @@
+package com.squareup.workflow1.internal
+
+import platform.Foundation.NSLock
+import platform.Foundation.NSThread
+import kotlin.concurrent.Volatile
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class ThreadLocalTest {
+
+  @Volatile
+  private var valueFromThread: Int = -1
+
+  @Test fun initialValue() {
+    val threadLocal = ThreadLocal(initialValue = { 42 })
+    assertEquals(42, threadLocal.get())
+  }
+
+  @Test fun settingValue() {
+    val threadLocal = ThreadLocal(initialValue = { 42 })
+    threadLocal.set(0)
+    assertEquals(0, threadLocal.get())
+  }
+
+  @Test fun initialValue_inSeparateThread_afterChanging() {
+    val threadLocal = ThreadLocal(initialValue = { 42 })
+    threadLocal.set(0)
+
+    val thread = NSThread {
+      valueFromThread = threadLocal.get()
+    }
+    thread.start()
+    thread.join()
+
+    assertEquals(42, valueFromThread)
+  }
+
+  @Test fun set_fromDifferentThreads_doNotConflict() {
+    val threadLocal = ThreadLocal(initialValue = { 0 })
+    val threadStartedLatch = NSLock().apply { lock() }
+    val firstReadLatch = NSLock().apply { lock() }
+    val firstReadDoneLatch = NSLock().apply { lock() }
+    val secondReadLatch = NSLock().apply { lock() }
+
+    val thread = NSThread {
+      threadStartedLatch.unlock()
+      firstReadLatch.lock()
+      threadLocal.set(1)
+      valueFromThread = threadLocal.get()
+      firstReadDoneLatch.unlock()
+      secondReadLatch.lock()
+      valueFromThread = threadLocal.get()
+    }
+    thread.start()
+
+    // Wait for the other thread to start, then both threads set the value to something different
+    // at the same time.
+    threadStartedLatch.lock()
+    firstReadLatch.unlock()
+    threadLocal.set(2)
+
+    // Wait for the background thread to finish setting value, then ensure that both threads see
+    // independent values.
+    firstReadDoneLatch.lock()
+    assertEquals(1, valueFromThread)
+    assertEquals(2, threadLocal.get())
+
+    // Change the value in this thread then read it again from the background thread.
+    threadLocal.set(3)
+    secondReadLatch.unlock()
+    thread.join()
+    assertEquals(1, valueFromThread)
+  }
+
+  private fun NSThread.join() {
+    while (!isFinished()) {
+      // Avoid being optimized out.
+      // Time interval is in seconds.
+      NSThread.sleepForTimeInterval(1.0 / 1000)
+    }
+  }
+}

workflow-runtime/src/jsMain/kotlin/com/squareup/workflow1/internal/Synchronization.js.kt

@@ -5,3 +5,13 @@ package com.squareup.workflow1.internal
 internal actual typealias Lock = Any
 
 internal actual inline fun <R> Lock.withLock(block: () -> R): R = block()
+
+internal actual class ThreadLocal<T>(private var value: T) {
+  actual fun get(): T = value
+  actual fun set(value: T) {
+    this.value = value
+  }
+}
+
+internal actual fun <T> threadLocalOf(initialValue: () -> T): ThreadLocal<T> =
+  ThreadLocal(initialValue())

workflow-runtime/src/jvmMain/kotlin/com/squareup/workflow1/internal/Synchronization.jvm.kt

@@ -3,3 +3,8 @@ package com.squareup.workflow1.internal
 internal actual typealias Lock = Any
 
 internal actual inline fun <R> Lock.withLock(block: () -> R): R = synchronized(this, block)
+
+internal actual typealias ThreadLocal<T> = java.lang.ThreadLocal<T>
+
+internal actual fun <T> threadLocalOf(initialValue: () -> T): ThreadLocal<T> =
+  ThreadLocal.withInitial(initialValue)