ylliX - Online Advertising Network
Weak self, a story about memory management and closure in Swift

Weak self, a story about memory management and closure in Swift


Memory management is a big topic in Swift and iOS development. If there are plenty of tutorials explaining when to use weak self with closure, here is a short story when memory leaks can still happen with it.

For the purpose of this blog post, let’s assume we have the following class with two functions. Each function execute something and finalize the execution with a closure executed.

UPDATE – 9 April 2022: I have revisited the examples to highlight when the reference counter is increased and when it can cause memory leaks.

class MyClass {

    func doSomething(_ completion: (() -> Void)?) {
        // do something
        completion?()
    }

    func doSomethingElse(_ completion: (() -> Void)?) {
        // do something else
        completion?()
    }
}

Now, here comes a new requirement, we want a new function doEverything that will execute both doSomething and doSomethingElse in this order. Along the way, we are changing the state of the class to follow the progression.

var didSomething: Bool = false
var didSomethingElse: Bool = false

func doEverything() {

    self.doSomething { 
        self.didSomething = true // <- strong reference to self
        print("did something")

        self.doSomethingElse {
            self.didSomethingElse = true // <- strong reference to self
            print("did something else")
        }
    }
}

Right off the bat, we can see self is strongly captured in the first and second closures: the closures keep a strong reference to self which internally increment the reference counter and can prevent to de-allocate the instance during the execution of doSomething.

It means if those functions were asynchronous and we want to release the instance before it completes the execution, the system would still have to wait to complete it before releasing the memory.

Of course, we know better and setup weak self for the closures:

func doEverything() {

    self.doSomething { [weak self] in 
        self?.didSomething = true
        print("did something")

        self?.doSomethingElse { [weak self] in 
            self?.didSomethingElse = true
            print("did something else")
        }
    }
}

Wait, do we actually need both [weak self] for each closure?

Actually, we don’t.

When we have nested closures like here, we should always set weak self to the first one, the outer closure. The inner closure, the one nested in the outer can reuse the same weak self.

func doEverything() {

    self.doSomething { [weak self] in 
        self?.didSomething = true
        print("did something")

        self?.doSomethingElse { in
            self?.didSomethingElse = true
            print("did something else")
        }
    }
}

However, if we did the other way around, having weak self only in the nested closure, the outer closure would still capture self strongly and increment the reference counter. So be careful where you set this one up.

func doEverything() {

    self.doSomething { in 
        self.didSomething = true // <- strong reference to self
        print("did something")

        self.doSomethingElse { [weak self] in 
            self?.didSomethingElse = true
            print("did something else")
        }
    }
}

So far so good.

Since we want to change other variables along the way, let’s clean the code with a guard let to make sure the instance is still available.

func doEverything() {

    self.doSomething { [weak self] in 
        guard let self = self else { 
            return 
        }
        self.didSomething = true
        print("did something")

        self.doSomethingElse { in
            self.didSomethingElse = true // <-- strong reference?
            print("did something else")
        }
    }
}

But now, here comes the question: since we have a strong reference called self in the outer closure, does the inner closure strongly capture it? How can we verify this?

This is the kind of questions that is worth diving into and Xcode Playground is perfect for this. I’ll include few logs to keep track of the steps as well as logging the reference counter.

For the first example, let’s keep it simple, so we can see how the reference counter is incremented along the way.

class MyClass {

    func doSomething(_ completion: (() -> Void)?) {
        // do something
        completion?()
    }

    func doSomethingElse(_ completion: (() -> Void)?) {
        // do something else
        completion?()
    }

    var didSomething: Bool = false
    var didSomethingElse: Bool = false

    deinit {
        print("Deinit")
    }

    func printCounter() {
        print(CFGetRetainCount(self))
    }

    func doEverything() {
        print("start")
        printCounter()
        self.doSomething {
            self.didSomething = true
            print("did something")
            self.printCounter()

            self.doSomethingElse {
                self.didSomethingElse = true
                print("did something else")
                self.printCounter()
            }
        }
        printCounter()
    }
}

do {
    let model = MyClass()
    model.doEverything()
}

Here is the output

# output
start
2
did something
4
did something else
6
2
Deinit

