Merge pull request #3 from MihaelIsaev/race-fix

🏎 Implement `knownEventLoop` to fix race condition
MihaelIsaev · Jun 28, 2021 · 0438b55 · 0438b55
2 parents 4910ff4 + 4781d86
commit 0438b55
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 33 deletions.
diff --git a/Sources/WS/Models/Configurator.swift b/Sources/WS/Models/Configurator.swift
@@ -95,4 +95,26 @@ public struct Configurator {
             application.storage[DefaultDecoderKey.self] = newValue
         }
     }
+
+    // MARK: - Default Decoder
+
+    struct KnownEventLoopKey: StorageKey {
+        typealias Value = EventLoop
+    }
+
+    /// Default encoder for all the observers, if `nil` then `JSONEncoder` is used.
+    public var knownEventLoop: EventLoop {
+        get {
+            if let eventLoop = application.storage[KnownEventLoopKey.self] {
+                return eventLoop
+            } else {
+                let eventLoop = application.eventLoopGroup.next()
+                application.storage[KnownEventLoopKey.self] = eventLoop
+                return eventLoop
+            }
+        }
+        nonmutating set {
+            application.storage[KnownEventLoopKey.self] = newValue
+        }
+    }
 }
diff --git a/Sources/WS/Protocols/AnyClient.swift b/Sources/WS/Protocols/AnyClient.swift
@@ -18,7 +18,7 @@ internal protocol _AnyClient: AnyClient, _Disconnectable, _Subscribable, _Sendab
 }
 
 extension AnyClient {
-    public var eventLoop: EventLoop { application.eventLoopGroup.next() }
+    public var eventLoop: EventLoop { application.ws.knownEventLoop }
     public var logger: Logger { application.logger }
 
     /// See `Broadcastable`

diff --git a/Sources/WS/Protocols/AnyObserver.swift b/Sources/WS/Protocols/AnyObserver.swift
@@ -41,7 +41,7 @@ internal protocol _AnyObserver: AnyObserver, _Disconnectable, _Sendable {
 // MARK: - Default implementation
 
 extension AnyObserver {
-    public var eventLoop: EventLoop { application.eventLoopGroup.next() }
+    public var eventLoop: EventLoop { application.ws.knownEventLoop }
 
     var _encoder: Encoder {
         if let encoder = self.encoder {

diff --git a/Sources/WS/Protocols/Subscribable.swift b/Sources/WS/Protocols/Subscribable.swift
@@ -3,22 +3,22 @@ import NIOWebSocket
 
 public protocol Subscribable {
     @discardableResult
-    func subscribe(to channels: String...) -> EventLoopFuture<Void>
+    func subscribe(to channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void>
     @discardableResult
-    func subscribe(to channels: [String]) -> EventLoopFuture<Void>
+    func subscribe(to channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void>
     @discardableResult
-    func unsubscribe(from channels: String...) -> EventLoopFuture<Void>
+    func unsubscribe(from channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void>
     @discardableResult
-    func unsubscribe(from channels: [String]) -> EventLoopFuture<Void>
+    func unsubscribe(from channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void>
 }
 
 extension Subscribable {
-    public func subscribe(to channels: String...) -> EventLoopFuture<Void> {
-        subscribe(to: channels)
+    public func subscribe(to channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        subscribe(to: channels, on: eventLoop)
     }
 
-    public func unsubscribe(from channels: String...) -> EventLoopFuture<Void> {
-        unsubscribe(from: channels)
+    public func unsubscribe(from channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        unsubscribe(from: channels, on: eventLoop)
     }
 }
 
@@ -28,41 +28,51 @@ internal protocol _Subscribable: class, Subscribable {
 }
 
 extension _Subscribable {
-    public func subscribe(to channels: String...) -> EventLoopFuture<Void> {
-        subscribe(to: channels)
+    public func subscribe(to channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        subscribe(to: channels, on: eventLoop)
     }
 
-    public func subscribe(to channels: [String]) -> EventLoopFuture<Void> {
-        channels.forEach { channel in
-            self.clients.forEach {
-                $0.channels.insert(channel)
+    public func subscribe(to channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        self.eventLoop.submit {
+            channels.forEach { channel in
+                self.clients.forEach {
+                    $0.channels.insert(channel)
+                }
             }
-        }
-        return eventLoop.future()
+        }.hop(to: eventLoop)
     }
 
-    public func unsubscribe(from channels: String...) -> EventLoopFuture<Void> {
-        unsubscribe(from: channels)
+    public func unsubscribe(from channels: String..., on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        unsubscribe(from: channels, on: eventLoop)
     }
 
-    public func unsubscribe(from channels: [String]) -> EventLoopFuture<Void> {
-        channels.forEach { channel in
-            self.clients.forEach {
-                $0.channels.remove(channel)
+    public func unsubscribe(from channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void> {
+        self.eventLoop.submit {
+            channels.forEach { channel in
+                self.clients.forEach {
+                    $0.channels.remove(channel)
+                }
             }
-        }
-        return eventLoop.future()
+        }.hop(to: eventLoop)
     }
 }
 
 // MARK: - EventLoopFuture
 
-extension EventLoopFuture: Subscribable where Value: Subscribable {
+extension EventLoopFuture where Value: Subscribable {
+    public func subscribe(to channels: String...) -> EventLoopFuture<Void> {
+        subscribe(to: channels)
+    }
+
     public func subscribe(to channels: [String]) -> EventLoopFuture<Void> {
-        flatMap { $0.subscribe(to: channels) }
+        flatMap { $0.subscribe(to: channels, on: self.eventLoop) }
+    }
+
+    public func unsubscribe(from channels: String...) -> EventLoopFuture<Void> {
+        unsubscribe(from: channels)
     }
 
     public func unsubscribe(from channels: [String]) -> EventLoopFuture<Void> {
-        flatMap { $0.unsubscribe(from: channels) }
+        flatMap { $0.unsubscribe(from: channels, on: self.eventLoop) }
     }
 }
diff --git a/Sources/WS/Utilities/Broadcaster.swift b/Sources/WS/Utilities/Broadcaster.swift
@@ -113,16 +113,16 @@ public class Broadcaster: Disconnectable, _Disconnectable, Sendable, _Sendable,
     /// See `Subscribable`
 
     /// Subscribe filtered clients to channels
-    public func subscribe(to channels: [String]) -> EventLoopFuture<Void> {
+    public func subscribe(to channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void> {
         clients.map {
-            $0.subscribe(to: channels)
+            $0.subscribe(to: channels, on: eventLoop)
         }.flatten(on: eventLoop)
     }
 
     /// Unsubscribe filtered clients from channels
-    public func unsubscribe(from channels: [String]) -> EventLoopFuture<Void> {
+    public func unsubscribe(from channels: [String], on eventLoop: EventLoop) -> EventLoopFuture<Void> {
         clients.map {
-            $0.unsubscribe(from: channels)
+            $0.unsubscribe(from: channels, on: eventLoop)
         }.flatten(on: eventLoop)
     }