KTOR-8989 Close or cancel engine only when the client reference count… (#5525)

Author: Stexxe

ktor-client/ktor-client-core/api/ktor-client-core.api

@@ -165,6 +165,8 @@ public final class io/ktor/client/engine/HttpClientEngine$DefaultImpls {
 }
 
 public abstract class io/ktor/client/engine/HttpClientEngineBase : io/ktor/client/engine/HttpClientEngine {
+	public static final synthetic field clientRefCount$FU$internal Ljava/util/concurrent/atomic/AtomicIntegerFieldUpdater;
+	public synthetic field clientRefCount$internal I
 	public fun <init> (Ljava/lang/String;)V
 	public fun close ()V
 	public fun getCoroutineContext ()Lkotlin/coroutines/CoroutineContext;

ktor-client/ktor-client-core/common/src/io/ktor/client/HttpClient.kt

@@ -15,6 +15,7 @@ import io.ktor.util.*
 import io.ktor.utils.io.*
 import io.ktor.utils.io.core.*
 import kotlinx.atomicfu.atomic
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableJob
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -647,15 +648,7 @@ public fun <T : HttpClientEngineConfig> HttpClient(
 ): HttpClient {
     val config: HttpClientConfig<T> = HttpClientConfig<T>().apply(block)
     val engine = engineFactory.create(config.engineConfig)
-    val client = HttpClient(engine, config, manageEngine = true)
-
-    // If the engine was created using factory Ktor is responsible for its lifecycle management. Otherwise user has to
-    // close engine by themself.
-    client.coroutineContext[Job]!!.invokeOnCompletion {
-        engine.close()
-    }
-
-    return client
+    return HttpClient(engine, config, manageEngine = true)
 }
 
 /**
@@ -1292,6 +1285,26 @@ public class HttpClient(
         manageEngine: Boolean
     ) : this(engine, userConfig) {
         this.manageEngine = manageEngine
+
+        if (this.manageEngine) {
+            if (engine is HttpClientEngineBase) {
+                engine.clientRefCount.incrementAndGet()
+            }
+
+            coroutineContext[Job]!!.invokeOnCompletion { cause ->
+                val shouldClose = engine !is HttpClientEngineBase || engine.clientRefCount.decrementAndGet() <= 0
+
+                if (shouldClose) {
+                    if (cause == null) {
+                        engine.close()
+                    } else {
+                        engine.cancel(
+                            cause as? CancellationException ?: CancellationException("Client scope is canceled", cause)
+                        )
+                    }
+                }
+            }
+        }
     }
 
     private val closed = atomic(false)
@@ -1352,14 +1365,6 @@ public class HttpClient(
     internal val config = HttpClientConfig<HttpClientEngineConfig>()
 
     init {
-        if (manageEngine) {
-            clientJob.invokeOnCompletion {
-                if (it != null) {
-                    engine.cancel()
-                }
-            }
-        }
-
         engine.install(this)
 
         sendPipeline.intercept(HttpSendPipeline.Receive) { call ->
@@ -1483,9 +1488,6 @@ public class HttpClient(
         }
 
         clientJob.complete()
-        if (manageEngine) {
-            engine.close()
-        }
     }
 
     override fun toString(): String = "HttpClient[$engine]"

ktor-client/ktor-client-core/common/src/io/ktor/client/engine/HttpClientEngineBase.kt

@@ -40,6 +40,9 @@ import kotlin.coroutines.*
 public abstract class HttpClientEngineBase(private val engineName: String) : HttpClientEngine {
     private val closed = atomic(false)
 
+    /** Number of [io.ktor.client.HttpClient] instances sharing this engine */
+    internal val clientRefCount: AtomicInt = atomic(0)
+
     override val dispatcher: CoroutineDispatcher by lazy { config.dispatcher ?: ioDispatcher() }
 
     override val coroutineContext: CoroutineContext by lazy {

ktor-client/ktor-client-core/common/test/HttpClientConfigTest.kt

@@ -3,9 +3,20 @@
  */
 
 import io.ktor.client.*
+import io.ktor.client.engine.*
 import io.ktor.client.plugins.api.*
+import io.ktor.client.request.*
+import io.ktor.http.*
+import io.ktor.util.date.*
+import io.ktor.utils.io.*
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.test.runTest
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertSame
+import kotlin.test.assertTrue
 
 class HttpClientConfigTest {
 
@@ -41,4 +52,88 @@ class HttpClientConfigTest {
         assertEquals(1, first)
         assertEquals(1, second)
     }
+
+    @Test
+    fun closingChildDoesNotCloseEngine() = runTest {
+        var engineClosed = false
+        val client = HttpClient(createTestEngineFactory { engineClosed = true })
+        val child = client.config {}
+
+        assertSame(client.engine, child.engine)
+
+        child.close()
+
+        assertFalse(engineClosed)
+        client.close()
+        assertTrue(engineClosed)
+    }
+
+    @Test
+    fun closingParentDoesNotCloseEngineWhileChildIsOpen() = runTest {
+        var engineClosed = false
+        val client = HttpClient(createTestEngineFactory { engineClosed = true })
+        val child = client.config {}
+
+        assertSame(client.engine, child.engine)
+
+        client.close()
+
+        assertFalse(engineClosed)
+        child.close()
+        assertTrue(engineClosed)
+    }
+
+    @Test
+    fun cancellingChildDoesNotCancelEngine() = runTest {
+        val client = HttpClient(createTestEngineFactory {})
+        val child = client.config {}
+
+        assertSame(client.engine, child.engine)
+
+        child.coroutineContext[Job]!!.cancel()
+
+        assertFalse(client.engine.coroutineContext[Job]!!.isCancelled)
+        client.coroutineContext[Job]!!.cancel()
+        assertTrue(client.engine.coroutineContext[Job]!!.isCancelled)
+    }
+
+    @Test
+    fun cancellingParentDoesNotCancelEngineWhileChildIsOpen() = runTest {
+        val client = HttpClient(createTestEngineFactory {})
+        val child = client.config {}
+
+        assertSame(client.engine, child.engine)
+
+        client.coroutineContext[Job]!!.cancel()
+
+        assertFalse(child.engine.coroutineContext[Job]!!.isCancelled)
+        child.coroutineContext[Job]!!.cancel()
+        assertTrue(child.engine.coroutineContext[Job]!!.isCancelled)
+    }
+
+    private fun createTestEngineFactory(onClose: () -> Unit): HttpClientEngineFactory<HttpClientEngineConfig> {
+        val engine = object : HttpClientEngineBase("test-engine") {
+            override val config: HttpClientEngineConfig = HttpClientEngineConfig()
+
+            @InternalAPI
+            override suspend fun execute(data: HttpRequestData): HttpResponseData {
+                return HttpResponseData(
+                    HttpStatusCode.OK,
+                    GMTDate.START,
+                    Headers.Empty,
+                    HttpProtocolVersion.HTTP_1_1,
+                    Unit,
+                    EmptyCoroutineContext
+                )
+            }
+
+            override fun close() {
+                super.close()
+                onClose()
+            }
+        }
+        return object : HttpClientEngineFactory<HttpClientEngineConfig> {
+            override fun create(block: HttpClientEngineConfig.() -> Unit) = engine
+        }
+    }
 }