KTOR-8409 & KTOR-7947 Fix closing of OkHttp SSE connection (#4811)

Author: bnorm

ktor-client/ktor-client-okhttp/jvm/src/io/ktor/client/engine/okhttp/OkHttpSSESession.kt

@@ -7,13 +7,13 @@ package io.ktor.client.engine.okhttp
 import io.ktor.client.plugins.sse.*
 import io.ktor.http.*
 import io.ktor.sse.*
-import kotlinx.coroutines.CancellationException
-import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.onFailure
 import kotlinx.coroutines.channels.trySendBlocking
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.consumeAsFlow
+import kotlinx.coroutines.flow.onCompletion
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.Response
@@ -22,18 +22,39 @@ import okhttp3.sse.EventSourceListener
 import okhttp3.sse.EventSources
 import kotlin.coroutines.CoroutineContext
 
-internal class OkHttpSSESession(
-    engine: OkHttpClient,
+internal class OkHttpSSESession private constructor(
+    factory: EventSource.Factory,
     engineRequest: Request,
     override val coroutineContext: CoroutineContext,
 ) : SSESession, EventSourceListener() {
-    private val serverSentEventsSource = EventSources.createFactory(engine).newEventSource(engineRequest, this)
+
+    constructor(
+        engine: OkHttpClient,
+        engineRequest: Request,
+        callContext: CoroutineContext,
+    ) : this(
+        factory = EventSources.createFactory(engine),
+        engineRequest = engineRequest,
+        coroutineContext = callContext + Job() + CoroutineName("OkHttpSSESession"),
+    )
+
+    private val serverSentEventsSource = factory.newEventSource(engineRequest, this)
 
     internal val originResponse: CompletableDeferred<Response> = CompletableDeferred()
 
     private val _incoming = Channel<ServerSentEvent>(8)
 
     override val incoming: Flow<ServerSentEvent> = _incoming.consumeAsFlow()
+        .onCompletion { cause ->
+            // Use onCompletion operator to handle CancellationExceptions which occur in downstream flow.
+            if (cause is CancellationException) close(cause = null)
+        }
+
+    init {
+        coroutineContext.job.invokeOnCompletion {
+            close(cause = null)
+        }
+    }
 
     override fun onOpen(eventSource: EventSource, response: Response) {
         originResponse.complete(response)
@@ -52,6 +73,7 @@ internal class OkHttpSSESession(
             (statusCode != HttpStatusCode.OK.value || contentType != ContentType.Text.EventStream.toString())
         ) {
             originResponse.complete(response)
+            close(cause = null)
         } else {
             val error = t?.let {
                 SSEClientException(
@@ -60,15 +82,19 @@ internal class OkHttpSSESession(
                 )
             } ?: mapException(response)
             originResponse.completeExceptionally(error)
+            close(cause = error)
         }
-
-        _incoming.close()
-        serverSentEventsSource.cancel()
     }
 
     override fun onClosed(eventSource: EventSource) {
-        _incoming.close()
+        close(cause = null)
+    }
+
+    private fun close(cause: Throwable?) {
+        _incoming.close(cause)
         serverSentEventsSource.cancel()
+        // Cancel context last so 'invokeOnCompletion' doesn't override 'cause' for closing '_incoming'.
+        coroutineContext.cancel()
     }
 
     private fun mapException(response: Response?): SSEClientException {

ktor-client/ktor-client-okhttp/jvm/test/io/ktor/client/engine/okhttp/OkHttpHttpClientTest.kt

@@ -4,6 +4,94 @@
 
 package io.ktor.client.engine.okhttp
 
+import io.ktor.client.*
+import io.ktor.client.plugins.sse.*
+import io.ktor.client.test.base.*
 import io.ktor.client.tests.*
+import io.ktor.network.sockets.*
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runTest
+import okhttp3.OkHttpClient
+import org.junit.jupiter.api.assertInstanceOf
+import java.util.concurrent.TimeUnit
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.fail
 
-class OkHttpHttpClientTest : HttpClientTest(OkHttp)
+class OkHttpHttpClientTest : HttpClientTest(OkHttp) {
+    @Test
+    fun testCancelSseRequestIncomingCollect() {
+        val okHttpClient = OkHttpClient()
+
+        HttpClient(OkHttp) {
+            engine { preconfigured = okHttpClient }
+            install(SSE)
+        }.use { client ->
+            runTest {
+                var request: Job? = null
+                request = launch {
+                    client.sse("${TEST_SERVER}/sse/hello?times=20&interval=100") {
+                        request?.cancel() // Cancel the request once the connection is open.
+                        incoming.collect() // Collect all messages.
+                    }
+                    fail("Request should be cancelled.")
+                }
+                request.join()
+            }
+        }
+
+        okHttpClient.connectionPool.evictAll() // Make sure idle connections are removed.
+        assertEquals(0, okHttpClient.connectionPool.connectionCount())
+    }
+
+    @Test
+    fun testCancelSseRequestWithDelay() {
+        val okHttpClient = OkHttpClient()
+
+        HttpClient(OkHttp) {
+            engine { preconfigured = okHttpClient }
+            install(SSE)
+        }.use { client ->
+            runTest {
+                var request: Job? = null
+                request = launch {
+                    client.sse("${TEST_SERVER}/sse/hello?times=20&interval=100") {
+                        request?.cancel() // Cancel the request once the connection is open.
+                        delay(1) // Never read from incoming.
+                    }
+                    fail("Request should be cancelled.")
+                }
+                request.join()
+            }
+        }
+
+        okHttpClient.connectionPool.evictAll() // Make sure idle connections are removed.
+        assertEquals(0, okHttpClient.connectionPool.connectionCount())
+    }
+
+    @Test
+    fun testSSESessionTimeout() {
+        val okHttpClient = OkHttpClient.Builder().apply {
+            readTimeout(1L, TimeUnit.SECONDS)
+        }.build()
+
+        HttpClient(OkHttp) {
+            engine { preconfigured = okHttpClient }
+            install(SSE)
+        }.use { client ->
+            runTest {
+                client.sse("$TEST_SERVER/sse/hello?delay=10000") {
+                    try {
+                        incoming.collect()
+                        fail("Request should error.")
+                    } catch (e: SSEClientException) {
+                        assertInstanceOf<SocketTimeoutException>(e.cause)
+                    }
+                }
+            }
+        }
+    }
+}