fix(rt): OkHttp request body retry failing due to Job completed (#832)

When OkHttp config has retry on connect failures enabled it will attempt
to re-use the request body in some scenarios. The coroutine Job was
setup to only be used once though and was completed when the original
request fails. This causes the (internal okhttp) retry to fail with an obscure
JobCompleted exception.

Author: aajtodd

.changes/a9e3a074-33fa-41da-8322-c83adc2e3637.json

@@ -0,0 +1,5 @@
+{
+    "id": "a9e3a074-33fa-41da-8322-c83adc2e3637",
+    "type": "bugfix",
+    "description": "Fix okhttp streaming body failing to retry"
+}

runtime/protocol/http-client-engines/http-client-engine-okhttp/jvm/src/aws/smithy/kotlin/runtime/http/engine/okhttp/StreamingRequestBody.kt

@@ -9,6 +9,7 @@ import aws.smithy.kotlin.runtime.http.HttpBody
 import aws.smithy.kotlin.runtime.io.internal.toOkio
 import aws.smithy.kotlin.runtime.io.internal.toSdk
 import aws.smithy.kotlin.runtime.io.readAll
+import aws.smithy.kotlin.runtime.tracing.trace
 import aws.smithy.kotlin.runtime.util.derivedName
 import kotlinx.coroutines.*
 import okhttp3.MediaType
@@ -20,44 +21,50 @@ import kotlin.coroutines.CoroutineContext
 /**
  * OkHttp [RequestBody] that reads from [body] channel or source
  */
+@OptIn(DelicateCoroutinesApi::class, ExperimentalStdlibApi::class)
 internal class StreamingRequestBody(
     private val body: HttpBody,
-    callContext: CoroutineContext,
-) : RequestBody(), CoroutineScope {
+    private val callContext: CoroutineContext,
+) : RequestBody() {
 
     init {
         require(body is HttpBody.ChannelContent || body is HttpBody.SourceContent) { "Invalid streaming body $body" }
     }
 
-    private val producerJob = Job(callContext[Job])
-    override val coroutineContext: CoroutineContext = callContext + producerJob + callContext.derivedName("send-request-body") + Dispatchers.IO
     override fun contentType(): MediaType? = null
     override fun contentLength(): Long = body.contentLength ?: -1
     override fun isOneShot(): Boolean = body.isOneShot
     override fun isDuplex(): Boolean = body.isDuplex
 
-    override fun writeTo(sink: BufferedSink) = try {
-        doWriteTo(sink)
-    } catch (ex: Exception) {
-        throw when (ex) {
-            is IOException -> ex
-            // wrap all exceptions thrown from inside `okhttp3.RequestBody#writeTo(..)` as an IOException
-            // see https://github.com/awslabs/aws-sdk-kotlin/issues/733
-            else -> IOException(ex)
+    override fun writeTo(sink: BufferedSink) {
+        try {
+            doWriteTo(sink)
+        } catch (ex: Exception) {
+            when (ex) {
+                is CancellationException -> {
+                    callContext.trace<StreamingRequestBody> { "request cancelled" }
+                    // shouldn't need to propagate the exception because okhttp is cancellation aware via executeAsync()
+                    return
+                }
+                is IOException -> throw ex
+                // wrap all exceptions thrown from inside `okhttp3.RequestBody#writeTo(..)` as an IOException
+                // see https://github.com/awslabs/aws-sdk-kotlin/issues/733
+                else -> throw IOException(ex)
+            }
         }
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
     private fun doWriteTo(sink: BufferedSink) {
+        val context = callContext + callContext.derivedName("send-request-body")
         if (isDuplex()) {
             // launch coroutine that writes to sink in the background
-            launch {
+            GlobalScope.launch(context + Dispatchers.IO) {
                 sink.use { transferBody(it) }
             }
         } else {
             // remove the current dispatcher (if it exists) and use the internal
             // runBlocking dispatcher that blocks the *current* thread
-            val blockingContext = coroutineContext.minusKey(CoroutineDispatcher)
+            val blockingContext = context.minusKey(CoroutineDispatcher)
 
             // Non-duplex (aka "normal") requests MUST write all of their request body
             // before this function returns. Requests are given a background thread to
@@ -68,13 +75,15 @@ internal class StreamingRequestBody(
             }
         }
     }
-    private suspend fun transferBody(sink: BufferedSink) = withJob(producerJob) {
+
+    private suspend fun transferBody(sink: BufferedSink) {
         when (body) {
             is HttpBody.ChannelContent -> {
                 val chan = body.readFrom()
                 val sdkSink = sink.toSdk()
                 chan.readAll(sdkSink)
             }
+
             is HttpBody.SourceContent -> {
                 val source = body.readFrom()
                 source.toOkio().use {
@@ -86,19 +95,3 @@ internal class StreamingRequestBody(
         }
     }
 }
-
-/**
- * Completes the given job when the block returns calling either `complete()` when the block runs
- * successfully or `completeExceptionally()` on exception.
- * @return the result of calling [block]
- */
-private inline fun <T> withJob(job: CompletableJob, block: () -> T): T {
-    try {
-        return block()
-    } catch (ex: Exception) {
-        job.completeExceptionally(ex)
-        throw ex
-    } finally {
-        job.complete()
-    }
-}

runtime/protocol/http-client-engines/http-client-engine-okhttp/jvm/test/aws/smithy/kotlin/runtime/http/engine/okhttp/StreamingRequestBodyTest.kt

@@ -15,7 +15,6 @@ import kotlinx.coroutines.test.runTest
 import okio.Buffer
 import okio.BufferedSink
 import org.junit.jupiter.api.Test
-import java.io.IOException
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.test.*
 import kotlin.time.Duration.Companion.milliseconds
@@ -110,13 +109,9 @@ class StreamingRequestBodyTest {
             val callContext = coroutineContext + Job(coroutineContext.job)
             val actual = StreamingRequestBody(body, callContext)
             val buffer = Buffer()
-            // see https://github.com/awslabs/aws-sdk-kotlin/issues/733 for why we expect
-            // this to be an IOException
-            val ex = assertFailsWith<IOException> {
-                actual.writeTo(buffer)
-            }
-            assertIs<CancellationException>(ex.cause)
+            actual.writeTo(buffer)
         }
+
         delay(100.milliseconds)
 
         job.cancel()
@@ -148,10 +143,9 @@ class StreamingRequestBodyTest {
 
         assertTrue(actual.isDuplex())
 
-        assertEquals(1, callJob.children.toList().size) // producerJob
-        assertEquals(0, callJob.children.toList()[0].children.toList().size)
+        assertEquals(0, callJob.children.toList().size)
         actual.writeTo(sink)
-        assertEquals(1, callJob.children.toList()[0].children.toList().size)
+        assertEquals(1, callJob.children.toList().size) // writer
         assertEquals(sink.buffer.size, 0)
         chan.writeAll(content.source())