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

@@ -41,8 +41,9 @@ internal fun <T> CoroutineScope.launchWorker(
   workerDiagnosticId: Long,
   workflowDiagnosticId: Long,
   diagnosticListener: WorkflowDiagnosticListener?
-): ReceiveChannel<T> = worker.run()
+): ReceiveChannel<ValueOrDone<T>> = worker.run()
     .wireUpDebugger(workerDiagnosticId, workflowDiagnosticId, diagnosticListener)
+    .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
@@ -74,3 +75,35 @@ private fun <T> Flow<T>.wireUpDebugger(
     }
   }
 }
+
+/**
+ * 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

@@ -64,7 +64,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
   )
 
@@ -182,14 +182,14 @@ 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.
-                return@onReceiveOrClosed null
+                return@onReceive null
               } else {
-                val update = case.acceptUpdate(valueOrClosed.value)
+                val update = case.acceptUpdate(valueOrDone.value)
                 acceptUpdate(update)
               }
             }

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

@@ -24,10 +24,12 @@ import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.emitAll
 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 {
@@ -58,11 +60,11 @@ 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())
 
@@ -94,6 +96,65 @@ class WorkersTest {
       yield()
 
       assertEquals("onWorkerStopped(0, 1)", listener.consumeNextEvent())
+      // Read the last event so the scope can complete.
+      assertTrue(outputs.receive().isDone)
     }
   }
+
+  @Test fun `emits done when complete immediately`() {
+    val channel = Channel<String>(capacity = 1)
+
+    runBlocking {
+      val workerOutputs = launchWorker(channel.asWorker(), 0, 0, null)
+      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(), 0, 0, null)
+      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(), 0, 0, null)
+        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(), 0, 0, null)
+      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>()