Skip to content

Commit f4f4b27

Browse files
IRusclaude
IRus
and
claude
authored
Add Store.awaitSnapshotCapturable for race-free snapshot coordination (#19)
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 <noreply@anthropic.com>
1 parent ddea5ee commit f4f4b27

2 files changed

Lines changed: 131 additions & 23 deletions

File tree

‎wasm-core/jvmTest/wasm/core/SnapshotSuspensionSafetyJvmTest.kt‎

Lines changed: 57 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@ import kotlinx.coroutines.CompletableDeferred
44
import kotlinx.coroutines.Dispatchers
55
import kotlinx.coroutines.asCoroutineDispatcher
66
import kotlinx.coroutines.async
7-
import kotlinx.coroutines.delay
8-
import kotlinx.coroutines.flow.first
97
import kotlinx.coroutines.runBlocking
108
import kotlinx.coroutines.withContext
119
import kotlinx.coroutines.withTimeout
@@ -56,17 +54,12 @@ class SnapshotSuspensionSafetyJvmTest {
5654
val invocation = async(Dispatchers.Unconfined) {
5755
instance.invoke("resume")
5856
}
59-
store.status.first { it == StoreStatus.InHostImport }
6057
// InHostImport is published before the import's continuation
61-
// parks. A capture probe only succeeds once the running segment
62-
// has suspended and released the execution gate, so poll it to
63-
// guarantee the completion below resumes a parked continuation
64-
// on the resumer thread instead of finishing the await inline.
65-
withTimeout(5_000) {
66-
while (runCatching { store.captureSnapshotState(instance) }.isFailure) {
67-
delay(1)
68-
}
69-
}
58+
// parks, so await the execution gate to guarantee the completion
59+
// below resumes a parked continuation on the resumer thread
60+
// instead of finishing the await inline.
61+
store.awaitSnapshotCapturable()
62+
assertEquals(StoreStatus.InHostImport, store.status.value)
7063
assertFalse(invocation.isCompleted)
7164

7265
withContext(resumerDispatcher) {
@@ -122,11 +115,10 @@ class SnapshotSuspensionSafetyJvmTest {
122115

123116
try {
124117
val invocation = async(invocationDispatcher) { instance.invoke("wait") }
125-
store.status.first { it == StoreStatus.InHostImport }
126-
// InHostImport is published before host code runs. Drain the
127-
// single-thread dispatcher to prove that the host continuation
128-
// has actually suspended and released the store execution gate.
129-
withContext(invocationDispatcher) { Unit }
118+
// InHostImport is published before the invocation continuation
119+
// parks, so waiting on status alone would race the capture below.
120+
store.awaitSnapshotCapturable()
121+
assertEquals(StoreStatus.InHostImport, store.status.value)
130122

131123
val captureStarted = CountDownLatch(1)
132124
val finishCapture = CountDownLatch(1)
@@ -159,6 +151,54 @@ class SnapshotSuspensionSafetyJvmTest {
159151
}
160152
}
161153

154+
@Test
155+
fun awaitSnapshotCapturableEnablesCrossThreadCaptureOfAParkedHostImport(): Unit = runBlocking {
156+
val type = FuncType(emptyList(), emptyList())
157+
val module = validatedModule {
158+
types += type
159+
imports += Import("host", "wait", ImportDesc.Function(0))
160+
exports += Export("wait", ExportDesc.Function(0))
161+
}
162+
val invocationDispatcher =
163+
Executors.newSingleThreadExecutor().asCoroutineDispatcher()
164+
165+
try {
166+
repeat(32) {
167+
val releaseImport = CompletableDeferred<Unit>()
168+
val store = Store()
169+
val instance = Instance(
170+
store,
171+
module,
172+
ResolvedImports(
173+
functions = listOf(
174+
HostImport(type) {
175+
releaseImport.await()
176+
emptyList()
177+
},
178+
),
179+
),
180+
)
181+
182+
val invocation = async(invocationDispatcher) { instance.invoke("wait") }
183+
store.awaitSnapshotCapturable()
184+
assertEquals(StoreStatus.InHostImport, store.status.value)
185+
186+
// The import is still blocked, so the guest cannot resume and
187+
// a capture from another thread must succeed deterministically.
188+
val snapshot = withContext(Dispatchers.Default) {
189+
store.captureSnapshotState(instance)
190+
}
191+
assertEquals(0, snapshot.pendingImport?.functionIndex)
192+
193+
releaseImport.complete(Unit)
194+
assertEquals(emptyList(), invocation.await())
195+
assertEquals(StoreStatus.Idle, store.status.value)
196+
}
197+
} finally {
198+
invocationDispatcher.close()
199+
}
200+
}
201+
162202
private fun validatedModule(configure: ModuleBuilder.() -> Unit): Module =
163203
ModuleBuilder()
164204
.apply(configure)

‎wasm-core/src/wasm/core/Store.kt‎

Lines changed: 74 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import kotlinx.coroutines.flow.first
1616
import kotlinx.coroutines.flow.update
1717
import kotlinx.coroutines.launch
1818
import kotlinx.coroutines.sync.Mutex
19+
import kotlinx.coroutines.sync.withLock
1920
import kotlin.concurrent.Volatile
2021
import kotlin.coroutines.Continuation
2122
import kotlin.coroutines.ContinuationInterceptor
@@ -84,7 +85,18 @@ public data class StoreConfig(
8485
}
8586
}
8687

87-
/** Coarse store state suitable for monitoring and snapshot coordination. */
88+
/**
89+
* Coarse store state suitable for monitoring and snapshot coordination.
90+
*
91+
* Suspension statuses are published from inside the store's gated execution
92+
* segments: [InHostImport], [Paused], and [WaitingForFuel] become externally
93+
* observable before the executing continuation actually parks and releases
94+
* the execution gate. A cross-thread observer that reacts to [Store.status]
95+
* alone can therefore call [Store.captureSnapshotState] while the segment is
96+
* still running and fail with "store execution has not parked". Await
97+
* [Store.awaitSnapshotCapturable] instead when the observation is meant to
98+
* precede a snapshot capture.
99+
*/
88100
@io.heapy.kwasm.ExperimentalKwasmApi
89101
public enum class StoreStatus {
90102
Idle,
@@ -957,6 +969,46 @@ public class Store(
957969
GuestStackFrame(it.functionIndex, it.functionName, it.currentInstructionIndex)
958970
}
959971

972+
/**
973+
* Suspend until the store parks at a snapshot-capturable suspension point.
974+
*
975+
* Suspension statuses are published before the executing continuation
976+
* parks and releases the execution gate, so reacting to [status] alone
977+
* can race [captureSnapshotState] and fail with "store execution has not
978+
* parked". This primitive returns only once a capturable status
979+
* ([StoreStatus.Paused], [StoreStatus.WaitingForFuel], or
980+
* [StoreStatus.InHostImport]) is published and the execution gate has
981+
* been released.
982+
*
983+
* After this returns, a [captureSnapshotState] call cannot fail with
984+
* "has not parked" unless the guest resumes in between. While the parked
985+
* host import or pause is still outstanding the guest cannot resume, so
986+
* awaiting this and then capturing is race-free; once the caller releases
987+
* the guest (completing the import, [PauseHandle.resume], [addFuel]) the
988+
* observed capturability is stale.
989+
*
990+
* Waits across [StoreStatus.Idle] and [StoreStatus.Running] for the next
991+
* capturable suspension point; if none is ever reached this suspends
992+
* until cancelled. Throws [SnapshotStateException] if the store is or
993+
* becomes [StoreStatus.Poisoned] while waiting.
994+
*/
995+
public suspend fun awaitSnapshotCapturable() {
996+
while (true) {
997+
val observed = statusState.first {
998+
it == StoreStatus.Poisoned || it.acceptsSnapshotCapture
999+
}
1000+
if (observed == StoreStatus.Poisoned) {
1001+
throw SnapshotStateException(
1002+
"store is poisoned and will not park at a snapshot-capturable suspension point",
1003+
)
1004+
}
1005+
val parkedCapturable = executionGate.withParkedExecution {
1006+
statusState.value.acceptsSnapshotCapture
1007+
}
1008+
if (parkedCapturable) return
1009+
}
1010+
}
1011+
9601012
/**
9611013
* Copy all state needed by the optional snapshot codec.
9621014
*
@@ -973,6 +1025,12 @@ public class Store(
9731025
* use this form because GC objects, host references, and registered host
9741026
* participants may need to be traversed before a fully detached byte
9751027
* representation exists.
1028+
*
1029+
* Capture requires the executing continuation to have parked, which
1030+
* happens strictly after the matching [StoreStatus] is published.
1031+
* Cross-thread callers coordinating through [status] must await
1032+
* [awaitSnapshotCapturable] first instead of capturing on the status
1033+
* observation alone.
9761034
*/
9771035
public fun <T> captureSnapshotState(
9781036
instance: Instance,
@@ -986,11 +1044,7 @@ public class Store(
9861044
}
9871045
try {
9881046
val currentStatus = statusState.value
989-
if (
990-
currentStatus != StoreStatus.Paused &&
991-
currentStatus != StoreStatus.WaitingForFuel &&
992-
currentStatus != StoreStatus.InHostImport
993-
) {
1047+
if (!currentStatus.acceptsSnapshotCapture) {
9941048
throw SnapshotStateException(
9951049
"store status is $currentStatus; " +
9961050
"snapshot requires Paused, WaitingForFuel, or a parked host import",
@@ -1357,6 +1411,11 @@ public class Store(
13571411
)
13581412
}
13591413

1414+
private val StoreStatus.acceptsSnapshotCapture: Boolean
1415+
get() = this == StoreStatus.Paused ||
1416+
this == StoreStatus.WaitingForFuel ||
1417+
this == StoreStatus.InHostImport
1418+
13601419
/** Excludes snapshot traversal from synchronous interpreter segments. */
13611420
private class StoreExecutionGate {
13621421
private val mutex = Mutex()
@@ -1367,6 +1426,15 @@ private class StoreExecutionGate {
13671426
mutex.unlock()
13681427
}
13691428

1429+
/**
1430+
* Run [block] under the gate, suspending until the current continuation
1431+
* segment (if any) parks and releases it. On return the gate is free
1432+
* again, so a follow-up [tryAcquireCapture] succeeds unless a new segment
1433+
* resumed in between.
1434+
*/
1435+
suspend fun <T> withParkedExecution(block: () -> T): T =
1436+
mutex.withLock { block() }
1437+
13701438
fun <T> resumeSegment(
13711439
continuation: Continuation<T>,
13721440
result: Result<T>,

0 commit comments

Comments
 (0)