Overview

On a recent project I was working on, I needed a key value store to persist some data which my app needed in various spots and wanted to somehow receive updates whenever values for certain keys were updated.

I started out with a simple in memory one (backed by a Dictionary), here’s what it looked like:

protocol KeyValueStore
{
    func valueForKey(key: String) -> String
    func setValue(value: String, forKey key: String)
    func removeValueForKey(key: String)
}

class DictionaryKeyValueStore : KeyValueStore
{
    private var store = [String : String]()
   
    // MARK: KeyValueStore
   
    func valueForKey(key: String) -> String?
    {
        return store[key]
    }
   
    func setValue(value: String, forKey key: String)
    {
        store[key] = value
    }
   
    func removeValueForKey(key: String)
    {
        store.removeValueForKey(key)
    }
}

In situations like this, a protocol is perfect! It removes any commitments to a particular implementation and allows an alternate one to be provided at a later time with minimal change to the users of that component.

A Simple Approach

My first (poor) stab at a solution looked like this:

protocol KeyValueStore
{
    func valueForKey(key: String) -> String
    func setValue(value: String, forKey key: String)
    func removeValueForKey(key: String)
    func registerForUpdates(key: String, callback: (key: String, newValue:String)->Void)
}

class DictionaryKeyValueStore : KeyValueStore
{
    typealias UpdatesCallback = (key: String, newValue:String?)->Void
    private var store = [String : String]()
    private var keyedCallbacks = [String: [UpdatesCallback]]()
   
    // MARK: KeyValueStore
    func valueForKey(key: String) -> String?
    {
        return store[key]
    }
   
    func setValue(value: String, forKey key: String)
    {
        store[key] = value
        notifyChanges(key, newValue: value)
    }
   
    func removeValueForKey(key: String)
    {
        store.removeValueForKey(key)
        notifyChanges(key, newValue: nil)
    }
   
    func registerForUpdates(key: String, callback: UpdatesCallback)
    {
        if var keyCallbacks = keyedCallbacks[key]
        {
            keyCallbacks.append(callback)
            keyedCallbacks[key] = keyCallbacks
        }
        else
        {
            keyedCallbacks[key] = [callback]
        }
    }
   
    // MARK: Private
    private func notifyChanges(key: String, newValue: String?)
    {
        if let keyCallbacks = keyedCallbacks[key]
        {
            for callback in keyCallbacks
            {
                callback(key: key, newValue: newValue)
            }
        }
    }
}

This worked as expected and allowed me to continue with my work but it did have its flaws. The main one being, there was no way of deregistering callbacks. This can be a problem especially if the store’s lifetime exceeds the lifetime of the classes that register callbacks.

Making it generic

Before solving the deregistration problem, I wanted to spend some time on making callback management a bit more generic as I found myself needing callbacks elsewhere in the app.

The callback management code was extracted to its own generic class.

class KeyedCallbacks<CallbackParameters>
{
    typealias Callback = (CallbackParameters) -> Void
    private var keyedCallbacks = [String: [Callback]]()
   
    func register(key: String, callback: Callback)
    {
        if var keyCallbacks = keyedCallbacks[key]
        {
            keyCallbacks.append(callback)
            keyedCallbacks[key] = keyCallbacks
        }
        else
        {
            keyedCallbacks[key] = [callback]
        }
    }
   
    func performCallbacksForKey(key:String, withParameters parameters: CallbackParameters)
    {
        if let keyCallbacks = keyedCallbacks[key]
        {
            for callback in keyCallbacks
            {
                callback(parameters)
            }
        }  
    }
}

class DictionaryKeyValueStore : KeyValueStore
{
    typealias CallbackParameters = (key: String, newValue:String?)
    typealias UpdatesCallback = CallbackParameters->Void
    private var store = [String : String]()
    private var keyedCallbacks = KeyedCallbacks<CallbackParameters>()
   
    // MARK: KeyValueStore
   
    func valueForKey(key: String) -> String?
    {
        return store[key]
    }
   
