From c85286beb3010866438bdc2aa8d7c877e91f2bfd Mon Sep 17 00:00:00 2001
From: Subeom Choi <subumm1@gmail.com>
Date: Wed, 17 Apr 2024 01:21:08 +0900
Subject: [PATCH] change: implement onCanceled without dispatchGroup

---
 Source/Concurrency/Contract/Contract.swift | 75 +++++++++++----------
 Source/Concurrency/Promise/Promise.swift   | 76 +++++++++++++---------
 2 files changed, 88 insertions(+), 63 deletions(-)

diff --git a/Source/Concurrency/Contract/Contract.swift b/Source/Concurrency/Contract/Contract.swift
index 712e9c1..019c24f 100644
--- a/Source/Concurrency/Contract/Contract.swift
+++ b/Source/Concurrency/Contract/Contract.swift
@@ -12,28 +12,18 @@ public final class Contract<Value, Failure: Error> {
     
     let queue: DispatchQueue
     var executeGroup: DispatchGroup
-    var cancelGroup: DispatchGroup
 
     var subscribers: [ContractSubscriber<Value, Failure>]
+    var cancelSubscribers: [ContractCancelSubscriber]
     
     init(queue: DispatchQueue = .global()) {
         self.state = Atomic(.executing)
         
         self.queue = queue
         self.executeGroup = DispatchGroup()
-        self.cancelGroup = DispatchGroup()
 
         self.subscribers = []
-        
-        cancelGroup.enter()
-    }
-    
-    deinit {
-        state.capture { state in
-            if case .canceled = state {} else {
-                cancelGroup.leave()
-            }
-        }
+        self.cancelSubscribers = []
     }
 }
 
