Add ability to erase storage, test deallocation of value with one thread

Sources/ThreadSpecific/Storage.swift

@@ -4,15 +4,23 @@
 //
 
 class Storage<T> {
-    private var value: T
+    private var value: T?
 
     init(value: T) {
         self.value = value
     }
 
-    func setValue(value: T){
-        self.value = value
+    func erase(){
+        self.value = nil
+    }
+
+    func getValue() throws -> T {
+        guard let value = value else {
+            throw ThreadSpecificStorageError.storageErased
+        }
+
+        return value
     }
 
-    func getValue() -> T { value }
+    func set(value: T) { self.value = value }
 }

Sources/ThreadSpecific/ThreadSpecific.swift

@@ -6,22 +6,37 @@ public class ThreadSpecific<T> {
     var allocatedStorages = [Storage<T>]()
     let valueClosure : () -> T
     var key = pthread_key_t()
-    let storageLock = NSLock()
+    var storageLock = pthread_rwlock_t()
 
     public init(wrappedValue: @autoclosure @escaping () -> T) {
         self.valueClosure = wrappedValue
         pthread_key_create(&key) {
             $0.deallocate()
         }
+        pthread_rwlock_init(&storageLock, nil)
     }
 
     deinit {
         pthread_key_delete(key)
+        pthread_rwlock_destroy(&storageLock)
+
+        for storage in allocatedStorages {
+            storage.erase()
+        }
     }
 
     public var wrappedValue: T {
-        get { return threadSpecificStorage.getValue() }
-        set (value) { threadSpecificStorage.setValue(value: value)}
+        get { return try! threadSpecificStorage.getValue() }
+        set (value) { threadSpecificStorage.set(value: value)}
+    }
+
+    public func erase() {
+        if (pthread_getspecific(key) == nil) {
+            return
+        }
+
+        threadSpecificStorage.erase()
+        pthread_setspecific(key, nil)      
     }
 
     private var threadSpecificStorage: Storage<T> {
@@ -44,10 +59,11 @@ public class ThreadSpecific<T> {
     func allocateStorage() -> Storage<T> {
         let defaultValue = valueClosure()
 
-        storageLock.lock()
+        pthread_rwlock_wrlock(&storageLock)
+        defer { pthread_rwlock_unlock(&storageLock) }
+
         let newStorage = Storage<T>(value: defaultValue)
         allocatedStorages.append(newStorage)
-        storageLock.unlock()
 
         return newStorage
     }

Sources/ThreadSpecific/ThreadSpecificWrapperError.swift

@@ -0,0 +1,8 @@
+//
+//  Storage.swift
+//  Created by Andrey Shavelev on 28/10/2019.
+//
+
+enum ThreadSpecificStorageError : Error {
+    case storageErased
+}

Tests/ThreadSpecificTests/ThreadSpecificTests.swift

@@ -12,9 +12,68 @@ final class ThreadSpecificTests: XCTestCase {
         termometer.degrees = 17
         XCTAssertEqual(17, termometer.degrees)
     }
+
+    func testDeallocatesOldValueWhenNewIsSet() {
+        let valueOwner = ValueOwner()
+        var oldValue: Value? = Value(42)
+        let newValue = Value(43)
+        weak var weakReferenceToOldValue = oldValue
+
+        valueOwner.value = oldValue!
+        oldValue = nil
+
+        XCTAssertNotNil(weakReferenceToOldValue)
+        valueOwner.value = newValue
+        XCTAssertNil(weakReferenceToOldValue)
+    }
+
+    func testDeallocatesValueWhenOwnerIsDeallocated() {
+        var valueOwner: ValueOwner? = ValueOwner()
+        var value: Value? = Value(42)
+        weak var weakReferenceToValue = value
+
+        valueOwner!.value = value!
+        value = nil
+
+        XCTAssertNotNil(weakReferenceToValue)
+        valueOwner = nil
+        XCTAssertNil(weakReferenceToValue)
+    }
+
+    func testDeallocatesValueAfterItWasErrased() {
+        let valueOwner = ValueOwner()
+        var value: Value? = Value(42)
+        weak var weakReferenceToValue = value
+
+        valueOwner.value = value!
+        value = nil
+
+        XCTAssertNotNil(weakReferenceToValue)
+        valueOwner.eraseValue()
+        XCTAssertNil(weakReferenceToValue)
+    }
+
+
 }
 
 public class Termometer {
     @ThreadSpecific
     var degrees: Int = 0
 }
+
+public class ValueOwner {
+    @ThreadSpecific
+    var value = Value(43)
+
+    func eraseValue() {
+        _value.erase()
+    }
+}
+
+public class Value {
+    let integer: Int
+
+    init(_ integer: Int) {
+        self.integer = integer
+    }
+}