Signed-off-by: arraykeys@gmail.com <arraykeys@gmail.com>

wangxian · May 21, 2018 · 8649bbc · 8649bbc
1 parent d775339
commit 8649bbc
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 15 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -8,6 +8,9 @@ v4.8
 3.重构了部分代码的日志部分,保证了日志按着预期输出.
 4.修复了sps\http代理初始化服务的时机不正确,导致nil异常的bug.
 5.优化了sps日志输出.
+6.--debug参数增加了Profiling功能,可以保存cpu,内存等多种调试数据到文件.
+7.优化了服务注册,避免了不必要的内存开销.
+
 
 v4.7
 1.增加了基于gomobile的sdk,对android/ios/windows/linux/mac提供SDK支持.

diff --git a/config.go b/config.go
@@ -7,6 +7,7 @@ import (
 	logger "log"
 	"os"
 	"os/exec"
+	"runtime/pprof"
 	"time"
 
 	"github.com/snail007/goproxy/services"
@@ -19,9 +20,11 @@ import (
 )
 
 var (
-	app     *kingpin.Application
-	service *services.ServiceItem
-	cmd     *exec.Cmd
+	app                                                                                                       *kingpin.Application
+	service                                                                                                   *services.ServiceItem
+	cmd                                                                                                       *exec.Cmd
+	cpuProfilingFile, memProfilingFile, blockProfilingFile, goroutineProfilingFile, threadcreateProfilingFile *os.File
+	isDebug                                                                                                   bool
 )
 
 func initConfig() (err error) {
@@ -227,7 +230,7 @@ func initConfig() (err error) {
 	spsArgs.ParentType = sps.Flag("parent-type", "parent protocol type <tls|tcp|kcp>").Short('T').Enum("tls", "tcp", "kcp")
 	spsArgs.LocalType = sps.Flag("local-type", "local protocol type <tls|tcp|kcp>").Default("tcp").Short('t').Enum("tls", "tcp", "kcp")
 	spsArgs.Local = sps.Flag("local", "local ip:port to listen,multiple address use comma split,such as: 0.0.0.0:80,0.0.0.0:443").Short('p').Default(":33080").String()
-	spsArgs.ParentServiceType = sps.Flag("parent-service-type", "parent service type <http|socks|ss>").Short('S').Enum("http", "socks", "ss")
+	spsArgs.ParentServiceType = sps.Flag("parent-service-type", "parent service type <http|socks>").Short('S').Enum("http", "socks", "ss")
 	spsArgs.DNSAddress = sps.Flag("dns-address", "if set this, proxy will use this dns for resolve doamin").Short('q').Default("").String()
 	spsArgs.DNSTTL = sps.Flag("dns-ttl", "caching seconds of dns query result").Short('e').Default("300").Int()
 	spsArgs.AuthFile = sps.Flag("auth-file", "http basic auth file,\"username:password\" each line in file").Short('F').String()
@@ -253,6 +256,8 @@ func initConfig() (err error) {
 	//parse args
 	serviceName := kingpin.MustParse(app.Parse(os.Args[1:]))
 
+	isDebug = *debug
+
 	//set kcp config
 
 	switch *kcpArgs.Mode {
@@ -310,6 +315,12 @@ func initConfig() (err error) {
 	flags := logger.Ldate
 	if *debug {
 		flags |= logger.Lshortfile | logger.Lmicroseconds
+		cpuProfilingFile, _ = os.Create("cpu.prof")
+		memProfilingFile, _ = os.Create("memory.prof")
+		blockProfilingFile, _ = os.Create("block.prof")
+		goroutineProfilingFile, _ = os.Create("goroutine.prof")
+		threadcreateProfilingFile, _ = os.Create("threadcreate.prof")
+		pprof.StartCPUProfile(cpuProfilingFile)
 	} else {
 		flags |= logger.Ltime
 	}
@@ -391,19 +402,40 @@ func initConfig() (err error) {
 	}
 	if *logfile == "" {
 		poster()
+		if *debug {
+			log.Println("[profiling] cpu profiling save to file : cpu.prof")
+			log.Println("[profiling] memory profiling save to file : memory.prof")
+			log.Println("[profiling] block profiling save to file : block.prof")
+			log.Println("[profiling] goroutine profiling save to file : goroutine.prof")
+			log.Println("[profiling] threadcreate profiling save to file : threadcreate.prof")
+		}
 	}
 	//regist services and run service
-	services.Regist("http", services.NewHTTP(), httpArgs, log)
-	services.Regist("tcp", services.NewTCP(), tcpArgs, log)
-	services.Regist("udp", services.NewUDP(), udpArgs, log)
-	services.Regist("tserver", services.NewTunnelServerManager(), tunnelServerArgs, log)
-	services.Regist("tclient", services.NewTunnelClient(), tunnelClientArgs, log)
-	services.Regist("tbridge", services.NewTunnelBridge(), tunnelBridgeArgs, log)
-	services.Regist("server", services.NewMuxServerManager(), muxServerArgs, log)
-	services.Regist("client", services.NewMuxClient(), muxClientArgs, log)
-	services.Regist("bridge", services.NewMuxBridge(), muxBridgeArgs, log)
-	services.Regist("socks", services.NewSocks(), socksArgs, log)
-	services.Regist("sps", services.NewSPS(), spsArgs, log)
+	//regist services and run service
+	switch serviceName {
+	case "http":
+		services.Regist(serviceName, services.NewHTTP(), httpArgs, log)
+	case "tcp":
+		services.Regist(serviceName, services.NewTCP(), tcpArgs, log)
+	case "udp":
+		services.Regist(serviceName, services.NewUDP(), udpArgs, log)
+	case "tserver":
+		services.Regist(serviceName, services.NewTunnelServerManager(), tunnelServerArgs, log)
+	case "tclient":
+		services.Regist(serviceName, services.NewTunnelClient(), tunnelClientArgs, log)
+	case "tbridge":
+		services.Regist(serviceName, services.NewTunnelBridge(), tunnelBridgeArgs, log)
+	case "server":
+		services.Regist(serviceName, services.NewMuxServerManager(), muxServerArgs, log)
+	case "client":
+		services.Regist(serviceName, services.NewMuxClient(), muxClientArgs, log)
+	case "bridge":
+		services.Regist(serviceName, services.NewMuxBridge(), muxBridgeArgs, log)
+	case "socks":
+		services.Regist(serviceName, services.NewSocks(), socksArgs, log)
+	case "sps":
+		services.Regist(serviceName, services.NewSPS(), spsArgs, log)
+	}
 	service, err = services.Run(serviceName, nil)
 	if err != nil {
 		log.Fatalf("run service [%s] fail, ERR:%s", serviceName, err)
@@ -423,3 +455,14 @@ func poster() {
 		
 		v%s`+" by snail , blog : http://www.host900.com/\n\n", APP_VERSION)
 }
+func saveProfiling() {
+	goroutine := pprof.Lookup("goroutine")
+	goroutine.WriteTo(goroutineProfilingFile, 1)
+	heap := pprof.Lookup("heap")
+	heap.WriteTo(memProfilingFile, 1)
+	block := pprof.Lookup("block")
+	block.WriteTo(blockProfilingFile, 1)
+	threadcreate := pprof.Lookup("threadcreate")
+	threadcreate.WriteTo(threadcreateProfilingFile, 1)
+	pprof.StopCPUProfile()
+}
diff --git a/main.go b/main.go
@@ -41,6 +41,9 @@ func Clean(s *services.Service) {
 				log.Printf("clean process %d", cmd.Process.Pid)
 				cmd.Process.Kill()
 			}
+			if isDebug {
+				saveProfiling()
+			}
 			cleanupDone <- true
 		}
 	}()

diff --git a/services/socks.go b/services/socks.go
@@ -309,6 +309,7 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 			s.udpSC.UDPListener.WriteToUDP(d, srcAddr)
 			s.udpSC.UDPListener.SetDeadline(time.Time{})
 			s.log.Printf("udp reply:%v", len(d))
+			d = nil
 		} else {
 			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(respBody, srcAddr)
@@ -363,6 +364,7 @@ func (s *Socks) udpCallback(b []byte, localAddr, srcAddr *net.UDPAddr) {
 			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(d, srcAddr)
 			s.udpSC.UDPListener.SetDeadline(time.Time{})
+			d = nil
 		} else {
 			s.udpSC.UDPListener.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout)))
 			s.udpSC.UDPListener.WriteToUDP(respPacket, srcAddr)