Allow VirtualTimeScheduler to run on any thread as long as critical m…

…ethods are all called on the same thread.
ReactiveX · Jul 10, 2024 · 16ca28f · 16ca28f
1 parent bfe5015
commit 16ca28f
Showing 1 changed file with 24 additions and 7 deletions.
diff --git a/RxSwift/Schedulers/VirtualTimeScheduler.swift b/RxSwift/Schedulers/VirtualTimeScheduler.swift
@@ -24,6 +24,8 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
 
     private var nextId = 0
 
+    private var thread: Thread!
+
     /// - returns: Current time.
     public var now: RxTime {
         self.converter.convertFromVirtualTime(self.clock)
@@ -106,7 +108,10 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
      - returns: The disposable object used to cancel the scheduled action (best effort).
      */
     public func scheduleAbsoluteVirtual<StateType>(_ state: StateType, time: VirtualTime, action: @escaping (StateType) -> Disposable) -> Disposable {
-        MainScheduler.ensureExecutingOnScheduler()
+        guard thread == nil || Thread.current == thread else {
+            rxFatalError("Executing on the wrong thread. Please ensure all work on the same thread.")
+        }
+        thread = Thread.current
 
         let compositeDisposable = CompositeDisposable()
 
@@ -130,12 +135,15 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
 
     /// Starts the virtual time scheduler.
     public func start() {
-        MainScheduler.ensureExecutingOnScheduler()
-
         if self.running {
             return
         }
 
+        guard thread == nil || Thread.current == thread else {
+            rxFatalError("Executing on the wrong thread. Please ensure all work on the same thread.")
+        }
+        thread = Thread.current
+
         self.running = true
         repeat {
             guard let next = self.findNext() else {
@@ -170,12 +178,15 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
     ///
     /// - parameter virtualTime: Absolute time to advance the scheduler's clock to.
     public func advanceTo(_ virtualTime: VirtualTime) {
-        MainScheduler.ensureExecutingOnScheduler()
-
         if self.running {
             fatalError("Scheduler is already running")
         }
 
+        guard thread == nil || Thread.current == thread else {
+            rxFatalError("Executing on the wrong thread. Please ensure all work on the same thread.")
+        }
+        thread = Thread.current
+
         self.running = true
         repeat {
             guard let next = self.findNext() else {
@@ -199,7 +210,10 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
 
     /// Advances the scheduler's clock by the specified relative time.
     public func sleep(_ virtualInterval: VirtualTimeInterval) {
-        MainScheduler.ensureExecutingOnScheduler()
+        guard thread == nil || Thread.current == thread else {
+            rxFatalError("Executing on the wrong thread. Please ensure all work on the same thread.")
+        }
+        thread = Thread.current
 
         let sleepTo = self.converter.offsetVirtualTime(self.clock, offset: virtualInterval)
         if self.converter.compareVirtualTime(sleepTo, self.clock).lessThen {
@@ -211,7 +225,10 @@ open class VirtualTimeScheduler<Converter: VirtualTimeConverterType>
 
     /// Stops the virtual time scheduler.
     public func stop() {
-        MainScheduler.ensureExecutingOnScheduler()
+        guard thread == nil || Thread.current == thread else {
+            rxFatalError("Executing on the wrong thread. Please ensure all work on the same thread.")
+        }
+        thread = Thread.current
 
         self.running = false
     }