Repository: Kotlin/kotlinx.coroutines
File: kotlinx-coroutines-core/jvm/test/ThreadLocalStressTest.kt (Kotlin)

Code before:
            assertEquals("initial", threadLocal.get())
        }
    }

    @Test
    fun testNonDispatcheableLeakWithContextSwitch() {
        repeat(100) {
            doTestWithPreparation(
                ::doTestWithContextSwitch,
                { threadLocal.set(null) }) { threadLocal.get() != null }
            assertNull(threadLocal.get())
        }
    }

    @Test
    fun testNonDispatcheableLeakWithInitialWithContextSwitch() {
        repeat(100) {
            doTestWithPreparation(
                ::doTestWithContextSwitch,
                { threadLocal.set("initial") }) { false /* can randomly wake up on the non-main thread */ }
            // Here we are always on the main thread
            assertEquals("initial", threadLocal.get())
        }
    }

    private fun doTestWithPreparation(testBody: suspend () -> Unit, setup: () -> Unit, isInvalid: () -> Boolean) {
        setup()
        val latch = CountDownLatch(1)
        testBody.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) {
            if (isInvalid()) {
                Thread.currentThread().uncaughtExceptionHandler.uncaughtException(
                    Thread.currentThread(),
                    IllegalStateException("Unexpected error: thread local was not cleaned")
                )
            }
            latch.countDown()
        })
        latch.await()
    }

    private suspend fun doTest() {
        withContext(threadLocal.asContextElement("foo")) {
            try {
                coroutineScope {
                    val semaphore = Semaphore(1, 1)
                    cancel()
                    semaphore.acquire()
                }
            } catch (e: CancellationException) {
                // Ignore cancellation
            }

Review comment (suggestion):
Minor, opinionated: `isInvalid` is a reverse of `isValid`, and most implementations here are of the form `!=` or `false`, which is a double negation, adding mental load to this and making code asymmetrical and, arguably, even non-idiomatic, as asserts are most of the time state their success conditions positively.

Code after:
            assertEquals("initial", threadLocal.get())
        }
    }

    @Test
    fun testNonDispatcheableLeakWithContextSwitch() {
        repeat(100) {
            doTestWithPreparation(
                ::doTestWithContextSwitch,
                { threadLocal.set(null) }) { threadLocal.get() == null }
            assertNull(threadLocal.get())
        }
    }

    @Test
    fun testNonDispatcheableLeakWithInitialWithContextSwitch() {
        repeat(100) {
            doTestWithPreparation(
                ::doTestWithContextSwitch,
                { threadLocal.set("initial") }) { true /* can randomly wake up on the non-main thread */ }
            // Here we are always on the main thread
            assertEquals("initial", threadLocal.get())
        }
    }

    private fun doTestWithPreparation(testBody: suspend () -> Unit, setup: () -> Unit, isValid: () -> Boolean) {
        setup()
        val latch = CountDownLatch(1)
        testBody.startCoroutineUninterceptedOrReturn(Continuation(EmptyCoroutineContext) {
            if (!isValid()) {
                Thread.currentThread().uncaughtExceptionHandler.uncaughtException(
                    Thread.currentThread(),
                    IllegalStateException("Unexpected error: thread local was not cleaned")
                )
            }
            latch.countDown()
        })
        latch.await()
    }

    private suspend fun doTest() {
        withContext(threadLocal.asContextElement("foo")) {
            try {
                coroutineScope {
                    val semaphore = Semaphore(1, 1)
                    cancel()
                    semaphore.acquire()
                }
            } catch (e: CancellationException) {
                // Ignore cancellation
            }