Fix schedule(after:interval:leeway:) disposal (#584)

* Fix `schedule(after:interval:leeway:)` disposal

Retain timers in `QueueScheduler` so that they will not be cancelled if the user doesn’t retain the returned `Disposable`.

* Add pull request # and handle

* Use `Atomic` instead of `Lock`

* Separate the tests for repeated action disposal

* Use `Set` instead of `Array` to store timers

* Use `ObjectIdentifier(value).hashValue`

* Weakify `self` in the returned `Disposable`

* Fix `QueueScheduler` tests

* Fix increment in test

Specify `pollInterval` on the asynchronous expectation to avoid false positive test success.

* Conform to timer wrapper to `Hashable` by self ref

`DispatchSourceTimer` doesn’t inherit from `NSObjectProtocol` on Linux. Basing hashing and equality on the wrapper is sufficient because the  same wrapper instance can be used insert and remove from the `Set` of retained timers.

Author: jjoelson

CHANGELOG.md

@@ -1,6 +1,8 @@
 # master
 *Please add new entries at the top.*
 
+1. Fixed `schedule(after:interval:leeway:)` being cancelled when the returned `Disposable` is not retained. (#584, kudos to @jjoelson)
+
 # 3.1.0-rc.1
 1. Fixed a scenario of downstream interruptions being dropped. (#577, kudos to @andersio)
 

Sources/Scheduler.swift

@@ -177,6 +177,26 @@ public final class UIScheduler: Scheduler {
 	}
 }
 
+/// A `Hashable` wrapper for `DispatchSourceTimer`. `Hashable` conformance is
+/// based on the identity of the wrapper object rather than the wrapped
+/// `DispatchSourceTimer`, so two wrappers wrapping the same timer will *not*
+/// be equal.
+private final class DispatchSourceTimerWrapper: Hashable {
+	private let value: DispatchSourceTimer
+	
+	fileprivate var hashValue: Int {
+		return ObjectIdentifier(self).hashValue
+	}
+	
+	fileprivate init(_ value: DispatchSourceTimer) {
+		self.value = value
+	}
+	
+	fileprivate static func ==(lhs: DispatchSourceTimerWrapper, rhs: DispatchSourceTimerWrapper) -> Bool {
+		return lhs === rhs
+	}
+}
+
 /// A scheduler backed by a serial GCD queue.
 public final class QueueScheduler: DateScheduler {
 	/// A singleton `QueueScheduler` that always targets the main thread's GCD
@@ -192,9 +212,12 @@ public final class QueueScheduler: DateScheduler {
 	}
 
 	public let queue: DispatchQueue
-
+	
+	private var timers: Atomic<Set<DispatchSourceTimerWrapper>>
+	
 	internal init(internalQueue: DispatchQueue) {
 		queue = internalQueue
+		timers = Atomic(Set())
 	}
 
 	/// Initializes a scheduler that will target the given queue with its
@@ -340,8 +363,20 @@ public final class QueueScheduler: DateScheduler {
 		timer.setEventHandler(handler: action)
 		timer.resume()
 
-		return AnyDisposable {
+		let wrappedTimer = DispatchSourceTimerWrapper(timer)
+		
+		timers.modify { timers in
+			timers.insert(wrappedTimer)
+		}
+
+		return AnyDisposable { [weak self] in
 			timer.cancel()
+			
+			if let scheduler = self {
+				scheduler.timers.modify { timers in
+					timers.remove(wrappedTimer)
+				}
+			}
 		}
 	}
 }

Tests/ReactiveSwiftTests/SchedulerSpec.swift

@@ -218,6 +218,67 @@ class SchedulerSpec: QuickSpec {
 					scheduler.queue.resume()
 					expect{count}.toEventually(equal(timesToRun))
 				}
+				
+				it("should repeatedly run actions after a given date when the disposable is not retained") {
+					let disposable = SerialDisposable()
+					
+					var count = 0
+					let timesToIncrement = 3
+					
+					// Schedule within a function so that the disposable is guaranteed to be deinitialised.
+					func scheduleAndDeinitDisposable() {
+						scheduler.schedule(after: Date(), interval: .milliseconds(10), leeway: .seconds(0)) {
+							expect(Thread.isMainThread) == false
+							
+							if count < timesToIncrement {
+								count += 1
+							}
+						}
+					}
+					
+					scheduleAndDeinitDisposable()
+					
+					expect(count) == 0
+					
+					scheduler.queue.resume()
+					expect(count).toEventually(equal(timesToIncrement), pollInterval: 0.1)
+				}
+				
+				it("should cancel repeatedly run actions on disposal") {
+					// Start two repeating timers, dispose the first, and ensure only the second runs.
+
+					let disposable1 = SerialDisposable()
+					let disposable2 = SerialDisposable()
+
+					var count = 0
+					let timesToRun = 3
+
+					let interval = DispatchTimeInterval.milliseconds(10)
+					
+					disposable1.inner = scheduler.schedule(after: Date(), interval: interval, leeway: .seconds(0)) {
+						fail("timer not cancelled on disposal")
+					}
+
+					disposable2.inner = scheduler.schedule(after: Date(), interval: interval, leeway: .seconds(0)) {
+						expect(Thread.isMainThread) == false
+						
+						count += 1
+						
+						if count == timesToRun {
+							disposable2.dispose()
+						}
+					}
+
+					disposable1.dispose()
+					
+					expect(count) == 0
+					
+					scheduler.queue.resume()
+					
+					// This expectation should take about 2.0 * interval to be fulfilled, and that's
+					// enough time to ensure that the first timer was actually cancelled.
+					expect(count).toEventually(equal(timesToRun))
+				}
 			}
 		}