Merge pull request #35 from Ehco1996/deadline

perf:  opt tcp/udp relay

Author: Ehco1996

README.md

@@ -131,4 +131,4 @@ iperf3 -c 0.0.0.0 -p 1234 -u -b 1G --length 1024
 | iperf | raw | relay(raw) | relay(ws) |relay(wss) | relay(mwss)|
 | ---- | ----  | ---- | ---- | ---- | ---- |
 | tcp  | 62.6 Gbits/sec | 23.9 Gbits/sec | 14.65 Gbits/sec | 4.22 Gbits/sec | 2.43 Gbits/sec |
-| udp  | 2.2 Gbits/sec | 2.2 Gbits/sec | 直接转发 | 直接转发 | 直接转发 |
+| udp  | 14.5 Gbits/sec | 3.3 Gbits/sec | 直接转发 | 直接转发 | 直接转发 |

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gobwas/pool v0.2.1 // indirect
 	github.com/gobwas/ws v1.0.4
 	github.com/gorilla/mux v1.8.0
+	github.com/prometheus/client_golang v1.9.0
 	github.com/urfave/cli/v2 v2.3.0
 	github.com/xtaci/smux v2.0.1+incompatible
 	go.uber.org/zap v1.16.0

go.sum

@@ -7,7 +7,7 @@ const (
 
 	MaxMWSSStreamCnt = 10
 	DialTimeOut      = 3 * time.Second
-	MaxConKeepAlive  = 10 * time.Minute
+	MaxConKeepAlive  = 3 * time.Second
 
 	Listen_RAW  = "raw"
 	Listen_WS   = "ws"

internal/constant/constant.go

@@ -99,12 +99,7 @@ func (r *Relay) RunLocalTCPServer() error {
 	for {
 		c, err := lis.AcceptTCP()
 		if err != nil {
-			logger.Logger.Infof("accept tcp con error: %s", err)
-			return err
-		}
-		if err := c.SetDeadline(time.Now().Add(constant.MaxConKeepAlive)); err != nil {
-			logger.Logger.Infof("set max deadline err: %s", err)
-			return err
+			logger.Logger.Fatal("accept tcp conn error: %s", err)
 		}
 		go func(c *net.TCPConn) {
 			if err := r.TP.HandleTCPConn(c); err != nil {
@@ -124,13 +119,12 @@ func (r *Relay) RunLocalUDPServer() error {
 	}
 	defer lis.Close()
 
-	buffers := transporter.NewBufferPool(constant.BUFFER_SIZE)
-	buf := buffers.Get().([]byte)
-	defer buffers.Put(buf)
+	buf := transporter.InboundBufferPool.Get().([]byte)
+	defer transporter.InboundBufferPool.Put(buf)
 	for {
 		n, addr, err := lis.ReadFromUDP(buf)
 		if err != nil {
-			return err
+			logger.Logger.Fatal("listen udp conn error: %s", err)
 		}
 		bc := r.TP.GetOrCreateBufferCh(addr)
 		bc.Ch <- buf[0:n]

internal/relay/relay.go

@@ -8,14 +8,15 @@ import (
 	"syscall"
 
 	"github.com/Ehco1996/ehco/internal/constant"
+	"github.com/Ehco1996/ehco/internal/web"
 )
 
 // 全局pool
-var inboundBufferPool, outboundBufferPool *sync.Pool
+var InboundBufferPool, OutboundBufferPool *sync.Pool
 
 func init() {
-	inboundBufferPool = NewBufferPool(constant.BUFFER_SIZE)
-	outboundBufferPool = NewBufferPool(constant.BUFFER_SIZE)
+	InboundBufferPool = NewBufferPool(constant.BUFFER_SIZE)
+	OutboundBufferPool = NewBufferPool(constant.BUFFER_SIZE)
 }
 
 func NewBufferPool(size int) *sync.Pool {
@@ -24,22 +25,62 @@ func NewBufferPool(size int) *sync.Pool {
 	}}
 }
 
-func copyBuffer(dst io.Writer, src io.Reader, bufferPool *sync.Pool) error {
+// NOTE adapeted from io.CopyBuffer
+func copyBuffer(dst io.Writer, src io.Reader, bufferPool *sync.Pool) (written int64, err error) {
 	buf := bufferPool.Get().([]byte)
 	defer bufferPool.Put(buf)
-	_, err := io.CopyBuffer(dst, src, buf)
-	return err
+
+	// If the reader has a WriteTo method, use it to do the copy.
+	// Avoids an allocation and a copy.
+	if wt, ok := src.(io.WriterTo); ok {
+		written, err = wt.WriteTo(dst)
+		web.NetWorkTransmitBytes.Add(float64(written))
+		return
+	}
+	// Similarly, if the writer has a ReadFrom method, use it to do the copy.
+	if rt, ok := dst.(io.ReaderFrom); ok {
+		written, err = rt.ReadFrom(src)
+		web.NetWorkTransmitBytes.Add(float64(written))
+		return
+	}
+	for {
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			nw, ew := dst.Write(buf[0:nr])
+			if nw > 0 {
+				written += int64(nw)
+			}
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+	web.NetWorkTransmitBytes.Add(float64(written * 2))
+	return
 }
 
 // NOTE must call setdeadline before use this func or may goroutine  leak
 func transport(rw1, rw2 io.ReadWriter) error {
 	errc := make(chan error, 1)
 	go func() {
-		errc <- copyBuffer(rw1, rw2, inboundBufferPool)
+		_, err := copyBuffer(rw1, rw2, InboundBufferPool)
+		errc <- err
 	}()
 
 	go func() {
-		errc <- copyBuffer(rw2, rw1, outboundBufferPool)
+		_, err := copyBuffer(rw2, rw1, OutboundBufferPool)
+		errc <- err
 	}()
 	err := <-errc
 	// NOTE 我们不关心operror 比如 eof/reset/broken pipe

internal/transporter/buffer.go

@@ -4,6 +4,7 @@ import (
 	"net"
 
 	"github.com/Ehco1996/ehco/internal/logger"
+	"github.com/Ehco1996/ehco/internal/web"
 )
 
 type Mwss struct {
@@ -21,6 +22,8 @@ func (s *Mwss) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
 
 func (s *Mwss) HandleTCPConn(c *net.TCPConn) error {
 	defer c.Close()
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 
 	node := s.raw.TCPNodes.PickMin()
 	defer s.raw.TCPNodes.DeferPick(node)

internal/transporter/mwss.go

@@ -4,9 +4,12 @@ import (
 	"net"
 	"net/http"
 	"sync"
+	"time"
 
+	"github.com/Ehco1996/ehco/internal/constant"
 	"github.com/Ehco1996/ehco/internal/lb"
 	"github.com/Ehco1996/ehco/internal/logger"
+	"github.com/Ehco1996/ehco/internal/web"
 	"github.com/gobwas/ws"
 )
 
@@ -28,6 +31,8 @@ func (raw *Raw) GetOrCreateBufferCh(uaddr *net.UDPAddr) *BufferCh {
 }
 
 func (raw *Raw) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
+	web.CurUDPNum.Inc()
+	defer web.CurUDPNum.Dec()
 
 	bc := raw.GetOrCreateBufferCh(uaddr)
 	node := raw.UDPNodes.PickMin()
@@ -39,10 +44,8 @@ func (raw *Raw) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
 		raw.UDPNodes.OnError(node)
 		return
 	}
-
 	defer func() {
 		rc.Close()
-		close(bc.Ch)
 		delete(raw.UDPBufferChMap, uaddr.String())
 	}()
 
@@ -52,34 +55,45 @@ func (raw *Raw) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
 	wg.Add(1)
 
 	go func() {
-		buf := outboundBufferPool.Get().([]byte)
+		buf := OutboundBufferPool.Get().([]byte)
+		wt := 0
 		for {
 			i, err := rc.Read(buf)
 			if err != nil {
 				logger.Logger.Info(err)
 				break
 			}
+			rc.SetReadDeadline(time.Now().Add(constant.MaxConKeepAlive))
 			if _, err := local.WriteToUDP(buf[0:i], uaddr); err != nil {
 				logger.Logger.Info(err)
 				break
 			}
+			wt += i
 		}
-		outboundBufferPool.Put(buf)
+		web.NetWorkTransmitBytes.Add(float64(wt * 2))
+		OutboundBufferPool.Put(buf)
 		wg.Done()
+		close(bc.Ch)
 	}()
 
+	wt := 0
 	for b := range bc.Ch {
+		wt += len(b)
+		rc.SetReadDeadline(time.Now().Add(constant.MaxConKeepAlive))
 		if _, err := rc.Write(b); err != nil {
 			logger.Logger.Info(err)
+			close(bc.Ch)
 			break
 		}
 	}
+	web.NetWorkTransmitBytes.Add(float64(wt * 2))
 	wg.Wait()
 }
 
 func (raw *Raw) HandleTCPConn(c *net.TCPConn) error {
 	defer c.Close()
-
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 	node := raw.TCPNodes.PickMin()
 	defer raw.TCPNodes.DeferPick(node)
 
@@ -89,18 +103,19 @@ func (raw *Raw) HandleTCPConn(c *net.TCPConn) error {
 		return err
 	}
 	logger.Logger.Infof("[raw] HandleTCPConn from %s to %s", c.LocalAddr().String(), node.Remote)
-
 	defer rc.Close()
+
 	return transport(c, rc)
 }
 
 func (raw *Raw) HandleWsRequset(w http.ResponseWriter, req *http.Request) {
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 	wsc, _, _, err := ws.UpgradeHTTP(req, w)
 	if err != nil {
 		return
 	}
 	defer wsc.Close()
-
 	node := raw.TCPNodes.PickMin()
 	defer raw.TCPNodes.DeferPick(node)
 
@@ -119,12 +134,13 @@ func (raw *Raw) HandleWsRequset(w http.ResponseWriter, req *http.Request) {
 }
 
 func (raw *Raw) HandleWssRequset(w http.ResponseWriter, req *http.Request) {
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 	wsc, _, _, err := ws.UpgradeHTTP(req, w)
 	if err != nil {
 		return
 	}
 	defer wsc.Close()
-
 	node := raw.TCPNodes.PickMin()
 	defer raw.TCPNodes.DeferPick(node)
 
@@ -143,8 +159,9 @@ func (raw *Raw) HandleWssRequset(w http.ResponseWriter, req *http.Request) {
 }
 
 func (raw *Raw) HandleMWssRequset(c net.Conn) {
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 	defer c.Close()
-
 	node := raw.TCPNodes.PickMin()
 	defer raw.TCPNodes.DeferPick(node)
 

internal/transporter/raw.go

@@ -5,6 +5,7 @@ import (
 	"net"
 
 	"github.com/Ehco1996/ehco/internal/logger"
+	"github.com/Ehco1996/ehco/internal/web"
 	"github.com/gobwas/ws"
 )
 
@@ -22,6 +23,8 @@ func (s *Ws) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
 
 func (s *Ws) HandleTCPConn(c *net.TCPConn) error {
 	defer c.Close()
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 
 	node := s.raw.TCPNodes.PickMin()
 	defer s.raw.TCPNodes.DeferPick(node)

internal/transporter/ws.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/Ehco1996/ehco/internal/logger"
 	"github.com/Ehco1996/ehco/internal/tls"
+	"github.com/Ehco1996/ehco/internal/web"
 	"github.com/gobwas/ws"
 )
 
@@ -23,6 +24,8 @@ func (s *Wss) HandleUDPConn(uaddr *net.UDPAddr, local *net.UDPConn) {
 
 func (s *Wss) HandleTCPConn(c *net.TCPConn) error {
 	defer c.Close()
+	web.CurTCPNum.Inc()
+	defer web.CurTCPNum.Dec()
 
 	node := s.raw.TCPNodes.PickMin()
 	defer s.raw.TCPNodes.DeferPick(node)

internal/transporter/wss.go

@@ -0,0 +1,43 @@
+package web
+
+import (
+	"os"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+var (
+	Hostname, _ = os.Hostname()
+
+	CurTCPNum = prometheus.NewGauge(prometheus.GaugeOpts{
+		Help: "当前tcp链接数",
+		Name: "current_tcp_num",
+		ConstLabels: map[string]string{
+			"hostname": Hostname,
+		},
+	})
+
+	CurUDPNum = prometheus.NewGauge(prometheus.GaugeOpts{
+		Help: "当前udp链接数",
+		Name: "current_udp_num",
+		ConstLabels: map[string]string{
+			"hostname": Hostname,
+		},
+	})
+
+	NetWorkTransmitBytes = prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "network_transmit_bytes",
+		Help: "传输流量总量bytes",
+		ConstLabels: map[string]string{
+			"hostname": Hostname,
+		},
+	})
+)
+
+func init() {
+
+	prometheus.MustRegister(CurTCPNum)
+	prometheus.MustRegister(CurUDPNum)
+	prometheus.MustRegister(NetWorkTransmitBytes)
+
+}

internal/web/metrics.go

@@ -9,6 +9,7 @@ import (
 	"github.com/Ehco1996/ehco/internal/constant"
 	"github.com/Ehco1996/ehco/internal/logger"
 	"github.com/gorilla/mux"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 func Index(w http.ResponseWriter, r *http.Request) {
@@ -21,16 +22,24 @@ func AttachProfiler(router *mux.Router) {
 	router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
 	router.HandleFunc("/debug/pprof/profile", pprof.Profile)
 	router.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+
+	// Manually add support for paths linked to by index page at /debug/pprof/
+	router.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
+	router.Handle("/debug/pprof/heap", pprof.Handler("heap"))
+	router.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
+	router.Handle("/debug/pprof/block", pprof.Handler("block"))
 }
 
 func StartWebServer(port string) {
 	addr := "0.0.0.0:" + port
 	logger.Logger.Infof("Start Web Server at http://%s/", addr)
 	r := mux.NewRouter()
 	AttachProfiler(r)
+	r.Handle("/metrics/", promhttp.Handler())
 	r.HandleFunc("/",
 		func(w http.ResponseWriter, r *http.Request) {
 			fmt.Fprintf(w, constant.IndexHTMLTMPL)
 		})
+
 	logger.Logger.Fatal(http.ListenAndServe(addr, r))
 }