36
\$\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\$
5
  • 3
    \$\begingroup\$ If the code works and you're looking to improve it, this is definitely on-topic here! \$\endgroup\$
    – Phrancis
    Apr 3, 2015 at 21:28
  • \$\begingroup\$ @sᴉɔuɐɹɥԀ Yup, that's what I thought. Thanks! \$\endgroup\$
    – Alex
    Apr 3, 2015 at 21:29
  • 4
    \$\begingroup\$ it has been three years. Do you want to review your own code, @Alex? \$\endgroup\$
    – Zeta
    Mar 30, 2018 at 15:46
  • 2
    \$\begingroup\$ It has been five years. Do you want to review your own code, @Alex? \$\endgroup\$
    – Linny
    Feb 17, 2020 at 0:15
  • 2
    \$\begingroup\$ It has been eight years. Do you want to review your own code, @Alex? \$\endgroup\$
    – J_H
    Jun 11, 2023 at 15:44

1 Answer 1

3
\$\begingroup\$

I appreciate that this is an old question, but a few thoughts. First, I might advise updating SynchronizedValue with modern GCD API. For example:

public class SynchronizedValue<Wrapped> {
    private let serialQueue = DispatchQueue(label: "…")
    private var _wrappedValue: Wrapped

    public init(wrappedValue v: Wrapped) { _wrappedValue = v }

    public var wrappedValue: Wrapped {
        get { serialQueue.sync { _wrappedValue } }
        set { serialQueue.sync { _wrappedValue = newValue } }
    }
}

Notes on the above:

  • The contemporary GCD syntax.
  • The sync method automatically returns whatever value the closure returns, so you do not need to capture the result in local variable, like you do in the get(action:) implementation. Just return the result of sync (or, because of SE-0255, you can even omit the return keyword).
  • I might move the get/set methods into accessors for the wrappedValue. (This convention of wrappedValue also means that you can make this a @propertyWrapper if you would like.)

Note, the GCD serial queue is a nice, simple pattern. Some might advise a reader-writer pattern (which can be a little faster, but often introduces its own issues). Generally, if performance was important, I might jump directly to a lock, e.g., NSLock:

public class SynchronizedValue<Wrapped> {
    private let lock = NSLock()
    private var _wrappedValue: Wrapped

    public init(wrappedValue v: Wrapped) { _wrappedValue = v }

    public var wrappedValue: Wrapped {
        get { lock.withLock { _wrappedValue } }
        set { lock.withLock { _wrappedValue = newValue } }
    }
}

Or an OSAllocatedUnfairLock:

import os.lock

public class SynchronizedValue<Wrapped: Sendable>: @unchecked Sendable {
    private let lock = OSAllocatedUnfairLock()
    private var _wrappedValue: Wrapped

    public init(wrappedValue v: Wrapped) { _wrappedValue = v }

    public var wrappedValue: Wrapped {
        get { lock.withLock { _wrappedValue } }
        set { lock.withLock { _wrappedValue = newValue } }
    }
}

You ask:

is this implementation genuinely thread-safe?

It is thread-safe insofar as you successfully prevent parallel access to the underlying dictionary.