    func setValue(value: String, forKey key: String)
    {
        store[key] = value
        notifyChanges(key, newValue: value)
    }
   
    func removeValueForKey(key: String)
    {
        store.removeValueForKey(key)
        notifyChanges(key, newValue: nil)
    }
   
    func registerForUpdates(key: String, callback: UpdatesCallback)
    {
        keyedCallbacks.register(key , callback: callback)
    }
   
    // MARK: Private
   
    private func notifyChanges(key: String, newValue: String?)
    {
        keyedCallbacks.performCallbacksForKey(key, withParameters: (key: key, newValue: newValue ))
    }
}

This relieved the key value store from the callback management responsibility making it cleaner and also allowed me to re-use the callback management class.

Deregistering callbacks

Now for the interesting part, how do we remove callbacks?

Since the callbacks are keyed in a dictionary we can remove all callbacks for a given key - but that is not desired, we want to remove a single callback for a given key. Had closures been equatable or had some means of comparison the following would have been suffice:

func deregister(key: String, callback: Callback)
{
    if var keyCallbacks = keyedCallbacks[key],
           index = keyCallbacks.indexOf(callback) // This does not work!
    {
        keyCallbacks.removeAtIndex(index)
        if keyCallbacks.count > 0
        {
            keyedCallbacks[key] = keyCallbacks
        }
        else
        {
            keyedCallbacks.removeValueForKey(key)
        }
    }
}

Sadly there is no reliable way to compare two closures in Swift. A few options have been mentioned on Stackoverflow - but as I say not reliable, nor nice for that matter.

Besides Even if there was, this solution may look good implementation wise, but its not so great for the users of the KeyedCallbacks class. They would now need to keep a reference to their callbacks and loose the ability to define them inline. What’s worse is if one of the callers instance methods is registered instead - that would result in a memory leak!

There are of course ways to mitigate the leaks, but all in all this is not a great solution.

A better approach

The issue we were facing was identifying closures in order to remove specific ones. Perhaps we could internally create some kind of token or handle that is equatable and hand it back to the caller. To deregister, the caller can hand back the handle and internally the KeyedCallbacks class will know exactly which callback to remove.

This is what that could look like:

protocol CallbackHandle
{
}

// ..
    func register(key: String, callback: Callback) -> CallbackHandle
    func deregister(handle: CallbackHandle)
// ..

public protocol CallbackHandle
{
   
}

private class KeyedCallbackHandle<CallbackParameters> : CallbackHandle, Equatable
{
    typealias Callback = (CallbackParameters) -> Void
    let key: String
    let callback : Callback
    init(key: String, callback: Callback)
    {
        self.key = key
        self.callback = callback
    }
}

private func ==<T>(lhs: KeyedCallbackHandle<T>, rhs: KeyedCallbackHandle<T>) -> Bool
{
    return lhs === rhs
}

public class KeyedCallbacks<CallbackParameters>
{
    typealias Callback = (CallbackParameters) -> Void
   
    private var keyedHandles = [String: [KeyedCallbackHandle<CallbackParameters>]]()
   
    public init()
    {
       
    }
   
    public func register(key: String, callback: Callback) -> CallbackHandle
    {
        let handle = KeyedCallbackHandle<CallbackParameters>(key:key, callback:callback)

        if var handles = keyedHandles[key]
        {
            handles.append(handle)
            keyedHandles[key] = handles
        }
        else
        {
            keyedHandles[key] = [handle]
        }
       
        return handle
    }
   
    public func deregister(handle: CallbackHandle)
    {
        if let keyedCallbackHandle = handle as? KeyedCallbackHandle<CallbackParameters>,
           var handles = keyedHandles[keyedCallbackHandle.key],
           let index = handles.indexOf(keyedCallbackHandle)
        {
            handles.removeAtIndex(index)
            if handles.count > 0
            {
                keyedHandles[keyedCallbackHandle.key] = handles
            }
            else
            {
                keyedHandles.removeValueForKey(keyedCallbackHandle.key)
            }
        }
    }
   
