29
\$\begingroup\$

This is a follow-up to this question.

While discussing some details about the code I posted there, I came upon a problem with thread-safety. After searching and trying different things, I reached a potential solution that I now present here. This worked without any problems during my tests, though since thread-safety is not easy, I'd like to ask: is this implementation genuinely thread-safe?

The thread-safety problem was the following: keys on the dictionary can be deallocated on any thread. That means, when a key disappears, the callback block can also be called on any thread. That in turn means the dict var can be changed while it is being accessed.

My solution was to use dispatch_sync to coordinate all reads/writes, including the one happening inside the callback block from DeallocWatcher.

public class WeakKeyDictionary<K: AnyObject, V where K: Hashable> {

    private var dict = SynchronizedValue(value: Dictionary<HashableWeakBox<K>, V>())
    public var block: (V)->() = { _ in }

    public init() {}
    public init(dictionary: Dictionary<K, V>) {
        for (k, v) in dictionary {
            setValue(v, forKey: k)
        }
    }

    public subscript(key: K) -> V? {
        get { return valueForKey(key) }
        set { setValue(newValue, forKey: key) }
    }

    public func valueForKey(key: K) -> V? {
        return dict.get { $0[HashableWeakBox(key)] }
    }
    public func setValue(newValue: V?, forKey key: K) {
        let hashableBox = HashableWeakBox(key)

        if let value = newValue {
            let watcher = DeallocWatcher { [weak self] in
                if let me = self {
                    if let v = me.syncedRemoveValueForKey(hashableBox) {
                        me.block(v)
                    }
                }
            }

            objc_setAssociatedObject(key, unsafeAddressOf(self), watcher, objc_AssociationPolicy(OBJC_ASSOCIATION_RETAIN_NONATOMIC))
            dict.access { $0[hashableBox] = value; return }
        }
        else {
            objc_setAssociatedObject(key, unsafeAddressOf(self), nil, 0)
        }
    }
    public func removeValueForKey(key: K) -> V? {
        objc_setAssociatedObject(key, unsafeAddressOf(self), nil, 0)
        return syncedRemoveValueForKey(HashableWeakBox(key))
    }

    private func syncedRemoveValueForKey(key: HashableWeakBox<K>) -> V? {
        var v: V?
        dict.access { v = $0.removeValueForKey(key); return }
        return v
    }

    public var count: Int { return dict.get { $0.count } }
    public var isEmpty: Bool { return dict.get { $0.isEmpty } }

    public var keyValues: [(K, V)] {
        return dict.get { dict in
            let v = dict.keys
                .filter { k in k.value != nil }
                .map { k -> (K, V) in (k.value!, dict[k]!) }
            return Array(v)
        }
    }
    public var keys: [K] {
        return keyValues.map { (k, v) in k }
    }
    public var values: [V] {
        return keyValues.map { (k, v) in v }
    }

    deinit {
        // Callback is not called when deallocing the helpers because in this case (inside deinit) 'self' is already nil
        dict.access {
            for box in $0.keys {
                objc_setAssociatedObject(box.value, unsafeAddressOf(self), nil, 0)
            }
        }
    }
}

extension WeakKeyDictionary: Printable {
    public var description: String  {
        let contents = dict.get { dict -> [String] in
            let v = dict.keys
                .filter { $0.value != nil }
                .map { "\($0.value!) : \(dict[$0]) " }
            return Array(v)
        }

        return "[ " + ", ".join(contents) + "]"
    }
}
extension WeakKeyDictionary: SequenceType {
    public func generate() -> IndexingGenerator<Array<(K, V)>> {
        return keyValues.generate()
    }
}

private class HashableWeakBox<T: AnyObject where T: Hashable>: Hashable {
    weak var value: T?
    let hashValueWhenNil: Int

    init(_ v: T) {
        value = v
        hashValueWhenNil = v.hashValue
    }

    var hashValue: Int { return value?.hashValue ?? hashValueWhenNil }
}

private class DeallocWatcher {
    let callback: ()->()
    init(_ c: ()->()) { callback = c }
    deinit { callback() }
}

This is the SynchronizedValue class (basically a wrapper around a GCD serial queue):

public class SynchronizedValue<T> {
    public let serialQueue = dispatch_queue_create("SynchronizedValue serial queue", DISPATCH_QUEUE_SERIAL)
    private var value: T

    public init(value v: T) { value = v }

    /// Should only return value types or thread-safe reference types
    public func get<V>(/*@noescape*/ action:  (T) -> V) -> V {
        var v: V?
        dispatch_sync(serialQueue) {
            v = action(self.value)
        }
        return v!
    }

    public func access(/*@noescape*/ action: (inout T) -> ()) {
        dispatch_sync(serialQueue) {
            action(&self.value)
        }
    }

    /// Should only be used for value types or for thread-safe reference types
    public func get() -> T {
        return get { $0 }
    }

    public func set(v: T) {
        access { $0 = v; return }
    }
}

Notes:

On the computed var keyValues, this line .filter { k in k.value != nil } is needed because a value can become nil while this dict.get { ... } block is running, and since the block is still running, the value hasn't been cleaned out yet (the corresponding DeallocWatcher callback is sitting on the queue waiting for its turn).

\$\endgroup\$
  • 2
    \$\begingroup\$ If the code works and you're looking to improve it, this is definitely on-topic here! \$\endgroup\$ – Phrancis Apr 3 '15 at 21:28
  • \$\begingroup\$ @sᴉɔuɐɹɥԀ Yup, that's what I thought. Thanks! \$\endgroup\$ – Alex Apr 3 '15 at 21:29
  • 2
    \$\begingroup\$ it has been three years. Do you want to review your own code, @Alex? \$\endgroup\$ – Zeta Mar 30 '18 at 15:46

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.