Merge pull request #102 from d-exclaimation/leak-fix

Fixing issues with ConcurrentEventStream from #101
GraphQLSwift · Jul 15, 2022 · 8f0a07e · 8f0a07e
2 parents 3071750 + a99f92a
commit 8f0a07e
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 11 deletions.
diff --git a/Sources/GraphQL/Subscription/EventStream.swift b/Sources/GraphQL/Subscription/EventStream.swift
@@ -31,24 +31,42 @@ public class ConcurrentEventStream<Element>: EventStream<Element> {
 extension AsyncThrowingStream {
     func mapStream<To>(_ closure: @escaping (Element) throws -> To) -> AsyncThrowingStream<To, Error> {
         return AsyncThrowingStream<To, Error> { continuation in
-            Task {
-                for try await event in self {
-                    let newEvent = try closure(event)
-                    continuation.yield(newEvent)
+            let task = Task {
+                do {
+                    for try await event in self {
+                        let newEvent = try closure(event)
+                        continuation.yield(newEvent)
+                    }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
                 }
             }
+
+            continuation.onTermination = { @Sendable reason in 
+                task.cancel()
+            }
         }
     }
 
     func filterStream(_ isIncluded: @escaping (Element) throws -> Bool) -> AsyncThrowingStream<Element, Error> {
         return AsyncThrowingStream<Element, Error> { continuation in
-            Task {
-                for try await event in self {
-                    if try isIncluded(event) {
-                        continuation.yield(event)
+            let task = Task {
+                do {
+                    for try await event in self {
+                        if try isIncluded(event) {
+                            continuation.yield(event)
+                        }
                     }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
                 }
             }
+
+            continuation.onTermination = { @Sendable _ in 
+                task.cancel()
+            }
         }
     }
 }

diff --git a/Tests/GraphQLTests/SubscriptionTests/SubscriptionTests.swift b/Tests/GraphQLTests/SubscriptionTests/SubscriptionTests.swift
@@ -605,12 +605,15 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
             }
         }
+
         var expected = [GraphQLResult]()
 
         db.trigger(email: Email(
@@ -675,6 +678,9 @@ class SubscriptionTests : XCTestCase {
         )
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should not trigger when subscription is already done'
@@ -701,7 +707,8 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
@@ -747,6 +754,9 @@ class SubscriptionTests : XCTestCase {
         // Ensure that the current result was the one before the db was stopped
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should not trigger when subscription is thrown'
@@ -861,7 +871,8 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
@@ -925,6 +936,9 @@ class SubscriptionTests : XCTestCase {
         )
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should pass through error thrown in source event stream'