    public func performCallbacksForKey(key:String, withParameters parameters: CallbackParameters)
    {
        if let handles = keyedHandles[key]
        {
            for handle in handles
            {
                handle.callback(parameters)
            }
        }
    }
}

Taking it a step further

Here’s the solution I landed on in the end which was inspired by a concept a few of my colleagues adopted in a different project.

We could make the callback automatically deregister as soon as the CallbackHandle goes out of scope. As such, callers don’t need to worry about deregistering, all they need to do is just keep the handle alive as long as they need callbacks.

To do so, we’re going to need an additional closure in KeyedCallbackHandle that gets executed when it gets deallocated.

private class KeyedCallbackHandle<CallbackParameters> : CallbackHandle, Equatable
{
    typealias Callback = (CallbackParameters) -> Void
    let key: String
    let callback : Callback
    var onDeinit : (()->Void)?
   
    init(key: String, callback: Callback)
    {
        self.key = key
        self.callback = callback
    }
   
    deinit
    {
        onDeinit?()
    }
}

The register method can now attach some additional code to that closure to remove the callback.

// ..   
    public func register(key: String, callback: Callback) -> CallbackHandle
    {
        let handle = KeyedCallbackHandle<CallbackParameters>(key:key, callback:callback)
        handle.onDeinit = { [weak self, weak handle] in
            if let s = self, let h = handle
            {
                s.deregister(h)
            }
        }
        if var handles = keyedHandles[key]
        {
            handles.append(handle)
            keyedHandles[key] = handles
        }
        else
        {
            keyedHandles[key] = [handle]
        }
       
        return handle
    }
// ..

The only thing left to do is to create weak wrapper for the KeyedCallbackHandle to store in the internal dictionary. That way the only strong reference to the handle will be the with the caller.

private class WeakKeyedCallbackHandle<CallbackParameters> : Equatable
{
    weak var handle : KeyedCallbackHandle<CallbackParameters>?
    init(_ handle: KeyedCallbackHandle<CallbackParameters>)
    {
        self.handle = handle
    }
}

private func ==<T>(lhs: WeakKeyedCallbackHandle<T>, rhs: WeakKeyedCallbackHandle<T>) -> Bool
{
    return lhs.handle == rhs.handle
}

public class KeyedCallbacks<CallbackParameters>
{
    typealias Callback = (CallbackParameters) -> Void
   
    private var keyedHandles = [String: [WeakKeyedCallbackHandle<CallbackParameters>]]()
   
    public init()
    {
       
    }
   
    public func register(key: String, callback: Callback) -> CallbackHandle
    {
        let handle = KeyedCallbackHandle<CallbackParameters>(key:key, callback:callback)
       
        handle.onDeinit = { [weak self, weak handle] in
            if let s = self, let h = handle
            {
                s.deregister(h)
            }
        }
        let weakHandle = WeakKeyedCallbackHandle(handle)
        if var handles = keyedHandles[key]
        {
            handles.append(weakHandle)
            keyedHandles[key] = handles
        }
        else
        {
            keyedHandles[key] = [weakHandle]
        }
       
        return handle
    }
   
    public func deregister(handle: CallbackHandle)
    {
        if let keyedCallbackHandle = handle as? KeyedCallbackHandle<CallbackParameters>,
           var handles = keyedHandles[keyedCallbackHandle.key],
           let index = handles.indexOf(WeakKeyedCallbackHandle(handle))
        {
            handles.removeAtIndex(index)
            if handles.count > 0
            {
                keyedHandles[keyedCallbackHandle.key] = handles
            }
            else
            {
                keyedHandles.removeValueForKey(keyedCallbackHandle.key)
            }
        }
    }
   
    public func performCallbacksForKey(key:String, withParameters parameters: CallbackParameters)
    {
        if let weakHandles = keyedHandles[key]
        {
            for weakHandle in weakHandles
            {
               weakHandle.handle?.callback(parameters)
            }
        }
    }
}

