Swift - Fighting memory leaks in closures

Swift is a great language. In fact, it's among my top 3 favorite languages. It is quite strict in terms of discipline and best practices which pay off in the long-term to offer performance and crash-free experience to end-users. There are also many ways where things can go wrong.

Today we are going to discuss one such area - Memory leaks.

What is the memory leak?

According to Wikipedia article,

In computer science, a memory leak is a type of resource leak that occurs when a computer program incorrectly manages memory allocations in such a way that memory which is no longer needed is not released

In iOS, when an application requests a memory resource but does not relinquish it when the task is done, a memory leak happens. Today we will look at similar scenarios with examples of code that cause a memory leak and code that looks like causing a leak but is actually not a source of the leak.

1. Closures capturing self instance

Consider following class named Sample. This class has a property named closure of type (() -> Void)? which does not accept or return any parameters. It also has a property named value of type Int

class Sample: UIViewController {

    var closure: (() -> Void)?
    var value: Int = 0
    ....
    ..
}

Now assume we write a code to define this closure which captures self within itself

closure = {
    print(self.value)
}

Is this a memory leak?

Yes, it is. Why? Because here's the retain cycle which prevents iOS from releasing this viewController when view goes out of scope. Here's how retain cycle looks like,

self owns closure and closure owns self

self -> closure
and
closure -> self

How can you fix this? Simple!

Just weakify self inside the closure and use optional value instead and retain cycle will be broken.

closure = { [weak self] in
    print(self?.value)
}

If you don't like optional values, you can also use optional chaining before starting using optional self

closure = { [weak self] in
    guard let strongSelf = self else { return }
    print(strongSelf.value)
}

After using [weak self] you've conveniently broken the existing retain cycle. However, be aware while using unowned instead of weak. [unowned self] will assume that self will always be a non-nil. If you use [unowned self] to break retain cycle and self turns out to be nil, application will crash

2. Capturing self inside closure owned by cell

Consider a slightly complex example with a tableViewCell. Similar to our previous example let's assume our viewController looks like this,

class Sample: UIViewController {
    var value: Int = 0
    let tableView = UITableView()
    
    override func viewDidLoad() {
        super.viewDidLoad()
        self.view.backgroundColor = .blue
        tableView.delegate = self
        tableView.dataSource = self
        tableView.register(SampleTableViewCell.self, forCellReuseIdentifier: "cell")
        self.view.addSubview(tableView)
    }    
}

extension Sample: UITableViewDataSource, UITableViewDelegate {
    func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
        return 1
    }

    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        let cell = tableView.dequeueReusableCell(withIdentifier: "cell", for: indexPath) as! SampleTableViewCell
        cell.myClosure = {
            print(self.value)
        }
        return cell
    }
}

And this is how the UITableViewCell subclass SampleTableViewCell looks like which also has an empty closure as a property

class SampleTableViewCell: UITableViewCell {

    var myClosure: (() -> Void)?
}

Above code in the ViewController named Sample is taking a cell and defining a closure. However, while defining a closure it's also causing a retain cycle.

How?

As we know tableView is a property on the Sample. So self, i.e. Sample instance owns the tableView which in turn owns SampleTableViewCell and cell owns a block. Now as it's clear from the code in tableView data source method cellForRowAt indexPath: IndexPath) that the closure owned by a cell also captures self which eventually causes a memory leak.

self -> tableView -> SampleTableViewCell -> myClosure -> self

How to break this retain cycle? Simple. We can do the same thing as we did in the first example. weakify the self inside myClosure definition

cell.myClosure = { [weak self] in
    print(self?.value)
}

weakifying the self instance will break the retain cycle and will no longer cause the memory leak

3. Closures associated with NotificationCenter

This is often trickier case compared to just normal interaction between self and associated closure.

Consider the following example where we set up an observer on current class and refer to self inside the associated closure

NotificationCenter.default.addObserver(forName: Notification.Name("ddd"), object: nil, queue: nil) { (not) in
    print(self.value)
}

This retain cycle is not obvious since it does not have a circular reference as we saw in first two examples. However, to add clear explanation self owns the added observer and observer owns the closure which gets called as soon as the notification with name ddd is received. As you can see inside the closure, it also owns the self which gives a birth to retain cycle.

It it important to note that this is an issue only with block based NotificationCenter. If you are using a method based NotificationCenter, you don't have to worry about memory leak. There is one more thing which strongly hints toward possibility of memory leak that the closure argument that this addObserver method takes is declared to be escaping

It's good that we recognized it as a closure. Now we can remove it using the same logic we used in first two examples - Use weak reference and retain cycle will be broken and no more memory leaks.

NotificationCenter.default.addObserver(forName: Notification.Name("ddd"), object: nil, queue: nil) { [weak self] (not) in
    print(self?.value)
}