@@ -80,8 +70,14 @@ extension Contract {
     
     func cancel() {
         state.mutate { state in
-            if case .canceled = state {} else {
-                cancelGroup.leave()
+            if case .executing = state {
+                for cancelSubscriber in self.cancelSubscribers {
+                    executeGroup.notify(queue: queue) {
+                        cancelSubscriber.onCanceled()
+                    }
+                }
+                subscribers = []
+                cancelSubscribers = []
                 return .canceled
             }
             
@@ -96,19 +92,21 @@ extension Contract {
         onCanceled: @escaping () -> Void
     ) {
         state.capture { state in
-            subscribers.append(ContractSubscriber(
-                queue: queue,
-                onResolved: onResolved,
-                onRejected: onRejected
-            ))
-        }
-        
-        let state = self.state
-        cancelGroup.notify(queue: queue) {
-            let state = state.capture { $0 }
-            
-            if case .canceled = state {
-                onCanceled()
+            if case .executing = state {
+                subscribers.append(ContractSubscriber(
+                    queue: queue,
+                    onResolved: onResolved,
+                    onRejected: onRejected
+                ))
+                cancelSubscribers.append(ContractCancelSubscriber(
+                    queue: queue,
+                    onCanceled: onCanceled
+                ))
+            }
+            else if case .canceled = state {
+                executeGroup.notify(queue: queue) {
+                    onCanceled()
+                }
             }
         }
     }
@@ -117,12 +115,17 @@ extension Contract {
         queue: DispatchQueue,
         onCanceled: @escaping () -> Void
     ) {
-        let state = self.state
-        cancelGroup.notify(queue: queue) {
-            let state = state.capture { $0 }
-            
-            if case .canceled = state {
-                onCanceled()
+        state.capture { state in
+            if case .executing = state {
+                cancelSubscribers.append(ContractCancelSubscriber(
+                    queue: queue,
+                    onCanceled: onCanceled
+                ))
+            }
+            else if case .canceled = state {
+                executeGroup.notify(queue: queue) {
+                    onCanceled()
+                }
             }
         }
     }
@@ -190,7 +193,6 @@ extension Contract {
     ) -> Contract<Value, Failure> {
         let contract = Contract(queue: queue)
         contract.state = Atomic(.canceled)
-        contract.cancelGroup.leave()
         
         return contract
     }
@@ -266,6 +268,11 @@ struct ContractSubscriber<Value, Failure: Error> {
     let onRejected: (Failure) -> Void
 }
 
+struct ContractCancelSubscriber {
+    let queue: DispatchQueue
+    let onCanceled: () -> Void
+}
+
 public final class ContractSchedule {
     private let mode: Mode
     
diff --git a/Source/Concurrency/Promise/Promise.swift b/Source/Concurrency/Promise/Promise.swift
index 269dc78..94c08f4 100644
--- a/Source/Concurrency/Promise/Promise.swift
+++ b/Source/Concurrency/Promise/Promise.swift
@@ -11,27 +11,25 @@ public final class Promise<Value, Failure: Error> {
     var state: Atomic<PromiseState<Value, Failure>>
     
     let queue: DispatchQueue
-    let executeGroup: DispatchGroup
-    let cancelGroup: DispatchGroup
+    let pendingGroup: DispatchGroup
+    
+    var cancelSubscribers: [PromiseCancelSubscriber]
     
     init(queue: DispatchQueue = .global()) {
         self.state = Atomic(.pending)
         
         self.queue = queue
-        self.executeGroup = DispatchGroup()
-        self.cancelGroup = DispatchGroup()
+        self.pendingGroup = DispatchGroup()
+        
+        self.cancelSubscribers = []
         
-        executeGroup.enter()
-        cancelGroup.enter()
+        pendingGroup.enter()
     }
     
     deinit {
         state.capture { state in
-            if case .canceled = state {} else {
-                cancelGroup.leave()
-                if case .pending = state {
-                    executeGroup.leave()
-                }
+            if case .pending = state {
+                pendingGroup.leave()
             }
         }
     }
@@ -222,7 +220,7 @@ extension Promise {
     func resolve(_ value: Value) {
         state.mutate { state in
             if case .pending = state {
-                executeGroup.leave()
+                pendingGroup.leave()
                 return .resolved(value)
             }
             
@@ -233,7 +231,7 @@ extension Promise {
     func reject(_ error: Failure) {
         state.mutate { state in
             if case .pending = state {
-                executeGroup.leave()
+                pendingGroup.leave()
                 return .rejected(error)
             }
             
@@ -244,10 +242,15 @@ extension Promise {
     func cancel() {
         state.mutate { state in
             if case .canceled = state {} else {
-                cancelGroup.leave()
                 if case .pending = state {
-                    executeGroup.leave()
+                    pendingGroup.leave()
                 }
+                for cancelSubscriber in self.cancelSubscribers {
+                    pendingGroup.notify(queue: queue) {
+                        cancelSubscriber.onCanceled()
+                    }
+                }
+                cancelSubscribers = []
                 return .canceled
             }
             
@@ -262,7 +265,7 @@ extension Promise {
         onCanceled: @escaping () -> Void
     ) {
         let state = self.state
-        executeGroup.notify(queue: queue) {
+        pendingGroup.notify(queue: queue) {
             let state = state.capture { $0 }
             
             if case .resolved(let value) = state {
@@ -272,11 +275,17 @@ extension Promise {
                 onRejected(error)
             }
         }
-        cancelGroup.notify(queue: queue) {
-            let state = state.capture { $0 }
-            
+        state.capture { state in
             if case .canceled = state {
-                onCanceled()
+                pendingGroup.notify(queue: queue) {
+                    onCanceled()
+                }
+            }
+            else {
+                cancelSubscribers.append(PromiseCancelSubscriber(
+                    queue: queue,
+                    onCanceled: onCanceled
+                ))
             }
         }
     }
@@ -285,12 +294,17 @@ extension Promise {
         queue: DispatchQueue,
         onCanceled: @escaping () -> Void
     ) {
-        let state = self.state
-        cancelGroup.notify(queue: queue) {
-            let state = state.capture { $0 }
-            
+        state.capture { state in
             if case .canceled = state {
-                onCanceled()
+                pendingGroup.notify(queue: queue) {
+                    onCanceled()
+                }
+            }
+            else {
+                cancelSubscribers.append(PromiseCancelSubscriber(
+                    queue: queue,
+                    onCanceled: onCanceled
+                ))
             }
         }
     }
@@ -373,7 +387,7 @@ extension Promise {
     ) -> Promise<Value, Failure> {
         let promise = Promise(queue: queue)
         promise.state = Atomic(.resolved(value))
-        promise.executeGroup.leave()
+        promise.pendingGroup.leave()
         
         return promise
     }
@@ -384,7 +398,7 @@ extension Promise {
     ) -> Promise<Value, Error> where Failure == Error {
         let promise = Promise<Value, Error>(queue: queue)
         promise.state = Atomic(.rejected(error))
-        promise.executeGroup.leave()
+        promise.pendingGroup.leave()
         
         return promise
     }
@@ -394,8 +408,7 @@ extension Promise {
     ) -> Promise<Value, Failure> {
         let promise = Promise(queue: queue)
         promise.state = Atomic(.canceled)
-        promise.executeGroup.leave()
-        promise.cancelGroup.leave()
+        promise.pendingGroup.leave()
         
         return promise
     }
@@ -444,6 +457,11 @@ enum PromiseState<Value, Failure: Error> {
     case canceled
 }
 
+struct PromiseCancelSubscriber {
+    let queue: DispatchQueue
+    let onCanceled: () -> Void
+}
+
 public enum PromiseError: Error {
     case timeout
 }