Finally some tidy ups can be made. The code that adds and removes callbacks from the dictionary can be extracted to its own class. This will keep the the main class shorter and easier to follow.

private class KeyedCallbackHandles<CallbackParameters>
{
    private var keyedHandles = [String: [WeakKeyedCallbackHandle<CallbackParameters>]]()
   
    func add(handle: KeyedCallbackHandle<CallbackParameters>)
    {
        handle.onDeinit = { [weak self, weak handle] in
            if let s = self, let h = handle
            {
                s.remove(h)
            }
        }
        let weakHandle = WeakKeyedCallbackHandle(handle)
        if var handles = keyedHandles[handle.key]
        {
            handles.append(weakHandle)
            keyedHandles[handle.key] = handles
        }
        else
        {
            keyedHandles[handle.key] = [weakHandle]
        }
    }
   
    func remove(handle: KeyedCallbackHandle<CallbackParameters>)
    {
         if var handles = keyedHandles[handle.key],
            let index = handles.indexOf(WeakKeyedCallbackHandle(handle))
        {
            handles.removeAtIndex(index)
            if handles.count > 0
            {
                keyedHandles[handle.key] = handles
            }
            else
            {
                keyedHandles.removeValueForKey(handle.key)
            }
        }
    }
   
    func handlesForKey(key: String) -> [KeyedCallbackHandle<CallbackParameters>]?
    {
        var handles : [KeyedCallbackHandle<CallbackParameters>]?
       
        if let weakHandles = keyedHandles[key]
        {
            handles = weakHandles.flatMap({ $0.handle })
        }
       
        return handles
    }
}

public class KeyedCallbacks<CallbackParameters>
{
    typealias Callback = (CallbackParameters) -> Void
   
    private var keyedHandles = KeyedCallbackHandles<CallbackParameters>()
   
    public init()
    {
       
    }
   
    public func register(key: String, callback: Callback) -> CallbackHandle
    {
        let handle = KeyedCallbackHandle<CallbackParameters>(key:key, callback:callback)
        keyedHandles.add(handle)
        return handle
    }
   
    public func deregister(handle: CallbackHandle)
    {
        if let keyedCallbackHandle = handle as? KeyedCallbackHandle<CallbackParameters>
        {
            keyedHandles.remove(keyedCallbackHandle)
        }
    }
   
    public func performCallbacksForKey(key:String, withParameters parameters: CallbackParameters)
    {
        if let handles = keyedHandles.handlesForKey(key)
        {
            for handle in handles
            {
                handle.callback(parameters)
            }
        }
    }
}

Usage

With KeyedCallbacks complete, the key value store needs one minor tweak to return the callback handle.

// ..   
class DictionaryKeyValueStore : KeyValueStore
{
// ..
    private var keyedCallbacks = KeyedCallbacks<CallbackParameters>()
// ..
    func registerForUpdates(key: String, callback: UpdatesCallback) -> CallbackHandle
    {
        return keyedCallbacks.register(key , callback: callback)
    }
// ..
}

// ..

Putting it all together - this is what registering for callbacks on the key value store looks like:

public func usage()
{
    let keyValueStore = DictionaryKeyValueStore()
    let handle1 = keyValueStore.registerForUpdates("a") { print(" ## \($0)") }
    unused(handle1)
    do
    {
        let handle2 = keyValueStore.registerForUpdates("a") { print(" >> \($0)") }
        unused(handle2)
        keyValueStore.setValue("cool!", forKey: "a")
        //prints
        // 
        //   ## ("a", Optional("cool!"))
        //   >> ("a", Optional("cool!"))
    }
    
    keyValueStore.setValue("warmer", forKey: "a")
    //prints
    //
    //   ## ("a", Optional("warmer"))
}

public func unused(_ : Any)
{
    // This is needed to supress the unused variable warning
}

… and we’re done! We now have a re-usable component for managing callbacks in a scoped fashion albeit not thread safe.

The complete example can be seen here. Also the complete class with further tidy ups and tests can be found on GitHub.