Skip to content

Add Store.awaitSnapshotCapturable for race-free snapshot coordination #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
IRus merged 2 commits into from
Jul 19, 2026
Merged

Add Store.awaitSnapshotCapturable for race-free snapshot coordination #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b494c6c
Add Store.awaitSnapshotCapturable for race-free snapshot coordination
IRus Jul 19, 2026
6627ccb
Merge main into fix/store-await-capturable
IRus Jul 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 57 additions & 17 deletions wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@ import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.async
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
Expand Down Expand Up @@ -56,17 +54,12 @@ class SnapshotSuspensionSafetyJvmTest {
val invocation = async(Dispatchers.Unconfined) {
instance.invoke("resume")
}
store.status.first { it == StoreStatus.InHostImport }
// InHostImport is published before the import's continuation
// parks. A capture probe only succeeds once the running segment
// has suspended and released the execution gate, so poll it to
// guarantee the completion below resumes a parked continuation
// on the resumer thread instead of finishing the await inline.
withTimeout(5_000) {
while (runCatching { store.captureSnapshotState(instance) }.isFailure) {
delay(1)
}
}
// parks, so await the execution gate to guarantee the completion
// below resumes a parked continuation on the resumer thread
// instead of finishing the await inline.
store.awaitSnapshotCapturable()
assertEquals(StoreStatus.InHostImport, store.status.value)
assertFalse(invocation.isCompleted)

withContext(resumerDispatcher) {
Expand Down Expand Up @@ -122,11 +115,10 @@ class SnapshotSuspensionSafetyJvmTest {

try {
val invocation = async(invocationDispatcher) { instance.invoke("wait") }
store.status.first { it == StoreStatus.InHostImport }
// InHostImport is published before host code runs. Drain the
// single-thread dispatcher to prove that the host continuation
// has actually suspended and released the store execution gate.
withContext(invocationDispatcher) { Unit }
// 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)
Expand Down Expand Up @@ -159,6 +151,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<Unit>()
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)
Expand Down
80 changes: 74 additions & 6 deletions wasm-core/src/wasm/core/Store.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
*
Expand All @@ -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 <T> captureSnapshotState(
instance: Instance,
Expand All @@ -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",
Expand Down Expand Up @@ -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()
Expand All @@ -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 <T> withParkedExecution(block: () -> T): T =
mutex.withLock { block() }

fun <T> resumeSegment(
continuation: Continuation<T>,
result: Result<T>,
Expand Down
Loading