fix: network connectivity speed

Signed-off-by: renxiangyu <[email protected]>
kosmos-io · Jan 8, 2024 · 62955f2 · 62955f2
1 parent 88a3efe
commit 62955f2
Showing 1 changed file with 63 additions and 29 deletions.
diff --git a/pkg/kosmosctl/floater/check.go b/pkg/kosmosctl/floater/check.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"sync"
 
 	"github.com/olekukonko/tablewriter"
 	"github.com/spf13/cobra"
@@ -45,6 +46,8 @@ type CommandCheckOptions struct {
 
 	SrcFloater *Floater
 	DstFloater *Floater
+
+	routinesMaxNum int
 }
 
 type PrintCheckData struct {
@@ -92,6 +95,7 @@ func NewCmdCheck() *cobra.Command {
 	flags.StringVar(&o.Port, "port", "8889", "Port used by floater.")
 	flags.IntVarP(&o.PodWaitTime, "pod-wait-time", "w", 30, "Time for wait pod(floater) launch.")
 	flags.StringVar(&o.Protocol, "protocol", string(TCP), "Protocol for the network problem.")
+	flags.IntVarP(&o.routinesMaxNum, "routines-max-number", "", 5, "Number of goroutines to use.")
 
 	return cmd
 }
@@ -199,60 +203,90 @@ func (o *CommandCheckOptions) Run() error {
 func (o *CommandCheckOptions) RunRange(iPodInfos []*FloatInfo, jPodInfos []*FloatInfo) []*PrintCheckData {
 	var resultData []*PrintCheckData
 
+	goroutinePool := make(chan struct{}, o.routinesMaxNum)
+	waitGroup := sync.WaitGroup{}
+
 	if len(iPodInfos) > 0 && len(jPodInfos) > 0 {
 		for _, iPodInfo := range iPodInfos {
 			for _, jPodInfo := range jPodInfos {
 				for _, ip := range jPodInfo.PodIPs {
-					var targetIP string
-					var err error
-					var cmdResult *command.Result
-					if o.DstFloater != nil {
-						targetIP, err = netmap.NetMap(ip, o.DstFloater.CIDRsMap)
-					} else {
-						targetIP = ip
-					}
-					if err != nil {
-						cmdResult = command.ParseError(err)
-					} else {
-						// ToDo RunRange && RunNative func support multiple commands, and the code needs to be optimized
-						cmdObj := &command.Ping{
-							TargetIP: targetIP,
+					goroutinePool <- struct{}{}
+					waitGroup.Add(1)
+					routineIPodInfo := iPodInfo
+					routineJPodInfo := jPodInfo
+					routineIp := ip
+					go func() {
+						var targetIP string
+						var err error
+						var cmdResult *command.Result
+						if o.DstFloater != nil {
+							targetIP, err = netmap.NetMap(routineIp, o.DstFloater.CIDRsMap)
+						} else {
+							targetIP = routineIp
+						}
+						if err != nil {
+							cmdResult = command.ParseError(err)
+						} else {
+							// ToDo RunRange && RunNative func support multiple commands, and the code needs to be optimized
+							cmdObj := &command.Ping{
+								TargetIP: targetIP,
+							}
+							cmdResult = o.SrcFloater.CommandExec(routineIPodInfo, cmdObj)
 						}
-						cmdResult = o.SrcFloater.CommandExec(iPodInfo, cmdObj)
-					}
-					resultData = append(resultData, &PrintCheckData{
-						*cmdResult,
-						iPodInfo.NodeName, jPodInfo.NodeName, targetIP,
-					})
+						resultData = append(resultData, &PrintCheckData{
+							*cmdResult,
+							routineIPodInfo.NodeName, routineJPodInfo.NodeName, targetIP,
+						})
+						defer func() {
+							<-goroutinePool
+							waitGroup.Done()
+						}()
+					}()
 				}
 			}
 		}
 	}
 
+	waitGroup.Wait()
 	return resultData
 }
 
 func (o *CommandCheckOptions) RunNative(iNodeInfos []*FloatInfo, jNodeInfos []*FloatInfo) []*PrintCheckData {
 	var resultData []*PrintCheckData
 
+	goroutinePool := make(chan struct{}, o.routinesMaxNum)
+	waitGroup := sync.WaitGroup{}
+
 	if len(iNodeInfos) > 0 && len(jNodeInfos) > 0 {
 		for _, iNodeInfo := range iNodeInfos {
 			for _, jNodeInfo := range jNodeInfos {
 				for _, ip := range jNodeInfo.NodeIPs {
-					// ToDo RunRange && RunNative func support multiple commands, and the code needs to be optimized
-					cmdObj := &command.Ping{
-						TargetIP: ip,
-					}
-					cmdResult := o.SrcFloater.CommandExec(iNodeInfo, cmdObj)
-					resultData = append(resultData, &PrintCheckData{
-						*cmdResult,
-						iNodeInfo.NodeName, jNodeInfo.NodeName, ip,
-					})
+					goroutinePool <- struct{}{}
+					waitGroup.Add(1)
+					routineINodeInfo := iNodeInfo
+					routineJNodeInfo := jNodeInfo
+					routineIp := ip
+					go func() {
+						// ToDo RunRange && RunNative func support multiple commands, and the code needs to be optimized
+						cmdObj := &command.Ping{
+							TargetIP: routineIp,
+						}
+						cmdResult := o.SrcFloater.CommandExec(routineINodeInfo, cmdObj)
+						resultData = append(resultData, &PrintCheckData{
+							*cmdResult,
+							routineINodeInfo.NodeName, routineJNodeInfo.NodeName, routineIp,
+						})
+						defer func() {
+							<-goroutinePool
+							waitGroup.Done()
+						}()
+					}()
 				}
 			}
 		}
 	}
 
+	waitGroup.Wait()
 	return resultData
 }