Improving Observer Pattern APIs in Swift With Weak Collections

Even if you don't know what the Observer pattern is, there's a good chance that you applied it somewhere in the past. This design pattern is used to create APIs that notify one or more subjects about changes in a certain object, with the NotificationCenter being the most popular use of this pattern in iOS.

One simple way to reproduce what the NotificationCenter does is to create a dictionary that maps a string (a notification) to an array of closures. Whenever that notification is "posted", all of the closures are executed.

final class NotificationCenter {

    var notifications = [String: [() -> Void]]()

    func register(_ closure: @escaping () -> Void, forNotification string: String) {
        notifications[string, default: []].append(closure)
    }

    func post(notification: String) {
        notifications[notification]?.forEach { $0() }
    }
}

However, the point of this article is not to attempt to reproduce the NotificationCenter, but to show you what this implementation implies. You must be aware that whenever you're using the basic Swift dictionaries, arrays or sets, all keys and values are retained! Additionally, you need to be aware that closures are reference types, so they are retained as well and can outlive their owners.

What this means is that you'll see that while this implementation works, it's going to be a huge memory issue. Because it's retaining the closures, they will never be unregistered. The notifications will attempt to execute them even if the object that registered it is long gone.

If you're been working with iOS for a long time, you might remember that iOS's own NotificationCenter had this issue! Prior to iOS 9, every observer had to be unregistered when being deallocated, because if you didn't, it would attempt to execute it when it shouldn't and crash your app.

deinit {
    NotificationCenter.default.removeObserver(self, ...)
}

In the case of our implementation, we could replicate this by adding the concept of "owners" to our closures, so that we are able to remove them if someone wishes to be unregistered. Fortunately, not only we don't need to go this far, but it's good if we don't. If you're developing an API, its usability should be one of your main priorities. In this case, let's take a look at how we can create an observer API that is memory safe while also not having to manually unregister the observers. The problem shown above that NotificationCenter had was fixed in iOS 9 (removing observers became an automatic process) when Apple started applying the same concept.

Weak Collections

Let's pretend we have a deeplink-based navigation system where "feature providers" can provide a feature (represented as an UIViewController) if they recognize the deeplink that the app wants to present:

final class FeaturePusher {

    typealias FeatureProvider = (URL) -> UIViewController?

    lazy var providers = [FeatureProvider]()

    func register(featureProvider: @escaping FeatureProvider) {
        providers.append(featureProvider)
    }

    func feature(forUrl url: URL) -> UIViewController? {
        return providers.lazy.compactMap { $0(url) }.first
    }
}

Like in the notification center example, this suffers from a memory issue. If whoever provided those closures ceases to exist, the FeaturePusher class will still be able to execute the closure and potentially crash the app. Fortunately, there are a few useful types in Foundation that can assist us in improving that.

As I've shown before in my Weak Dictionary Values article, Foundation offers a series of lower-level Obj-C collection types that are more powerful than the basic Swift ones. Two of them specifically are NSMapTable and NSHashTable, which are Obj-C versions of Dictionary and Set, respectively. Both of them allow a higher range of memory management options, which include weak references for both values and keys. If instead of using a base Swift array we used a NSMapTable that has our closures as values and weak references to whoever provided that block as a key, our navigation system would automatically evict and deallocate the closures whenever the related providers are deallocated. That's because in weak collections if the weak component is deallocated, the entire entry will be evicted from the collection.

Creating Weak Collections is just a matter of using the correct initializer. A dictionary with weak keys can be initialized with NSMapTable.weakToStrongObjects(), while one with weak values can be initialized with NSMapTable.strongToWeakObjects(). If we want our navigation system's closures to be automatically unregistered if the object that registered them was deallocated, we can create a weak-keyed dictionary that maps an object to an array of closures:

lazy var providers = NSMapTable<AnyObject, NSHashTable<FeatureProviderBox>>.weakToStrongObjects()

Because the keys are weak, the closures will automatically be evicted from the dictionary if the key ceases to exist.

Note that NSMapTable is an Obj-C API, so all keys and values must be class objects. That's why we have to use a NSHashTable as a value instead of a regular Set or Array.

