From 354163cef7f2e931d9677d9f50918d4b0bade1c8 Mon Sep 17 00:00:00 2001
From: Subeom Choi <subumm1@gmail.com>
Date: Wed, 17 Apr 2024 20:11:17 +0900
Subject: [PATCH] feature: add cancel when deinit to client

---
 Source/Concurrency/Promise/Promise.swift      | 16 +++----
 .../Network/Implement/Client/DataClient.swift | 46 ++++++++++++++-----
 .../Network/Implement/Client/JSONClient.swift |  7 ++-
 Test/Network/DataClientTest.swift             | 12 ++---
 Test/Network/JSONClientTest.swift             | 18 ++++----
 5 files changed, 63 insertions(+), 36 deletions(-)

diff --git a/Source/Concurrency/Promise/Promise.swift b/Source/Concurrency/Promise/Promise.swift
index fe20d73..e023dff 100644
--- a/Source/Concurrency/Promise/Promise.swift
+++ b/Source/Concurrency/Promise/Promise.swift
@@ -346,7 +346,7 @@ extension Promise {
 extension Promise {
     public static func pending(
         on queue: DispatchQueue = .global(),
-        cancelWhen: PromisePending<Value, Failure>.CancelWhen = .none
+        cancelWhen: PromisePendingCancelWhen = .none
     ) -> PromisePending<Value, Failure> {
         PromisePending(
             queue: queue,
@@ -394,11 +394,11 @@ public final class PromisePending<Value, Failure: Error> {
     public let cancel: () -> Void
     public let onCancel: (@escaping () -> Void) -> Void
     
-    let cancelWhen: CancelWhen
+    let cancelWhen: PromisePendingCancelWhen
     
     init(
         queue: DispatchQueue,
-        cancelWhen: CancelWhen
+        cancelWhen: PromisePendingCancelWhen
     ) {
         let promise = Promise<Value, Failure>(queue: queue)
         
@@ -416,11 +416,11 @@ public final class PromisePending<Value, Failure: Error> {
             cancel()
         }
     }
-    
-    public enum CancelWhen {
-        case `deinit`
-        case none
-    }
+}
+
+public enum PromisePendingCancelWhen {
+    case `deinit`
+    case none
 }
 
 enum PromiseState<Value, Failure: Error> {
diff --git a/Source/Network/Implement/Client/DataClient.swift b/Source/Network/Implement/Client/DataClient.swift
index 4a5199c..13434dc 100644
--- a/Source/Network/Implement/Client/DataClient.swift
+++ b/Source/Network/Implement/Client/DataClient.swift
@@ -15,20 +15,33 @@ public final class DataClient: Client {
     public typealias Request = Data?
     public typealias Response = Data?
     
+    var storage = [ObjectIdentifier: PromisePending<ClientResult<Data?>, Error>]()
+    
     let session: URLSession
+    let cancelWhen: PromisePendingCancelWhen
     
-    init(session: URLSession) {
+    init(
+        session: URLSession,
+        cancelWhen: PromisePendingCancelWhen
+    ) {
         self.session = session
+        self.cancelWhen = cancelWhen
     }
 }
 
 extension DataClient {
-    public convenience init(optionBlock: (inout URLSessionConfiguration) -> Void = { _ in }) {
+    public convenience init(
+        cancelWhen: PromisePendingCancelWhen,
+        optionBlock: (inout URLSessionConfiguration) -> Void = { _ in }
+    ) {
         var configuration = URLSessionConfiguration.default
         optionBlock(&configuration)
     
         let session = URLSession(configuration: configuration)
-        self.init(session: session)
+        self.init(
+            session: session,
+            cancelWhen: cancelWhen
+        )
     }
 }
 
@@ -41,19 +54,19 @@ extension DataClient {
         timeout: Interval? = nil,
         optionBlock: @escaping (inout URLRequest) -> Void = { _ in }
     ) -> Promise<ClientResult<Data?>, Error> {
-        let pending = Promise<ClientResult<Data?>, Error>.pending()
+        let pending = Promise<ClientResult<Data?>, Error>.pending(cancelWhen: cancelWhen)
         
         var request = URLRequest(url: url)
-        optionBlock(&request)
-
-        request.url = url
         request.httpMethod = method.rawValue
         header.forEach { (key, value) in
             request.setValue(value, forHTTPHeaderField: key)
         }
         request.httpBody = body
+        optionBlock(&request)
         
-        let task = session.dataTask(with: request) { data, response, error in
+        let task = session.dataTask(with: request) { [weak pending] data, response, error in
+            guard let pending else { return }
+            
             if let error = error {
                 pending.reject(error)
                 return
@@ -76,9 +89,9 @@ extension DataClient {
             task.cancel()
         }
         if let timeout {
-            let item = DispatchWorkItem {
-                pending.reject(DataClientError.timeout)
-                task.cancel()
+            let item = DispatchWorkItem { [weak task, weak pending] in
+                task?.cancel()
+                pending?.reject(DataClientError.timeout)
             }
             DispatchQueue.global().asyncAfter(
                 deadline: .now() + timeout.dispatchTime,
@@ -94,6 +107,17 @@ extension DataClient {
         
         task.resume()
         
+        storage[ObjectIdentifier(pending)] = pending
+        let remove = { [weak self, weak pending] in
+            guard let self, let pending else { return }
+            self.storage[ObjectIdentifier(pending)] = nil
+        }
+        pending.promise.subscribe(
+            onResolved: { _ in remove() },
+            onRejected: { _ in remove() },
+            onCanceled: { remove() }
+        )
+        
         return pending.promise
     }
 }
diff --git a/Source/Network/Implement/Client/JSONClient.swift b/Source/Network/Implement/Client/JSONClient.swift
index 2201c0a..673c463 100644
--- a/Source/Network/Implement/Client/JSONClient.swift
+++ b/Source/Network/Implement/Client/JSONClient.swift
@@ -23,8 +23,11 @@ public final class JSONClient: Client {
 }
 
 extension JSONClient {
-    public convenience init(optionBlock: (inout URLSessionConfiguration) -> Void = { _ in }) {
-        let client = DataClient(optionBlock: optionBlock)
+    public convenience init(
+        cancelWhen: PromisePendingCancelWhen,
+        optionBlock: (inout URLSessionConfiguration) -> Void = { _ in }
+    ) {
+        let client = DataClient(cancelWhen: cancelWhen, optionBlock: optionBlock)
         self.init(client: client)
     }
 }
diff --git a/Test/Network/DataClientTest.swift b/Test/Network/DataClientTest.swift
index d562821..4845ad2 100644
--- a/Test/Network/DataClientTest.swift
+++ b/Test/Network/DataClientTest.swift
@@ -14,14 +14,14 @@ import SabyConcurrency
 
 final class DataClientTest: XCTestCase {
     func test__init() {
-        let client = DataClient()
+        let client = DataClient(cancelWhen: .deinit)
         let configuration = URLSessionConfiguration.default
         
         XCTAssertEqual(client.session.configuration, configuration)
     }
     
     func test__init_option_block() {
-        let client = DataClient() {
+        let client = DataClient(cancelWhen: .deinit) {
             $0.timeoutIntervalForRequest = 3000
         }
         let configuration = URLSessionConfiguration.default
@@ -40,7 +40,7 @@ final class DataClientTest: XCTestCase {
                 )
             ]
         }
-        let client = DataClient() {
+        let client = DataClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -59,7 +59,7 @@ final class DataClientTest: XCTestCase {
                 )
             ]
         }
-        let client = DataClient() {
+        let client = DataClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -81,7 +81,7 @@ final class DataClientTest: XCTestCase {
                 )
             ]
         }
-        let client = DataClient() {
+        let client = DataClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -100,7 +100,7 @@ final class DataClientTest: XCTestCase {
                 )
             ]
         }
-        let client = DataClient() {
+        let client = DataClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
diff --git a/Test/Network/JSONClientTest.swift b/Test/Network/JSONClientTest.swift
index b32a46c..f374acd 100644
--- a/Test/Network/JSONClientTest.swift
+++ b/Test/Network/JSONClientTest.swift
@@ -15,14 +15,14 @@ import SabyConcurrency
 
 final class JSONClientTest: XCTestCase {
     func test__init() {
-        let client = JSONClient()
+        let client = JSONClient(cancelWhen: .deinit)
         let configuration = URLSessionConfiguration.default
         
         XCTAssertEqual(client.client.session.configuration, configuration)
     }
     
     func test__init_option_block() {
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.timeoutIntervalForRequest = 3000
         }
         let configuration = URLSessionConfiguration.default
@@ -41,7 +41,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -67,7 +67,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -93,7 +93,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -118,7 +118,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -144,7 +144,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -169,7 +169,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }
         
@@ -194,7 +194,7 @@ final class JSONClientTest: XCTestCase {
                 )
             ]
         }
-        let client = JSONClient() {
+        let client = JSONClient(cancelWhen: .deinit) {
             $0.protocolClasses = [MockURLProtocol<MockURLResultStorage>.self]
         }