With only strong references to self, we can see the counter going up to 6. However, as expected, once both functions are executed, the instance is de-allocated.

Now let’s introduce weak self in the outer closure.

func doEverything() {
    print("start")
    printCounter()
    self.doSomething { [weak self] in
        self?.didSomething = true
        print("did something")
        self?.printCounter()

        self?.doSomethingElse {
            self?.didSomethingElse = true
            print("did something else")
            self?.printCounter()
        }
    }
    printCounter()
}

With the first weak self, the instance is still de-allocated, and the counter goes only up to 4.

# output
start
2
did something
3
did something else
4
2
Deinit

So what happens with guard let self?

func doEverything() {
    print("start")
    printCounter()
    self.doSomething { [weak self] in
        guard let self = self else { return }
        self.didSomething = true
        print("did something")
        self.printCounter()

        self.doSomethingElse {
            self.didSomethingElse = true
            print("did something else")
            self.printCounter()
        }
    }
    printCounter()
}

Here is the output

# output
start
2
did something
3
did something else
5
2
Deinit

If the instance is successfully de-initialized, we can see that the counter is actually increased from 4 to 5 when we execute doSomethingElse, which means the nester closure capture strongly our temporary self.

It looks already suspicious, but let’s try with a different example. What if, instead of functions, doSomething and doSomethingElse are closure properties of the class. Let’s adapt the code for a similar execution.

class MyClass {

    var doSomething: (() -> Void)?
    var doSomethingElse: (() -> Void)?

    var didSomething: Bool = false
    var didSomethingElse: Bool = false

    deinit {
        print("Deinit")
    }

    func printCounter() {
        print(CFGetRetainCount(self))
    }

    func doEverything() {

        print("start")
        printCounter()
        doSomething = { [weak self] in
            guard let self = self else { return }
            self.didSomething = true
            print("did something")
            self.printCounter()

            self.doSomethingElse = {
                self.didSomethingElse = true
                print("did something else")
                self.printCounter()
            }

            self.doSomethingElse?()
        }
        doSomething?()
        printCounter()
    }
}

do {
    let model = MyClass()
    model.doEverything()
}

Here is the output

# output
start
2
did something
3
did something else
5
3

This time, the class is not even de-allocated 🤯. To fix it back, we have to keep a weak instance.

func doEverything() {

    print("start")
    printCounter()
    doSomething = { [weak self] in
        self?.didSomething = true
        print("did something")
        self?.printCounter()

        self?.doSomethingElse = {
            self?.didSomethingElse = true
            print("did something else")
            self?.printCounter()
        }

        self?.doSomethingElse?()
    }
    doSomething?()
    printCounter()
}
# output
start
2
did something
3
did something else
3
2
Deinit

Yay, the instance is properly de-allocated this time. It’s confirmed, the inner closure was making a strong reference to guard let self.


So, what does it mean for my code?

When we face a closure, we tend to write weak self followed by a guard let to quickly go around without thinking too much about the execution further down. This is where we still need to be careful. It’s easy to miss this kind of memory leaks, so here are few takeaways:

First, regarding the format, I personally use guard let strongSelf in closure instead of guard let self. The reason is that during code review, it can be tricky to know which “self” we are referring to further down in the code.

Second, if there is nested closure, I would prefer keeping reference to the weak (and optional) self? and never point back to strongSelf, so I have the insurance to avoid any strong reference to it.

func doEverything() {
    doSomething = { [weak self] in
        guard let strongSelf = self else { return }
        strongSelf.didSomething = true
        print("did something")

        strongSelf.doSomethingElse = {
            self?.didSomethingElse = true
            print("did something else")
        }

        strongSelf.doSomethingElse?()
    }
    doSomething?()
}

Last but not least, if we have too many closures to handles, the best is still to refactor it as a separate function, or use newer API to avoid those mistakes. I’m thinking about functional reactive programming like RxSwift or Combine, but also swift Async to get by.

Of course, the code shared today might be a bit farfetched and might not reflect your daily usage of closures, but in my opinion, it’s still important to keep in mind the memory management and reference we make of our instances. Also, this question came right into the middle of a peer review, so we never too careful 😉

Hope you enjoyed this post, happy coding!



Source link

Leave a Reply

Your email address will not be published. Required fields are marked *