Fix a flaky check and a TSan violation in StateFlowStressTest. (#4483)

There's no guarantee that any particular collector will ever encounter
a value emitted by a particular emitter. It's entirely possible for the
value to be overwritten by a different emitter before the collector gets
a chance to collect it. It's very unlikely for a collector to miss the
second half of all values emitted by a particular emitter, but it is
still possible and this causes the test to be flaky.

We can instead check if the collector has collected a recent enough
value from *any* emitter. This should be sufficient to verify that the
collector was still running near the end of the test.

Also fixed a race condition in the test when printing test progress.
The race condition is benign, but it causes TSan to fail for the test,
which could prevent it from finding other concurrency bugs.

Author: chaoren

kotlinx-coroutines-core/jvm/test/flow/StateFlowStressTest.kt

@@ -1,9 +1,10 @@
 package kotlinx.coroutines.flow
 
-import kotlinx.coroutines.testing.*
+import java.util.concurrent.atomic.AtomicLongArray
+import kotlin.random.*
 import kotlinx.coroutines.*
+import kotlinx.coroutines.testing.*
 import org.junit.*
-import kotlin.random.*
 
 class StateFlowStressTest : TestBase() {
     private val nSeconds = 3 * stressTestMultiplier
@@ -17,7 +18,7 @@ class StateFlowStressTest : TestBase() {
 
     fun stress(nEmitters: Int, nCollectors: Int) = runTest {
         pool = newFixedThreadPoolContext(nEmitters + nCollectors, "StateFlowStressTest")
-        val collected = Array(nCollectors) { LongArray(nEmitters) }
+        val collected = Array(nCollectors) { AtomicLongArray(nEmitters) }
         val collectors = launch {
             repeat(nCollectors) { collector ->
                 launch(pool) {
@@ -37,38 +38,39 @@ class StateFlowStressTest : TestBase() {
                             }
                             c[emitter] = current
 
-                        }.take(batchSize).map { 1 }.sum()
+                        }.take(batchSize).count()
                     } while (cnt == batchSize)
                 }
             }
         }
-        val emitted = LongArray(nEmitters)
+        val emitted = AtomicLongArray(nEmitters)
         val emitters = launch {
             repeat(nEmitters) { emitter ->
                 launch(pool) {
-                    var current = 1L
                     while (true) {
-                        state.value = current * nEmitters + emitter
-                        emitted[emitter] = current
-                        current++
-                        if (current % 1000 == 0L) yield() // make it cancellable
+                        state.value = emitted.incrementAndGet(emitter) * nEmitters + emitter
+                        if (emitted[emitter] % 1000 == 0L) yield() // make it cancellable
                     }
                 }
             }
         }
         for (second in 1..nSeconds) {
             delay(1000)
             val cs = collected.map { it.sum() }
-            println("$second: emitted=${emitted.sum()}, collected=${cs.minOrNull()}..${cs.maxOrNull()}")
+            println("$second: emitted=${emitted.sum()}, collected=${cs.min()}..${cs.max()}")
         }
         emitters.cancelAndJoin()
         collectors.cancelAndJoin()
         // make sure nothing hanged up
-        require(collected.all { c ->
-            c.withIndex().all { (emitter, current) -> current > emitted[emitter] / 2 }
-        })
+        for (i in 0..<nCollectors) {
+            check((0..<nEmitters).any { j -> collected[i][j] > emitted[j] * 0.9 }) {
+                "collector #$i failed to collect any of the most recently emitted values"
+            }
+        }
     }
 
+    private fun AtomicLongArray.sum() = (0..<length()).sumOf(::get)
+
     @Test
     fun testSingleEmitterAndCollector() = stress(1, 1)