Skip to content

Commit e9abba3

Browse files
NeedleInAJayStackNeedleInAJayStack
committed
Fixes memory reference issues
1 parent 87d28d9 commit e9abba3

File tree

5 files changed

+66
-48
lines changed

5 files changed

+66
-48
lines changed

‎README.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,12 @@ Features:
77
- Server implementation that implements defined protocol conversations
88
- Client and Server types that wrap messengers
99
- Codable Server and Client message structures
10+
11+
## Memory Management
12+
13+
Memory ownership among the Server, Client, and Messager may seem a little backwards. This is because the Swift/Vapor WebSocket
14+
implementation persists WebSocket objects long after their callback and they are expected to retain strong memory references to the
15+
objects required for responses. In order to align cleanly and avoid memory cycles, Server and Client are injected strongly into Messager
16+
callbacks, and only hold weak references to their Messager. This means that Messager objects (or their enclosing WebSocket) must
17+
be persisted to have the connected Server or Client objects function. That is, if a Server's Messager falls out of scope and deinitializes,
18+
the Server will no longer respond to messages.

‎Sources/GraphQLWS/Client.swift

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import GraphQL
55

66
/// Client is an open-ended implementation of the client side of the protocol. It parses and adds callbacks for each type of server respose.
77
public class Client {
8-
let messenger: Messenger
8+
// We keep this weak because we strongly inject this object into the messenger callback
9+
weak var messenger: Messenger?
910

1011
var onConnectionError: (ConnectionErrorResponse, Client) -> Void = { _, _ in }
1112
var onConnectionAck: (ConnectionAckResponse, Client) -> Void = { _, _ in }
@@ -26,9 +27,8 @@ public class Client {
2627
messenger: Messenger
2728
) {
2829
self.messenger = messenger
29-
30-
self.messenger.onRecieve { [weak self] message in
31-
guard let self = self else { return }
30+
messenger.onRecieve { message in
31+
guard let messenger = self.messenger else { return }
3232

3333
self.onMessage(message, self)
3434

@@ -40,7 +40,7 @@ public class Client {
4040

4141
guard let json = message.data(using: .utf8) else {
4242
let error = GraphQLWSError.invalidEncoding()
43-
self.messenger.error(error.message, code: error.code)
43+
messenger.error(error.message, code: error.code)
4444
return
4545
}
4646

@@ -50,56 +50,56 @@ public class Client {
5050
}
5151
catch {
5252
let error = GraphQLWSError.noType()
53-
self.messenger.error(error.message, code: error.code)
53+
messenger.error(error.message, code: error.code)
5454
return
5555
}
5656

5757
switch response.type {
5858
case .GQL_CONNECTION_ERROR:
5959
guard let connectionErrorResponse = try? self.decoder.decode(ConnectionErrorResponse.self, from: json) else {
6060
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_CONNECTION_ERROR)
61-
self.messenger.error(error.message, code: error.code)
61+
messenger.error(error.message, code: error.code)
6262
return
6363
}
6464
self.onConnectionError(connectionErrorResponse, self)
6565
case .GQL_CONNECTION_ACK:
6666
guard let connectionAckResponse = try? self.decoder.decode(ConnectionAckResponse.self, from: json) else {
6767
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_CONNECTION_ACK)
68-
self.messenger.error(error.message, code: error.code)
68+
messenger.error(error.message, code: error.code)
6969
return
7070
}
7171
self.onConnectionAck(connectionAckResponse, self)
7272
case .GQL_CONNECTION_KEEP_ALIVE:
7373
guard let connectionKeepAliveResponse = try? self.decoder.decode(ConnectionKeepAliveResponse.self, from: json) else {
7474
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_CONNECTION_KEEP_ALIVE)
75-
self.messenger.error(error.message, code: error.code)
75+
messenger.error(error.message, code: error.code)
7676
return
7777
}
7878
self.onConnectionKeepAlive(connectionKeepAliveResponse, self)
7979
case .GQL_DATA:
8080
guard let nextResponse = try? self.decoder.decode(DataResponse.self, from: json) else {
8181
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_DATA)
82-
self.messenger.error(error.message, code: error.code)
82+
messenger.error(error.message, code: error.code)
8383
return
8484
}
8585
self.onData(nextResponse, self)
8686
case .GQL_ERROR:
8787
guard let errorResponse = try? self.decoder.decode(ErrorResponse.self, from: json) else {
8888
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_ERROR)
89-
self.messenger.error(error.message, code: error.code)
89+
messenger.error(error.message, code: error.code)
9090
return
9191
}
9292
self.onError(errorResponse, self)
9393
case .GQL_COMPLETE:
9494
guard let completeResponse = try? self.decoder.decode(CompleteResponse.self, from: json) else {
9595
let error = GraphQLWSError.invalidResponseFormat(messageType: .GQL_COMPLETE)
96-
self.messenger.error(error.message, code: error.code)
96+
messenger.error(error.message, code: error.code)
9797
return
9898
}
9999
self.onComplete(completeResponse, self)
100100
case .unknown:
101101
let error = GraphQLWSError.invalidType()
102-
self.messenger.error(error.message, code: error.code)
102+
messenger.error(error.message, code: error.code)
103103
}
104104
}
105105
}
@@ -148,6 +148,7 @@ public class Client {
148148

149149
/// Send a `connection_init` request through the messenger
150150
public func sendConnectionInit(payload: ConnectionInitAuth?) {
151+
guard let messenger = messenger else { return }
151152
messenger.send(
152153
ConnectionInitRequest(
153154
payload: payload
@@ -157,6 +158,7 @@ public class Client {
157158

158159
/// Send a `start` request through the messenger
159160
public func sendStart(payload: GraphQLRequest, id: String) {
161+
guard let messenger = messenger else { return }
160162
messenger.send(
161163
StartRequest(
162164
payload: payload,
@@ -167,6 +169,7 @@ public class Client {
167169

168170
/// Send a `stop` request through the messenger
169171
public func sendStop(id: String) {
172+
guard let messenger = messenger else { return }
170173
messenger.send(
171174
StopRequest(
172175
id: id
@@ -176,6 +179,7 @@ public class Client {
176179

177180
/// Send a `connection_terminate` request through the messenger
178181
public func sendConnectionTerminate() {
182+
guard let messenger = messenger else { return }
179183
messenger.send(
180184
ConnectionTerminateRequest().toJSON(encoder)
181185
)

‎Sources/GraphQLWS/Messenger.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@ import Foundation
44
import NIO
55

66
/// Protocol for an object that can send and recieve messages
7-
public protocol Messenger {
7+
public protocol Messenger: AnyObject {
8+
// AnyObject compliance requires that the implementing object is a class and we can reference it weakly
89
func send<S>(_ message: S) -> Void where S: Collection, S.Element == Character
910
func onRecieve(callback: @escaping (String) -> Void) -> Void
1011
func close() -> Void

‎Sources/GraphQLWS/Server.swift

Lines changed: 36 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@ import RxSwift
88

99
/// Server implements the server-side portion of the protocol, allowing a few callbacks for customization.
1010
public class Server {
11-
let messenger: Messenger
11+
// We keep this weak because we strongly inject this object into the messenger callback
12+
weak var messenger: Messenger?
1213

1314
let onExecute: (GraphQLRequest) -> EventLoopFuture<GraphQLResult>
1415
let onSubscribe: (GraphQLRequest) -> EventLoopFuture<SubscriptionResult>
@@ -38,8 +39,8 @@ public class Server {
3839
self.onExecute = onExecute
3940
self.onSubscribe = onSubscribe
4041

41-
self.messenger.onRecieve { [weak self] message in
42-
guard let self = self else { return }
42+
messenger.onRecieve { message in
43+
guard let messenger = self.messenger else { return }
4344

4445
self.onMessage(message)
4546

@@ -51,7 +52,7 @@ public class Server {
5152

5253
guard let json = message.data(using: .utf8) else {
5354
let error = GraphQLWSError.invalidEncoding()
54-
self.messenger.error(error.message, code: error.code)
55+
messenger.error(error.message, code: error.code)
5556
return
5657
}
5758

@@ -61,42 +62,42 @@ public class Server {
6162
}
6263
catch {
6364
let error = GraphQLWSError.noType()
64-
self.messenger.error(error.message, code: error.code)
65+
messenger.error(error.message, code: error.code)
6566
return
6667
}
6768

6869
switch request.type {
6970
case .GQL_CONNECTION_INIT:
7071
guard let connectionInitRequest = try? self.decoder.decode(ConnectionInitRequest.self, from: json) else {
7172
let error = GraphQLWSError.invalidRequestFormat(messageType: .GQL_CONNECTION_INIT)
72-
self.messenger.error(error.message, code: error.code)
73+
messenger.error(error.message, code: error.code)
7374
return
7475
}
75-
self.onConnectionInit(connectionInitRequest)
76+
self.onConnectionInit(connectionInitRequest, messenger)
7677
case .GQL_START:
7778
guard let startRequest = try? self.decoder.decode(StartRequest.self, from: json) else {
7879
let error = GraphQLWSError.invalidRequestFormat(messageType: .GQL_START)
79-
self.messenger.error(error.message, code: error.code)
80+
messenger.error(error.message, code: error.code)
8081
return
8182
}
82-
self.onStart(startRequest)
83+
self.onStart(startRequest, messenger)
8384
case .GQL_STOP:
8485
guard let stopRequest = try? self.decoder.decode(StopRequest.self, from: json) else {
8586
let error = GraphQLWSError.invalidRequestFormat(messageType: .GQL_STOP)
86-
self.messenger.error(error.message, code: error.code)
87+
messenger.error(error.message, code: error.code)
8788
return
8889
}
89-
self.onStop(stopRequest, self.messenger)
90+
self.onStop(stopRequest, messenger)
9091
case .GQL_CONNECTION_TERMINATE:
9192
guard let connectionTerminateRequest = try? self.decoder.decode(ConnectionTerminateRequest.self, from: json) else {
9293
let error = GraphQLWSError.invalidRequestFormat(messageType: .GQL_CONNECTION_TERMINATE)
93-
self.messenger.error(error.message, code: error.code)
94+
messenger.error(error.message, code: error.code)
9495
return
9596
}
96-
self.onConnectionTerminate(connectionTerminateRequest)
97+
self.onConnectionTerminate(connectionTerminateRequest, messenger)
9798
case .unknown:
9899
let error = GraphQLWSError.invalidType()
99-
self.messenger.error(error.message, code: error.code)
100+
messenger.error(error.message, code: error.code)
100101
}
101102
}
102103

@@ -126,7 +127,7 @@ public class Server {
126127
self.onMessage = callback
127128
}
128129

129-
private func onConnectionInit(_ connectionInitRequest: ConnectionInitRequest) {
130+
private func onConnectionInit(_ connectionInitRequest: ConnectionInitRequest, _ messenger: Messenger) {
130131
guard !initialized else {
131132
let error = GraphQLWSError.tooManyInitializations()
132133
messenger.error(error.message, code: error.code)
@@ -148,7 +149,7 @@ public class Server {
148149
// TODO: Should we send the `ka` message?
149150
}
150151

151-
private func onStart(_ startRequest: StartRequest) {
152+
private func onStart(_ startRequest: StartRequest, _ messenger: Messenger) {
152153
guard initialized else {
153154
let error = GraphQLWSError.notInitialized()
154155
messenger.error(error.message, code: error.code)
@@ -169,48 +170,50 @@ public class Server {
169170

170171
if isStreaming {
171172
let subscribeFuture = onSubscribe(graphQLRequest)
172-
subscribeFuture.whenSuccess { [weak self] result in
173-
guard let self = self else { return }
173+
subscribeFuture.whenSuccess { result in
174174
guard let streamOpt = result.stream else {
175175
// API issue - subscribe resolver isn't stream
176176
let error = GraphQLWSError.internalAPIStreamIssue()
177-
self.messenger.error(error.message, code: error.code)
177+
messenger.error(error.message, code: error.code)
178178
return
179179
}
180180
let stream = streamOpt as! ObservableSubscriptionEventStream
181181
let observable = stream.observable
182182
observable.subscribe(
183-
onNext: { resultFuture in
183+
onNext: { [weak self] resultFuture in
184+
guard let self = self, let messenger = self.messenger else { return }
184185
resultFuture.whenSuccess { result in
185-
self.messenger.send(DataResponse(result, id: id).toJSON(self.encoder))
186+
messenger.send(DataResponse(result, id: id).toJSON(self.encoder))
186187
}
187188
resultFuture.whenFailure { error in
188-
self.messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
189+
messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
189190
}
190191
},
191-
onError: { error in
192-
self.messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
192+
onError: { [weak self] error in
193+
guard let self = self, let messenger = self.messenger else { return }
194+
messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
193195
},
194-
onCompleted: {
195-
self.messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
196-
_ = self.messenger.close()
196+
onCompleted: { [weak self] in
197+
guard let self = self, let messenger = self.messenger else { return }
198+
messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
199+
_ = messenger.close()
197200
}
198201
).disposed(by: self.disposeBag)
199202
}
200203
subscribeFuture.whenFailure { error in
201204
let error = GraphQLWSError.graphQLError(error)
202-
_ = self.messenger.error(error.message, code: error.code)
205+
_ = messenger.error(error.message, code: error.code)
203206
}
204207
}
205208
else {
206209
let executeFuture = onExecute(graphQLRequest)
207210
executeFuture.whenSuccess { result in
208-
self.messenger.send(DataResponse(result, id: id).toJSON(self.encoder))
209-
self.messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
211+
messenger.send(DataResponse(result, id: id).toJSON(self.encoder))
212+
messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
210213
}
211214
executeFuture.whenFailure { error in
212-
self.messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
213-
self.messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
215+
messenger.send(ErrorResponse(error, id: id).toJSON(self.encoder))
216+
messenger.send(CompleteResponse(id: id).toJSON(self.encoder))
214217
}
215218
}
216219
}
@@ -224,7 +227,7 @@ public class Server {
224227
onExit()
225228
}
226229

227-
private func onConnectionTerminate(_: ConnectionTerminateRequest) {
230+
private func onConnectionTerminate(_: ConnectionTerminateRequest, _ messenger: Messenger) {
228231
onExit()
229232
_ = messenger.close()
230233
}

‎Tests/GraphQLWSTests/GraphQLWSTests.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,12 @@ import XCTest
1010

1111
class GraphqlWsTests: XCTestCase {
1212
var clientMessenger: TestMessenger!
13+
var serverMessenger: TestMessenger!
1314
var server: Server!
1415

1516
override func setUp() {
1617
clientMessenger = TestMessenger()
17-
let serverMessenger = TestMessenger()
18+
serverMessenger = TestMessenger()
1819

1920
clientMessenger.other = serverMessenger
2021
serverMessenger.other = clientMessenger

0 commit comments

Comments
 (0)