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!!)))
               }
             }
           }