There is one caveat though. Since you have added an entry to the notification center's dispatch table that includes a notification queue and a block, you will have to manually unregister observation. To do this, you pass the object returned by this method to removeObserver method. You must invoke removeObserver or removeObserver:name:object: before any object specified by addObserverForName:object:queue:usingBlock: is deallocated

This can be done by keeping a local reference to observer returned by addObserver method while creating an observation for any object. This local reference can then be used to manually unregister the observation as follows

// Adding an observer
observer = NotificationCenter.default.addObserver(forName: Notification.Name("ddd"), object: nil, queue: nil) { [weak self] (not) in
    print(self?.value)
}

// Removing an observer
override func viewWillDisappear(_ animated: Bool) {
    NotificationCenter.default.removeObserver(observer)
    super.viewWillDisappear(animated)
}

4. Escaping closure as method argument

Escaping closures are dangerous in terms of memory leak as they can give rise to hidden retain cycles and memory leaks. (As we saw in the previous example) Let's look at the example.

class CustomManager {
    static let sharedInstance = CustomManager()
    var simpleClosure: (() -> Void)?
    var anotherClosure: (() -> Void)?

    func doSomething(closure: @escaping () -> Void) {
        self.anotherClosure = closure
        self.anotherClosure?()
    }
}

// In another class
class Sample: UIViewController {

    let customManager = CustomManager.sharedInstance
    var value: Int = 0
    
    override func viewDidLoad() {
        super.viewDidLoad()
        self.customManager.doSomething {
            print(self.value)
        }
    }

In this case, the closure is declared as escaping in the CustomManager class. We also store the passed closure as a property on the CustomManager and this definitely gives rise to retain cycle as follows:

Here, self owns the customManager and when doSomething method is called, customManager will own the closure (Which is escaping here) as it stores it to its own anotherClosure property, and as this closure captures the self within itself, it will give rise to the memory leak unless we weakify the self instance in the closure definition of the second example,

// Memory leak
self.customManager.doSomething {
    print(self.value)
}

// Memory leak fixed
self.customManager.doSomething { [weak self] in
    print(self?.value)
}

So every time you write the closure definition which has been passed as escaping, pay close attention if you are not retaining the instance (In this case self) which itself owns the closure or its owner

5. Lazy properties as non-closures

What happens when you declare a lazy closure property that uses self inside closure?

var value: Int = 0
lazy var myVariable: () -> Int = {
    print(self.value)
    return 100
}

Here, self owns the myVariable which is retained for the lifetime of the class. myVariable also retains self which gives birth to retain cycle. So myVariable and value live and die together. So as long as one exists, other lives.

Given the mutual relationship between both variables, we can safely use unowned here since self will always be in memory when this closure is called. So we use [unowned self] here to break the retain cycle

var value: Int = 0
lazy var myVariable: () -> Int = { [unowned self] in
    print(self.value)
    return 100
}

6. Using blocks inside closure associated with singleton object

Let's look at another example involving Singleton and Closure.

Here we have our classic singleton class named CustomManager

class CustomManager {
    static let sharedInstance = CustomManager()
    var simpleClosure: (() -> Void)?
}

And this is another class where we are using this singleton and associated closure, simpleClosure

class Sample: UIViewController {

