优化主wg

forward/forward.go

@@ -37,8 +37,7 @@ var bufPool = sync.Pool{
 }
 
 // 开启转发，负责分发具体转发
-func Run(stats *ConnectionStats, wg *sync.WaitGroup) {
+func Run(stats *ConnectionStats) {
-	defer wg.Done()
 	defer releaseResources(stats) // 在函数返回时释放资源
 	var ctx, cancel = context.WithCancel(context.Background())
 	var innerWg sync.WaitGroup
@@ -165,14 +164,12 @@ func (cs *ConnectionStats) handleTCPConnection(wg *sync.WaitGroup, clientConn ne
 			}
 		}
 	}()
-
 	copyWG.Wait()
 }
 
 // UDP转发
 func (cs *ConnectionStats) handleUDPConnection(wg *sync.WaitGroup, localConn *net.UDPConn, remoteAddr *net.UDPAddr, ctx context.Context) {
 	defer wg.Done()
-
 	for {
 		select {
 		case <-ctx.Done():

main.go

@@ -49,10 +49,12 @@ func main() {
 	}
 	// 设置 WaitGroup 计数为连接数
 	conf.Wg.Add(len(largeStats.Connections))
-
 	// 并发执行多个转发
 	for _, stats := range largeStats.Connections {
-		go forward.Run(stats, &conf.Wg)
+		go func(s *forward.ConnectionStats) {
+			forward.Run(s)
+			conf.Wg.Done()
+		}(stats)
 	}
 	conf.Wg.Wait()
 	defer close(conf.Ch)

utils/utils.go

@@ -28,7 +28,10 @@ func AddForward(newF conf.ConnectionStats) bool {
 			TotalBytesLock: sync.Mutex{},
 		}
 		conf.Wg.Add(1)
-		go forward.Run(stats, &conf.Wg)
+		go func() {
+			forward.Run(stats)
+			conf.Wg.Done()
+		}()
 		return true
 	}
 	return false
@@ -62,7 +65,10 @@ func ExStatus(f conf.ConnectionStats) bool {
 				TotalBytesLock: sync.Mutex{},
 			}
 			conf.Wg.Add(1)
-			go forward.Run(stats, &conf.Wg)
+			go func() {
+				forward.Run(stats)
+				conf.Wg.Done()
+			}()
 			return true
 		} else {
 			conf.Ch <- f.LocalPort + f.Protocol