KTOR-3504 KTOR-2430 Fix WS timeout

Timeout should close incoming and outgoing channels, after sending close frame https://youtrack.jetbrains.com/issue/KTOR-2430 https://youtrack.jetbrains.com/issue/KTOR-3504
ktorio · Oct 13, 2022 · 22cc34f · 22cc34f
1 parent 5ec968c
commit 22cc34f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/ktor-shared/ktor-websockets/common/src/io/ktor/websocket/DefaultWebSocketSession.kt b/ktor-shared/ktor-websockets/common/src/io/ktor/websocket/DefaultWebSocketSession.kt
@@ -10,6 +10,7 @@ import io.ktor.utils.io.core.*
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
+import io.ktor.utils.io.errors.*
 import kotlin.coroutines.*
 
 /**
@@ -210,7 +211,6 @@ internal class DefaultWebSocketSessionImpl(
         }
     }
 
-    @OptIn(ExperimentalCoroutinesApi::class)
     private fun runOutgoingProcessor(): Job = launch(
         OutgoingProcessorCoroutineName + Dispatchers.Unconfined,
         start = CoroutineStart.UNDISPATCHED
@@ -247,7 +247,7 @@ internal class DefaultWebSocketSessionImpl(
     }
 
     @OptIn(InternalAPI::class)
-    private suspend fun sendCloseSequence(reason: CloseReason?) {
+    private suspend fun sendCloseSequence(reason: CloseReason?, exception: Throwable? = null) {
         if (!tryClose()) return
         context.complete()
 
@@ -260,6 +260,11 @@ internal class DefaultWebSocketSessionImpl(
             }
         } finally {
             closeReasonRef.complete(reasonToSend)
+
+            if (exception != null) {
+                outgoingToBeProcessed.close(exception)
+                filtered.close(exception)
+            }
         }
     }
 
@@ -270,7 +275,9 @@ internal class DefaultWebSocketSessionImpl(
 
         val newPinger: SendChannel<Frame.Pong>? = when {
             closed.value -> null
-            interval > 0L -> pinger(raw.outgoing, interval, timeoutMillis)
+            interval > 0L -> pinger(raw.outgoing, interval, timeoutMillis) {
+                sendCloseSequence(it, IOException("Ping timeout"))
+            }
             else -> null
         }
 

diff --git a/ktor-shared/ktor-websockets/common/src/io/ktor/websocket/PingPong.kt b/ktor-shared/ktor-websockets/common/src/io/ktor/websocket/PingPong.kt
@@ -45,6 +45,7 @@ internal fun CoroutineScope.pinger(
     outgoing: SendChannel<Frame>,
     periodMillis: Long,
     timeoutMillis: Long,
+    onTimeout: suspend (CloseReason) -> Unit
 ): SendChannel<Frame.Pong> {
     val actorJob = Job()
 
@@ -82,8 +83,7 @@ internal fun CoroutineScope.pinger(
                     // we were unable to send the ping or hadn't got a valid pong message in time,
                     // so we are triggering close sequence (if already started then the following close frame could be ignored)
 
-                    val closeFrame = Frame.Close(CloseReason(CloseReason.Codes.INTERNAL_ERROR, "Ping timeout"))
-                    outgoing.send(closeFrame)
+                    onTimeout(CloseReason(CloseReason.Codes.INTERNAL_ERROR, "Ping timeout"))
                     break
                 }
             }