调整生命周期的管理

Tun/Punchnet/Actors/SDLContextActor.swift

@@ -166,6 +166,10 @@ actor SDLContextActor {
         self.quicClient = quicClient
         await quicClient.start()
         
+        defer {
+            self.quicClient = nil
+        }
+        
         // 这里必须等待quic的协商完成
         try await Task.sleep(for: .seconds(0.5))
         SDLLogger.log("[SDLContext] start quic client: \(self.config.serverHost)")
@@ -284,6 +288,11 @@ actor SDLContextActor {
         let dnsClient = DNSCloudClient(host: self.config.serverHost, port: 15353)
         self.dnsClient = dnsClient
         dnsClient.start()
+        
+        defer {
+            self.dnsClient = nil
+        }
+        
         do {
             try await withTaskCancellationHandler {
                 for try await packet in dnsClient.packetFlow {
@@ -319,6 +328,10 @@ actor SDLContextActor {
         SDLLogger.log("[SDLContext] dnsLocalClient started")
         self.dnsLocalClient = dnsLocalClient
         
+        defer {
+            self.dnsLocalClient = nil
+        }
+        
         do {
             try await withTaskCancellationHandler {
                 // 处理事件流
@@ -352,24 +365,24 @@ actor SDLContextActor {
     private func startUDPHole() async throws {
         // 启动udp服务器
         let udpHole = try SDLUDPHole()
-        let localAddress = try udpHole.start()
+        let localAddress = try await udpHole.start()
         SDLLogger.log("[SDLContext] udpHole started, on address: \(localAddress)")
         self.udpHole = udpHole
         self.udpHoleLocalAddress = localAddress
         
         defer {
-            self.udpHole?.stop()
             self.udpHole = nil
-            self.udpHoleLocalAddress = nil
         }
         
-        try await withThrowingTaskGroup { group in
+        // 开始探测nat的类型
-            defer {
+        Task {
-                group.cancelAll()
+            await self.probeNatType()
-            }
+        }
+        SDLLogger.log("[SDLContext] udpHole ready")
         
-            group.addTask {
+        do {
-                for await (remoteAddress, message) in udpHole.messageStream {
+            try await withTaskCancellationHandler {
+                for try await (remoteAddress, message) in await udpHole.messageStream() {
                     try Task.checkCancellation()
                     
                     switch message.inboundMessage {
@@ -388,25 +401,14 @@ actor SDLContextActor {
                         try? await self.handleHoleData(data: data)
                     }
                 }
-            }
+            } onCancel: {
-            
+                Task {
-            group.addTask {
+                    await udpHole.stop()
-                for await event in udpHole.eventStream {
-                    try Task.checkCancellation()
-                    switch event {
-                    case .ready:
-                        // 开始探测nat的类型
-                        Task {
-                            await self.probeNatType()
-                        }
-                        SDLLogger.log("[SDLContext] udpHole ready")
-                    case .closed, .errorCaught:
-                        throw SDLContextError.udpHoleClosed
-                    }
                 }
             }
-            
+        } catch let err {
-            try await group.next()
+            await udpHole.stop()
+            throw err
         }
     }
     
@@ -564,7 +566,7 @@ actor SDLContextActor {
         }
     }
     
-    private func sendStunRequest(probeReply: SDLV6AssistProbeReply?) {
+    private func sendStunRequest(probeReply: SDLV6AssistProbeReply?) async {
         guard let sessionToken = self.sessionToken else {
             return
         }
@@ -582,7 +584,7 @@ actor SDLContextActor {
         }
         
         if let stunData = try? stunRequest.serializedData() {
-            self.sendSuperPacket(type: .stunRequest, data: stunData)
+            await self.sendSuperPacket(type: .stunRequest, data: stunData)
         }
     }
     
@@ -661,23 +663,23 @@ actor SDLContextActor {
     }
     
     // 发送给super/stun节点的数据
-    private func sendSuperPacket(type: SDLPacketType, data: Data) {
+    private func sendSuperPacket(type: SDLPacketType, data: Data) async {
-        self.sendPacket(type: type, data: data, remoteAddress: self.config.stunSocketAddress)
+        await self.sendPacket(type: type, data: data, remoteAddress: self.config.stunSocketAddress)
     }
     
     // 发送给peer的数据
-    private func sendPeerPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) {
+    private func sendPeerPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
-        self.sendPacket(type: type, data: data, remoteAddress: remoteAddress)
+        await self.sendPacket(type: type, data: data, remoteAddress: remoteAddress)
     }
     
-    private func sendPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) {
+    private func sendPacket(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
         switch remoteAddress {
         case .v4:
             guard let udpHole = self.udpHole else {
                 SDLLogger.log("[SDLContext] udpHole is nil for remoteAddress: \(remoteAddress)", for: .debug)
                 return
             }
-            udpHole.send(type: type, data: data, remoteAddress: remoteAddress)
+            await udpHole.send(type: type, data: data, remoteAddress: remoteAddress)
         case .v6:
             guard let udpHoleV6 = self.udpHoleV6 else {
                 SDLLogger.log("[SDLContext] udpHoleV6 is nil for remoteAddress: \(remoteAddress)", for: .debug)
@@ -821,8 +823,8 @@ extension SDLContextActor {
         case .removeSession(let dstMac):
             await self.sessionManager.removeSession(dstMac: dstMac)
         case .sendRegister(let registerData, let remoteAddresses):
-            remoteAddresses.forEach { remoteAddress in
+            for remoteAddress in remoteAddresses {
-                self.sendPeerPacket(type: .register, data: registerData, remoteAddress: remoteAddress)
+                await self.sendPeerPacket(type: .register, data: registerData, remoteAddress: remoteAddress)
             }
         case .shutdown(let message):
             self.publishTunnelEvent(message: message)
@@ -870,7 +872,7 @@ extension SDLContextActor {
             registerAck.srcMac = networkAddr.mac
             registerAck.dstMac = register.srcMac
             
-            self.sendPeerPacket(type: .registerAck, data: try registerAck.serializedData(), remoteAddress: remoteAddress)
+            await self.sendPeerPacket(type: .registerAck, data: try registerAck.serializedData(), remoteAddress: remoteAddress)
             // 这里需要建立到来源的会话, 在复杂网络下，通过super-node查询到的nat地址不一定靠谱，需要通过udp包的来源地址作为nat地址
             if let session = Session(dstMac: register.srcMac, natAddress: remoteAddress, addressType: source.convertAddressType()) {
                 await self.sessionManager.addSession(session: session)
@@ -1024,15 +1026,15 @@ extension SDLContextActor {
         switch plan {
         case .superNode(let payload):
             // 通过super_node进行转发
-            self.sendSuperPacket(type: .data, data: payload)
+            await self.sendSuperPacket(type: .data, data: payload)
         case .peer(let payload, let session):
             // 通过session发送到对端
             SDLLogger.log("[SDLContext] step 5 send packet by session: \(session)", for: .trace)
-            self.sendPeerPacket(type: .data, data: payload, remoteAddress: session.natAddress)
+            await self.sendPeerPacket(type: .data, data: payload, remoteAddress: session.natAddress)
             self.flowTracer.inc(num: payload.count, type: .p2p)
         case .superNodeAndPunch(let payload, let request):
             // 通过super_node进行转发
-            self.sendSuperPacket(type: .data, data: payload)
+            await self.sendSuperPacket(type: .data, data: payload)
             SDLLogger.log("[SDLContext] step 5 send packet by super: \(self.config.stunSocketAddress)", for: .trace)
             // 流量统计
             self.flowTracer.inc(num: payload.count, type: .forward)

Tun/Punchnet/Actors/SDLNATProberActor.swift

@@ -158,10 +158,10 @@ actor SDLNATProberActor {
     // MARK: - Internal helpers
 
     private func sendProbe(using udpHole: SDLUDPHole, cookie: UInt32) async {
-        udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 1, attr: .none), remoteAddress: addressArray[0][0])
+        await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 1, attr: .none), remoteAddress: addressArray[0][0])
-        udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 2, attr: .none), remoteAddress: addressArray[1][1])
+        await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 2, attr: .none), remoteAddress: addressArray[1][1])
-        udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 3, attr: .peer), remoteAddress: addressArray[0][0])
+        await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 3, attr: .peer), remoteAddress: addressArray[0][0])
-        udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 4, attr: .port), remoteAddress: addressArray[0][0])
+        await udpHole.send(type: .stunProbe, data: makeProbePacket(cookieId: cookie, step: 4, attr: .port), remoteAddress: addressArray[0][0])
     }
     
     private func makeProbePacket(cookieId: UInt32, step: UInt32, attr: SDLProbeAttr) -> Data {

Tun/Punchnet/Actors/SDLPuncherActor.swift

@@ -127,7 +127,7 @@ actor SDLPuncherActor {
         if peerInfo.hasV4Info {
             if let remoteAddress = try? await peerInfo.v4Info.socketAddress() {
                 SDLLogger.log("[SDLContext] hole sock address: \(remoteAddress)", for: .debug)
-                self.sendRegister(using: udpHole, udpHoleV6: udpHoleV6, registerData: registerData, remoteAddress: remoteAddress)
+                await self.sendRegister(using: udpHole, udpHoleV6: udpHoleV6, registerData: registerData, remoteAddress: remoteAddress)
             } else {
                 SDLLogger.log("[SDLPuncherActor] failed to resolve peerInfo.v4Info", for: .debug)
             }
@@ -136,7 +136,7 @@ actor SDLPuncherActor {
         if peerInfo.hasV6Info {
             if let remoteAddress = try? await peerInfo.v6Info.socketAddress() {
                 SDLLogger.log("[SDLContext] hole sock address v6: \(remoteAddress)", for: .debug)
-                self.sendRegister(using: udpHole, udpHoleV6: udpHoleV6, registerData: registerData, remoteAddress: remoteAddress)
+                await self.sendRegister(using: udpHole, udpHoleV6: udpHoleV6, registerData: registerData, remoteAddress: remoteAddress)
             } else {
                 SDLLogger.log("[SDLPuncherActor] failed to resolve peerInfo.v6Info", for: .debug)
             }
@@ -156,14 +156,14 @@ actor SDLPuncherActor {
         }
     }
     
-    private func sendRegister(using udpHole: SDLUDPHole?, udpHoleV6: SDLUDPHoleV6?, registerData: Data, remoteAddress: SocketAddress) {
+    private func sendRegister(using udpHole: SDLUDPHole?, udpHoleV6: SDLUDPHoleV6?, registerData: Data, remoteAddress: SocketAddress) async {
         switch remoteAddress {
         case .v4:
             guard let udpHole else {
                 SDLLogger.log("[SDLPuncherActor] udpHole is nil when v4 peerInfo arrived", for: .debug)
                 return
             }
-            udpHole.send(type: .register, data: registerData, remoteAddress: remoteAddress)
+            await udpHole.send(type: .register, data: registerData, remoteAddress: remoteAddress)
         case .v6:
             guard let udpHoleV6 else {
                 SDLLogger.log("[SDLPuncherActor] udpHoleV6 is nil when v6 peerInfo arrived", for: .debug)

Tun/Punchnet/UDPHole/SDLUDPHole.swift

@@ -11,40 +11,73 @@ import SwiftProtobuf
 
 enum SDLUDPHoleError: Error {
     case invalidLocalAddress
+    case closed
+    case errorCaught
+    case sendFaied(Error)
+}
+
+actor SDLUDPHole {
+    enum State {
+        case idle
+        case running
+        case stopped
+    }
+    
+    private var state: State = .idle
+    
+    private let udpHoleHandler: SDLUDPHoleHandler
+    
+    init() throws {
+        self.udpHoleHandler = try SDLUDPHoleHandler()
+    }
+    
+    func start() throws -> SocketAddress {
+        let localAddress = try self.udpHoleHandler.start()
+        self.state = .running
+        
+        return localAddress
+    }
+    
+    func messageStream() -> AsyncThrowingStream<(SocketAddress, SDLHoleMessage), Error> {
+        return self.udpHoleHandler.messageStream
+    }
+    
+    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) {
+        guard self.state == .running else {
+            return
+        }
+        
+        self.udpHoleHandler.send(type: type, data: data, remoteAddress: remoteAddress)
+    }
+    
+    func stop() {
+        guard self.state != .stopped else {
+            return
+        }
+        self.state = .stopped
+        self.udpHoleHandler.stop()
+    }
+    
 }
 
 // 处理和sn-server服务器之间的通讯
-final class SDLUDPHole: ChannelInboundHandler {
+private final class SDLUDPHoleHandler: ChannelInboundHandler {
     typealias InboundIn = AddressedEnvelope<ByteBuffer>
     
-    // 事件
-    enum HoleEvent {
-        case ready
-        case closed
-        case errorCaught
-    }
-    
-    private var isStopped: Bool = false
-    
     private let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
     private var channel: Channel?
     
-    public let messageStream: AsyncStream<(SocketAddress, SDLHoleMessage)>
+    private let locker = NSLock()
-    private let messageContinuation: AsyncStream<(SocketAddress, SDLHoleMessage)>.Continuation
     
-    // 事件相关逻辑
+    public let messageStream: AsyncThrowingStream<(SocketAddress, SDLHoleMessage), Error>
-    public let eventStream: AsyncStream<HoleEvent>
+    private let messageContinuation: AsyncThrowingStream<(SocketAddress, SDLHoleMessage), Error>.Continuation
-    private let eventContinuation: AsyncStream<HoleEvent>.Continuation
+    private var isMessageContinuationFinished: Bool = false
     
     // 启动函数
     init() throws {
-        let (stream, continuation) = AsyncStream.makeStream(of: (SocketAddress, SDLHoleMessage).self, bufferingPolicy: .bufferingNewest(2048))
+        let (stream, continuation) = AsyncThrowingStream.makeStream(of: (SocketAddress, SDLHoleMessage).self, bufferingPolicy: .bufferingNewest(2048))
         self.messageStream = stream
         self.messageContinuation = continuation
-        
-        let eventPair = AsyncStream.makeStream(of: HoleEvent.self)
-        self.eventStream = eventPair.stream
-        self.eventContinuation = eventPair.continuation
     }
     
     func start() throws -> SocketAddress {
@@ -61,7 +94,6 @@ final class SDLUDPHole: ChannelInboundHandler {
         }
         
         self.channel = channel
-        eventContinuation.yield(.ready)
         
         return localAddress
     }
@@ -89,13 +121,13 @@ final class SDLUDPHole: ChannelInboundHandler {
     
     func channelInactive(context: ChannelHandlerContext) {
         SDLLogger.log("[SDLUDPHole] channelInactive", for: .debug)
-        self.eventContinuation.yield(.closed)
+        self.finishMessageContinuationIfNeed(throwing: .closed)
     }
     
     func errorCaught(context: ChannelHandlerContext, error: any Error) {
         SDLLogger.log("[SDLUDPHole] channel error: \(error)", for: .debug)
         context.close(promise: nil)
-        self.eventContinuation.yield(.errorCaught)
+        self.finishMessageContinuationIfNeed(throwing: .errorCaught)
     }
     
     // MARK: 处理写入逻辑
@@ -109,24 +141,37 @@ final class SDLUDPHole: ChannelInboundHandler {
         buffer.writeBytes(data)
         
         let envelope = AddressedEnvelope<ByteBuffer>(remoteAddress: remoteAddress, data: buffer)
-        _ = channel.eventLoop.submit {
+        
-            channel.writeAndFlush(envelope, promise: nil)
+        let promise = channel.eventLoop.makePromise(of: Void.self)
+        channel.eventLoop.execute {
+            channel.writeAndFlush(envelope, promise: promise)
+        }
+        
+        promise.futureResult.whenFailure { err in
+            self.finishMessageContinuationIfNeed(throwing: .sendFaied(err))
         }
     }
         
     func stop() {
-        guard !self.isStopped else {
-            return
-        }
-        
-        self.isStopped = true
-        
         SDLLogger.log("[SDLUDPHole] stop", for: .debug)
-        self.messageContinuation.finish()
+        self.finishMessageContinuationIfNeed(throwing: nil)
-        self.eventContinuation.finish()
+        try? self.channel?.close().wait()
         self.channel = nil
-        
         try? self.group.syncShutdownGracefully()
     }
     
+    private func finishMessageContinuationIfNeed(throwing error: SDLUDPHoleError?) {
+        locker.lock()
+        defer {
+            locker.unlock()
+        }
+        
+        guard !self.isMessageContinuationFinished else {
+            return
+        }
+        
+        self.isMessageContinuationFinished = true
+        self.messageContinuation.finish(throwing: error)
+    }
+    
 }