Merge pull request #33 from grahamburgsma/websocket-retain-cycle

Fixed websocket reference cycle
vapor · Mar 26, 2019 · d85e5b6 · d85e5b6
2 parents 21eb477 + bb2b082
commit d85e5b6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/Sources/WebSocket/WebSocket.swift b/Sources/WebSocket/WebSocket.swift
@@ -83,21 +83,21 @@ public final class WebSocket: BasicWorker {
         self.onErrorCallback = { _, _ in }
         self.onCloseCodeCallback = { _ in }
 
-        self.onTextCallback = { _, text in
+        self.onTextCallback = { [unowned self] _, text in
             guard self.textBuffer != nil else {
                 self.textBuffer = [text]
                 return
             }
             self.textBuffer?.append(text)
         }
-        self.onBinaryCallback = { _, binary in
+        self.onBinaryCallback = { [unowned self] _, binary in
             guard self.binaryBuffer != nil else {
                 self.binaryBuffer = [binary]
                 return
             }
             self.binaryBuffer?.append(binary)
         }
-        self.onErrorCallback = { _, error in
+        self.onErrorCallback = { [unowned self] _, error in
             guard self.errorBuffer != nil else {
                 self.errorBuffer = [error]
                 return

diff --git a/Tests/WebSocketTests/WebSocketTests.swift b/Tests/WebSocketTests/WebSocketTests.swift
@@ -112,11 +112,22 @@ class WebSocketTests: XCTestCase {
         try server.onClose.wait()
     }
 
+    func testDeallocation() throws {
+        let worker = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+
+        weak var ws = try HTTPClient.webSocket(hostname: "echo.websocket.org", on: worker).wait()
+        ws?.close()
+        try ws?.onClose.wait()
+
+        XCTAssertNil(ws, "Websocket not deallocated")
+    }
+
     static let allTests = [
         ("testClient", testClient),
         ("testClientTLS", testClientTLS),
         ("testServer", testServer),
         ("testServerContinuation", testServerContinuation),
+        ("testDeallocation", testDeallocation),
     ]
 }