diff --git a/runtime/src/main/java/com/dylibso/chicory/runtime/ByteArrayMemory.java b/runtime/src/main/java/com/dylibso/chicory/runtime/ByteArrayMemory.java index 12de4dedd..2509cac70 100644 --- a/runtime/src/main/java/com/dylibso/chicory/runtime/ByteArrayMemory.java +++ b/runtime/src/main/java/com/dylibso/chicory/runtime/ByteArrayMemory.java @@ -114,6 +114,14 @@ private static final class WaitState { int waiterCount; // The number of waiting threads that have been scheduled to wake up int pendingWakeups; + // Bumped by every notify that issues wakeups. A waiter may only consume + // a pending wakeup minted while it was waiting (generation moved since + // it registered). Without this, a newly-arrived waiter could steal a + // wakeup meant for an already-parked thread. That thread would then + // re-check, find none, and sleep forever (lost wakeup). + // Only ever compared with ==/!=, never ordered, so wraparound is not + // a practical concern. + long generation; } private final Map waitStates; @@ -146,9 +154,10 @@ private int waitOn(int address, BooleanSupplier condition, long timeout) { } state.waiterCount++; + final long arrivalGeneration = state.generation; try { - while (state.pendingWakeups == 0) { + while (state.pendingWakeups == 0 || state.generation == arrivalGeneration) { long remaining = deadline - System.nanoTime(); if (remaining <= 0) { return 2; // timeout @@ -165,11 +174,12 @@ private int waitOn(int address, BooleanSupplier condition, long timeout) { return 0; // woken } finally { - if (state.pendingWakeups > 0) { - // any thread leaving the try block is correct to consume - // a pending wakeup IF available: + if (state.pendingWakeups > 0 && state.generation != arrivalGeneration) { + // any thread leaving the try block is correct to consume a + // pending wakeup IF one was minted while it was waiting: // ret 0 - woken thread correctly consumes a wakeup - // ret 2 - timeout can only occur with pendingWakeups == 0 so will not consume + // ret 2 - timeout cannot coincide with a consumable wakeup + // (the while loop would have exited instead) // throw - isn't part of the wasm runtime but is semantically correct to consume state.pendingWakeups--; } @@ -220,8 +230,11 @@ public int notify(int address, int maxThreads) { toWake = Math.min(actualWaiters, maxThreads); } - // Add pending wakeups - waiters will consume these + // Add pending wakeups - waiters will consume these. The generation + // bump marks every currently-registered waiter as eligible; threads + // arriving later must wait for the next notify. state.pendingWakeups += toWake; + state.generation++; assert (state.pendingWakeups <= state.waiterCount); // It's always safe to notify all. In the wait routine we consume pendingWakeups diff --git a/runtime/src/main/java/com/dylibso/chicory/runtime/ByteBufferMemory.java b/runtime/src/main/java/com/dylibso/chicory/runtime/ByteBufferMemory.java index dbcef602c..a6a695289 100644 --- a/runtime/src/main/java/com/dylibso/chicory/runtime/ByteBufferMemory.java +++ b/runtime/src/main/java/com/dylibso/chicory/runtime/ByteBufferMemory.java @@ -85,6 +85,14 @@ private static final class WaitState { int waiterCount; // The number of waiting threads that have been scheduled to wake up int pendingWakeups; + // Bumped by every notify that issues wakeups. A waiter may only consume + // a pending wakeup minted while it was waiting (generation moved since + // it registered). Without this, a newly-arrived waiter could steal a + // wakeup meant for an already-parked thread. That thread would then + // re-check, find none, and sleep forever (lost wakeup). + // Only ever compared with ==/!=, never ordered, so wraparound is not + // a practical concern. + long generation; } private final Map waitStates; @@ -117,9 +125,10 @@ private int waitOn(int address, BooleanSupplier condition, long timeout) { } state.waiterCount++; + final long arrivalGeneration = state.generation; try { - while (state.pendingWakeups == 0) { + while (state.pendingWakeups == 0 || state.generation == arrivalGeneration) { long remaining = deadline - System.nanoTime(); if (remaining <= 0) { return 2; // timeout @@ -136,11 +145,12 @@ private int waitOn(int address, BooleanSupplier condition, long timeout) { return 0; // woken } finally { - if (state.pendingWakeups > 0) { - // any thread leaving the try block is correct to consume - // a pending wakeup IF available: + if (state.pendingWakeups > 0 && state.generation != arrivalGeneration) { + // any thread leaving the try block is correct to consume a + // pending wakeup IF one was minted while it was waiting: // ret 0 - woken thread correctly consumes a wakeup - // ret 2 - timeout can only occur with pendingWakeups == 0 so will not consume + // ret 2 - timeout cannot coincide with a consumable wakeup + // (the while loop would have exited instead) // throw - isn't part of the wasm runtime but is semantically correct to consume state.pendingWakeups--; } @@ -191,8 +201,11 @@ public int notify(int address, int maxThreads) { toWake = Math.min(actualWaiters, maxThreads); } - // Add pending wakeups - waiters will consume these + // Add pending wakeups - waiters will consume these. The generation + // bump marks every currently-registered waiter as eligible; threads + // arriving later must wait for the next notify. state.pendingWakeups += toWake; + state.generation++; assert (state.pendingWakeups <= state.waiterCount); // It's always safe to notify all. In the wait routine we consume pendingWakeups diff --git a/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryStressTest.java b/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryStressTest.java new file mode 100644 index 000000000..120bfe5c5 --- /dev/null +++ b/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryStressTest.java @@ -0,0 +1,127 @@ +package com.dylibso.chicory.runtime; + +import static org.junit.jupiter.api.Assertions.assertFalse; + +import com.dylibso.chicory.wasm.types.MemoryLimits; +import java.util.ArrayList; +import java.util.List; +import java.util.function.Supplier; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +public class MemoryStressTest { + + static final class RwLock { + private static final int WRITER_BIT = 0x80000000; + private static final int READER_MASK = 0x7FFFFFFF; + + private final Memory memory; + private final int addr; + + public RwLock(Memory memory, int addr) { + this.memory = memory; + this.addr = addr; + } + + public void lock() { + while (true) { + int s = memory.atomicReadInt(addr); + if ((s & WRITER_BIT) != 0) { + memory.atomicWait(addr, s, -1L); + continue; + } + if (memory.atomicCmpxchgInt(addr, s, s | WRITER_BIT) == s) { + break; + } + } + while (true) { + int s = memory.atomicReadInt(addr); + if ((s & READER_MASK) == 0) { + return; + } + memory.atomicWait(addr, s, -1L); + } + } + + public void unlock() { + memory.atomicWriteInt(addr, 0); + memory.atomicNotify(addr, -1); + } + + public void lockShared() { + while (true) { + int s = memory.atomicReadInt(addr); + if ((s & WRITER_BIT) != 0) { + memory.atomicWait(addr, s, -1L); + continue; + } + if (memory.atomicCmpxchgInt(addr, s, s + 1) == s) { + return; + } + } + } + + public void unlockShared() { + int prev = memory.atomicAddInt(addr, -1); + if ((prev & READER_MASK) == 1 && (prev & WRITER_BIT) != 0) { + memory.atomicNotify(addr, -1); + } + } + } + + private static Stream sharedMemoryImplementations() { + return Stream.of( + Arguments.of( + "ByteArrayMemory", + (Supplier) () -> new ByteArrayMemory(new MemoryLimits(1, 1, true))), + Arguments.of( + "ByteBufferMemory", + (Supplier) + () -> new ByteBufferMemory(new MemoryLimits(1, 1, true)))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("sharedMemoryImplementations") + public void rwLockSurvivesContention(String name, Supplier memorySupplier) + throws InterruptedException { + final RwLock lock = new RwLock(memorySupplier.get(), 0); + final long deadline = System.currentTimeMillis() + 3_000; + + final List threads = new ArrayList(); + for (int r = 0; r < 6; r++) { + threads.add( + new Thread( + () -> { + while (System.currentTimeMillis() < deadline) { + lock.lockShared(); + lock.unlockShared(); + } + }, + "reader-" + r)); + } + threads.add( + new Thread( + () -> { + while (System.currentTimeMillis() < deadline) { + lock.lock(); + lock.unlock(); + } + }, + "writer")); + + for (Thread t : threads) { + t.setDaemon(true); + t.start(); + } + final long joinDeadline = deadline + 1000; + for (Thread t : threads) { + final long remaining = Long.max(10, joinDeadline - System.currentTimeMillis()); + t.join(remaining); + } + + final boolean anyThreadBlocked = threads.stream().anyMatch(Thread::isAlive); + assertFalse(anyThreadBlocked); + } +} diff --git a/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryTest.java b/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryTest.java index a82d9eda1..c6ebf0a43 100644 --- a/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryTest.java +++ b/runtime/src/test/java/com/dylibso/chicory/runtime/MemoryTest.java @@ -1,9 +1,11 @@ package com.dylibso.chicory.runtime; import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; import com.dylibso.chicory.wasm.types.MemoryLimits; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import java.util.stream.Stream; @@ -165,4 +167,53 @@ public void nonOverlappingCopyCrossPage(String name, Supplier memorySupp // Source should be unchanged assertArrayEquals(data, memory.readBytes(100, data.length)); } + + /** + * This test verifies a permit stealing bug that existed in the memory implementations. + * An atomicWait could avoid waiting if it found a "notify permit". This is wrong by definition + * since atomicNotify can only wake threads in a waiting state. + * + * The test creates a thread that calls (A) atomicWait, while the main thread calls atomicNotify and + * then (B) atomicWait. If the notify call returns one, it means it should have woken (A) since + * (B) has not yet happened. This means that (B) must timeout since there are no other notify calls + * involved. + */ + @ParameterizedTest(name = "{0}") + @MethodSource("growableMemoryImplementations") + public void notifyWakesParkedThreadNotLateArrival(String name, Supplier memorySupplier) + throws Exception { + final int addr = 0; + final int sentinel = 42; + var memory = memorySupplier.get(); + + for (int round = 0; round < 200; round++) { + memory.atomicWriteInt(addr, sentinel); + + var parkedReady = new CountDownLatch(1); + + var parked = + new Thread( + () -> { + parkedReady.countDown(); + memory.atomicWait(addr, sentinel, 100_000_000L); + }); + parked.setDaemon(true); + parked.start(); + parkedReady.await(); + // This sleep gives the other thread a higher chance of reaching the atomicWait, + // but setting it too high makes a passing test slow (200 x sleep) + Thread.sleep(1); + + int woken = memory.atomicNotify(addr, 1); + if (woken == 0) { + // we didn't wake up the other thread. This case is not interesting, retry + continue; + } + + // Late arrival: this wait MUST timeout as we have only woken ONE thread. + final int result = memory.atomicWait(addr, sentinel, 0L); + + assertEquals(2, result, "round " + round + ": main thread should have timed out (2)"); + } + } }