Can we define an async/await function that returns one value instantly, as well as another value asynchronously (as you'd normally expect)?

Question

Picture an image loading function with a closure completion. Let's say it returns a token ID that you can use to cancel the asynchronous operation if needed.

@discardableResult
func loadImage(url: URL, completion: @escaping (Result<UIImage, Error>) -> Void) -> UUID? {
    
    if let image = loadedImages[url] {
        completion(.success(image))
        return nil
    }
    
    let id = UUID()
    
    let task = URLSession.shared.dataTask(with: url) { data, response, error in
        
        defer {
            self.requests.removeValue(forKey: id)
        }
        
        if let data, let image = UIImage(data: data) {
            DispatchQueue.main.async {
                self.loadedImages[url] = image
                completion(.success(image))
            }
            return
        }
        
        if let error = error as? NSError, error.code == NSURLErrorCancelled {
            return
        }
        
        //TODO: Handle response errors
        print(response as Any)
        completion(.failure(.loadingError))
    }
    task.resume()
    
    requests[id] = task
    return id
}

func cancelRequest(id: UUID) {
    requests[id]?.cancel()
    requests.removeValue(forKey: id)
    
    print("ImageLoader: cancelling request")
}

How would we accomplish this (elegantly) with swift concurrency? Is it even possible or practical?

Answer 1

I haven't done much testing on this, but I believe this is what you're looking for. It allows you to simply await an image load, but you can cancel using the URL from somewhere else. It also merges near-simultaneous requests for the same URL so you don't re-download something you're in the middle of.

actor Loader {
    private var tasks: [URL: Task<UIImage, Error>] = [:]

    func loadImage(url: URL) async throws -> UIImage {
        if let imageTask = tasks[url] {
            return try await imageTask.value
        }

        let task = Task {
            // Rather than removing here, you could skip that and this would become a
            // cache of results. Making that correct would take more work than the
            // question asks for, so I won't go into it
            defer { tasks.removeValue(forKey: url) }

            let data = try await URLSession.shared.data(from: url).0
            guard let image = UIImage(data: data) else {
                throw DecodingError.dataCorrupted(.init(codingPath: [],
                                                        debugDescription: "Invalid image"))
            }
            return image
        }
        tasks[url] = task

        return try await task.value
    }

    func cancelRequest(url: URL) {
        // Remove, and cancel if it's removed
        tasks.removeValue(forKey: url)?.cancel()
        print("ImageLoader: cancelling request")
    }
}

Calling it looks like:

let image = try await loader.loadImage(url: url)

And you can cancel a request if it's still pending using:

loader.cancelRequest(url: url)

A key lesson here is that it is very natural to access task.value multiple times. If the task has already completed, then it will just return immediately.

Answer 2

Return a task in a tuple or other structure.

In the cases where you don't care about the ID, do this:

try await imageTask(url: url).task.value

private var requests: [UUID: Task<UIImage, Swift.Error>] = [:]

func imageTask(url: URL) -> (id: UUID?, task: Task<UIImage, Swift.Error>) {
  switch loadedImages[url] {
  case let image?: return (id: nil, task: .init { image } )
  case nil:
    let id = UUID()
    let task = Task {
      defer { requests[id] = nil }

      guard let image = UIImage(data: try await URLSession.shared.data(from: url).0)
      else { throw Error.loadingError }

      try Task.checkCancellation()
      Task { @MainActor in loadedImages[url] = image }
      return image
    }

    requests[id] = task
    return (id: id, task: task)
  }
}

Answer 3

Is it even possible or practical?

Yes to both.

As I say in a comment, I think you may be missing the fact that a Task is an object you can retain and later cancel. Thus, if you create an architecture where you apply an ID to a task as you ask for the task to start, you can use that same ID to cancel that task before it has returned.

Here's a simple demonstration. I've deliberately written it as Playground code (though I actually developed it in an iOS project).

First, here is a general TimeConsumer class that wraps a single time-consuming Task. We can ask for the task to be created and started, but because we retain the task, we can also cancel it midstream. It happens that my task doesn't return a value, but that's neither here nor there; it could if we wanted.

class TimeConsumer {
    var current: Task<(), Error>?
    func consume(seconds: Int) async throws {
        let task = Task {
            try await Task.sleep(for: .seconds(seconds))
        }
        current = task
        _ = await task.result
    }
    func cancel() {
        current?.cancel()
    }
}

Now then. In front of my TimeConsumer I'll put a TaskVendor actor. A TimeConsumer represents just one time-consuming task, but a TaskVendor has the ability to maintain multiple time-consuming tasks, identifying each task with an identifier.

actor TaskVendor {
    private var tasks = [UUID: TimeConsumer?]()
    func giveMeATokenPlease() -> UUID {
        let uuid = UUID()
        tasks[uuid] = nil
        return uuid
    }
    func beginTheTask(uuid: UUID) async throws {
        let consumer = TimeConsumer()
        tasks[uuid] = consumer
        try await consumer.consume(seconds: 10)
        tasks[uuid] = nil
    }
    func cancel(uuid: UUID) {
        tasks[uuid]??.cancel()
        tasks[uuid] = nil
    }
}

That's all there is to it! Observe how TaskVendor is configured. I can do three things: I can ask for a token (really my actual TaskVendor needn't bother doing this, but I wanted to centralize everything for generality); I can start the task with that token; and, optionally, I can cancel the task with that token.

So here's a simple test harness. Here we go!

let vendor = TaskVendor()
func test() async throws {
    let uuid = await vendor.giveMeATokenPlease()
    print("start")
    Task {
        try await Task.sleep(for: .seconds(2))
        print("cancel?")
        // await vendor.cancel(uuid: uuid)
    }
    try await vendor.beginTheTask(uuid: uuid)
    print("finish")
}
Task {
    try await test()
}

What you will see in the console is:

start
[two seconds later] cancel?
[eight seconds after that] finish

We didn't cancel anything; the word "cancel?" signals the place where our test might cancel, but we didn't, because I wanted to prove to you that this is working as we expect: it takes a total of 10 seconds between "start" and "finish", so sure enough, we are consuming the expected time fully.

Now uncomment the await vendor.cancel line. What you will see now is:

start
[two seconds later] cancel?
[immediately!] finish

We did it! We made a cancellable task vendor.

Answer 4

I'm including one possible answer to the question, for the benefit of others. I'll leave the question in place in case someone has another take on it.

The only way that I know of having a 'one-shot' async method that would return a token before returning the async result is by adding an inout argument:

func loadImage(url: URL, token: inout UUID?) async -> Result<UIImage, Error> {
    
    token = UUID()
    
    //...
}

Which we may call like this:

var requestToken: UUID? = nil
let result = await imageLoader.loadImage(url: url, token: &requestToken)

However, this approach and the two-shot solution by @matt both seem fussy, from the api design standpoint. Of course, as he suggests, this leads to a bigger question: How do we implement cancellation with swift concurrency (ideally without too much overhead)? And indeed, using tasks and wrapper objects seems unavoidable, but it certainly seems like a lot of legwork for a fairly simple pattern.