But, two caveats:

  1. There is a potential race between the deinit of DeallocWatcher and the adding of new values. Consider:

    // remove strong reference to a key of something already in the map table; this will result it the map table’s weak key to be (eventually) deallocated
    let firstKey = savedKeys.removeFirst()
    
    // create a new key with the same identifier
    let newKey = SampleKey(id: firstKey.id)
    
    // add new value to the weak-to-strong map table
    let value = SampleValue(…)
    mapTable.setValue(value, forKey: newKey)
    
    // make sure to save a strong reference to this new key
    savedKeys.append(newKey)
    
    // report results
    print("Removed ", firstKey)
    print("Added ", newKey, value)
    

    If I remove the last strong reference to a key and then create a new key instance that is equal to the first one, the DeallocWatcher will be called after the new value was added, it might be unable to differentiate these two separate key instances, resulting in both being removed from the dictionary.

    This is admittedly a contrived example, but it illustrates the subtle race between keys being deallocated and the eviction of the items from the broader dictionary when the last strong reference to this weak property is removed.

  2. In a broader thread-safety observation, the underlying dictionary is not necessarily the right level for the synchronization. E.g., what will happen if you are iterating through the keys and values of the dictionary and another thread comes in and removes a key or adds a key. You might see the same value emitted twice, or see a value omitted, or get a nil value, etc.

    If you are using this as a cache (where you are just looking for hits or misses), for example, that might not be a (serious) issue, but it is more problematic if you are iterating through the collection (which the presence of keys and values arrays suggest was your intent).

    I might suggest eliminating the separate values array (as that is inconsistent with the concept of a collection that might be mutated in parallel). Or, personally, I would entirely eliminate these arrays and simply expose an interface like the access method that is currently buried in the SynchronizedValue object. The synchronization is not well suited at the subscript accessor level, but at a higher level of abstraction. You generally want to synchronize at the “ok, let me do a bunch of some stuff with the dictionary” level. (Yes, we see people share “thread-safe dictionaries” with all the dictionary-like interfaces on Stack Overflow, but, that is generally a mistake. I understand the intuitive appeal of the approach, but it is generally misguided, IMHO.)


Now, you only asked about the thread-safety question, but here are a few observations on the whole weak-to-strong map table concept:

  1. It should be noted that the proposed implementation only handles one type of map table (weak keys, strong values, aka weak-to-strong map table). But, NSMapTable handles other permutations (e.g., strong keys with weak values, etc.).

  2. Note that reference types for keys are somewhat unusual in Swift. Often keys are strings, integers, UUIDs, all of which are value types in Swift. Model objects are sometimes reference types (though we often prefer value types nowadays), but rarely keys.

    Even if you try to be sneaky and introduce a NSString as a reference type key, note that NSString is class cluster with highly optimized (i.e., very atypical) reference semantics. It may not release them when you might otherwise suspect with standard ARC rules.

  3. Also note, weak-to-strong NSMapTable has some surprising/unobvious behaviors (see NSMapTable and Zeroing Weak References), so it should be noted that this implementation has very different memory behaviors. As Luke noted in How NSMapTable works, the Mountain Lion release notes discourage the use of weak-to-strong NSMapTable entirely:

    However, weak-to-strong NSMapTables are not currently recommended, as the strong values for weak keys which get zero'd out do not get cleared away (and released) until/unless the map table resizes itself.

    This weak-to-strong of NSMapTable is not so much a “evict as soon as the last strong reference is removed” pattern as a “the next time the table is resized, any keys without any more strong references will be evicted.”

  4. The sleight of hand with hashes in HashableWeakBox is, IMHO, highly suspect (i.e., that you allow the weak reference to change the value to nil when there are no more strong references, but you want it to keep the same hash).

    Do not get me wrong: I completely understand why this trick was employed (i.e., dictionaries achieve O(1) performance via hashed keys, so it would be problematic to change the hash behind the scenes), but does not seem copacetic to my eye (without some explicit assurances re mutating keys; if a key changes, we would generally remove the old key and add a new key/value pair).

    Frankly, I think it is a moot issue, as I think the whole idea of weak-to-strong collections is distinctly unswifty (see point 4, above). And having played around with various attempts to implement a native weak-to-strong map table implementation, I now have an appreciation as to why Apple implemented a lazy removal pattern (see point 5) in the first place. At first glance, I blithely dismissed their implementation as just poorly designed, but in retrospect, I now have an appreciation of the language constraints that they were dealing with.


In case you are interested, here is my modernization of your example:

public class MapTable<Key: AnyObject & Hashable, Value> {
    private let dictionary = SynchronizedValue(wrappedValue: Dictionary<HashableWeakBox<Key>, Value>())

    public init() {}

    public init(dictionary: Dictionary<Key, Value>) {
        for (k, v) in dictionary {
            setValue(v, forKey: k)                      // as aside, this might be a tad inefficient
        }
    }

    public subscript(key: Key) -> Value? {
        get { value(forKey: key) }
        set { setValue(newValue, forKey: key) }
    }

    public func value(forKey key: Key) -> Value? {
        dictionary.get { $0[HashableWeakBox(key)] }
    }

