Dispose on deinit

Use NSRecursiveLock for subscriber reads/writes
UrbanCompass · Sep 10, 2021 · e506ef5 · e506ef5
1 parent 2ce2f5d
commit e506ef5
Show file tree

Hide file tree

Showing 4 changed files with 108 additions and 25 deletions.
diff --git a/Snail/Disposer.swift b/Snail/Disposer.swift
@@ -7,16 +7,23 @@ public protocol DisposableType {
 }
 
 public class Disposer {
-    private(set) public var disposables: [DisposableType] = []
+    private(set) var disposables: [DisposableType] = []
+    private let recursiveLock = NSRecursiveLock()
 
     public init() {}
 
+    deinit {
+        disposeAll()
+    }
+
     public func disposeAll() {
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
         disposables.forEach { $0.dispose() }
         disposables.removeAll()
     }
 
     public func add(disposable: DisposableType) {
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
         disposables.append(disposable)
     }
 }
diff --git a/Snail/Observable.swift b/Snail/Observable.swift
@@ -6,8 +6,8 @@ import Dispatch
 public class Observable<T>: ObservableType {
     private var isStopped: Int32 = 0
     private var stoppedEvent: Event<T>?
-    private var subscribers: [Subscriber<T>] = []
-    private let subscribersQueue = DispatchQueue(label: "snail-observable-queue", attributes: .concurrent)
+    private(set) var subscribers: [Subscriber<T>] = []
+    private let recursiveLock = NSRecursiveLock()
 
     public init() {}
 
@@ -28,31 +28,27 @@ public class Observable<T>: ObservableType {
             return subscriber
         }
 
-        subscribersQueue.async(flags: .barrier) {
-            self.subscribers.append(subscriber)
-        }
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
+        subscribers.append(subscriber)
 
         return subscriber
     }
 
     public func on(_ event: Event<T>) {
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
         switch event {
         case .next:
             guard isStopped == 0 else {
                 return
             }
 
-            subscribersQueue.sync {
-                self.subscribers.forEach {
-                    notify(subscriber: $0, event: event)
-                }
+            subscribers.forEach {
+                notify(subscriber: $0, event: event)
             }
         case .error, .done:
             if OSAtomicCompareAndSwap32Barrier(0, 1, &isStopped) {
-                subscribersQueue.sync {
-                    self.subscribers.forEach {
-                        notify(subscriber: $0, event: event)
-                    }
+                subscribers.forEach {
+                    notify(subscriber: $0, event: event)
                 }
                 stoppedEvent = event
             }
@@ -69,21 +65,18 @@ public class Observable<T>: ObservableType {
     }
 
     public func removeSubscribers() {
-        subscribersQueue.async(flags: .barrier) {
-            self.subscribers.removeAll()
-        }
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
+        subscribers.removeAll()
     }
 
     public func removeSubscriber(subscriber: Subscriber<T>) {
-        subscribersQueue.sync {
-            guard let index = self.subscribers.enumerated().first(where: { $0.element === subscriber })?.offset else {
-                return
-            }
-
-            subscribersQueue.async(flags: .barrier) {
-                self.subscribers.remove(at: index)
-            }
+        recursiveLock.lock(); defer { recursiveLock.unlock() }
+
+        guard let index = subscribers.enumerated().first(where: { $0.element === subscriber })?.offset else {
+            return
         }
+
+        subscribers.remove(at: index)
     }
 
     public func map<U>(_ transform: @escaping (T) -> U) -> Observable<U> {

diff --git a/SnailTests/DisposerTests.swift b/SnailTests/DisposerTests.swift
@@ -20,4 +20,28 @@ class DisposerTests: XCTestCase {
         subject.disposeAll()
         XCTAssertTrue(subject.disposables.count == 0)
     }
+
+    func testDisposeOnDeinitRemovesSubscribers() {
+        var subject: Disposer? = Disposer()
+        let exp = expectation(description: "test")
+        var sum = 0
+
+        let obs = Observable<Void>()
+        for _ in 0..<10 {
+            if subject == nil {
+                subject = Disposer()
+            }
+            obs.subscribe(onNext: {
+                sum += 1
+            }).add(to: subject!)
+
+        }
+        DispatchQueue.global(qos: .utility).asyncAfter(deadline: .now() + 0.5) {
+            subject = nil
+            obs.on(.next(()))
+            XCTAssertEqual(sum, 0)
+            exp.fulfill()
+        }
+        waitForExpectations(timeout: 2, handler: nil)
+    }
 }
diff --git a/SnailTests/ObservableTests.swift b/SnailTests/ObservableTests.swift
@@ -1054,4 +1054,63 @@ class ObservableTests: XCTestCase {
 
         subject.value = true
     }
+
+    func testRemovingFirstSubscriberWithinSubscription_NoDeadlock() {
+        class TestObject {
+            let disposer = Disposer()
+            private let variable: Variable<Bool>
+
+            init(variable: Variable<Bool>) {
+                self.variable = variable
+
+                self.variable.asObservable().subscribe(onNext: { _ in
+                    if let firstSubscriber = variable.subject.subscribers.first {
+                        self.variable.subject.removeSubscriber(subscriber: firstSubscriber)
+                    }
+                })
+                .add(to: disposer)
+            }
+        }
+
+        let subject = Variable(true)
+        let disposer = Disposer()
+
+        subject.asObservable().subscribe(onNext: { _ in
+            _ = TestObject(variable: subject)
+        })
+        .add(to: disposer)
+
+        subject.value = true
+    }
+
+    func testDisposerRemovesSubscriptionFromObserverOnDeinit() {
+        class TestOwner {
+            let disposer = Disposer()
+
+            func bind(observable: Observable<Int>) {
+                observable
+                    .subscribe(queue: .main, onNext: { _ in })
+                    .add(to: disposer)
+            }
+        }
+
+        let observable = Observable<Int>()
+
+        var owner: TestOwner! = TestOwner()
+        owner.bind(observable: observable)
+
+        let isDone = expectation(description: "Disposes")
+
+        DispatchQueue.global().asyncAfter(deadline: .now() + 0.1) {
+            owner = nil
+        }
+
+        DispatchQueue.global().asyncAfter(deadline: .now() + 0.5) {
+            isDone.fulfill()
+        }
+
+        waitForExpectations(timeout: 1) { _ in
+            XCTAssertEqual(observable.subscribers.count, 0)
+        }
+    }
 }