nick/amethyst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

Merge pull request #2808 from vitorpamplona/claude/review-quic-blocking-code-autOK

Browse Source 
perf(quic): lock-free hot paths — ThreadLocal Cipher, AtomicReference close, @Volatile getters
This commit is contained in:
Vitor Pamplona
2026-05-08 23:19:37 -04:00
committed by GitHub
parent 9ec4f04bd7 2bb55ff9ec
commit 2a9b5bc17d
7 changed files with 155 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnection.kt
+14
View File
@@ -324,6 +324,16 @@ class QuicConnection(
*/
private val closeStateMonitor = Any()
/**
* Single-shot signal completed when [status] transitions to CLOSED —
* either by the writer flushing CONNECTION_CLOSE in `drainOutbound`,
* or by [markClosedExternally] forcing the state. Replaces an earlier
* 1 ms polling loop in `QuicConnectionDriver.close()`. Idempotent:
* `complete(Unit)` returns false on subsequent firers, so racing
* teardown paths are safe.
*/
internal val closingDrainSignal: CompletableDeferred<Unit> = CompletableDeferred()
/** App-level error code for graceful close. */
var closeReason: String? = null
private set
@@ -1522,6 +1532,10 @@ class QuicConnection(
closeReason = reason
true
}
// Wake any teardown coroutine waiting for the close to land. Safe
// to call after the monitor — complete() is idempotent and the
// CLOSED transition has already been published via @Volatile.
if (firstClose) closingDrainSignal.complete(Unit)
if (firstClose) {
// "remote" covers both peer-initiated CONNECTION_CLOSE and
// local invariant violations (CID mismatch, frame decode
quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionDriver.kt
+48 -36
View File
@@ -31,6 +31,8 @@ import kotlinx.coroutines.joinAll
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withTimeoutOrNull
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.ExperimentalAtomicApi
/**
* Owns the UDP socket and runs the read + send loops for a [QuicConnection].
@@ -48,6 +50,7 @@ import kotlinx.coroutines.withTimeoutOrNull
* and [QuicConnection.openBidiStream]/[com.vitorpamplona.quic.stream.SendBuffer.enqueue])
* call [wakeup] to nudge the send loop.
*/
@OptIn(ExperimentalAtomicApi::class)
class QuicConnectionDriver(
val connection: QuicConnection,
private val socket: UdpSocket,
@@ -94,17 +97,20 @@ class QuicConnectionDriver(
* to poll `connection.status == CLOSED` and trust that the rest of
* the cleanup eventually settled.
*/
internal val closeTeardownJob: Job? get() = closeJob
internal val closeTeardownJob: Job? get() = closeJob.load()
/**
* Round-5 concurrency #5: close() guard. A second concurrent invocation
* (e.g. session close + read-loop death close racing) used to launch a
* parallel teardown that called scope.cancel() and socket.close() while
* the first close was mid-joinAll. We now memoize the teardown Job so
* the second caller awaits the first's completion instead.
* the first close was mid-joinAll. We memoize the teardown Job so the
* second caller awaits the first's completion instead.
*
* Lock-free CAS replaces the previous `synchronized(this)` double-checked
* init: the close path is single-shot ("first writer wins"), which is
* exactly what `compareAndSet` expresses.
*/
@Volatile
private var closeJob: Job? = null
private val closeJob: AtomicReference<Job?> = AtomicReference(null)
fun start() {
connection.start()
@@ -365,38 +371,44 @@ class QuicConnectionDriver(
// second concurrent caller (which is common: session.close() and
// read-loop death both race to close()) awaits the same Job rather
// than launching a parallel teardown.
if (closeJob != null) return
synchronized(this) {
if (closeJob != null) return
closeJob =
parentScope.launch {
connection.close(0L, "")
wakeup()
val send = sendJob
// Bounded wait for the send loop to flush CONNECTION_CLOSE.
// We don't want to hang forever if the writer is wedged —
// the timeout is the upper bound on how long close() blocks.
withTimeoutOrNull(CLOSE_FLUSH_TIMEOUT_MILLIS) {
// Spin until the writer has actually drained the queued
// close. The CLOSING-status check transitions to CLOSED
// once drainOutbound builds the CONNECTION_CLOSE packet.
while (connection.status == QuicConnection.Status.CLOSING) {
kotlinx.coroutines.delay(1)
}
}
// Now flip to CLOSED so both loops exit their while-guards.
connection.markClosedExternally("driver close requested")
wakeup()
// Wait for both loops to actually exit — joinAll won't
// return until the in-flight socket.send() completes.
withTimeoutOrNull(CLOSE_FLUSH_TIMEOUT_MILLIS) {
listOfNotNull(readJob, send).joinAll()
}
// Final teardown — cancel guarantees both jobs are done
// before we close the socket.
scope.cancel()
socket.close()
if (closeJob.load() != null) return
// Build the teardown coroutine LAZY so we can race-test the CAS
// without paying for a launched-and-cancelled Job on the loser.
val teardown =
parentScope.launch(start = kotlinx.coroutines.CoroutineStart.LAZY) {
connection.close(0L, "")
wakeup()
val send = sendJob
// Bounded wait for the send loop to flush CONNECTION_CLOSE.
// Event-driven via [QuicConnection.closingDrainSignal] —
// both `drainOutbound` (after building the close datagram)
// and `markClosedExternally` (forced transition) complete
// the deferred. Replaces an earlier 1 ms polling loop.
// The timeout is the upper bound on how long close() blocks
// if the writer is wedged.
withTimeoutOrNull(CLOSE_FLUSH_TIMEOUT_MILLIS) {
connection.closingDrainSignal.await()
}
// Now flip to CLOSED so both loops exit their while-guards.
connection.markClosedExternally("driver close requested")
wakeup()
// Wait for both loops to actually exit — joinAll won't
// return until the in-flight socket.send() completes.
withTimeoutOrNull(CLOSE_FLUSH_TIMEOUT_MILLIS) {
listOfNotNull(readJob, send).joinAll()
}
// Final teardown — cancel guarantees both jobs are done
// before we close the socket.
scope.cancel()
socket.close()
}
if (closeJob.compareAndSet(null, teardown)) {
teardown.start()
} else {
// A concurrent close() already installed the teardown Job; drop
// ours without ever starting it. The winner's Job runs, this
// call is a no-op (matching the original idempotent contract).
teardown.cancel()
}
}
quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionWriter.kt
+7 -2
View File
@@ -18,6 +18,8 @@
* 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.
*/
@file:OptIn(kotlin.concurrent.atomics.ExperimentalAtomicApi::class)
package com.vitorpamplona.quic.connection
import com.vitorpamplona.quartz.utils.Log
@@ -86,6 +88,9 @@ fun drainOutbound(
if (conn.status == QuicConnection.Status.CLOSING) {
val datagram = buildClosingDatagram(conn, nowMillis)
conn.status = QuicConnection.Status.CLOSED
// Signal the teardown coroutine that the close datagram is built;
// the driver awaits this instead of polling status.
conn.closingDrainSignal.complete(Unit)
return datagram
}
@@ -952,7 +957,7 @@ private fun appendFlowControlUpdates(
// resetAcked / stopSendingAcked so subsequent stale loss tokens
// are dropped.
for (stream in conn.streamsListLocked()) {
val resetState = stream.resetState
val resetState = stream.resetState.load()
if (resetState != null && stream.resetEmitPending && !stream.resetAcked) {
frames +=
ResetStreamFrame(
@@ -968,7 +973,7 @@ private fun appendFlowControlUpdates(
)
stream.resetEmitPending = false
}
val stopSendingState = stream.stopSendingState
val stopSendingState = stream.stopSendingState.load()
if (stopSendingState != null && stream.stopSendingEmitPending && !stream.stopSendingAcked) {
frames +=
StopSendingFrame(
quic/src/commonMain/kotlin/com/vitorpamplona/quic/stream/QuicStream.kt
+39 -25
View File
@@ -23,6 +23,8 @@ package com.vitorpamplona.quic.stream
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.ExperimentalAtomicApi
/**
* One QUIC stream (bidirectional or unidirectional). Application code
@@ -30,6 +32,7 @@ import kotlinx.coroutines.flow.flow
* APIs; the [QuicConnection] owns the underlying buffers and drains them
* into STREAM frames on the wire.
*/
@OptIn(ExperimentalAtomicApi::class)
class QuicStream(
val streamId: Long,
val direction: Direction,
@@ -255,24 +258,28 @@ class QuicStream(
* (two app threads racing the writer's clear-after-emit).
*/
fun resetStream(errorCode: Long) {
// Synchronized atomic compare-and-set: pre-fix the
// `if (resetState != null) return` plus the assignment was
// racy. Two concurrent callers (e.g. the application aborting
// a request while STOP_SENDING from the peer triggers our own
// resetStream from the parser) could both observe null and
// both write — the second write would clobber the first
// errorCode while [resetEmitPending] was already set. The
// writer would then emit a RESET_STREAM with whichever
// errorCode landed last, possibly different from what the
// application asked for. The synchronized block makes
// first-call-wins genuinely first-call-wins.
synchronized(this) {
if (resetState != null) return
resetState =
ResetState(
errorCode = errorCode,
finalSize = send.nextOffset,
)
// Lock-free first-call-wins. The previous synchronized block
// existed because the naive `if (resetState != null) return;
// resetState = …` had a write-write race: two concurrent
// callers (e.g. the application aborting a request while
// STOP_SENDING from the peer triggers our own resetStream
// from the parser) could both observe null and both write,
// letting the second clobber the first errorCode while
// resetEmitPending was already set. The writer would then
// emit a RESET_STREAM with whichever errorCode landed last.
//
// compareAndSet collapses that to a single CAS: only the
// first writer succeeds, all others observe non-null and
// bail out. `resetEmitPending = true` happens only on the
// winning path, so its @Volatile write happens-after the
// resetState publication — the writer reading the flag sees
// the populated state.
val newState =
ResetState(
errorCode = errorCode,
finalSize = send.nextOffset,
)
if (resetState.compareAndSet(null, newState)) {
resetEmitPending = true
}
}
@@ -292,10 +299,8 @@ class QuicStream(
* original frame already on the wire).
*/
fun stopSending(errorCode: Long) {
// Same atomic-CAS rationale as [resetStream].
synchronized(this) {
if (stopSendingState != null) return
stopSendingState = StopSendingState(errorCode = errorCode)
// Same lock-free first-call-wins rationale as [resetStream].
if (stopSendingState.compareAndSet(null, StopSendingState(errorCode = errorCode))) {
stopSendingEmitPending = true
}
}
@@ -304,8 +309,13 @@ class QuicStream(
* RFC 9000 §3.5 send-side reset state. Set once by [resetStream]
* (subsequent calls no-op); read by the writer + loss/ACK
* dispatchers. Once set, contents are immutable.
*
* Held in an [AtomicReference] so [resetStream] can use
* `compareAndSet(null, …)` to win the first-call-wins race
* without acquiring a lock. Readers use `.load()` — the published
* state is immutable after the CAS so a single load is enough.
*/
internal var resetState: ResetState? = null
internal val resetState: AtomicReference<ResetState?> = AtomicReference(null)
/**
* True while a RESET_STREAM emit is pending. Cleared after the
@@ -331,8 +341,12 @@ class QuicStream(
@Volatile
internal var resetAcked: Boolean = false
/** Receive-side stop-sending state. Set by [stopSending]. */
internal var stopSendingState: StopSendingState? = null
/**
* Receive-side stop-sending state. Set once by [stopSending]
* (subsequent calls no-op); read by the writer + loss/ACK
* dispatchers. Atomic for the same reason as [resetState].
*/
internal val stopSendingState: AtomicReference<StopSendingState?> = AtomicReference(null)
@Volatile
internal var stopSendingEmitPending: Boolean = false
quic/src/commonMain/kotlin/com/vitorpamplona/quic/stream/SendBuffer.kt
+20 -8
View File
@@ -60,9 +60,14 @@ package com.vitorpamplona.quic.stream
* send loop under the connection mutex; [markAcked] / [markLost] run on
* the parser path also under the connection mutex. The two execution
* paths are NOT serialised by a shared lock, so all internal state is
* mutated under `synchronized(this)`. Even the cheap getters
* ([readableBytes], [sentOffset], [finPending], [finSent]) take the
* monitor so a writer pre-flight check can't observe torn state.
* mutated under `synchronized(this)`. Single-field reads
* ([nextOffset], [sentOffset], [finPending], [finSent], [finAcked])
* use `@Volatile` backing fields and bypass the monitor — they cannot
* tear (Boolean is single-byte; Long writes happen inside the
* synchronized block on JVM/Android, where `@Volatile Long` is
* atomic). [readableBytes] still synchronizes because its formula
* combines two fields and would otherwise observe a transient
* negative value if read mid-`takeChunk`.
*
* # FIN
*
@@ -105,6 +110,7 @@ class SendBuffer(
private var flushedFloor: Long = 0L
/** Logical offset just past the last byte. Advances on [enqueue]. */
@Volatile
private var _nextOffset: Long = 0L
/**
@@ -115,6 +121,7 @@ class SendBuffer(
*
* Invariant: `flushedFloor <= nextSendOffset <= nextOffset`.
*/
@Volatile
private var nextSendOffset: Long = 0L
/**
@@ -151,14 +158,19 @@ class SendBuffer(
*/
private var retransmitTotalBytes: Long = 0L
@Volatile
private var _finPending: Boolean = false
@Volatile
private var _finSent: Boolean = false
@Volatile
private var _finAcked: Boolean = false
val nextOffset: Long get() = synchronized(this) { _nextOffset }
val finPending: Boolean get() = synchronized(this) { _finPending }
val finSent: Boolean get() = synchronized(this) { _finSent }
val finAcked: Boolean get() = synchronized(this) { _finAcked }
val nextOffset: Long get() = _nextOffset
val finPending: Boolean get() = _finPending
val finSent: Boolean get() = _finSent
val finAcked: Boolean get() = _finAcked
/**
* Bytes the writer would emit on the next [takeChunk] before any
@@ -183,7 +195,7 @@ class SendBuffer(
* the high-water mark of *fresh* sends only, not the cumulative
* retransmit volume.
*/
val sentOffset: Long get() = synchronized(this) { nextSendOffset }
val sentOffset: Long get() = nextSendOffset
fun enqueue(bytes: ByteArray) {
if (bytes.isEmpty()) return
quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/crypto/PlatformCrypto.kt
+14 -2
View File
@@ -26,11 +26,23 @@ import javax.crypto.spec.SecretKeySpec
/**
* One-block AES-ECB encryption via JCA. Used only by QUIC header protection
* (one block per packet, so no need for a more elaborate API).
* (one block per packet). The `Cipher` instance is cached per-thread so the
* read/send loops avoid `Cipher.getInstance(...)` provider lookup on every
* packet — which is the dominant cost for a 16-byte one-shot AEAD-less call.
*
* `ThreadLocal` is safe here because the call is stateless: every invocation
* re-`init`s with the caller-supplied key before `doFinal`, so a coroutine
* that hops dispatchers between calls just lands on whichever thread's cached
* `Cipher` it ends up on. No state leaks across calls.
*/
private val aesEcbCipher: ThreadLocal<Cipher> =
ThreadLocal.withInitial { Cipher.getInstance("AES/ECB/NoPadding") }
actual val PlatformAesOneBlock: AesOneBlockEncrypt =
AesOneBlockEncrypt { key, block ->
val cipher = Cipher.getInstance("AES/ECB/NoPadding")
// .get() is non-null because withInitial supplies a Cipher, but
// Kotlin sees the Java return type as platform-nullable.
val cipher = aesEcbCipher.get()!!
cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"))
cipher.doFinal(block)
}
quic/src/jvmAndroid/kotlin/com/vitorpamplona/quic/tls/JdkCertificateValidator.kt
+13 -3
View File
@@ -191,10 +191,20 @@ class JdkCertificateValidator(
// GeneralName type 2 = dNSName, type 7 = iPAddress.
if (type == 2 && dnsMatches(idnAscii(value), normalizedHost)) return true
if (type == 7 && hostAsIp != null) {
// Defense-in-depth: a malformed cert could put a hostname in
// a type 7 SAN. Without this gate, InetAddress.getByName(value)
// would perform a DNS A/AAAA lookup on the validation path,
// both leaking the name in plaintext and blocking the read
// loop on the system resolver. Forcing a literal check keeps
// the JDK call to pure parsing (no I/O, no name service).
val sanIp =
try {
InetAddress.getByName(value).hostAddress
} catch (_: Throwable) {
if (looksLikeIpLiteral(value)) {
try {
InetAddress.getByName(value).hostAddress
} catch (_: Throwable) {
null
}
} else {
null
}
if (sanIp != null && sanIp.equals(hostAsIp, ignoreCase = true)) return true