withoutActuallyEscaping + async functions

Using Swift
1

Is the behavior of the following code undefined? A strict interpretation of the documentation of withoutActuallyEscaping would conclude that it is. However, it's an interesting question because neither block nor escapingBlock escape offload.

func offload<R>(to queue: DispatchQueue, _ block: () -> R) async -> R {
    await withUnsafeContinuation { continuation in
        withoutActuallyEscaping(block) { escapingBlock in
            queue.async {
                continuation.resume(returning: escapingBlock())
            }
        }
    }
}

Thanks.

2

Yes, this is unsafe. queue.async absolutely does escape its closure argument, and so block is incorrectly escaped here.

2 Likes
3

To elaborate on @lukasa's answer, the rule listed in the docs for withoutActuallyEscaping(_:do:) is:

The escapable copy of closure passed to body is only valid during the call to withoutActuallyEscaping(_:do:). It is undefined behavior for the escapable closure to be stored, referenced, or executed after the function returns.

This differs subtly from what you've said above:

It's clearly the case (as you've noted) that escapingBlock does indeed escape the call to withoutActuallyEscaping, and if you actually attempt to use the code you've written, you'll get a runtime error thanks to checking that withoutActuallyEscaping does.

% ./test
0    test                               0x00000001001ff08c closure #1 in offload<A>(to:_:) + 188
1    test                               0x00000001001ff8b4 closure #1 in withUnsafeContinuation<A>(_:) + 64
2    test                               0x00000001001ff578 withUnsafeContinuation<A>(_:) + 172
3    libswift_Concurrency.dylib         0x000000023785a34c swift::runJobInEstablishedExecutorContext(swift::Job*) + 376
4    libswift_Concurrency.dylib         0x000000023785b3c0 swift_job_runImpl(swift::Job*, swift::ExecutorRef) + 72
5    libdispatch.dylib                  0x00000001ad49fe08 _dispatch_root_queue_drain + 396
6    libdispatch.dylib                  0x00000001ad4a071c _dispatch_worker_thread2 + 164
7    libsystem_pthread.dylib            0x00000001ad611fe0 _pthread_wqthread + 228
8    libsystem_pthread.dylib            0x00000001ad610e18 start_wqthread + 8
closure argument was escaped in withoutActuallyEscaping block: file test.swift, line 5, column 9 
zsh: trace trap  ./test

I believe the correct way to perform this is to invert the calls to withoutActuallyEscaping and withUnsafeContinuation, which guarantees that we won't actually return from withoutActuallyEscaping until after the closure passed to queue.async has completed execution:

func offload<R>(to queue: DispatchQueue, _ block: () -> R) async -> R {
    await withoutActuallyEscaping(block) { escapingBlock in
        await withUnsafeContinuation { continuation in
            queue.async {
                continuation.resume(returning: escapingBlock())
            }
        }
    }
}
2 Likes
4

It is remarkable that inverting the two with* calls compiles! withoutActuallyEscaping is declared as a synchronous function that takes a synchronous closure:

@_transparent
@_semantics("typechecker.withoutActuallyEscaping(_:do:)")
public func withoutActuallyEscaping<ClosureType, ResultType>(
  _ closure: ClosureType,
  do body: (_ escapingClosure: ClosureType) throws -> ResultType
) rethrows -> ResultType {
  // This implementation is never used, since calls to
  // `Swift.withoutActuallyEscaping(_:do:)` are resolved as a special case by
  // the type checker.
  Builtin.staticReport(_trueAfterDiagnostics(), true._value,
    ("internal consistency error: 'withoutActuallyEscaping(_:do:)' operation failed to resolve"
     as StaticString).utf8Start._rawValue)
  Builtin.unreachable()
}

It appears that the compiler's special handling of this function also allows async usage. However, I do wonder if this is an accident and it's unsafe to use it this way.

1 Like
5

What about with structured concurrency?

func offload(block: () async -> Void) async {
  await withoutActuallyEscaping(block) { escapingBlock in
    await withTaskGroup(of: Void.self) { group in
      group.addTask { await escapingBlock() }
      await group.waitForAll()
    }
  }
}

TaskGroup.addTask is escaping but is that an implementation detail? From the surface it’s not clear how it could escape.

SE-0304 noted the decision to use @_implicitSelfCapture which achieves my end goal but it’s a mouthful.

6

I've since concluded that inverting the calls is still unsafe. escapingBlock isn't guaranteed to be released by the time the task resumes and swift_isEscapingClosureAtFileLocation is called. The consume operator in Swift 5.9 may help here, but only if it guarantees the immediate release of escapingBlock.

In the mean time, I've created a safe version of this function using dispatch_async_f and locally allocated contexts, which this margin is too narrow to contain.

1 Like
7

Interesting, are you able to determine what's still holding onto escapingBlock by the time the error is triggered?

8

The error never triggered for me, but it's possible. To encourage the error to trigger, as a demonstration, call sleep right after resuming the continuation.

1 Like
9

