[*] Use conc.WaitGroup, it's safer than before

LoCyan-Team · Aug 19, 2023 · d8d47d2 · d8d47d2
1 parent 1555a3c
commit d8d47d2
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 49 deletions.
diff --git a/cmd/frpc/sub/root.go b/cmd/frpc/sub/root.go
@@ -24,12 +24,13 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
-	"sync"
 	"syscall"
 	"time"
 
 	"github.com/spf13/cobra"
 
+	"github.com/sourcegraph/conc"
+
 	"github.com/fatedier/frp/client"
 	"github.com/fatedier/frp/pkg/api"
 	"github.com/fatedier/frp/pkg/auth"
@@ -116,8 +117,9 @@ var rootCmd = &cobra.Command{
 			return nil
 		}
 
-		log.Info("欢迎使用LoCyanFrp映射客户端! v0.51.0 #20230710001")
-		var wg sync.WaitGroup
+		log.Info("欢迎使用LoCyanFrp映射客户端! v0.51.3 #20230710001")
+		var wg conc.WaitGroup
+		defer wg.Wait()
 
 		// If cfgDir is not empty, run multiple frpc service for each config file in cfgDir.
 		// Note that it's only designed for testing. It's not guaranteed to be stable.
@@ -143,44 +145,42 @@ var rootCmd = &cobra.Command{
 			}
 
 			if len(ids) > 1 {
-				for i := 0; i < len(ids); i++ {
-					// 每一个都是新的协程
-					wg.Add(1)
-					id := ids[i]
-					go func(id string) {
-						defer wg.Done()
-						// != -1 -> 存在多开
-						if strings.Contains(cfgProxyid, ",") {
-							cfg, err := s.EZStartGetCfg(cfgToken, id)
-							if err != nil {
-								log.Warn("Get Config File Failed, Please Check Your Args, err: %s", err)
-								// 无法获取配置文件，直接关闭软件，防止启动上一个配置文件导致二次报错
-								os.Exit(1)
-							}
-
-							file, err := os.OpenFile("./ini/"+id+".ini", os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0777)
-							if err != nil {
-								log.Warn("Open File Failed, Err: %s", err)
-								os.Exit(1)
-							}
-							defer file.Close()
-							str := cfg
-							_, err = file.WriteString(str) //直接写入字符串数据
-							// 写入文件是否成功检测
-							if err != nil {
-								log.Warn("文本在写入的过程中发生了致命错误, Err: %s", err)
-								os.Exit(1)
-							}
-
-							// 内容写入后直接启动
-							fliePath := "./ini/" + id + ".ini"
-							err4 := runClient(fliePath, &wg)
-							if err4 != nil {
-								log.Warn("启动的过程中发生致命错误, Err: %s", err4)
-								os.Exit(1)
-							}
+				// 每一个都是新的协程
+				for _, id := range ids {
+
+					// 将循环变量赋值给局部变量
+					idCopy := id
+
+					wg.Go(func() {
+						fliePath := "./ini/" + idCopy + ".ini"
+						cfg, err := s.EZStartGetCfg(cfgToken, idCopy)
+						if err != nil {
+							log.Warn("Get Config File Failed, Please Check Your Args, err: %s", err)
+							// 无法获取配置文件，直接关闭软件，防止启动上一个配置文件导致二次报错
+							os.Exit(1)
+						}
+
+						file, err := os.OpenFile(fliePath, os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0777)
+						if err != nil {
+							log.Warn("Open File Failed, Err: %s", err)
+							os.Exit(1)
+						}
+						defer file.Close()
+						str := cfg
+						_, err = file.WriteString(str) //直接写入字符串数据
+						// 写入文件是否成功检测
+						if err != nil {
+							log.Warn("文本在写入的过程中发生了致命错误, Err: %s", err)
+							os.Exit(1)
+						}
+
+						// 内容写入后直接启动
+						err4 := runClient(fliePath, &wg)
+						if err4 != nil {
+							log.Warn("启动的过程中发生致命错误, Err: %s", err4)
+							os.Exit(1)
 						}
-					}(id)
+					})
 				}
 				wg.Wait()
 				return nil
@@ -232,23 +232,21 @@ var rootCmd = &cobra.Command{
 }
 
 func runMultipleClients(cfgDir string) error {
-	var wg sync.WaitGroup
+	var wg conc.WaitGroup
+	defer wg.Wait()
 	err := filepath.WalkDir(cfgDir, func(path string, d fs.DirEntry, err error) error {
 		if err != nil || d.IsDir() {
 			return nil
 		}
-		wg.Add(1)
 		time.Sleep(time.Millisecond)
-		go func() {
-			defer wg.Done()
+		wg.Go(func() {
 			err := runClient(path, &wg)
 			if err != nil {
 				fmt.Printf("frpc service error for config file [%s]\n", path)
 			}
-		}()
+		})
 		return nil
 	})
-	wg.Wait()
 	return err
 }
 
@@ -303,8 +301,7 @@ func parseClientCommonCfgFromCmd() (cfg config.ClientCommonConf, err error) {
 	return
 }
 
-func runClient(cfgFilePath string, wg *sync.WaitGroup) error {
-	defer wg.Done()
+func runClient(cfgFilePath string, wg *conc.WaitGroup) error {
 	cfg, pxyCfgs, visitorCfgs, err := config.ParseClientConfig(cfgFilePath)
 	if err != nil {
 		fmt.Println(err)

diff --git a/go.mod b/go.mod
@@ -64,6 +64,7 @@ require (
 	github.com/prometheus/procfs v0.10.1 // indirect
 	github.com/quic-go/qtls-go1-20 v0.3.1 // indirect
 	github.com/rogpeppe/go-internal v1.11.0 // indirect
+	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect
 	github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b // indirect

diff --git a/go.sum b/go.sum
@@ -137,6 +137,8 @@ github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUz
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/samber/lo v1.38.1 h1:j2XEAqXKb09Am4ebOg31SpvzUTTs6EN3VfgeLUhPdXM=
 github.com/samber/lo v1.38.1/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA=
+github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
+github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
 github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
@@ -285,4 +287,4 @@ k8s.io/apimachinery v0.27.4/go.mod h1:XNfZ6xklnMCOGGFNqXG7bUrQCoR04dh/E7FprV6pb+
 k8s.io/client-go v0.27.4 h1:vj2YTtSJ6J4KxaC88P4pMPEQECWMY8gqPqsTgUKzvjk=
 k8s.io/client-go v0.27.4/go.mod h1:ragcly7lUlN0SRPk5/ZkGnDjPknzb37TICq07WhI6Xc=
 k8s.io/utils v0.0.0-20230209194617-a36077c30491 h1:r0BAOLElQnnFhE/ApUsg3iHdVYYPBjNSSOMowRZxxsY=
-k8s.io/utils v0.0.0-20230209194617-a36077c30491/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
+k8s.io/utils v0.0.0-20230209194617-a36077c30491/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=