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

@@ -52,7 +52,7 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    *  - prevent modifications to this object after [freeze] is called.
    *  - prevent sending to sinks before render returns.
    */
-  private var frozen = false
+  private var frozen by threadLocalOf { true }
 
   override val actionSink: Sink<WorkflowAction<PropsT, StateT, OutputT>> get() = this
 
@@ -100,7 +100,7 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    * Freezes this context so that any further calls to this context will throw.
    */
   fun freeze() {
-    checkNotFrozen("freeze") { "freeze" }
+    // checkNotFrozen("freeze") { "freeze" }
     frozen = true
   }
 
@@ -117,7 +117,10 @@ internal class RealRenderContext<PropsT, StateT, OutputT>(
    *
    * @see checkWithKey
    */
-  private inline fun checkNotFrozen(stackTraceKey: Any, lazyMessage: () -> Any) =
+  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/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)