URLSession Implicit Cancellation using Async/Await Helper

We have a wrapper around a custom URLSession instance that is supposed to be deallocated in certain situations and then implicitly cancel all running network requests.

When starting a request using async/await however, this instance gets retained until the request finishes.

We start our requests like this:

// Somewhere in unstructured UIKit
Task { [weak self] in
    do {
        ...
        let thing = try await self?.urlSessionWrapper.loadSomething()
        ...
    } catch { ... }
}

With an implementation that looks like this:

class URLSessionWrapper {
    let urlSession: URLSession

    deinit {
        urlSession.invalidateAndCancel()
    }
}
...
extension URLSessionWrapper {
    func loadSomething() async throws -> Decoded {
        let resp = try await urlSession.data(from: URL(string: "https://httpbin.io/delay/3")!)
        // decode
        return decoded
    }
}

After being very confused and annoyed, I experimented with this for a few days until I finally found out that when moving loadSomething from URLSessionWrapper to URLSession everything suddenly works fine: The URLSessionWrapper is immediately deallocated and all running requests are cancelled as expected.

Does someone have an explanation for this? Could this be a bug in the concurrency implementation, or URLSession, or am I just doing something wrong?

I have created a sample project at GitHub - iteracticman/AsyncLeak to make it easy for everyone to reproduce this. With this UIKit app you can easily start multiple requests using either URLSessionWrapper or URLSession from a UIViewController. When the UIViewController gets deallocated while requests are in progress, the URLSessionWrapper gets deallocated immediately and all requests cancelled only when all requests were started via urlSessionWrapper.urlSession.loadSomething(), but not when one was started via urlSessionWrapper.loadSomething().

1 Like

i believe the reason this occurs is that, in the formulation where the wrapper is used, the wrapper cannot be deallocated until any awaited calls to loadSomething() have returned. if this were not the case, it would presumably break a lot of expected (and desirable) behaviors (e.g. self could become nil across a suspension point within an instance method). things work out in the desired manner when loading via the URLSession instance because the wrapper can be deallocated in that case since it doesn't have any async methods still executing that are keeping its retain count above zero.

2 Likes

Right. Ditto for normal (sync) calls:

var a: A? = A()

class A {
    deinit {
        print("A deinit")
    }
    func foo() {
        a = nil
        print("exiting foo")
    }
}

a!.foo()
print("done")

produces:

exiting foo
A deinit
done

Note that "A deinit" happens after returning from foo, for the duration of "foo" call the class is kept retained.

2 Likes

Sorry, I do not understand this.

When I have an implementation that looks like this:

extension URLSessionWrapper {
    func loadSomething() async throws -> Data {
        try await self.urlSession.loadSomething()
    }
}

extension URLSession {
    func loadSomething() async throws -> Data {
        try await self.data(from: URL(string: "https://httpbin.io/delay/3")!).0
    }
}

And then call

try await self?.urlSessionWrapper.urlSession.loadSomething()

and everything works as expected, why exactly does calling

try await self?.urlSessionWrapper.loadSomething()

instead behave differently?

Is this documented somewhere? Am I the only one who finds this unintuitive and surprising?

And let's suppose this is perfectly fine: Shouldn't there be a way to work around this?

The first:

is equivalent to:

if let urlSession = self?.urlSessionWrapper.urlSession {
    try await urlSession.loadSomething()
}

Note that the "wrapper" object is not used (doesn't need to be retained) while the control is on the second line → so it will be deallocated → so it's deinit will run → so it will call "invalidateAndCancel", the thing you want.

The second:

is equivalent to:

if let urlSessionWrapper = self?.urlSessionWrapper {
    try await urlSessionWrapper.loadSomething()
}

In this one the wrapper object is being used for the whole duration of wrapper's "loadSomething" call → so it has to be retained → so it won't be deallocated → so it's deinit won't be called → so it won't call "invalidateAndCancel".

One workaround that immediately springs to mind is to use a closure in your wrapper object - a closure that will fire upon the internal async call completion.

self?.urlSessionWrapper.loadSomething { [weak self] result, error in
    // TODO
}

Note that in this case the wrapper's loadSomething call is neither throwing nor async. You can still use a throwing async for the URLSession callSomething call, although you'd need to cross the async → sync bridge to bring the result back to the completion closure.

I don't immediately see a workaround for the throwing+async wrapper's version of your code, hopefully others can suggest it.

Edit: another workaround is to move "invalidateAndCancel" from deinit to some explicit "invalidate" call and call it appropriately (e.g. from viewController didDisappear or something).

Yet another workaround is to not expose "loadSomething" on the wrapper object. Potentially even removing the wrapper object and implementing your "wrapper" functionality in an extension methods of URLSession. Or use a subclass of URLSession (this is good if you need to keep an extra state in your wrapper).

2 Likes

Thank you for the detailed explanation @tera. I understand the mechanism now.

It makes sense, but it's not exactly obvious and could maybe also have been implemented very differently.

I spent a few days investigating this and think proper documentation would have saved me a lot of time and nerves.

Regarding the workaround: I will add this explicit "invalidate" method as you suggested. Thanks again!

Glad to be of help and glad you asked: your example highlights a quite important difference between closure based code and seemingly "equivalent" async/await code. Distilling your example to get rid URLSession and making it standalone:

class A {
    static func run() {
        var a: A! = A()
        DispatchQueue.main.asyncAfter(deadline: .now() + 1) {
            a = nil
        }
        a.foo()
        RunLoop.current.run(until: .distantFuture)
    }

    deinit {
        print("A deinit")
    }
    
    private func load_foo(execute: @escaping (Int?, Error?) -> Void) {
        DispatchQueue.main.asyncAfter(deadline: .now() + 10) {
            execute(42, nil)
        }
    }
    func foo() {
        print("foo started")
        load_foo { [weak self] result, error in
            guard self != nil else {
                print("no more self, returning")
                return
            }
            if let result { print("result:", result) }
            else { print("error:", error!) }
        }
    }
}
A.run()

Then you want to modernise it and switch from closures to async await:

class A {
    static func run() {
        ...
        a.bar() // was foo()
    }

    deinit {
        print("A deinit")
        continuation?.resume(throwing: NSError(domain: "cancelled", code: -128))
        continuation = nil
    }
    ... get rid of foo's

    private var continuation: CheckedContinuation<Int, Error>?
    
    private func load_bar() async throws -> Int {
        try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation<Int, Error>) in
            guard let self else {
                print("no more self (1), returning")
                return
            }
            self.continuation = continuation
            DispatchQueue.main.asyncAfter(deadline: .now() + 10) { [weak self] in
                guard let self else {
                    print("no more self (2), returning")
                    return
                }
                self.continuation = nil
                continuation.resume(returning: 42)
            }
        }
    }
    func bar() {
        print("bar started")
        Task {
            do {
                let result = try await load_bar()
                print("result:", result)
            } catch {
                print("error: \(error)")
            }
        }
    }
}

The first closure-based "foo" version gives this output:

foo started
... 1 second later:
A deinit
... 9 seconds later:
no more self, returning

The seemingly equivalent async/await "bar" version gives a totally different output:

bar started
... 10 seconds later:
result: 42
A deinit

Beware.

2 Likes

Thanks for that great example.

Using withCheckedThrowingContinuation with a block that tries to weakly capture self illustrates the issue much better.