Skip to content

Commit 81708e8

Browse files
AzoyAzoy
committed
Add a WasmMutex
1 parent 3b7b2dd commit 81708e8

File tree

3 files changed

+164
-3
lines changed

3 files changed

+164
-3
lines changed

‎stdlib/public/Synchronization/CMakeLists.txt

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,13 @@ set(SWIFT_SYNCHRONIZATION_LINUX_SOURCES
2828
Mutex/LinuxImpl.swift
2929
)
3030

31-
# Windows depdencies and sources
31+
# Wasm dependencies and sources
32+
33+
set(SWIFT_SYNCHRONIZATION_WASM_SOURCES
34+
Mutex/WasmImpl.swift
35+
)
36+
37+
# Windows dependencies and sources
3238

3339
set(SWIFT_SYNCHRONIZATION_WINDOWS_DEPENDENCIES)
3440

@@ -68,6 +74,8 @@ add_swift_target_library(swiftSynchronization ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES
6874
${SWIFT_SYNCHRONIZATION_DARWIN_SOURCES}
6975
SWIFT_SOURCES_DEPENDS_LINUX
7076
${SWIFT_SYNCHRONIZATION_LINUX_SOURCES}
77+
SWIFT_SOURCES_DEPENDS_WASM
78+
${SWIFT_SYNCHRONIZATION_WASM_SOURCES}
7179
SWIFT_SOURCES_DEPENDS_WINDOWS
7280
${SWIFT_SYNCHRONIZATION_WINDOWS_SOURCES}
7381

‎stdlib/public/Synchronization/Mutex/LinuxImpl.swift

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import SynchronizationShims
1414

1515
extension Atomic where Value == UInt32 {
16-
borrowing func wait() {
16+
borrowing func wait() -> Bool {
1717
_swift_stdlib_wait(.init(rawAddress))
1818
}
1919

@@ -75,7 +75,15 @@ internal struct _MutexHandle: ~Copyable {
7575
// Block until unlock has been called. This will return early if the call
7676
// to unlock happened between attempting to acquire and attempting to
7777
// wait while nobody else managed to acquire it yet.
78-
storage.wait()
78+
//
79+
// Once we return from this, if the return value is 0 or false, then it
80+
// means the kernel was able to successfully acquire the lock. This can
81+
// occur when we fail the initial compare and exchange, but the mutex was
82+
// unlocked before we called 'wait'.
83+
if !storage.wait() {
84+
// Locked!
85+
return
86+
}
7987

8088
(exchanged, state) = storage.weakCompareExchange(
8189
expected: 0,
@@ -136,5 +144,7 @@ internal struct _MutexHandle: ~Copyable {
136144

137145
// Wake up the next highest priority waiter.
138146
storage.wake()
147+
148+
// Unlocked!
139149
}
140150
}

‎stdlib/public/Synchronization/Mutex/WasmImpl.swift

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the Swift Atomics open source project
4+
//
5+
// Copyright (c) 2024 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
@_extern(c, "llvm.wasm32.memory.atomic.wait32")
14+
func _swift_stdlib_wait(
15+
on: UnsafePointer<UInt32>,
16+
expected: UInt32,
17+
timeout: Int64
18+
) -> UInt32
19+
20+
@_extern(c, "llvm.wasm32.memory.atomic.notify")
21+
func _swift_stdlib_wake(on: UnsafePointer<UInt32>, count: UInt32)
22+
23+
extension Atomic where Value == UInt32 {
24+
borrowing func wait(expected: _MutexHandle.State) {
25+
_swift_stdlib_wait(
26+
on: .init(rawAddress),
27+
expected: expected.rawValue,
28+
29+
// A timeout of < 0 means indefinitely.
30+
timeout: -1
31+
)
32+
}
33+
34+
borrowing func wake() {
35+
// Only wake up 1 thread
36+
_swift_stdlib_wake(on: .init(rawAddress), count: 1)
37+
}
38+
}
39+
40+
@available(SwiftStdlib 5.11, *)
41+
@frozen
42+
@usableFromInline
43+
@_staticExclusiveOnly
44+
internal struct _MutexHandle: ~Copyable {
45+
@usableFromInline
46+
let storage: Atomic<State>
47+
48+
@available(SwiftStdlib 5.11, *)
49+
@_alwaysEmitIntoClient
50+
@_transparent
51+
init() {
52+
storage = Atomic(.unlocked)
53+
}
54+
55+
@available(SwiftStdlib 5.11, *)
56+
@usableFromInline
57+
borrowing func lock() {
58+
// Note: We could probably merge this cas into a do/while style loop, but we
59+
// really want to perform the strong variant before attempting to do weak
60+
// ones in the loop.
61+
62+
var (exchanged, state) = storage.compareExchange(
63+
expected: .unlocked,
64+
desired: .locked,
65+
successOrdering: .acquiring,
66+
failureOrdering: .relaxed
67+
)
68+
69+
if _fastPath(exchanged) {
70+
// Locked!
71+
return
72+
}
73+
74+
while !exchanged {
75+
// If we're not already contended, go ahead and transition the mutex state
76+
// into being contended. If when we do this that the value stored there
77+
// was unlocked, then we know we unintentionally acquired the lock. A
78+
// weird quirk is that
79+
if state != .contended, storage.exchange(
80+
.contended,
81+
ordering: .acquiring
82+
) == .unlocked {
83+
// Locked!
84+
return
85+
}
86+
87+
// Block until unlock has been called. This will return early if the call
88+
// to unlock happened between attempting to acquire and attempting to
89+
// wait while nobody else managed to acquire it yet.
90+
storage.wait(expected: .contended)
91+
92+
(exchanged, state) = storage.weakCompareExchange(
93+
expected: .unlocked,
94+
desired: .locked,
95+
successOrdering: .acquiring,
96+
failureOrdering: .relaxed
97+
)
98+
}
99+
100+
// Locked!
101+
}
102+
103+
@available(SwiftStdlib 5.11, *)
104+
@usableFromInline
105+
borrowing func tryLock() -> Bool {
106+
storage.compareExchange(
107+
expected: .unlocked,
108+
desired: .locked,
109+
successOrdering: .acquiring,
110+
failureOrdering: .relaxed
111+
).exchanged
112+
}
113+
114+
@available(SwiftStdlib 5.11, *)
115+
@usableFromInline
116+
borrowing func unlock() {
117+
// Transition our state from being either .locked or .contended to .unlocked.
118+
// At this point the mutex is freely acquirable. If the value that was
119+
// stored in the mutex was .locked, then no one else was waiting on this
120+
// mutex so we can just skip trying to wake up a thread.
121+
guard storage.exchange(.unlocked, ordering: .releasing) == .contended else {
122+
// Unlocked!
123+
return
124+
}
125+
126+
// Otherwise, wake up our next lucky random thread to acquire the mutex.
127+
// (Assuming no new thread acquires the lock before it does)
128+
storage.wake()
129+
130+
// Unlocked!
131+
}
132+
}
133+
134+
extension _MutexHandle {
135+
@available(SwiftStdlib 5.11, *)
136+
@frozen
137+
@usableFromInline
138+
internal enum State: UInt32, AtomicRepresentable {
139+
case unlocked
140+
case locked
141+
case contended
142+
}
143+
}

0 commit comments

Comments
 (0)