KTOR-8194 Fix SelectorManager dispatcher ignored on Native

ktorio · Feb 13, 2025 · 75459e6 · 75459e6
1 parent bd652e4
commit 75459e6
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 15 deletions.
diff --git a/ktor-client/ktor-client-cio/common/src/io/ktor/client/engine/cio/CIOEngine.kt b/ktor-client/ktor-client-cio/common/src/io/ktor/client/engine/cio/CIOEngine.kt
@@ -66,7 +66,6 @@ internal class CIOEngine(
                 requestJob.join()
             } finally {
                 selector.close()
-                selector.coroutineContext[Job]!!.join()
             }
         }
     }

diff --git a/ktor-network/nix/src/io/ktor/network/selector/SelectUtilsNix.kt b/ktor-network/nix/src/io/ktor/network/selector/SelectUtilsNix.kt
@@ -67,7 +67,7 @@ internal actual class SelectorHelper {
     }
 
     @OptIn(ExperimentalForeignApi::class, InternalAPI::class)
-    private fun selectionLoop() {
+    private suspend fun selectionLoop() {
         val completed = mutableSetOf<EventInfo>()
         val watchSet = mutableSetOf<EventInfo>()
         val closeSet = mutableSetOf<Int>()
@@ -84,6 +84,8 @@ internal actual class SelectorHelper {
 
                 maxDescriptor = max(maxDescriptor + 1, wakeupSignalEvent.descriptor + 1)
 
+                yield()
+
                 try {
                     selector_pselect(maxDescriptor + 1, readSet, writeSet, errorSet).check()
                 } catch (_: PosixException.BadFileDescriptorException) {
@@ -263,4 +265,8 @@ internal actual class SelectorHelper {
     private fun isDescriptorValid(descriptor: Int): Boolean {
         return fcntl(descriptor, F_GETFL) != -1 || errno != EBADF
     }
+
+    private suspend inline fun dispatchIfNeeded() {
+        yield() // it will always redispatch it to the right thread
+    }
 }
diff --git a/ktor-network/nix/test/io/ktor/network/sockets/tests/TcpSocketTestNix.kt b/ktor-network/nix/test/io/ktor/network/sockets/tests/TcpSocketTestNix.kt
@@ -26,7 +26,7 @@ class TcpSocketTestNix {
         socket.close()
         selector.close()
 
-        selector.coroutineContext[Job]?.join()
+        socket.awaitClosed()
 
         val isDescriptorValid = fcntl(descriptor, F_GETFL) != -1 || errno != EBADF
         check(!isDescriptorValid) { "Descriptor was not closed" }
@@ -59,7 +59,9 @@ class TcpSocketTestNix {
         server.close()
         selector.close()
 
-        selector.coroutineContext[Job]?.join()
+        serverConnection.awaitClosed()
+        clientConnection.awaitClosed()
+        server.awaitClosed()
 
         val isServerDescriptorValid = fcntl(serverDescriptor, F_GETFL) != -1 || errno != EBADF
         check(!isServerDescriptorValid) { "Server descriptor was not closed" }

diff --git a/ktor-network/nix/test/io/ktor/network/sockets/tests/UdpSocketTestNix.kt b/ktor-network/nix/test/io/ktor/network/sockets/tests/UdpSocketTestNix.kt
@@ -24,7 +24,7 @@ class UdpSocketTestNix {
         socket.close()
         selector.close()
 
-        selector.coroutineContext[Job]?.join()
+        socket.awaitClosed()
 
         val isDescriptorValid = fcntl(descriptor, F_GETFL) != -1 || errno != EBADF
         check(!isDescriptorValid) { "Descriptor was not closed" }

diff --git a/ktor-network/posix/src/io/ktor/network/selector/SelectorManager.kt b/ktor-network/posix/src/io/ktor/network/selector/SelectorManager.kt
@@ -11,7 +11,7 @@ import kotlin.coroutines.*
 
 public actual fun SelectorManager(
     dispatcher: CoroutineContext
-): SelectorManager = WorkerSelectorManager()
+): SelectorManager = WorkerSelectorManager(dispatcher)
 
 public actual interface SelectorManager : CoroutineScope, Closeable {
     /**

diff --git a/ktor-network/posix/src/io/ktor/network/selector/WorkerSelectorManager.kt b/ktor-network/posix/src/io/ktor/network/selector/WorkerSelectorManager.kt
@@ -8,12 +8,8 @@ import kotlinx.coroutines.*
 import kotlinx.io.*
 import kotlin.coroutines.*
 
-@OptIn(ExperimentalCoroutinesApi::class)
-internal class WorkerSelectorManager : SelectorManager {
-    @OptIn(DelicateCoroutinesApi::class)
-    private val selectorContext = newSingleThreadContext("WorkerSelectorManager")
-    private val job = Job()
-    override val coroutineContext: CoroutineContext = selectorContext + job
+internal class WorkerSelectorManager(context: CoroutineContext) : SelectorManager {
+    override val coroutineContext: CoroutineContext = context + CoroutineName("selector")
 
     private val selector = SelectorHelper()
 
@@ -39,7 +35,5 @@ internal class WorkerSelectorManager : SelectorManager {
 
     override fun close() {
         selector.requestTermination()
-        selectorContext.close()
-        job.cancel()
     }
 }
diff --git a/ktor-network/windows/src/io/ktor/network/selector/SelectUtilsWindows.kt b/ktor-network/windows/src/io/ktor/network/selector/SelectUtilsWindows.kt
@@ -49,14 +49,17 @@ internal actual class SelectorHelper {
     }
 
     @OptIn(ExperimentalForeignApi::class, InternalAPI::class)
-    private fun selectionLoop() {
+    private suspend fun selectionLoop() {
         val completed = mutableSetOf<EventInfo>()
         val watchSet = mutableSetOf<EventInfo>()
         val closeSet = mutableSetOf<Int>()
 
         try {
             while (!interestQueue.isClosed) {
                 val wsaEvents = fillHandlersOrClose(watchSet, completed, closeSet)
+
+                yield()
+
                 val index = memScoped {
                     val length = wsaEvents.size + 1
                     val wsaEventsWithWake = allocArray<CPointerVarOf<COpaquePointer>>(length).apply {