    let customManager = CustomManager.sharedInstance
    override func viewDidLoad() {
        super.viewDidLoad()
        self.customManager.simpleClosure = {
            print(self.value)
        }
    }

Inside class Sample, we are defining the simpleClosure which captures the self instance. Here's an interesting thing.

simpleClosure is a property on CustomManager. In other words, CustomManager owns simpleClosure. As it's clear from above code inside Sample class, self owns the customManager and also, simpleClosure owns the self. So there is a clear retain cycle in this case. Let's rewrite above closure definition with weak so as to break this cycle,

self -> customManager -> simpleClosure -> self

// No retain cycle
self.customManager.simpleClosure = { [weak self] in
    print(self?.value)
}

Wowww! So just using [weak self] inside closure we effectively broke the retain cycle and hence avoided the memory leak.

Things which look like a memory leak, but in fact aren't

Next we will look at examples of code which looks like a memory leak, but in fact, they are not. This will help you avoid writing redundant code which weakifies the self instance.

1. UIView animation blocks

UIView.animate(withDuration: 2.0) {
    print(self.value)
}

This is not a retain cycle. Because self does not own static method animate. It's true that animation block owns the self. But the entity which owns this animation block is other than self. So there is an ownership in one direction, but being absent in other direction makes sure there is no retain cycle.

Granted that the animation block is escaping and is caputring self, but in this case self does not own the entity which owns the animation block which also removes any possibility of retain cycle.

2. Dispatch queue blocks used for a delay

DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
    print(self.value)
}

This is not a retain cycle either. Because self does not own the static method asyncAfter. It's true that async block owns the self. But the entity which owns this async block is other than self. So there is an ownership in one direction, but being absent in other direction makes sure there is no retain cycle.

3. Dismiss blocks

self.dismiss(animated: true) {
    print(self.value)
}

Do you think there is a retain cycle cause self owns the dismiss method which in turn owns the closure and the closure eventually owns the self inside?

First, even though we are calling dismiss method using self instance, it's not the self of current class it refers to. It refers to dismiss method defined on UIViewController superclass. So the current class is not an owner of dismiss method. Second, the block which captures the self is not owned by self itself and block is non-escaping which means it is not stored by the parent class as well. This eliminates the possibility of reverse ownership and thus a retain cycle.

So even though dismiss block captures self, self (Current class instance) is not capturing a block - So there is no retain cycle.

Let's look at another example of calling a dismiss method.

self.navigationController?.dismiss(animated: true, completion: {
    print(self.value)
})

We can apply the same reasoning here. It's clear that dismiss is not owned by your subclass of UIViewController. So there is no ownership claim from self (a current instance of your UIViewController subclass) to dismiss. Also, the completion block is altogether owned by Apple's UIKit framework which has no ownership association with self either. It's true that completion block owns the self, but there is no such indication that self owns the block. So there is not a retain cycle.

4. Non-escaping closures as method argument

But do we have to be careful dealing with defining closures passed as a method argument? Let's look how it works -

For the sake of an example, let's make a Singleton class named CustomManager

class CustomManager {
    static let sharedInstance = CustomManager()
    var simpleClosure: (() -> Void)?
    
    func doSomething(closure: () -> Void) {
        closure()
    }
}

Above closure has a method named doSomething which takes the empty closure as an argument. Now, let's assume another arbitrary class keeping reference to this singleton and invoking a method

class Sample: UIViewController {
    let customManager = CustomManager.sharedInstance    
    var value: Int = 0
    .....
    ...
    ..
    override func viewDidLoad() {
        super.viewDidLoad()
        self.customManager.doSomething {
            print(self.value)
        }
    }

The question is when we called a doSomething in above code and used self inside closure, were we supposed to weakify it? As it turns out in this case, it wasn't necessary.

It's true that self owns the customManager, and closure captures the self, but self does not own the closure which is passed as an argument to doSomething method. (Latter is the necessary condition to have a retain cycle). This is a non-escaping closure and takes care of memory management, since it's not retained by the customManager singleton. Thus we can call doSomething method on customManager singleton instance and use self inside without causing a memory leak.

So as long as you are using self inside the closure passed as a method argument, and closure is not retained by its owner (That is, it is non-escaping), you should be free from the memory leak. But as usual, run the memory graph and make sure app calls the deinit method as soon as viewController is popped or dismissed. Meanwhile, you can use the above-mentioned examples with escaping and non-escaping closures to get the idea of how they affect the creation of retain cycle and memory leaks

5. Computed properties

   • As simple variables
     When computed property is declared as a simple variable, you do not need to use weakify as there is no creation of retain cycle in this case

   var myVal: Int {
       return 100
   }
   
   // This is similar to
   func myVal() -> Int {
       return 100
   }
   

   Just like you do not have to weakify inside of equivalent method, you don't have to weakify inside computed property closure.
   In fact, Swift makes it really easy. If you try to use [weak self] inside computed variable, it will throw you a nice compiler error

   

   • As Closures
     We can use the similar reasoning for computed properties stored as a closure. But here is the tricky situation.

    var value: Int = 0
    var myVar: () -> Void {
        return {
            self.value = 0
        }
    }
   

   Should we use the [weak self] inside the retuning closure? It's a confusing situation since self owns the myVar and returned closure seems to be capturing self. But returned closure is not stored anywhere as it's simply returned. Plus similar to the previous example, we can simplify it as a method as follows which makes it even more clear that there is no retain cycle,

   func myVar() -> (() ->Void) {
       return {
           self.value = 0
       }
   }
   

6. Lazy properties as non-closures
   Let's look at the following example involving lazy property where the value is computed using closure

   lazy var myVariable: Int = {
       print(self.value)
       return 100
   }()
   

   In spite of the fact that we are using self inside the closure associated with lazy variable, it is not retained for the lifetime of the controller. As soon as the myVariable is called for the first time, closure is executed, it's value is assigned to myVariable (Which is then used subsequently for optimization) and closure is immediately deallocated since as it's the case for lazy variable, it will be executed only once.

References:

Apple documentation - NSNotifications
Retain cycles weak and unowned in Swift
Weak and unowned references in Swift