Stop using onReceiveOrClosed because it is extremely broken.

I think this fixes #626.

Likely caused by Kotlin/kotlinx.coroutines#1584.

Author: zach-klippenstein

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/Workers.kt

@@ -21,8 +21,11 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.channels.Channel.Factory.RENDEZVOUS
 import kotlinx.coroutines.channels.ReceiveChannel
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.catch
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.produceIn
 
 /**
@@ -31,8 +34,9 @@ import kotlinx.coroutines.flow.produceIn
  * will emit everything from the worker. The channel will be closed when the flow completes.
  */
 @UseExperimental(FlowPreview::class, ExperimentalCoroutinesApi::class)
-internal fun <T> CoroutineScope.launchWorker(worker: Worker<T>): ReceiveChannel<T> =
+internal fun <T> CoroutineScope.launchWorker(worker: Worker<T>): ReceiveChannel<ValueOrDone<T>> =
   worker.run()
+      .transformToValueOrDone()
       .catch { e ->
         // Workers that failed (as opposed to just cancelled) should have their failure reason
         // re-thrown from the workflow runtime. If we don't unwrap the cause here, they'll just
@@ -42,3 +46,35 @@ internal fun <T> CoroutineScope.launchWorker(worker: Worker<T>): ReceiveChannel<
       }
       .buffer(RENDEZVOUS)
       .produceIn(this)
+
+/**
+ * Pretend we can use ReceiveChannel.onReceiveOrClosed.
+ *
+ * See https://github.com/Kotlin/kotlinx.coroutines/issues/1584 and
+ * https://github.com/square/workflow/issues/626.
+ */
+internal class ValueOrDone<out T> private constructor(private val _value: Any?) {
+
+  val isDone: Boolean get() = this === Done
+
+  @Suppress("UNCHECKED_CAST")
+  val value: T
+    get() {
+      check(!isDone)
+      return _value as T
+    }
+
+  companion object {
+    private val Done = ValueOrDone<Nothing>(null)
+
+    fun <T> value(value: T): ValueOrDone<T> = ValueOrDone(value)
+    fun done(): ValueOrDone<Nothing> = Done
+  }
+}
+
+private fun <T> Flow<T>.transformToValueOrDone(): Flow<ValueOrDone<T>> = flow {
+  collect {
+    emit(ValueOrDone.value(it))
+  }
+  emit(ValueOrDone.done())
+}

kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt

@@ -60,7 +60,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
    * continues to listen to the worker after it finishes.
    */
   private class WorkerSession(
-    val channel: ReceiveChannel<*>,
+    val channel: ReceiveChannel<ValueOrDone<*>>,
     var tombstone: Boolean = false
   )
 
@@ -143,8 +143,8 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
         .filter { (_, session) -> !session.tombstone }
         .forEach { (case, session) ->
           with(selector) {
-            session.channel.onReceiveOrClosed { valueOrClosed ->
-              if (valueOrClosed.isClosed) {
+            session.channel.onReceive { valueOrDone ->
+              if (valueOrDone.isDone) {
                 // Set the tombstone flag so we don't continue to listen to the subscription.
                 session.tombstone = true
                 // Nothing to do on close other than update the session, so don't emit any output.
@@ -156,10 +156,10 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
                         )
                     )
                 )
-                return@onReceiveOrClosed OutputEnvelope(null, debugInfo)
+                return@onReceive OutputEnvelope(null, debugInfo)
               } else {
-                val update = case.acceptUpdate(valueOrClosed.value)
-                acceptUpdate(update, Kind.Updated(Source.Worker(case.key, valueOrClosed.value!!)))
+                val update = case.acceptUpdate(valueOrDone.value)
+                acceptUpdate(update, Kind.Updated(Source.Worker(case.key, valueOrDone.value!!)))
               }
             }
           }

kotlin/workflow-runtime/src/test/java/com/squareup/workflow/internal/WorkersTest.kt

@@ -22,10 +22,13 @@ import kotlinx.coroutines.CoroutineStart.UNDISPATCHED
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.supervisorScope
 import kotlinx.coroutines.yield
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
 
 class WorkersTest {
 
@@ -50,16 +53,73 @@ class WorkersTest {
       yield()
       assertEquals(1, counter.getAndIncrement())
 
-      assertEquals("a", workerOutputs.poll())
+      assertEquals("a", workerOutputs.poll()!!.value)
       yield()
       assertEquals(3, counter.getAndIncrement())
 
-      assertEquals("b", workerOutputs.poll())
+      assertEquals("b", workerOutputs.poll()!!.value)
       yield()
       assertEquals(6, counter.getAndIncrement())
 
       // Cancel the worker so we can exit this loop.
       workerOutputs.cancel()
     }
   }
+
+  @Test fun `emits done when complete immediately`() {
+    val channel = Channel<String>(capacity = 1)
+
+    runBlocking {
+      val workerOutputs = launchWorker(channel.asWorker())
+      assertTrue(workerOutputs.isEmpty)
+
+      channel.close()
+      assertTrue(workerOutputs.receive().isDone)
+    }
+  }
+
+  @Test fun `emits done when complete after sending`() {
+    val channel = Channel<String>(capacity = 1)
+
+    runBlocking {
+      val workerOutputs = launchWorker(channel.asWorker())
+      assertTrue(workerOutputs.isEmpty)
+
+      channel.send("foo")
+      assertEquals("foo", workerOutputs.receive().value)
+
+      channel.close()
+      assertTrue(workerOutputs.receive().isDone)
+    }
+  }
+
+  @Test fun `does not emit done when failed`() {
+    val channel = Channel<String>(capacity = 1)
+
+    runBlocking {
+      // Needed so that cancelling the channel doesn't cancel our job, which means receive will
+      // throw the JobCancellationException instead of the actual channel failure.
+      supervisorScope {
+        val workerOutputs = launchWorker(channel.asWorker())
+        assertTrue(workerOutputs.isEmpty)
+
+        channel.close(ExpectedException())
+        assertFailsWith<ExpectedException> { workerOutputs.receive() }
+      }
+    }
+  }
+
+  @Test fun `completes after emitting done`() {
+    val channel = Channel<String>(capacity = 1)
+
+    runBlocking {
+      val workerOutputs = launchWorker(channel.asWorker())
+      channel.close()
+      assertTrue(workerOutputs.receive().isDone)
+
+      assertTrue(channel.isClosedForReceive)
+    }
+  }
+
+  private class ExpectedException : RuntimeException()
 }

kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerStressTest.kt

@@ -12,13 +12,11 @@ import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.reduce
 import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.runBlocking
-import kotlin.test.Ignore
 import kotlin.test.Test
 import kotlin.test.assertEquals
 
 class WorkerStressTest {
 
-  @Ignore("https://github.com/square/workflow/issues/626")
   @UseExperimental(ExperimentalCoroutinesApi::class)
   @Test fun `multiple subscriptions to single channel when closed`() {
     val channel = Channel<Unit>()