How to use withTaskCancellationHandler properly?

What does this look like to people here? A goal is to try to cut down on thread context switching for the purposes of wrapping any cancellable API in a non-async/await context.

import Foundation

private class CancellableCheckedThrowingContinuation<T> {

    private enum State {
        case ready
        case executing(cancel: () -> Void, continuation: CheckedContinuation<T, Error>)
        case cancelled
    }

    private var state: State = .ready
    private var lock: UnsafeMutablePointer<os_unfair_lock>

    init() {
        lock = UnsafeMutablePointer<os_unfair_lock>.allocate(capacity: 1)
        lock.initialize(to: os_unfair_lock())
    }

    deinit {
        lock.deallocate()
    }

    func operation(_ body: @escaping (CheckedContinuation<T, Error>) -> () -> Void) async throws -> T {
        try Task.checkCancellation()
        return try await withTaskCancellationHandler(operation: {
            try Task.checkCancellation()
            return try await withCheckedThrowingContinuation({ continuation in
                let cancel = body(continuation)
                os_unfair_lock_lock(lock)
                if case .ready = state {
                    state = .executing(cancel: cancel, continuation: continuation)
                    os_unfair_lock_unlock(lock)
                } else {
                    os_unfair_lock_unlock(lock)
                    cancel()
                    continuation.resume(throwing: CancellationError())
                }
            })
        }, onCancel: {
            os_unfair_lock_lock(lock)
            if case .executing(let cancel, let continuation) = state {
                state = .cancelled
                os_unfair_lock_unlock(lock)
                cancel()
                continuation.resume(throwing: CancellationError())
            } else {
                state = .cancelled
                os_unfair_lock_unlock(lock)
            }
        })
    }
}

public func withCancellableCheckedThrowingContinuation<T>(_ body: @escaping (CheckedContinuation<T, Error>) -> () -> Void) async throws -> T {
    try Task.checkCancellation()
    return try await CancellableCheckedThrowingContinuation().operation(body)
}

It seems to be working, but I question a few things.

  1. Is that many calls to try Task.checkCancellation() actually worth anything?
  2. Could I actually safely invoke the body function within the lock?
    Something like this?
os_unfair_lock_lock(lock)
if case .ready = state {
    state = .executing(cancel: body(continuation), continuation: continuation)
    os_unfair_lock_unlock(lock)
} else {
    os_unfair_lock_unlock(lock)
    continuation.resume(throwing: CancellationError())
}

Usage:

class Something {
    init(_ completion: (Result<Data, Error>) -> Void) {
        // ... non-async/await stuff
    }
    func cancel() {
        // ... non-async/await stuff
    }
}

public func doSomething() async throws -> Data {
    return try await withCancellableCheckedThrowingContinuation { continuation in
        let something = Something { result in
            switch result {
            case .success(let data):
                continuation.resume(returning: data)
            case .failure(let error):
                continuation.resume(throwing: error)
            }
        }
        return something.cancel
    }
}
2 Likes