Skip to content

Commit 7db041f

Browse files
RecouseRecouse
committed
Fixed DataTask memory leak
1 parent 8761379 commit 7db041f

File tree

1 file changed

+95
-89
lines changed

1 file changed

+95
-89
lines changed

‎Sources/EventSource/EventSource.swift‎

Lines changed: 95 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ public extension EventSource {
6868

6969
private let timeoutInterval: TimeInterval
7070

71+
private var continuation: AsyncStream<EventType>.Continuation?
72+
7173
private var urlSession: URLSession?
7274

7375
private var sessionDelegate = SessionDelegate()
@@ -102,17 +104,21 @@ public extension EventSource {
102104

103105
public func events() -> AsyncStream<EventType> {
104106
AsyncStream { continuation in
105-
continuation.onTermination = { @Sendable _ in
106-
close()
107+
continuation.onTermination = { @Sendable [weak self] _ in
108+
self?.close()
107109
}
108110

111+
self.continuation = continuation
112+
109113
urlSession = URLSession(
110114
configuration: urlSessionConfiguration,
111115
delegate: sessionDelegate,
112116
delegateQueue: nil
113117
)
114118

115-
sessionDelegate.onEvent = { event in
119+
sessionDelegate.onEvent = { [weak self] event in
120+
guard let self else { return }
121+
116122
switch event {
117123
case let .didCompleteWithError(error):
118124
handleSessionError(error)
@@ -126,96 +132,96 @@ public extension EventSource {
126132
urlSessionDataTask = urlSession!.dataTask(with: urlRequest)
127133
urlSessionDataTask!.resume()
128134
readyState = .connecting
129-
130-
func handleSessionError(_ error: Error?) {
131-
guard readyState != .closed else {
132-
close()
133-
return
134-
}
135-
136-
// Send error event
137-
if let error {
138-
sendErrorEvent(with: error)
139-
}
140-
141-
// Close connection
142-
close()
143-
}
144-
145-
func handleSessionResponse(
146-
_ response: URLResponse,
147-
completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
148-
) {
149-
guard readyState != .closed else {
150-
completionHandler(.cancel)
151-
return
152-
}
153-
154-
guard let httpResponse = response as? HTTPURLResponse else {
155-
completionHandler(.cancel)
156-
return
157-
}
158-
159-
// Stop connection when 204 response code, otherwise keep open
160-
guard httpResponse.statusCode != 204 else {
161-
completionHandler(.cancel)
162-
close()
163-
return
164-
}
165-
166-
if 200...299 ~= httpResponse.statusCode {
167-
if readyState != .open {
168-
setOpen()
169-
}
170-
} else {
171-
httpResponseErrorStatusCode = httpResponse.statusCode
172-
}
173-
174-
completionHandler(.allow)
175-
}
176-
177-
/// Closes the connection, if one was made,
178-
/// and sets the `readyState` property to `.closed`.
179-
/// - Returns: State before closing.
180-
@Sendable func close() {
181-
let previousState = self.readyState
182-
if previousState != .closed {
183-
continuation.yield(.closed)
184-
continuation.finish()
185-
}
186-
cancel()
187-
}
188-
189-
func parseMessages(from data: Data) {
190-
if let httpResponseErrorStatusCode {
191-
self.httpResponseErrorStatusCode = nil
192-
handleSessionError(
193-
EventSourceError.connectionError(statusCode: httpResponseErrorStatusCode, response: data)
194-
)
195-
return
196-
}
197-
198-
let messages = messageParser.parse(data)
135+
}
136+
}
137+
138+
private func handleSessionError(_ error: Error?) {
139+
guard readyState != .closed else {
140+
close()
141+
return
142+
}
143+
144+
// Send error event
145+
if let error {
146+
sendErrorEvent(with: error)
147+
}
199148

200-
// Update last message ID
201-
if let lastMessageWithId = messages.last(where: { $0.id != nil }) {
202-
lastMessageId = lastMessageWithId.id ?? ""
203-
}
149+
// Close connection
150+
close()
151+
}
152+
153+
private func handleSessionResponse(
154+
_ response: URLResponse,
155+
completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
156+
) {
157+
guard readyState != .closed else {
158+
completionHandler(.cancel)
159+
return
160+
}
161+
162+
guard let httpResponse = response as? HTTPURLResponse else {
163+
completionHandler(.cancel)
164+
return
165+
}
204166

205-
messages.forEach {
206-
continuation.yield(.message($0))
207-
}
208-
}
209-
210-
func setOpen() {
211-
readyState = .open
212-
continuation.yield(.open)
213-
}
214-
215-
func sendErrorEvent(with error: Error) {
216-
continuation.yield(.error(error))
167+
// Stop connection when 204 response code, otherwise keep open
168+
guard httpResponse.statusCode != 204 else {
169+
completionHandler(.cancel)
170+
close()
171+
return
172+
}
173+
174+
if 200...299 ~= httpResponse.statusCode {
175+
if readyState != .open {
176+
setOpen()
217177
}
178+
} else {
179+
httpResponseErrorStatusCode = httpResponse.statusCode
218180
}
181+
182+
completionHandler(.allow)
183+
}
184+
185+
/// Closes the connection, if one was made,
186+
/// and sets the `readyState` property to `.closed`.
187+
/// - Returns: State before closing.
188+
@Sendable private func close() {
189+
let previousState = self.readyState
190+
if previousState != .closed {
191+
continuation?.yield(.closed)
192+
continuation?.finish()
193+
}
194+
cancel()
195+
}
196+
197+
private func parseMessages(from data: Data) {
198+
if let httpResponseErrorStatusCode {
199+
self.httpResponseErrorStatusCode = nil
200+
handleSessionError(
201+
EventSourceError.connectionError(statusCode: httpResponseErrorStatusCode, response: data)
202+
)
203+
return
204+
}
205+
206+
let messages = messageParser.parse(data)
207+
208+
// Update last message ID
209+
if let lastMessageWithId = messages.last(where: { $0.id != nil }) {
210+
lastMessageId = lastMessageWithId.id ?? ""
211+
}
212+
213+
messages.forEach {
214+
continuation?.yield(.message($0))
215+
}
216+
}
217+
218+
private func setOpen() {
219+
readyState = .open
220+
continuation?.yield(.open)
221+
}
222+
223+
private func sendErrorEvent(with error: Error) {
224+
continuation?.yield(.error(error))
219225
}
220226

221227
public func cancel() {

0 commit comments

Comments
 (0)