Ah, I suppose because the task awaiting the continuation.resume call is permitted to resume (and then exit the withoutActuallyEscaping body) before the block passed to queue.async actually returns, especially when there's nontrivial work following the call to continuation.resume(returning:).

Interestingly, I can reproduce the failure when building in debug mode but it appears that when building with -O the compiler is able to determine that the last use of escapingBlock must occur before we exit withoutActuallyEscaping, and so omits any retains/releases for escapingBlock, averting the crash. This does seem like something that consume would help guarantee though would love to hear @Joe_Groff's thoughts on that.

10

I don't think the consume operator will release closure captures. We need to prevent the queue.async closure from strongly capturing escapingBlock, by simulating an unowned(unsafe) capture via unsafeBitCasts to and from (Int, Int).

1 Like
11

Yeah, in the general case the compiler would have to assume that the closure could be called again and so the captures might be needed, but it appears that it's possible for the optimizer to eliminate the lifetime extension for the captures of queue.async here, so perhaps it has some special knowledge about how DispatchQueue behaves?

12

I've been using something similar: https://github.com/troughton/Substrate/blob/968a0ea06e42c34c1b6cee1b93f61e9243a8e7b8/Sources/SubstrateUtilities/TaskStream.swift#L44. The key is to store the closure in a class instance and use unowned and withExtendedLifetime on that class instance, e.g.:

@usableFromInline final class ClosureHolder<R> {
    @usableFromInline let closure: @Sendable () -> R
    
    @inlinable
    init(closure: @escaping @Sendable  ()  -> R) {
        self.closure = closure
    }
}

extension DispatchQueue {
    @inlinable @_unsafeInheritExecutor
    public func offload<T>(@_implicitSelfCapture _ perform: @Sendable () -> T) async -> T {
        let taskStreamContinuation = self.taskStreamContinuation
        return await withoutActuallyEscaping(perform) { perform in
            let closureHolder = ClosureHolder<T>(closure: perform)
            let result: T = await withUnsafeContinuation { continuation in
                self.async { [unowned closureHolder, continuation] in
                    let result = closureHolder.closure()
                    continuation.resume(returning: result)
                }
            }
            withExtendedLifetime(closureHolder) {}
            return result
        }
    }
}
2 Likes
13

I realize I'm reviving an old topic, but I wanted to shout out some thanks to @Torust for this non-obvious solution.

I'll admit I don't fully understand the dynamics at play, but (to enable a RunLoop-backed actor executor) I am sending the block from withoutActuallyEscaping to another thread via an escaping block and then using locks to guarantee the requirement of withoutActuallyEscaping is met, but was still running into occasional closure argument was escaped in withoutActuallyEscaping - wrapping the non-escaping block with the ClosureHolder, passing the reference unowned, and extending the lifetime manually addresses the issue, presumably caused by the captured block getting released late otherwise.

I will say that I am surprised by this behaviour, especially in light of Apple's example usage in docs (though I suppose here escapableF isn't captured by an async block):

///     func perform(_ f: () -> Void, simultaneouslyWith g: () -> Void) {
///         withoutActuallyEscaping(f) { escapableF in
///             withoutActuallyEscaping(g) { escapableG in
///                 let queue = DispatchQueue(label: "perform", attributes: .concurrent)
///                 queue.async(execute: escapableF)
///                 queue.async(execute: escapableG)
///                 queue.sync(flags: .barrier) {}
///             }
///         }
///     }

(as hinted at by @Jumhyn)

Ideally, this case ought to be made to succeed trivially by the compiler/language and not require special wrangling or knowledge of internals.

14

For some reason, even after using ClosureHolder + unowned +withExtendedLifetime, I still see the error occasionally.

Currently I'm trying something like this:

final class ClosureHolder<R>: @unchecked Sendable {
    @usableFromInline var closure: (() throws -> R)?
    init(closure: (() throws -> R)?) {
        self.closure = closure
    }
}

struct SendableWrapper<T>: @unchecked Sendable {
    @usableFromInline var value: T
    init(value: T) {
        self.value = value
    }
}

func launch1<R>(_ closure: () throws -> R) async throws -> R {

    try await withoutActuallyEscaping(closure) { escapingClosure in

        let holderPtr = UnsafeMutablePointer<ClosureHolder<R>>.allocate(capacity: 1)
        holderPtr.initialize(to: .init(closure: nil))
        holderPtr.pointee.closure = escapingClosure

        let wrapped = SendableWrapper(value: holderPtr)
    
        let result = try await withCheckedThrowingContinuation { continuation in

            DispatchQueue.global().async {
                do {
                    let result = try wrapped.value.pointee.closure!()
                    wrapped.value.deinitialize(count: 1)
                    wrapped.value.deallocate()
                    continuation.resume(returning: result)
                } catch {
                    wrapped.value.deinitialize(count: 1)
                    wrapped.value.deallocate()
                    continuation.resume(throwing: error)
                }
            }

        }

        return result

    }

}

Not sure whether this guarantees eliminating the problem, but currently it works fine after hundreds of trials