Fix workflow freezing thread safety. #1375

zach-klippenstein · 2025-07-18T17:12:16Z

Two commits:

First one adds a stress test that shows the bug.
Second one fixes the bug and makes the test pass.

zach-klippenstein · 2025-07-18T18:18:37Z

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

@@ -100,7 +100,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" }


This was just a defensive programming check, since this method is only called by internal code. It's more useful for it to be idempotent now.

rjrjr

You have my LG for what it's worth. Not my strong suit.

rjrjr · 2025-07-18T18:31:58Z

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

+ */
+internal actual class ThreadLocal<T>(
+  private val initialValue: () -> T
+) : NSObject(), NSCopyingProtocol {


@jamieQ maybe you could give this a sanity check?

yes please.

cc @mjohnson12 , @square-tomb , @watt

rjrjr · 2025-07-18T18:35:07Z

workflow-runtime/src/jvmTest/kotlin/com/squareup/workflow1/StressTestHelpers.kt

+import java.util.concurrent.CountDownLatch
+
+/**
+ * Returns the maximum number of threads that can be ran in parallel on the host system, rounded


Suggested change

* Returns the maximum number of threads that can be ran in parallel on the host system, rounded

* Returns the maximum number of threads that can be run in parallel on the host system, rounded

:D

After August 1 there's gonna be so many more "ran" in the codebases :P

Bonus points if one of the last things you do is write a lint check that catches this.

Oh my God.

Oh no.

rjrjr · 2025-07-18T18:37:29Z

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


And we're confident that we only ever call freeze() from the render thread? I can't think of any violations.

Ya I can't think of RenderContext methods that would be called off main thread? will keep thinking.

I can add another comment to freeze saying that's a requirement, but yea we only call it from one place.

And to @steve-the-edwards reply to this comment (which doesn't show up in the thread for some reason, fucking github), it doesn't matter if we call it on the main thread or not, or even (to Ray's question) if they're called only from render, we just need to make sure freeze and unfreeze are always called symmetrically from the same thread (i.e. when one thread unfreezes, the same thread later freezes).

If arbitrary RenderContext methods are called from various threads, and that code isn't already thread-safe, then it's already broken since (1) frozen is not a thread synchronization mechanism and (2) even if it were, we mutate internal data structures outside of the unfrozen section. actionSink.send is thread-safe since it just sends to a channel. But i think we were catching that case with the default value of frozen before, which we're not now, so I need to rethink that. Maybe a separate flag for whether sending to a sink is allowed.

steve-the-edwards · 2025-07-18T20:44:09Z

...runtime/src/jvmTest/kotlin/com/squareup/workflow1/WorkflowRuntimeMultithreadingStressTest.kt

+      })
+
+    val testDispatcher = newFixedThreadPoolContext(nThreads = testThreadCount, name = "test")
+    testDispatcher.use {


What is this API and why do we need it? is this like a resource API for the threads underneath this disaptcher?

It's a generic helper on Closeable: https://kotlinlang.org/api/core/kotlin-stdlib/kotlin.io/use.html

This thread pool is such, and the kdoc is extremely insistent about the fact that there are resources owned by this dispatcher and it must be closed to avoid leaks.

steve-the-edwards · 2025-07-18T20:49:42Z

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

+ */
+internal actual class ThreadLocal<T>(
+  private val initialValue: () -> T
+) : NSObject(), NSCopyingProtocol {


yes please.

steve-the-edwards · 2025-07-18T20:50:30Z

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


Ya I can't think of RenderContext methods that would be called off main thread? will keep thinking.

steve-the-edwards · 2025-07-18T20:51:28Z

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

+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class ThreadLocalTest {


@jamieQ and whoever else to review this as well then

jamieQ · 2025-07-18T21:19:45Z

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

+
+  @Test fun set_fromDifferentThreads_doNotConflict() {
+    val threadLocal = ThreadLocal(initialValue = { 0 })
+    val threadStartedLatch = NSLock().apply { lock() }


what does apply { ... } do?

edit: the AI tells me it runs the closure on self and returns the instance – so, effectively is a configuration block

yup, config block.

Equivalent to this:

val threadStartedLatch = NSLock() threadStartedLatch.lock()

jamieQ · 2025-07-18T22:05:39Z

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

+    val secondReadLatch = NSLock().apply { lock() }
+
+    val thread = NSThread {
+      threadStartedLatch.unlock()


i'm a bit surprised this works without some assertion failure, as the lock API typically require that unlocks happen on the same thread that acquired the lock.

i'd like to offer some possible alternatives, but i'm having trouble figuring out where these 'platform' imports are defined, so am not sure what's available (e.g. semaphores might be the way to go here). anyone know where i might try looking?

platform is just the prefix used for what is bundled by default as a prebuilt framework with the kotlin native interop.

So this is the Foundation framework accessed through those prebuilts. (https://www.jetbrains.com/help/kotlin-multiplatform-dev/multiplatform-ios-dependencies.html)

There are a bunch of others available, e.g.:

Though your point about Semaphore does make me think - why don't we just use the kotlinx.coroutines Semaphore or Mutex that are already in common kotlin code shared on platforms.

I looked those up and they use kotlinx.atomicfu for synchronization - e.g., https://github.com/Kotlin/kotlinx-atomicfu?tab=readme-ov-file#locks for Locks.

I think that's likely better tested and supported tthen rolling our own?

Was there another reason to roll our own Zach?

Should we pair on this? Anything in Cocoa should be available I think, but I'm not sure what the exact rules are. I will also make this test fail to ensure that it's actually running, although when I run it from the IDE it says it runs and passes.

This isn't coroutines code, we need actual thread-level synchronization here. I will try using atomicfu but it was a headache last time I tried, it's a compiler plugin so it's a bit more complex but maybe they've ironed out the kinks.

Also AFAIK atomicfu doesn't have anything like ThreadLocal, so we'll have to roll our own for that anyway — I think. @jamieQ from what I found, Cocoa doesn't have ThreadLocal either, just this threadDictionary thing, but I would be very happy to be wrong about this.

good poitn about not beign able to use the coroutine synchronization 🤦🏻 .

But yes, something lower level like atomicfu might be good. if not, and you come up with better actuals with @jamieQ then that's great!

If there's no native ThreadLocal, then TBH i'm not too worried about using (the correct) explicit Apple type for a lock over atomicfu - a lock is a lock, ultimately it's not that complicated. The tests here are for the ThreadLocal implementation, I didn't write (and don't intend to) write any tests for the lock expect/actuals.

I looked at Compose and they use every one of these techniques – KMP uses atomicfu SynchronizedObject, but there's unix-specific code in AOSP that even manually builds a lock type on atomics. I really don't want to do that, ever.

Add broken stress test for multithreaded runtime behavior.

daff58b

zach-klippenstein force-pushed the zachklipp/workflow-multithreading branch 4 times, most recently from a1d2b3b to 0fe019b Compare July 18, 2025 18:16

zach-klippenstein marked this pull request as ready for review July 18, 2025 18:16

zach-klippenstein requested a review from a team as a code owner July 18, 2025 18:16

zach-klippenstein force-pushed the zachklipp/workflow-multithreading branch from 0fe019b to 58c2441 Compare July 18, 2025 18:17

zach-klippenstein commented Jul 18, 2025

View reviewed changes

Fix workflow freezing thread safety.

f3c38bb

zach-klippenstein force-pushed the zachklipp/workflow-multithreading branch from 58c2441 to f3c38bb Compare July 18, 2025 18:25

rjrjr approved these changes Jul 18, 2025

View reviewed changes

steve-the-edwards approved these changes Jul 18, 2025

View reviewed changes

jamieQ reviewed Jul 18, 2025

View reviewed changes

	* Returns the maximum number of threads that can be ran in parallel on the host system, rounded
	* Returns the maximum number of threads that can be run in parallel on the host system, rounded

Fix workflow freezing thread safety. #1375

Are you sure you want to change the base?

Fix workflow freezing thread safety. #1375

Conversation

zach-klippenstein commented Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rjrjr left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zach-klippenstein Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zach-klippenstein Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jamieQ Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zach-klippenstein Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zach-klippenstein Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zach-klippenstein Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zach-klippenstein commented Jul 18, 2025 •

edited

Loading

zach-klippenstein Jul 21, 2025 •

edited

Loading

zach-klippenstein Jul 21, 2025 •

edited

Loading

jamieQ Jul 18, 2025 •

edited

Loading

zach-klippenstein Jul 21, 2025 •

edited

Loading

zach-klippenstein Jul 21, 2025 •

edited

Loading

zach-klippenstein Jul 21, 2025 •

edited

Loading