Fix Jetty upgrade cancellation on connection close

ktorio · Feb 5, 2020 · 1734027 · 1734027
1 parent 614f637
commit 1734027
Showing 1 changed file with 8 additions and 13 deletions.
diff --git a/ktor-server/ktor-server-jetty/jvm/src/io/ktor/server/jetty/internal/JettyUpgradeImpl.kt b/ktor-server/ktor-server-jetty/jvm/src/io/ktor/server/jetty/internal/JettyUpgradeImpl.kt
@@ -31,20 +31,9 @@ object JettyUpgradeImpl : ServletUpgrade {
         // for upgraded connections IDLE timeout should be significantly increased
         endPoint.idleTimeout = TimeUnit.MINUTES.toMillis(60L)
 
-        withContext(engineContext) {
+        withContext(engineContext + CoroutineName("upgrade-scope")) {
             try {
                 coroutineScope {
-                    endPoint.connection.addListener(
-                        object : Connection.Listener {
-                            override fun onOpened(connection: Connection?) {
-                            }
-
-                            override fun onClosed(connection: Connection?) {
-                                cancel()
-                            }
-                        }
-                    )
-
                     val inputChannel = ByteChannel(autoFlush = true)
                     val reader = EndPointReader(endPoint, coroutineContext, inputChannel)
                     val writer = endPointWriter(endPoint)
@@ -54,10 +43,16 @@ object JettyUpgradeImpl : ServletUpgrade {
                     if (endPoint is AbstractEndPoint) {
                         endPoint.upgrade(reader)
                     }
-                    upgrade.upgrade(
+                    val upgradeJob = upgrade.upgrade(
                         inputChannel, outputChannel, coroutineContext,
                         coroutineContext + userContext
                     )
+
+                    upgradeJob.invokeOnCompletion {
+                        inputChannel.cancel()
+                        outputChannel.close()
+                        cancel()
+                    }
                 }
             } finally {
                 connection.close()