From b494c6c01421c7c97adce5db4d58319325ca7605 Mon Sep 17 00:00:00 2001 From: Ruslan Ibrahimau Date: Sun, 19 Jul 2026 11:53:15 +0300 Subject: [PATCH] Add Store.awaitSnapshotCapturable for race-free snapshot coordination StoreStatus.InHostImport (and Paused/WaitingForFuel) is published from inside a gated execution segment, before the executing continuation parks and StoreExecutionGate releases its mutex. An embedder that observes status.first { it == InHostImport } on another thread and then calls captureSnapshotState intermittently hits the fail-fast tryAcquireCapture and gets "store execution has not parked". This exact race made SnapshotSuspensionSafetyJvmTest flaky. Reordering the publication is not an option: the gate release is implicit in the ContinuationInterceptor machinery and only happens once the segment actually suspends. Instead, expose a public awaitable primitive: awaitSnapshotCapturable() suspends until a capturable status is published AND the gate has been released (by briefly taking the gate mutex), so a follow-up captureSnapshotState cannot fail with "has not parked" unless the guest resumes in between - which it cannot while the parked host import is still blocked. Also document the publish-before-park ordering on StoreStatus and captureSnapshotState, switch the previously flaky test to the new primitive, and add a focused stress test for capture of a parked host import from another thread. Supersedes the test-level dispatcher drain on branch fix/snapshot-suspension-safety-jvm-test. Co-Authored-By: Claude Fable 5 --- .../core/SnapshotSuspensionSafetyJvmTest.kt | 53 +++++++++++- wasm-core/src/wasm/core/Store.kt | 80 +++++++++++++++++-- 2 files changed, 126 insertions(+), 7 deletions(-) diff --git a/wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt b/wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt index 1e356cf..23cd321 100644 --- a/wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt +++ b/wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt @@ -111,7 +111,10 @@ class SnapshotSuspensionSafetyJvmTest { try { val invocation = async(invocationDispatcher) { instance.invoke("wait") } - store.status.first { it == StoreStatus.InHostImport } + // InHostImport is published before the invocation continuation + // parks, so waiting on status alone would race the capture below. + store.awaitSnapshotCapturable() + assertEquals(StoreStatus.InHostImport, store.status.value) val captureStarted = CountDownLatch(1) val finishCapture = CountDownLatch(1) @@ -144,6 +147,54 @@ class SnapshotSuspensionSafetyJvmTest { } } + @Test + fun awaitSnapshotCapturableEnablesCrossThreadCaptureOfAParkedHostImport(): Unit = runBlocking { + val type = FuncType(emptyList(), emptyList()) + val module = validatedModule { + types += type + imports += Import("host", "wait", ImportDesc.Function(0)) + exports += Export("wait", ExportDesc.Function(0)) + } + val invocationDispatcher = + Executors.newSingleThreadExecutor().asCoroutineDispatcher() + + try { + repeat(32) { + val releaseImport = CompletableDeferred() + val store = Store() + val instance = Instance( + store, + module, + ResolvedImports( + functions = listOf( + HostImport(type) { + releaseImport.await() + emptyList() + }, + ), + ), + ) + + val invocation = async(invocationDispatcher) { instance.invoke("wait") } + store.awaitSnapshotCapturable() + assertEquals(StoreStatus.InHostImport, store.status.value) + + // The import is still blocked, so the guest cannot resume and + // a capture from another thread must succeed deterministically. + val snapshot = withContext(Dispatchers.Default) { + store.captureSnapshotState(instance) + } + assertEquals(0, snapshot.pendingImport?.functionIndex) + + releaseImport.complete(Unit) + assertEquals(emptyList(), invocation.await()) + assertEquals(StoreStatus.Idle, store.status.value) + } + } finally { + invocationDispatcher.close() + } + } + private fun validatedModule(configure: ModuleBuilder.() -> Unit): Module = ModuleBuilder() .apply(configure) diff --git a/wasm-core/src/wasm/core/Store.kt b/wasm-core/src/wasm/core/Store.kt index 9da15b6..97e950c 100644 --- a/wasm-core/src/wasm/core/Store.kt +++ b/wasm-core/src/wasm/core/Store.kt @@ -16,6 +16,7 @@ import kotlinx.coroutines.flow.first import kotlinx.coroutines.flow.update import kotlinx.coroutines.launch import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock import kotlin.concurrent.Volatile import kotlin.coroutines.Continuation import kotlin.coroutines.ContinuationInterceptor @@ -84,7 +85,18 @@ public data class StoreConfig( } } -/** Coarse store state suitable for monitoring and snapshot coordination. */ +/** + * Coarse store state suitable for monitoring and snapshot coordination. + * + * Suspension statuses are published from inside the store's gated execution + * segments: [InHostImport], [Paused], and [WaitingForFuel] become externally + * observable before the executing continuation actually parks and releases + * the execution gate. A cross-thread observer that reacts to [Store.status] + * alone can therefore call [Store.captureSnapshotState] while the segment is + * still running and fail with "store execution has not parked". Await + * [Store.awaitSnapshotCapturable] instead when the observation is meant to + * precede a snapshot capture. + */ @io.heapy.kwasm.ExperimentalKwasmApi public enum class StoreStatus { Idle, @@ -957,6 +969,46 @@ public class Store( GuestStackFrame(it.functionIndex, it.functionName, it.currentInstructionIndex) } + /** + * Suspend until the store parks at a snapshot-capturable suspension point. + * + * Suspension statuses are published before the executing continuation + * parks and releases the execution gate, so reacting to [status] alone + * can race [captureSnapshotState] and fail with "store execution has not + * parked". This primitive returns only once a capturable status + * ([StoreStatus.Paused], [StoreStatus.WaitingForFuel], or + * [StoreStatus.InHostImport]) is published and the execution gate has + * been released. + * + * After this returns, a [captureSnapshotState] call cannot fail with + * "has not parked" unless the guest resumes in between. While the parked + * host import or pause is still outstanding the guest cannot resume, so + * awaiting this and then capturing is race-free; once the caller releases + * the guest (completing the import, [PauseHandle.resume], [addFuel]) the + * observed capturability is stale. + * + * Waits across [StoreStatus.Idle] and [StoreStatus.Running] for the next + * capturable suspension point; if none is ever reached this suspends + * until cancelled. Throws [SnapshotStateException] if the store is or + * becomes [StoreStatus.Poisoned] while waiting. + */ + public suspend fun awaitSnapshotCapturable() { + while (true) { + val observed = statusState.first { + it == StoreStatus.Poisoned || it.acceptsSnapshotCapture + } + if (observed == StoreStatus.Poisoned) { + throw SnapshotStateException( + "store is poisoned and will not park at a snapshot-capturable suspension point", + ) + } + val parkedCapturable = executionGate.withParkedExecution { + statusState.value.acceptsSnapshotCapture + } + if (parkedCapturable) return + } + } + /** * Copy all state needed by the optional snapshot codec. * @@ -973,6 +1025,12 @@ public class Store( * use this form because GC objects, host references, and registered host * participants may need to be traversed before a fully detached byte * representation exists. + * + * Capture requires the executing continuation to have parked, which + * happens strictly after the matching [StoreStatus] is published. + * Cross-thread callers coordinating through [status] must await + * [awaitSnapshotCapturable] first instead of capturing on the status + * observation alone. */ public fun captureSnapshotState( instance: Instance, @@ -986,11 +1044,7 @@ public class Store( } try { val currentStatus = statusState.value - if ( - currentStatus != StoreStatus.Paused && - currentStatus != StoreStatus.WaitingForFuel && - currentStatus != StoreStatus.InHostImport - ) { + if (!currentStatus.acceptsSnapshotCapture) { throw SnapshotStateException( "store status is $currentStatus; " + "snapshot requires Paused, WaitingForFuel, or a parked host import", @@ -1357,6 +1411,11 @@ public class Store( ) } +private val StoreStatus.acceptsSnapshotCapture: Boolean + get() = this == StoreStatus.Paused || + this == StoreStatus.WaitingForFuel || + this == StoreStatus.InHostImport + /** Excludes snapshot traversal from synchronous interpreter segments. */ private class StoreExecutionGate { private val mutex = Mutex() @@ -1367,6 +1426,15 @@ private class StoreExecutionGate { mutex.unlock() } + /** + * Run [block] under the gate, suspending until the current continuation + * segment (if any) parks and releases it. On return the gate is free + * again, so a follow-up [tryAcquireCapture] succeeds unless a new segment + * resumed in between. + */ + suspend fun withParkedExecution(block: () -> T): T = + mutex.withLock { block() } + fun resumeSegment( continuation: Continuation, result: Result,