    public func setValue(_ value: Value?, forKey key: Key) {
        let hashableBox = HashableWeakBox(key)

        if let value {
            let watcher = DeallocWatcher { [weak self] in
                guard let self else { return }
                syncedRemoveValue(forKey: hashableBox)
            }

            setAssociatedObject(key, watcher)
            dictionary.access { $0[hashableBox] = value }
        } else {
            setAssociatedObject(key, nil, .OBJC_ASSOCIATION_ASSIGN)
        }
    }

    func setAssociatedObject(
        _ object: Any,
        _ value: Any?,
        _ policy: objc_AssociationPolicy = .OBJC_ASSOCIATION_RETAIN_NONATOMIC
    ) {
        withUnsafePointer(to: self) { pointer in
            objc_setAssociatedObject(object, pointer, value, policy)
        }
    }

    @discardableResult
    public func removeValue(forKey key: Key) -> Value? {
        setAssociatedObject(key, nil, .OBJC_ASSOCIATION_ASSIGN)
        return syncedRemoveValue(forKey: HashableWeakBox(key))
    }

    @discardableResult
    private func syncedRemoveValue(forKey key: HashableWeakBox<Key>) -> Value? {
        dictionary.access { $0.removeValue(forKey: key) }
    }

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

    public var keyValues: [(Key, Value)] {
        dictionary.get { dict in
            dict.keys
                .filter { k in k.value != nil }
                .map { k -> (Key, Value) in (k.value!, dict[k]!) }
        }
    }

    public var keys: [Key] {
        dictionary.get { dict in
            dict.keys
                .filter { $0.value != nil }
                .map { $0.value! }
        }
    }

    public var values: [Value] {
        dictionary.get { dict in
            dict.keys
                .filter { $0.value != nil }
                .map { dict[$0]! }
        }
    }

    deinit {
        // Callback is not called when deallocing the helpers because in this case (inside deinit) 'self' is already nil
        dictionary.access {
            $0.keys
                .compactMap { $0.value }
                .forEach { setAssociatedObject($0, nil, .OBJC_ASSOCIATION_ASSIGN) }
        }
    }
}

extension MapTable: CustomStringConvertible {
    public var description: String  {
        let string = dictionary.get { dict in
            dict.keys
                .filter { $0.value != nil }
                .map {
                    if let value = dict[$0] {
                        "\($0.value!) : \(value)"
                    } else {
                        "\($0.value!) : null"
                    }
                }
                .joined(separator: ",\n    ")
        }

        return "[\n    " + string + "\n]"
    }
}

// I have attempted to preserve this “original hash” concept, but this implementation is incorrect, IMHO; do not use this

private class HashableWeakBox<T: AnyObject & Hashable>: Hashable {
    private(set) weak var value: T?
    var originalHashValue: Int

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

    static func == (lhs: HashableWeakBox, rhs: HashableWeakBox) -> Bool {
        lhs.value == rhs.value && lhs.originalHashValue == rhs.originalHashValue
    }

    func hash(into hasher: inout Hasher) {
        hasher.combine(originalHashValue)
    }
}

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

public class SynchronizedValue<Wrapped> {
    fileprivate let serialQueue = DispatchQueue(label: "SynchronizedValue serial queue")
    fileprivate var _wrappedValue: Wrapped

    var wrappedValue: Wrapped {
        get { serialQueue.sync { _wrappedValue } }
        set { serialQueue.sync { _wrappedValue = newValue } }
    }

    public init(wrappedValue v: Wrapped) { _wrappedValue = v }

    /// Should only return value types or thread-safe reference types
    public func get<T>(execute work: (Wrapped) throws -> T) rethrows -> T {
        try serialQueue.sync {
            try work(_wrappedValue)
        }
    }

    public func access<T>(execute work: (inout Wrapped) throws -> T) rethrows -> T {
        try serialQueue.sync {
            try work(&_wrappedValue)
        }
    }
}

I am not advising the above (notably, the HashableWeakBox is incorrect), but have merely updated the original code snippet for modern API conventions and language features.

\$\endgroup\$

Not the answer you're looking for? Browse other questions tagged or ask your own question.