You can make Obj-C types like NSMapTable able to hold Swift structs by creating a generic Box class wrapper type. Here, we create one to be able to represent our feature closure as a class object (FeatureProviderBox) in order to be able to store it inside the NSHashTable.

final class Box<T> {
    let obj: T
    init(obj: T) {
        self.obj = obj
    }
}

final class FeaturePusher {

    typealias FeatureProvider = (URL) -> UIViewController?
    typealias FeatureProviderBox = Box<FeatureProvider>

    lazy var providers = NSMapTable<AnyObject, NSHashTable<FeatureProviderBox>>.weakToStrongObjects()

    func register(featureProvider: @escaping FeatureProvider, forObject object: AnyObject) {
        if providers.object(forKey: object) == nil {
            providers.setObject(NSHashTable(), forKey: object)
        }
        let box = FeatureProviderBox(obj: featureProvider)
        providers.object(forKey: object)?.add(box)
    }

    func feature(forUrl url: URL) -> UIViewController? {
        let allValues = providers.objectEnumerator()
        while let table = allValues?.nextObject() as? NSHashTable<FeatureProviderBox> {
            if let feature = table.allObjects.lazy.compactMap { $0.obj(url) }.first {
                return feature
            }
        }
        return nil
    }
}

Unit Testing Weak Collections (and Reference Cycles)

To check if our improvement worked, we can create a unit test that checks if the correct view controllers are returned:

func test_observerReturnsTheCorrectFeature() {

    let pusher = FeaturePusher()
    let swiftRocksUrl = URL(string: "myApp://swiftRocks")!
    let swiftRocksVC = SwiftRocksViewController()

    let observerObject: UIView = UIView()

    pusher.register(featureProvider: { url in
        return url == swiftRocksUrl ? swiftRocksVC : nil
    }, forObject: observerObject)

    XCTAssertTrue(pusher.feature(forUrl: swiftRocksUrl) === swiftRocksVC)

    let someOtherURL = URL(string: "myApp://notSwiftRocks")!
    XCTAssertNil(pusher.feature(forUrl: someOtherURL))
}

However, we are mostly interested in seeing if the automatic eviction is working. To test that the observers are being deallocated and the closures are being evicted, we can use an autoreleasepool. As described in my autoreleasepool article, you can use a pool whenever you want something to be deallocated as soon as possible:

func test_observerIsDeallocated() {

    let pusher = FeaturePusher()
    let swiftRocksUrl = URL(string: "myApp://swiftRocks")!
    let swiftRocksVC = SwiftRocksViewController()

    autoreleasepool {
        let observerObject: UIView = UIView()

        pusher.register(featureProvider: { url in
            return url == swiftRocksUrl ? swiftRocksVC : nil
        }, forObject: observerObject)

        XCTAssertTrue(pusher.feature(forUrl: swiftRocksUrl) === swiftRocksVC)

        let someOtherURL = URL(string: "myApp://notSwiftRocks")!
        XCTAssertNil(pusher.feature(forUrl: someOtherURL))
    }

    XCTAssertNil(pusher.feature(forUrl: swiftRocksUrl))
}

You'll see that this test will pass, but if you're not sure why, try removing the pool to see what happens. The test will fail, and the reason is that objects aren't deallocated as soon as they go out of scope in iOS (that will usually happen at the end of a RunLoop). In this case, the pool is simply a way to force it to deallocate immediately for unit testing purposes. This same trick can be applied to unit test any type of reference cycle situation :)

Conclusion: Final considerations for Weak Collections

Weak Collections are a great way to build better APIs, but you must be aware of their possible limitations. While types like NSHashTable and NSPointerArray are all-around great tools, you may see that NSMapTable's documentation tells you to be careful with weakToStrongObjects(). In that configuration, although the values are ejected from the table as expected, they still may be held in memory for a larger period of time. That's why this article didn't attempt to fully reproduce the NotificationCenter, as it took me a while to realize that doing so would require a pretty ugly workaround. However, you'll find that NSHashTable is good to go under any configuration.

References and Good Reads