实现全局的TCP最大连接数

internal/nodes/client_conn.go

@@ -8,8 +8,10 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/events"
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/monitor"
+	"github.com/TeaOSLab/EdgeNode/internal/ratelimit"
 	"github.com/iwind/TeaGo/maps"
 	"net"
+	"sync"
 	"sync/atomic"
 	"time"
 )
@@ -44,18 +46,22 @@ func init() {
 type ClientConn struct {
 	rawConn  net.Conn
 	isClosed bool
+
+	once    sync.Once
+	limiter *ratelimit.Counter
 }
 
-func NewClientConn(conn net.Conn, quickClose bool) net.Conn {
+func NewClientConn(conn net.Conn, quickClose bool, limiter *ratelimit.Counter) net.Conn {
 	if quickClose {
+		// TCP
 		tcpConn, ok := conn.(*net.TCPConn)
 		if ok {
 			// TODO 可以设置此值
-			_ = tcpConn.SetLinger(3)
+			_ = tcpConn.SetLinger(nodeconfigs.DefaultTCPLinger)
 		}
 	}
 
-	return &ClientConn{rawConn: conn}
+	return &ClientConn{rawConn: conn, limiter: limiter}
 }
 
 func (this *ClientConn) Read(b []byte) (n int, err error) {
@@ -76,6 +82,11 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 
 func (this *ClientConn) Close() error {
 	this.isClosed = true
+	this.once.Do(func() {
+		if this.limiter != nil {
+			this.limiter.Release()
+		}
+	})
 	return this.rawConn.Close()
 }
 

internal/nodes/client_listener.go

@@ -3,12 +3,16 @@
 package nodes
 
 import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
 	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
 	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
+	"github.com/TeaOSLab/EdgeNode/internal/ratelimit"
 	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"net"
 )
 
+var sharedConnectionsLimiter = ratelimit.NewCounter(nodeconfigs.DefaultTCPMaxConnections)
+
 // ClientListener 客户端网络监听
 type ClientListener struct {
 	rawListener net.Listener
@@ -23,10 +27,21 @@ func NewClientListener(listener net.Listener, quickClose bool) net.Listener {
 }
 
 func (this *ClientListener) Accept() (net.Conn, error) {
+	// 限制并发连接数
+	var isOk = false
+	var limiter = sharedConnectionsLimiter
+	limiter.Ack()
+	defer func() {
+		if !isOk {
+			limiter.Release()
+		}
+	}()
+
 	conn, err := this.rawListener.Accept()
 	if err != nil {
 		return nil, err
 	}
+
 	// 是否在WAF名单中
 	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
 	if err == nil {
@@ -42,7 +57,8 @@ func (this *ClientListener) Accept() (net.Conn, error) {
 		}
 	}
 
-	return NewClientConn(conn, this.quickClose), nil
+	isOk = true
+	return NewClientConn(conn, this.quickClose, limiter), nil
 }
 
 func (this *ClientListener) Close() error {

internal/nodes/node.go

@@ -14,6 +14,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/goman"
 	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
 	"github.com/TeaOSLab/EdgeNode/internal/metrics"
+	"github.com/TeaOSLab/EdgeNode/internal/ratelimit"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
 	"github.com/TeaOSLab/EdgeNode/internal/stats"
@@ -51,12 +52,15 @@ type Node struct {
 	sock     *gosock.Sock
 	locker   sync.Mutex
 
-	timezone string
+	maxCPU     int32
+	maxThreads int
+	timezone   string
 }
 
 func NewNode() *Node {
 	return &Node{
-		sock: gosock.NewTmpSock(teaconst.ProcessName),
+		sock:       gosock.NewTmpSock(teaconst.ProcessName),
+		maxThreads: -1,
 	}
 }
 
@@ -637,18 +641,39 @@ func (this *Node) listenSock() error {
 // 重载配置调用
 func (this *Node) onReload(config *nodeconfigs.NodeConfig) {
 	// max cpu
-	if config.MaxCPU > 0 && config.MaxCPU < int32(runtime.NumCPU()) {
-		runtime.GOMAXPROCS(int(config.MaxCPU))
-	} else {
-		runtime.GOMAXPROCS(runtime.NumCPU())
+	if config.MaxCPU != this.maxCPU {
+		if config.MaxCPU > 0 && config.MaxCPU < int32(runtime.NumCPU()) {
+			runtime.GOMAXPROCS(int(config.MaxCPU))
+			remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(config.MaxCPU)+"'")
+		} else {
+			runtime.GOMAXPROCS(runtime.NumCPU())
+			remotelogs.Println("NODE", "[CPU]set max cpu to '"+types.String(runtime.NumCPU())+"'")
+		}
+
+		this.maxCPU = config.MaxCPU
 	}
 
 	// max threads
-	if config.MaxThreads > 0 {
-		debug.SetMaxThreads(config.MaxThreads)
-		remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(config.MaxThreads)+"'")
-	} else {
-		debug.SetMaxThreads(nodeconfigs.DefaultMaxThreads)
+	if config.MaxThreads != this.maxThreads {
+		if config.MaxThreads > 0 {
+			debug.SetMaxThreads(config.MaxThreads)
+			remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(config.MaxThreads)+"'")
+		} else {
+			debug.SetMaxThreads(nodeconfigs.DefaultMaxThreads)
+			remotelogs.Println("NODE", "[THREADS]set max threads to '"+types.String(nodeconfigs.DefaultMaxThreads)+"'")
+		}
+		this.maxThreads = config.MaxThreads
+	}
+
+	// max tcp connections
+	if config.TCPMaxConnections <= 0 {
+		config.TCPMaxConnections = nodeconfigs.DefaultTCPMaxConnections
+	}
+	if config.TCPMaxConnections != sharedConnectionsLimiter.Count() {
+		remotelogs.Println("NODE", "[TCP]changed tcp max connections to '"+types.String(config.TCPMaxConnections)+"'")
+
+		sharedConnectionsLimiter.Close()
+		sharedConnectionsLimiter = ratelimit.NewCounter(config.TCPMaxConnections)
 	}
 
 	// timezone