fix: correctly handle async cancellation of call context in OkHttp engine

.changes/289fe06b-6d65-47da-a007-360398c39244.json

@@ -0,0 +1,8 @@
+{
+    "id": "289fe06b-6d65-47da-a007-360398c39244",
+    "type": "bugfix",
+    "description": "Correctly handle async cancellation of call context in OkHttp engine",
+    "issues": [
+        "awslabs/smithy-kotlin#1061"
+    ]
+}

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

@@ -50,15 +50,16 @@ public class OkHttpEngine(
         val engineCall = client.newCall(engineRequest)
         val engineResponse = mapOkHttpExceptions { engineCall.executeAsync() }
 
-        callContext.job.invokeOnCompletion {
-            engineResponse.body.close()
-        }
-
         val response = engineResponse.toSdkResponse()
         val requestTime = Instant.fromEpochMilliseconds(engineResponse.sentRequestAtMillis)
         val responseTime = Instant.fromEpochMilliseconds(engineResponse.receivedResponseAtMillis)
 
-        return OkHttpCall(request, response, requestTime, responseTime, callContext, engineCall)
+        return OkHttpCall(request, response, requestTime, responseTime, callContext, engineCall).also { call ->
+            callContext.job.invokeOnCompletion { error ->
+                if (error != null) call.cancelInFlight()
+                engineResponse.body.close()
+            }
+        }
     }
 
     override fun shutdown() {

runtime/protocol/http-client-engines/test-suite/common/test/aws/smithy/kotlin/runtime/http/test/AsyncStressTest.kt

@@ -87,4 +87,45 @@ class AsyncStressTest : AbstractEngineTest() {
             assertEquals(engineJobsBefore.size, engineJobsAfter.size, message)
         }
     }
+
+    @Test
+    fun testJobCancellation() = testEngines {
+        // https://github.com/smithy-lang/smithy-kotlin/issues/1061
+
+        test { _, client ->
+            val req = HttpRequest {
+                testSetup()
+                url.path.decoded = "slow"
+            }
+
+            // Expect CancellationException because we're cancelling
+            assertFailsWith<CancellationException> {
+                coroutineScope {
+                    val parentScope = this
+
+                    println("Invoking call on ctx $coroutineContext")
+                    val call = client.call(req)
+
+                    val bytes = async {
+                        delay(100.milliseconds)
+                        println("Body of type ${call.response.body} on ctx $coroutineContext")
+                        try {
+                            call.response.body.readAll()
+                        } catch (e: Throwable) {
+                            // IllegalStateException: "Unbalanced enter/exit" will be thrown if body closed improperly
+                            assertIsNot<IllegalStateException>(e)
+                            null
+                        }
+                    }
+
+                    val cancellation = async {
+                        delay(400.milliseconds)
+                        parentScope.cancel("Cancelling!")
+                    }
+
+                    awaitAll(bytes, cancellation)
+                }
+            }
+        }
+    }
 }

runtime/protocol/http-client-engines/test-suite/jvm/src/aws/smithy/kotlin/runtime/http/test/suite/Concurrency.kt

@@ -8,6 +8,8 @@ package aws.smithy.kotlin.runtime.http.test.suite
 import io.ktor.server.application.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
+import kotlinx.coroutines.delay
+import kotlin.time.Duration.Companion.milliseconds
 
 internal fun Application.concurrentTests() {
     routing {
@@ -18,5 +20,17 @@ internal fun Application.concurrentTests() {
                 call.respondText(text.repeat(respSize / text.length))
             }
         }
+
+        route("slow") {
+            get {
+                val chunk = ByteArray(256) { it.toByte() }
+                call.respondOutputStream {
+                    repeat(10) {
+                        delay(200.milliseconds)
+                        write(chunk)
+                    }
+                }
+            }
+        }
     }
 }