Cooperative task cancellation and NSManagedObjectContext

I am trying to understand how (or if) cooperative task cancellation behavior can interact with a closure being executed on the queue associated with an instance of NSManagedObjectContext. I crafted a relatively simple test case that captures what I thought would work. On both iOS 14 and 15, however, Task.isCancelled always returns false in the closure supplied to either perform(_:) or perform(schedule:_:).

Have I misunderstood something here, or is checking task cancellation state via Task.isCancelled not available in this circumstance? I have crafted a workaround for this scenario using withTaskCancellationHandler(operation:onCancel:) that basically mimics Task.isCancelled and Task.checkCancellation(), but it's not as nice.

import XCTest
@testable import TaskCancellation
import CoreData

class TaskCancellationTests: XCTestCase {
    func testCancellation() {
        let expectation = expectation(description: "...")
        let task = Task {
            do {
                let context = PersistenceController.shared.container.newBackgroundContext()
                try await doWork(context: context)
            }
            catch is CancellationError {
            }
            catch {
                XCTFail("Unexpected error: \(error)")
            }

            expectation.fulfill()
        }

        task.cancel()
        waitForExpectations(timeout: 2)
    }

    @discardableResult
    private func doWork(context: NSManagedObjectContext) async throws -> Bool {
        // Artificial delay to ensure that the parent task is cancelled before trying to use the NSManagedObjectContext
        Thread.sleep(forTimeInterval: 0.5)

        XCTAssertTrue(Task.isCancelled)
        await anotherMethod()

        return try await context.performAsync {
            if Task.isCancelled {
                return true
            }
            XCTFail("Should not get here")
            throw NSError(domain: "Test", code: 0, userInfo: [NSLocalizedDescriptionKey: "Task.isCancelled returned false"])
        }
    }

    private func anotherMethod() async {
        XCTAssertTrue(Task.isCancelled)
    }
}

extension NSManagedObjectContext {
    func performAsync<T>(_ block: @escaping () throws -> T) async throws -> T {
        if #available(iOS 15.0, *) {
            return try await perform(block)
        }
        else {
            return try await withCheckedThrowingContinuation { continuation in
                perform {
                    do {
                        let result = try block()
                        continuation.resume(returning: result)
                    }
                    catch {
                        continuation.resume(throwing: error)
                    }
                }
            }
        }
    }
}

Since the block handed to performAsync is @escaping it isn't guaranteed to be executed in the same task context as performAsync or doWork, and so it can lose its ability to see if the parent task was cancelled.

In particular, if NSManagedObjectContext.perform does anything that is non-synchronous with block (which it certainly does) such as dispatching to a queue, then it will lose the task context.

Here's a simplified code sample to show the behavior:

func perform(block: @escaping () -> Void) async {
  // Wait enough time for cancellation to happen
  try? await Task.sleep(nanoseconds: 500_000_000)
  print("outer", "Task.isCancelled", Task.isCancelled)

  // Does not lose cancellation context
  block()

  // Loses cancellation context
  Task {
    block()
  }

  // Loses cancellation context
  DispatchQueue.main.async {
    block()
  }
}

let task = Task {
  await perform {
    print("inner", "Task.isCancelled", Task.isCancelled)
  }
}

Task {
  // Wait a little bit of time before cancelling
  try await Task.sleep(nanoseconds: 100_000_000)
  task.cancel()
  print("!!!")
}

As you can see, if perform invokes its block in another task or on a dispatch queue then the block no longer can check if the task was cancelled.

If the block argument was not @escaping then you could guarantee that it will be able to see the cancellation.

2 Likes

Thanks for the explanation. I was wondering about the case of dispatching to another queue specifically, and your code captures the situation quite clearly.