fix: ensure CRT doesn't starve the dispatcher (#294)

Author: aajtodd

aws-runtime/crt-util/common/src/aws/sdk/kotlin/runtime/crt/ReadChannelBodyStream.kt

@@ -35,15 +35,12 @@ public class ReadChannelBodyStream(
     override val coroutineContext: CoroutineContext = callContext + producerJob
 
     private val currBuffer = atomic<SdkBuffer?>(null)
-    private val bufferChan = Channel<SdkBuffer>(1)
+    private val bufferChan = Channel<SdkBuffer>(Channel.UNLIMITED)
 
     init {
         producerJob.invokeOnCompletion { cause ->
             bodyChan.cancel(cause)
         }
-
-        // launch a coroutine to fill the buffer channel
-        proxyRequestBody()
     }
 
     // lie - CRT tries to control this via normal seek operations (e.g. when they calculate a hash for signing
@@ -52,16 +49,51 @@ public class ReadChannelBodyStream(
     // and handle these concerns.
     override fun resetPosition(): Boolean = true
 
-    @OptIn(ExperimentalCoroutinesApi::class)
     override fun sendRequestBody(buffer: MutableBuffer): Boolean {
+        return doSendRequestBody(buffer).also { if (it) producerJob.complete() }
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    private fun doSendRequestBody(buffer: MutableBuffer): Boolean {
+        // ensure the request context hasn't been cancelled
+        callContext.ensureActive()
         var outgoing = currBuffer.getAndSet(null) ?: bufferChan.tryReceive().getOrNull()
 
+        if (bodyChan.availableForRead > 0 && outgoing == null) {
+            // NOTE: It is critical that the coroutine launched doesn't actually suspend because it will never
+            // get a chance to resume. The CRT will consume the dispatcher thread until the data has been read
+            // completely. We could launch one of the coroutines into a different dispatcher but this won't work
+            // on platforms (e.g. JS) that don't have multiple threads. Essentially the CRT will starve
+            // the dispatcher and not allow other coroutines to make progress.
+            // see: https://github.com/awslabs/aws-sdk-kotlin/issues/282
+            //
+            // TODO - we could get rid of this extra copy + coroutine if readAvailable() had a non-suspend version
+            // see: https://youtrack.jetbrains.com/issue/KTOR-2772
+            //
+            // To get around this, if there is data to read we launch a coroutine UNDISPATCHED so that it runs
+            // immediately in the current thread. The coroutine will fill the buffer but won't suspend because
+            // we know data is available.
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                val sdkBuffer = SdkBuffer(bodyChan.availableForRead)
+                bodyChan.readAvailable(sdkBuffer)
+                bufferChan.send(sdkBuffer)
+            }.invokeOnCompletion { cause ->
+                if (cause != null) {
+                    producerJob.completeExceptionally(cause)
+                    bufferChan.close(cause)
+                }
+            }
+        }
+
+        if (bodyChan.availableForRead == 0 && bodyChan.isClosedForRead) {
+            bufferChan.close()
+        }
+
         if (outgoing == null) {
             if (bufferChan.isClosedForReceive) {
                 return true
             }
-            // ensure the request context hasn't been cancelled
-            callContext.ensureActive()
+
             outgoing = bufferChan.tryReceive().getOrNull() ?: return false
         }
 
@@ -73,29 +105,4 @@ public class ReadChannelBodyStream(
 
         return bufferChan.isClosedForReceive && currBuffer.value == null
     }
-
-    private fun proxyRequestBody() {
-        // TODO - we could get rid of this extra copy + coroutine if readAvailable() had a non-suspend version
-        // see: https://youtrack.jetbrains.com/issue/KTOR-2772
-        val job = launch(start = CoroutineStart.UNDISPATCHED) {
-            while (!bodyChan.isClosedForRead) {
-                bodyChan.awaitContent()
-                if (bodyChan.isClosedForRead) return@launch
-
-                // TODO - we could pool these
-                val buffer = SdkBuffer(bodyChan.availableForRead)
-                bodyChan.readAvailable(buffer)
-                bufferChan.send(buffer)
-            }
-        }
-
-        job.invokeOnCompletion { cause ->
-            bufferChan.close(cause)
-            if (cause != null) {
-                producerJob.completeExceptionally(cause)
-            } else {
-                producerJob.complete()
-            }
-        }
-    }
 }

aws-runtime/crt-util/common/test/aws/sdk/kotlin/runtime/crt/ReadChannelBodyStreamTest.kt

@@ -34,17 +34,17 @@ class ReadChannelBodyStreamTest {
     }
 
     @Test
-    fun testCancellation() = runSuspendTest {
+    fun testCancellation(): Unit = runSuspendTest {
         val chan = SdkByteChannel()
         val job = Job()
         val stream = ReadChannelBodyStream(chan, coroutineContext + job)
-        yield()
 
         job.cancel()
-        yield()
 
         val (sendBuffer, _) = mutableBuffer(16)
-        assertTrue(stream.sendRequestBody(sendBuffer))
+        assertFailsWith<CancellationException> {
+            stream.sendRequestBody(sendBuffer)
+        }
     }
 
     @Test

services/build.gradle.kts

@@ -104,6 +104,10 @@ subprojects {
                     classpath = compileDependencyFiles + runtimeDependencyFiles
                     testClassesDirs = output.classesDirs
                     useJUnitPlatform()
+                    testLogging {
+                        events("passed", "skipped", "failed")
+                        showStandardStreams = true
+                    }
                 }
             }
         }