From f601f2f6238e444cdb9745188fb21e29448bedd1 Mon Sep 17 00:00:00 2001
From: Esonhugh <esonhughoutside@gmail.com>
Date: Mon, 15 Apr 2024 16:57:19 +0800
Subject: [PATCH] update: multi threading mode update, supported -t and -n to
 specific threading

---
 cmd/all/all.go        |  8 ++++----
 cmd/root.go           |  6 ++++--
 cmd/subnet/subnet.go  |  8 ++++----
 pkg/mutli/executor.go |  4 ++--
 pkg/mutli/subnet.go   | 21 ++++++++++++++++-----
 5 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/cmd/all/all.go b/cmd/all/all.go
index af23219..b166ca2 100644
--- a/cmd/all/all.go
+++ b/cmd/all/all.go
@@ -37,8 +37,8 @@ var AllCmd = &cobra.Command{
 			log.Warnf("ParseStringToIPNet failed: %v", err)
 			return
 		}
-		if command.Opts.BatchMode {
-			RunBatch(ipNets)
+		if command.Opts.MultiThreadingMode {
+			RunMultiThread(ipNets, command.Opts.ThreadingNum)
 		} else {
 			Run(ipNets)
 		}
@@ -55,8 +55,8 @@ func Run(net *net.IPNet) {
 	printResult(records)
 }
 
-func RunBatch(net *net.IPNet) {
-	scan := mutli.ScanAll(net)
+func RunMultiThread(net *net.IPNet, count int) {
+	scan := mutli.ScanAll(net, count)
 	var finalRecord []define.Record
 	for r := range scan {
 		finalRecord = append(finalRecord, r...)
diff --git a/cmd/root.go b/cmd/root.go
index 3de8515..2f4fef6 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -19,7 +19,8 @@ var Opts = struct {
 	OutputFile string
 	Verbose    string
 
-	BatchMode bool
+	MultiThreadingMode bool
+	ThreadingNum       int
 }{}
 
 func init() {
@@ -29,7 +30,8 @@ func init() {
 	RootCmd.PersistentFlags().StringVarP(&Opts.Zone, "zone", "z", "cluster.local", "zone")
 	RootCmd.PersistentFlags().StringVarP(&Opts.OutputFile, "output-file", "o", "", "output file")
 	RootCmd.PersistentFlags().StringVarP(&Opts.Verbose, "verbose", "v", "info", "log level (debug,info,trace,warn,error,fatal,panic)")
-	RootCmd.PersistentFlags().BoolVarP(&Opts.BatchMode, "batch-mode", "b", false, "batch mode")
+	RootCmd.PersistentFlags().BoolVarP(&Opts.MultiThreadingMode, "thread", "t", false, "multi threading mode, work pair with -n")
+	RootCmd.PersistentFlags().IntVarP(&Opts.ThreadingNum, "thread-num", "n", 16, "threading num, default 16")
 }
 
 var RootCmd = &cobra.Command{
diff --git a/cmd/subnet/subnet.go b/cmd/subnet/subnet.go
index 4e9064d..07be242 100644
--- a/cmd/subnet/subnet.go
+++ b/cmd/subnet/subnet.go
@@ -30,8 +30,8 @@ var SubNetCmd = &cobra.Command{
 			log.Warnf("ParseStringToIPNet failed: %v", err)
 			return
 		}
-		if command.Opts.BatchMode {
-			BatchRun(ipNets)
+		if command.Opts.MultiThreadingMode {
+			RunMultiThread(ipNets, command.Opts.ThreadingNum)
 		} else {
 			Run(ipNets)
 		}
@@ -47,8 +47,8 @@ func Run(net *net.IPNet) {
 	printResult(records)
 }
 
-func BatchRun(net *net.IPNet) {
-	scan := mutli.NewSubnetScanner()
+func RunMultiThread(net *net.IPNet, num int) {
+	scan := mutli.NewSubnetScanner(num)
 	var finalRecord []define.Record
 	for r := range scan.ScanSubnet(net) {
 		finalRecord = append(finalRecord, r...)
diff --git a/pkg/mutli/executor.go b/pkg/mutli/executor.go
index 082faf9..5ddb91f 100644
--- a/pkg/mutli/executor.go
+++ b/pkg/mutli/executor.go
@@ -6,8 +6,8 @@ import (
 	"github.com/esonhugh/k8spider/define"
 )
 
-func ScanAll(subnet *net.IPNet) (result <-chan []define.Record) {
-	subs := NewSubnetScanner()
+func ScanAll(subnet *net.IPNet, num int) (result <-chan []define.Record) {
+	subs := NewSubnetScanner(num)
 	result = ScanServiceWithChan(subs.ScanSubnet(subnet))
 	return result
 }
diff --git a/pkg/mutli/subnet.go b/pkg/mutli/subnet.go
index 3e3cc0a..4acf048 100644
--- a/pkg/mutli/subnet.go
+++ b/pkg/mutli/subnet.go
@@ -12,12 +12,20 @@ import (
 )
 
 type SubnetScanner struct {
-	wg *sync.WaitGroup
+	wg    *sync.WaitGroup
+	count int
 }
 
-func NewSubnetScanner() *SubnetScanner {
-	return &SubnetScanner{
-		wg: new(sync.WaitGroup),
+func NewSubnetScanner(threading ...int) *SubnetScanner {
+	if len(threading) == 0 {
+		return &SubnetScanner{
+			wg: new(sync.WaitGroup),
+		}
+	} else {
+		return &SubnetScanner{
+			wg:    new(sync.WaitGroup),
+			count: threading[0],
+		}
 	}
 }
 
@@ -29,9 +37,12 @@ func (s *SubnetScanner) ScanSubnet(subnet *net.IPNet) <-chan []define.Record {
 	out := make(chan []define.Record, 100)
 	go func() {
 		log.Debugf("splitting subnet into 16 pices")
-		if subnets, err := pkg.SubnetShift(subnet, 4); err != nil {
+		// if subnets, err := pkg.SubnetShift(subnet, 4); err != nil {
+		if subnets, err := pkg.SubnetInto(subnet, s.count); err != nil {
+			log.Errorf("Subnet split into %v failed, fallback to single mode, reason: %v", s.count, err)
 			go s.scan(subnet, out)
 		} else {
+			log.Debugf("Subnet split into %v success", len(subnets))
 			for _, sn := range subnets {
 				go s.scan(sn, out)
 			}