diff --git a/quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/transport/UdpSocket.kt b/quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/transport/UdpSocket.kt index 72a1bd4568..6a68360b5e 100644 --- a/quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/transport/UdpSocket.kt +++ b/quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/transport/UdpSocket.kt @@ -21,6 +21,8 @@ package com.vitorpamplona.quic.transport import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.ExecutorCoroutineDispatcher +import kotlinx.coroutines.asCoroutineDispatcher import kotlinx.coroutines.withContext import java.net.InetAddress import java.net.InetSocketAddress @@ -28,17 +30,31 @@ import java.net.StandardSocketOptions import java.nio.ByteBuffer import java.nio.channels.ClosedChannelException import java.nio.channels.DatagramChannel +import java.util.concurrent.Executors import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicLong /** - * JVM/Android UDP socket using blocking [DatagramChannel] dispatched onto - * [Dispatchers.IO]. We don't use NIO selectors because each QUIC connection - * has exactly one socket and one receive loop — Selector doesn't pay for - * itself at this scale. + * JVM/Android UDP socket using a blocking [DatagramChannel]. We don't use NIO + * selectors because each QUIC connection has exactly one socket and one receive + * loop — a Selector doesn't pay for itself at this scale. * - * The receive buffer is sized to 64 KiB (max IPv4/IPv6 datagram); QUIC packets - * cap at MTU (~1500 in practice). + * Threading: the blocking `recvfrom` parks its thread for the *entire* life of + * the connection (it only returns when a datagram arrives or the socket + * closes). If that ran on the shared [Dispatchers.IO] pool it would pin one + * pool thread per connection, and past ~64 concurrent connections it would + * starve *all* other `Dispatchers.IO` work in the process — this module's and + * the host app's alike. So each socket owns two dedicated daemon threads: + * [recvDispatcher] for the perpetually-blocked receive, and [sendDispatcher] + * for the (rarely-blocking, but still-blocking) send. The receive can't share a + * thread with send — it would monopolise it — hence two. Both are shut down in + * [close]. QUIC's blocking socket I/O therefore never touches the shared pool. + * + * [connect] still resolves DNS + binds on [Dispatchers.IO]: that's a one-shot + * setup cost, not a lifetime parker, so it doesn't need isolation. + * + * The receive buffer is sized to typical Ethernet MTU; QUIC packets cap at MTU + * (~1500 in practice). */ actual class UdpSocket private constructor( private val channel: DatagramChannel, @@ -46,6 +62,19 @@ actual class UdpSocket private constructor( ) { private val closed = AtomicBoolean(false) + // Dedicated single-thread executors so the blocking socket calls never + // occupy the shared Dispatchers.IO pool. Daemon threads so a leaked socket + // can't keep the JVM alive. Separate recv/send threads because the receive + // parks continuously and would otherwise block sends behind it. We keep the + // ExecutorService handles (not just the dispatchers) so close() can call + // shutdownNow() — an interrupt that breaks the parked recvfrom immediately + // (ClosedByInterruptException) rather than the graceful shutdown() that + // dispatcher.close() would do. + private val recvExecutor = Executors.newSingleThreadExecutor { r -> Thread(r, "quic-udp-recv").apply { isDaemon = true } } + private val sendExecutor = Executors.newSingleThreadExecutor { r -> Thread(r, "quic-udp-send").apply { isDaemon = true } } + private val recvDispatcher: ExecutorCoroutineDispatcher = recvExecutor.asCoroutineDispatcher() + private val sendDispatcher: ExecutorCoroutineDispatcher = sendExecutor.asCoroutineDispatcher() + // Sized to typical Ethernet MTU + a bit; QUIC tops out at ~1500 in practice // and any larger inbound frame is dropped as malformed anyway. The previous // 64 KiB buffer was wasteful per connection. @@ -76,15 +105,19 @@ actual class UdpSocket private constructor( actual val receiveBufferSizeBytes: Int get() = channel.getOption(StandardSocketOptions.SO_RCVBUF) - actual suspend fun send(payload: ByteArray): Int = - withContext(Dispatchers.IO) { + actual suspend fun send(payload: ByteArray): Int { + // Fail fast without dispatching onto a possibly shut-down executor. + if (closed.get()) throw ClosedChannelException() + return withContext(sendDispatcher) { if (closed.get()) throw ClosedChannelException() val buf = ByteBuffer.wrap(payload) channel.send(buf, remote) } + } - actual suspend fun receive(): ByteArray? = - withContext(Dispatchers.IO) { + actual suspend fun receive(): ByteArray? { + if (closed.get()) return null + return withContext(recvDispatcher) { if (closed.get()) return@withContext null try { // No synchronized — only the read loop touches readBuf, by @@ -101,6 +134,7 @@ actual class UdpSocket private constructor( null } } + } actual fun close() { if (closed.compareAndSet(false, true)) { @@ -109,6 +143,15 @@ actual class UdpSocket private constructor( } catch (_: Throwable) { // already closed } + // shutdownNow() interrupts the dedicated workers: a thread parked in + // a blocking recvfrom throws ClosedByInterruptException (a + // ClosedChannelException, caught below), so receive() returns null + // and both threads exit promptly instead of leaking per closed + // connection. channel.close() above would also unblock it + // (AsynchronousCloseException), but the interrupt is immediate and + // guarantees the executor terminates. + recvExecutor.shutdownNow() + sendExecutor.shutdownNow() } } diff --git a/quic/src/jvmTest/kotlin/com/vitorpamplona/quic/transport/UdpSocketTest.kt b/quic/src/jvmTest/kotlin/com/vitorpamplona/quic/transport/UdpSocketTest.kt new file mode 100644 index 0000000000..64acd1d0b2 --- /dev/null +++ b/quic/src/jvmTest/kotlin/com/vitorpamplona/quic/transport/UdpSocketTest.kt @@ -0,0 +1,146 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quic.transport + +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.withTimeoutOrNull +import java.net.DatagramPacket +import java.net.DatagramSocket +import java.net.InetSocketAddress +import java.nio.channels.ClosedChannelException +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertNull +import kotlin.test.assertTrue + +class UdpSocketTest { + // A plain UDP peer on loopback that echoes one datagram back. + private val peer = DatagramSocket(InetSocketAddress("127.0.0.1", 0)) + + @AfterTest + fun tearDown() { + runCatching { peer.close() } + } + + private fun threadNames(): List = Thread.getAllStackTraces().keys.map { it.name } + + private fun hasThreadPrefixed(prefix: String) = threadNames().any { it.startsWith(prefix) } + + @Test + fun `round trips a datagram over loopback`() { + runBlocking { + val socket = UdpSocket.connect("127.0.0.1", peer.localPort) + try { + // Peer thread: receive one packet and echo it back to the sender. + val echo = + Thread { + val buf = ByteArray(2048) + val incoming = DatagramPacket(buf, buf.size) + peer.receive(incoming) + peer.send(DatagramPacket(incoming.data, incoming.length, incoming.socketAddress)) + }.apply { + isDaemon = true + start() + } + + socket.send(byteArrayOf(1, 2, 3, 4)) + val reply = withTimeoutOrNull(3_000) { socket.receive() } + echo.join(1_000) + + assertTrue(reply != null && reply.contentEquals(byteArrayOf(1, 2, 3, 4)), "should echo the datagram back") + assertEquals(1, socket.receivedDatagramCount) + } finally { + socket.close() + } + } + } + + @Test + fun `blocking receive runs on a dedicated thread, not Dispatchers-IO`() { + runBlocking { + val socket = UdpSocket.connect("127.0.0.1", peer.localPort) + try { + // Park a receive with no incoming datagram — it blocks in recvfrom + // on the socket's dedicated recv thread, not a Dispatchers.IO worker. + val pending = async(Dispatchers.IO) { socket.receive() } + // Give the receive time to reach the blocking call on its own thread. + delay(200) + + assertTrue(hasThreadPrefixed("quic-udp-recv"), "a dedicated recv thread must carry the blocking receive") + + // Closing unblocks the parked receive (returns null), proving the + // blocking call was on the dedicated thread and is released on close. + socket.close() + val result = withTimeoutOrNull(2_000) { pending.await() } + assertNull(result, "receive() must return null once the socket closes") + } finally { + socket.close() + } + } + } + + @Test + fun `close shuts down the dedicated threads`() { + runBlocking { + val socket = UdpSocket.connect("127.0.0.1", peer.localPort) + try { + // The single-thread executors spawn their thread lazily on first + // use, so touch BOTH directions to bring both threads up. + socket.send(byteArrayOf(0)) + val recv = launch(Dispatchers.IO) { socket.receive() } + delay(200) + assertTrue( + hasThreadPrefixed("quic-udp-recv") && hasThreadPrefixed("quic-udp-send"), + "both dedicated threads must be up while open", + ) + + socket.close() + recv.join() + + // The executors shut down on close; their threads must exit promptly. + val gone = + withTimeoutOrNull(2_000) { + while (hasThreadPrefixed("quic-udp-recv") || hasThreadPrefixed("quic-udp-send")) delay(25) + true + } + assertTrue(gone == true, "dedicated recv/send threads must be gone after close() — else they leak per connection") + } finally { + socket.close() + } + } + } + + @Test + fun `after close receive returns null and send throws`() { + runBlocking { + val socket = UdpSocket.connect("127.0.0.1", peer.localPort) + socket.close() + assertNull(socket.receive(), "receive() returns null after close") + assertFailsWith { socket.send(byteArrayOf(9)) } + } + } +}