From 6cb8641869eabbb73564a6a8dc3d2ddaabd60eb5 Mon Sep 17 00:00:00 2001
From: anlicheng <244108715@qq.com>
Date: Sun, 3 Aug 2025 12:26:54 +0800
Subject: [PATCH] fix task group

---
 Sources/Punchnet/SDLSuperClient.swift | 28 +++++++++++++++------------
 Sources/Punchnet/SDLUDPHole.swift     |  6 +++++-
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/Sources/Punchnet/SDLSuperClient.swift b/Sources/Punchnet/SDLSuperClient.swift
index 4700fd0..9607631 100644
--- a/Sources/Punchnet/SDLSuperClient.swift
+++ b/Sources/Punchnet/SDLSuperClient.swift
@@ -63,13 +63,14 @@ actor SDLSuperClient {
                 group.addTask {
                     try await self.asyncChannel.channel.closeFuture.get()
                     NSLog("[SDLSuperClient] socket closed")
-                    self.writeContinuation.finish()
-                    self.inboundContinuation.finish()
-                    
                     throw SDLError.socketClosed
                 }
                 
                 group.addTask {
+                    defer {
+                        self.inboundContinuation.finish()
+                    }
+                    
                     for try await var packet in inbound {
                         if let message = SDLSuperClientDecoder.decode(buffer: &packet) {
                             SDLLogger.log("[SDLSuperTransport] read message: \(message)", level: .warning)
@@ -87,6 +88,10 @@ actor SDLSuperClient {
                 }
                 
                 group.addTask {
+                    defer {
+                        self.writeContinuation.finish()
+                    }
+                    
                     for try await message in self.writeStream {
                         var buffer = self.asyncChannel.channel.allocator.buffer(capacity: message.data.count + 5)
                         buffer.writeInteger(message.packetId, as: UInt32.self)
@@ -100,22 +105,21 @@ actor SDLSuperClient {
                 // --MARK: 心跳机制
                 group.addTask {
                     while true {
-                        
                         do {
-                            
-                            try Task.checkCancellation()
-                            //await self.ping()
-                            try Task.checkCancellation()
+                            await self.ping()
                             try await Task.sleep(nanoseconds: 5 * 1_000_000_000)
-                            try Task.checkCancellation()
                         } catch let err {
-                            NSLog("[SDLSuperClient] heartbeat cancelled")
-                            throw err
+                            NSLog("[SDLSuperClient] heartbeat cancelled with error: \(err)")
+                            break
                         }
                     }
                 }
                 
-                try await group.waitForAll()
+                // 迭代等待所有任务的退出, 第一个异常会被抛出
+                for try await _ in group {
+                    
+                }
+                
                 NSLog("[SDLSuperClient] group closed")
             }
         }
diff --git a/Sources/Punchnet/SDLUDPHole.swift b/Sources/Punchnet/SDLUDPHole.swift
index 08a649f..c68c6a4 100644
--- a/Sources/Punchnet/SDLUDPHole.swift
+++ b/Sources/Punchnet/SDLUDPHole.swift
@@ -100,7 +100,11 @@ actor SDLUDPHole {
                     }
                 }
                 
-                try await group.waitForAll()
+                for try await _ in group {
+                    
+                }
+                
+                SDLLogger.log("[SDLUDPHole] group closed", level: .debug)
             }
         }
     }