Skip to content

Commit f2d28cf

Browse files
jckarterjckarter
authored
Merge pull request #35188 from jckarter/checked-continuation
Concurrency: Introduce a `CheckedContinuation` adapter.
2 parents 7e7ab2f + 5f09745 commit f2d28cf

File tree

5 files changed

+292
-0
lines changed

5 files changed

+292
-0
lines changed

‎include/swift/Runtime/Concurrency.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,10 @@ void swift_continuation_throwingResumeWithError(/* +1 */ SwiftError *error,
345345
void *continuation,
346346
const Metadata *resumeType);
347347

348+
/// SPI helper to log a misuse of a `CheckedContinuation` to the appropriate places in the OS.
349+
extern "C" SWIFT_CC(swift)
350+
void swift_continuation_logFailedCheck(const char *message);
351+
348352
}
349353

350354
#endif

‎stdlib/public/Concurrency/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ endif()
3636
add_swift_target_library(swift_Concurrency ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS_STDLIB
3737
Actor.cpp
3838
Actor.swift
39+
CheckedContinuation.swift
3940
GlobalExecutor.cpp
4041
PartialAsyncTask.swift
4142
Task.cpp

‎stdlib/public/Concurrency/CheckedContinuation.swift

Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
import Swift
2+
3+
@_silgen_name("swift_continuation_logFailedCheck")
4+
internal func logFailedCheck(_ message: UnsafeRawPointer)
5+
6+
/// Implementation class that holds the `UnsafeContinuation` instance for
7+
/// a `CheckedContinuation`.
8+
internal final class CheckedContinuationCanary {
9+
// The instance state is stored in tail-allocated raw memory, so that
10+
// we can atomically check the continuation state.
11+
12+
private init() { fatalError("must use create") }
13+
14+
private static func _create(continuation: UnsafeRawPointer, function: String)
15+
-> Self {
16+
let instance = Builtin.allocWithTailElems_1(self,
17+
1._builtinWordValue,
18+
(UnsafeRawPointer?, String).self)
19+
20+
instance._continuationPtr.initialize(to: continuation)
21+
instance._functionPtr.initialize(to: function)
22+
return instance
23+
}
24+
25+
private var _continuationPtr: UnsafeMutablePointer<UnsafeRawPointer?> {
26+
return UnsafeMutablePointer<UnsafeRawPointer?>(
27+
Builtin.projectTailElems(self, (UnsafeRawPointer?, String).self))
28+
}
29+
private var _functionPtr: UnsafeMutablePointer<String> {
30+
let tailPtr = UnsafeMutableRawPointer(
31+
Builtin.projectTailElems(self, (UnsafeRawPointer?, String).self))
32+
33+
let functionPtr = tailPtr
34+
+ MemoryLayout<(UnsafeRawPointer?, String)>.offset(of: \(UnsafeRawPointer?, String).1)!
35+
36+
return functionPtr.assumingMemoryBound(to: String.self)
37+
}
38+
39+
internal static func create<T>(continuation: UnsafeContinuation<T>,
40+
function: String) -> Self {
41+
return _create(
42+
continuation: unsafeBitCast(continuation, to: UnsafeRawPointer.self),
43+
function: function)
44+
}
45+
46+
internal static func create<T>(continuation: UnsafeThrowingContinuation<T>,
47+
function: String) -> Self {
48+
return _create(
49+
continuation: unsafeBitCast(continuation, to: UnsafeRawPointer.self),
50+
function: function)
51+
}
52+
53+
internal var function: String {
54+
return _functionPtr.pointee
55+
}
56+
57+
// Take the continuation away from the container, or return nil if it's
58+
// already been taken.
59+
private func _takeContinuation() -> UnsafeRawPointer? {
60+
// Atomically exchange the current continuation value with a null pointer.
61+
let rawContinuationPtr = unsafeBitCast(_continuationPtr,
62+
to: Builtin.RawPointer.self)
63+
let rawOld = Builtin.atomicrmw_xchg_seqcst_Word(rawContinuationPtr,
64+
0._builtinWordValue)
65+
66+
return unsafeBitCast(rawOld, to: UnsafeRawPointer?.self)
67+
}
68+
69+
internal func takeContinuation<T>() -> UnsafeContinuation<T>? {
70+
return unsafeBitCast(_takeContinuation(),
71+
to: UnsafeContinuation<T>.self)
72+
}
73+
internal func takeThrowingContinuation<T>() -> UnsafeThrowingContinuation<T>? {
74+
return unsafeBitCast(_takeContinuation(),
75+
to: UnsafeThrowingContinuation<T>.self)
76+
}
77+
78+
deinit {
79+
_functionPtr.deinitialize(count: 1)
80+
// Log if the continuation was never consumed before the instance was
81+
// destructed.
82+
if _continuationPtr.pointee != nil {
83+
logFailedCheck("SWIFT TASK CONTINUATION MISUSE: \(function) leaked its continuation!\n")
84+
}
85+
}
86+
}
87+
88+
/// A wrapper class for `UnsafeContinuation` that logs misuses of the
89+
/// continuation, logging a message if the continuation is resumed
90+
/// multiple times, or if an object is destroyed without its continuation
91+
/// ever being resumed.
92+
///
93+
/// Raw `UnsafeContinuation`, like other unsafe constructs, requires the
94+
/// user to apply it correctly in order to maintain invariants. The key
95+
/// invariant is that the continuation must be resumed exactly once,
96+
/// and bad things happen if this invariant is not upheld--if a continuation
97+
/// is abandoned without resuming the task, then the task will be stuck in
98+
/// the suspended state forever, and conversely, if the same continuation is
99+
/// resumed multiple times, it will put the task in an undefined state.
100+
/// `UnsafeContinuation` avoids enforcing these invariants at runtime because
101+
/// it aims to be a low-overhead mechanism for interfacing Swift tasks with
102+
/// event loops, delegate methods, callbacks, and other non-`async` scheduling
103+
/// mechanisms. However, during development, being able to verify that the
104+
/// invariants are being upheld in testing is important.
105+
///
106+
/// `CheckedContinuation` is designed to be a drop-in API replacement for
107+
/// `UnsafeContinuation` that can be used for testing purposes, at the cost
108+
/// of an extra allocation and indirection for the wrapper object.
109+
/// Changing a call of `withUnsafeContinuation` into a call of
110+
/// `withCheckedContinuation` should be enough to obtain the extra checking
111+
/// without further source modification in most circumstances.
112+
public struct CheckedContinuation<T> {
113+
let canary: CheckedContinuationCanary
114+
115+
/// Initialize a `CheckedContinuation` wrapper around an
116+
/// `UnsafeContinuation`.
117+
///
118+
/// In most cases, you should use `withCheckedContinuation` instead.
119+
/// You only need to initialize your own `CheckedContinuation<T>` if you
120+
/// already have an `UnsafeContinuation` you want to impose checking on.
121+
///
122+
/// - Parameters:
123+
/// - continuation: a fresh `UnsafeContinuation` that has not yet
124+
/// been resumed. The `UnsafeContinuation` must not be used outside of
125+
/// this object once it's been given to the new object.
126+
/// - function: a string identifying the declaration that is the notional
127+
/// source for the continuation, used to identify the continuation in
128+
/// runtime diagnostics related to misuse of this continuation.
129+
public init(continuation: UnsafeContinuation<T>, function: String = #function) {
130+
canary = CheckedContinuationCanary.create(
131+
continuation: continuation,
132+
function: function)
133+
}
134+
135+
/// Resume the task awaiting the continuation by having it return normally
136+
/// from its suspension point.
137+
///
138+
/// A continuation must be resumed exactly once. If the continuation has
139+
/// already been resumed through this object, then the attempt to resume
140+
/// the continuation again will be logged, but otherwise have no effect.
141+
public func resume(returning x: __owned T) {
142+
if let c: UnsafeContinuation<T> = canary.takeContinuation() {
143+
c.resume(returning: x)
144+
} else {
145+
logFailedCheck("SWIFT TASK CONTINUATION MISUSE: \(canary.function) tried to resume its continuation more than once, returning \(x)!\n")
146+
}
147+
}
148+
}
149+
150+
public func withCheckedContinuation<T>(
151+
function: String = #function,
152+
_ body: (CheckedContinuation<T>) -> Void
153+
) async -> T {
154+
return await withUnsafeContinuation {
155+
body(CheckedContinuation(continuation: $0, function: function))
156+
}
157+
}
158+
159+
/// A wrapper class for `UnsafeThrowingContinuation` that logs misuses of the
160+
/// continuation, logging a message if the continuation is resumed
161+
/// multiple times, or if an object is destroyed without its continuation
162+
/// ever being resumed.
163+
///
164+
/// Raw `UnsafeThrowingContinuation`, like other unsafe constructs, requires the
165+
/// user to apply it correctly in order to maintain invariants. The key
166+
/// invariant is that the continuation must be resumed exactly once,
167+
/// and bad things happen if this invariant is not upheld--if a continuation
168+
/// is abandoned without resuming the task, then the task will be stuck in
169+
/// the suspended state forever, and conversely, if the same continuation is
170+
/// resumed multiple times, it will put the task in an undefined state.
171+
/// `UnsafeThrowingContinuation` avoids enforcing these invariants at runtime because
172+
/// it aims to be a low-overhead mechanism for interfacing Swift tasks with
173+
/// event loops, delegate methods, callbacks, and other non-`async` scheduling
174+
/// mechanisms. However, during development, being able to verify that the
175+
/// invariants are being upheld in testing is important.
176+
///
177+
/// `CheckedThrowingContinuation` is designed to be a drop-in API replacement for
178+
/// `UnsafeThrowingContinuation` that can be used for testing purposes, at the cost
179+
/// of an extra allocation and indirection for the wrapper object.
180+
/// Changing a call of `withUnsafeThrowingContinuation` into a call of
181+
/// `withCheckedThrowingContinuation` should be enough to obtain the extra checking
182+
/// without further source modification in most circumstances.
183+
public struct CheckedThrowingContinuation<T> {
184+
let canary: CheckedContinuationCanary
185+
186+
/// Initialize a `CheckedThrowingContinuation` wrapper around an
187+
/// `UnsafeThrowingContinuation`.
188+
///
189+
/// In most cases, you should use `withCheckedThrowingContinuation` instead.
190+
/// You only need to initialize your own `CheckedThrowingContinuation<T>` if you
191+
/// already have an `UnsafeThrowingContinuation` you want to impose checking on.
192+
///
193+
/// - Parameters:
194+
/// - continuation: a fresh `UnsafeThrowingContinuation` that has not yet
195+
/// been resumed. The `UnsafeThrowingContinuation` must not be used outside of
196+
/// this object once it's been given to the new object.
197+
/// - function: a string identifying the declaration that is the notional
198+
/// source for the continuation, used to identify the continuation in
199+
/// runtime diagnostics related to misuse of this continuation.
200+
public init(continuation: UnsafeThrowingContinuation<T>, function: String = #function) {
201+
canary = CheckedContinuationCanary.create(
202+
continuation: continuation,
203+
function: function)
204+
}
205+
206+
/// Resume the task awaiting the continuation by having it return normally
207+
/// from its suspension point.
208+
///
209+
/// A continuation must be resumed exactly once. If the continuation has
210+
/// already been resumed through this object, whether by `resume(returning:)`
211+
/// or by `resume(throwing:)`, then the attempt to resume
212+
/// the continuation again will be logged, but otherwise have no effect.
213+
public func resume(returning x: __owned T) {
214+
if let c: UnsafeThrowingContinuation<T> = canary.takeThrowingContinuation() {
215+
c.resume(returning: x)
216+
} else {
217+
logFailedCheck("SWIFT TASK CONTINUATION MISUSE: \(canary.function) tried to resume its continuation more than once, returning \(x)!\n")
218+
}
219+
}
220+
221+
/// Resume the task awaiting the continuation by having it throw an error
222+
/// from its suspension point.
223+
///
224+
/// A continuation must be resumed exactly once. If the continuation has
225+
/// already been resumed through this object, whether by `resume(returning:)`
226+
/// or by `resume(throwing:)`, then the attempt to resume
227+
/// the continuation again will be logged, but otherwise have no effect.
228+
public func resume(throwing x: __owned Error) {
229+
if let c: UnsafeThrowingContinuation<T> = canary.takeThrowingContinuation() {
230+
c.resume(throwing: x)
231+
} else {
232+
logFailedCheck("SWIFT TASK CONTINUATION MISUSE: \(canary.function) tried to resume its continuation more than once, throwing \(x)!\n")
233+
}
234+
}
235+
}
236+
237+
public func withCheckedThrowingContinuation<T>(
238+
function: String = #function,
239+
_ body: (CheckedThrowingContinuation<T>) -> Void
240+
) async throws -> T {
241+
return try await withUnsafeThrowingContinuation {
242+
body(CheckedThrowingContinuation(continuation: $0, function: function))
243+
}
244+
}
245+

‎stdlib/public/Concurrency/Task.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,3 +523,8 @@ void swift::swift_continuation_throwingResumeWithError(/* +1 */ SwiftError *erro
523523
bool swift::swift_task_isCancelled(AsyncTask *task) {
524524
return task->isCancelled();
525525
}
526+
527+
SWIFT_CC(swift)
528+
void swift::swift_continuation_logFailedCheck(const char *message) {
529+
swift_reportError(0, message);
530+
}

‎test/Concurrency/Runtime/checked_continuation.swift

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
// RUN: %target-run-simple-swift(-Xfrontend -enable-experimental-concurrency) 2>&1 | %FileCheck %s
2+
3+
// REQUIRES: executable_test
4+
// REQUIRES: concurrency
5+
6+
import _Concurrency
7+
8+
struct TestError: Error {}
9+
10+
runAsyncAndBlock {
11+
let x: Int = await withCheckedContinuation { c in
12+
c.resume(returning: 17)
13+
c.resume(returning: 38)
14+
}
15+
// CHECK: main tried to resume its continuation more than once, returning 38!
16+
17+
do {
18+
let x: Int = await try withCheckedThrowingContinuation { c in
19+
c.resume(throwing: TestError())
20+
c.resume(returning: 679)
21+
c.resume(throwing: TestError())
22+
}
23+
} catch {
24+
// CHECK-NEXT: main tried to resume its continuation more than once, returning 679!
25+
// CHECK-NEXT: main tried to resume its continuation more than once, throwing TestError()!
26+
}
27+
28+
_ = Task.runDetached {
29+
let _: Void = await withCheckedContinuation { _ in
30+
/*do nothing, leaking the task*/
31+
}
32+
// TODO: Whether the detached task gets a chance to run or not before
33+
// the process exits is currently platform-dependent, and we don't yet
34+
// have the API for yielding to the task runtime.
35+
// C/HECK-NEXT: main leaked its continuation!
36+
}
37+
}

0 commit comments

Comments
 (0)