From 3d14f70c11110a0b09abd5016f0afc3edd4773a0 Mon Sep 17 00:00:00 2001
From: Andrey Shavelev <shavelev.andrey@gmail.com>
Date: Sat, 2 Nov 2019 19:37:59 +0200
Subject: [PATCH] Add ability to erase storage, test deallocation of value with
 one thread

---
 Sources/ThreadSpecific/Storage.swift          | 16 +++--
 Sources/ThreadSpecific/ThreadSpecific.swift   | 26 ++++++--
 .../ThreadSpecificWrapperError.swift          |  8 +++
 .../ThreadSpecificTests.swift                 | 59 +++++++++++++++++++
 4 files changed, 100 insertions(+), 9 deletions(-)
 create mode 100644 Sources/ThreadSpecific/ThreadSpecificWrapperError.swift

diff --git a/Sources/ThreadSpecific/Storage.swift b/Sources/ThreadSpecific/Storage.swift
index 517cf80..d5cada4 100644
--- a/Sources/ThreadSpecific/Storage.swift
+++ b/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 }
 }
diff --git a/Sources/ThreadSpecific/ThreadSpecific.swift b/Sources/ThreadSpecific/ThreadSpecific.swift
index 013f052..9269723 100644
--- a/Sources/ThreadSpecific/ThreadSpecific.swift
+++ b/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
     }
diff --git a/Sources/ThreadSpecific/ThreadSpecificWrapperError.swift b/Sources/ThreadSpecific/ThreadSpecificWrapperError.swift
new file mode 100644
index 0000000..d16add4
--- /dev/null
+++ b/Sources/ThreadSpecific/ThreadSpecificWrapperError.swift
@@ -0,0 +1,8 @@
+//
+//  Storage.swift
+//  Created by Andrey Shavelev on 28/10/2019.
+//
+
+enum ThreadSpecificStorageError : Error {
+    case storageErased
+}
diff --git a/Tests/ThreadSpecificTests/ThreadSpecificTests.swift b/Tests/ThreadSpecificTests/ThreadSpecificTests.swift
index a63bfd7..e56ccf5 100644
--- a/Tests/ThreadSpecificTests/ThreadSpecificTests.swift
+++ b/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
+    }
+}