[统计]记录流量（以5分钟作为间隔）

internal/const/const.go

@@ -1,7 +1,7 @@
 package teaconst
 
 const (
-	Version = "0.0.5"
+	Version = "0.0.6"
 
 	ProductName = "Edge Node"
 	ProcessName = "edge-node"

internal/nodes/http_request.go

@@ -199,7 +199,14 @@ func (this *HTTPRequest) doBegin() {
 
 // 结束调用
 func (this *HTTPRequest) doEnd() {
+	// 记录日志
 	this.log()
+
+	// 流量统计
+	// TODO 增加是否开启开关
+	if this.Server != nil {
+		SharedTrafficStatManager.Add(this.Server.Id, this.writer.sentBodyBytes)
+	}
 }
 
 // 原始的请求URI

internal/nodes/http_request_websocket.go

@@ -69,7 +69,20 @@ func (this *HTTPRequest) doWebsocket() {
 	}()
 
 	go func() {
-		_, _ = io.Copy(clientConn, originConn)
+		buf := make([]byte, 4*1024) // TODO 使用内存池
+		for {
+			n, err := originConn.Read(buf)
+			if n > 0 {
+				this.writer.sentBodyBytes += int64(n)
+				_, err = clientConn.Write(buf[:n])
+				if err != nil {
+					break
+				}
+			}
+			if err != nil {
+				break
+			}
+		}
 		_ = clientConn.Close()
 		_ = originConn.Close()
 	}()

internal/nodes/listener_tcp.go

@@ -65,7 +65,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 	}
 
 	go func() {
-		originBuffer := make([]byte, 32*1024) // TODO 需要可以设置，并可以使用Pool
+		originBuffer := make([]byte, 4*1024) // TODO 需要可以设置，并可以使用Pool
 		for {
 			n, err := originConn.Read(originBuffer)
 			if n > 0 {
@@ -74,6 +74,9 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 					closer()
 					break
 				}
+
+				// 记录流量
+				SharedTrafficStatManager.Add(firstServer.Id, int64(n))
 			}
 			if err != nil {
 				closer()
@@ -82,7 +85,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 		}
 	}()
 
-	clientBuffer := make([]byte, 32*1024) // TODO 需要可以设置，并可以使用Pool
+	clientBuffer := make([]byte, 4*1024) // TODO 需要可以设置，并可以使用Pool
 	for {
 		n, err := conn.Read(clientBuffer)
 		if n > 0 {

internal/nodes/traffic_stat_manager.go

@@ -0,0 +1,102 @@
+package nodes
+
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeNode/internal/logs"
+	"github.com/TeaOSLab/EdgeNode/internal/rpc"
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"github.com/iwind/TeaGo/Tea"
+	"strconv"
+	"sync"
+	"time"
+)
+
+var SharedTrafficStatManager = NewTrafficStatManager()
+
+// 流量统计
+type TrafficStatManager struct {
+	m      map[string]int64 // [timestamp serverId] => bytes
+	locker sync.Mutex
+}
+
+// 获取新对象
+func NewTrafficStatManager() *TrafficStatManager {
+	manager := &TrafficStatManager{
+		m: map[string]int64{},
+	}
+
+	go manager.Start()
+
+	return manager
+}
+
+// 启动自动任务
+func (this *TrafficStatManager) Start() {
+	duration := 5 * time.Minute
+	if Tea.IsTesting() {
+		// 测试环境缩短上传时间，方便我们调试
+		duration = 30 * time.Second
+	}
+	ticker := time.NewTicker(duration)
+	for range ticker.C {
+		err := this.Upload()
+		if err != nil {
+			logs.Error("TRAFFIC_STAT_MANAGER", "upload stats failed: "+err.Error())
+		}
+	}
+}
+
+// 添加流量
+func (this *TrafficStatManager) Add(serverId int64, bytes int64) {
+	if bytes == 0 {
+		return
+	}
+
+	timestamp := utils.UnixTime() / 300 * 300
+
+	key := strconv.FormatInt(timestamp, 10) + strconv.FormatInt(serverId, 10)
+	this.locker.Lock()
+	this.m[key] += bytes
+	this.locker.Unlock()
+}
+
+// 上传流量
+func (this *TrafficStatManager) Upload() error {
+	if sharedNodeConfig == nil {
+		return nil
+	}
+
+	client, err := rpc.SharedRPC()
+	if err != nil {
+		return err
+	}
+
+	this.locker.Lock()
+	m := this.m
+	this.m = map[string]int64{}
+	this.locker.Unlock()
+
+	pbStats := []*pb.ServerDailyStat{}
+	for key, bytes := range m {
+		timestamp, err := strconv.ParseInt(key[:10], 10, 64)
+		if err != nil {
+			return err
+		}
+		serverId, err := strconv.ParseInt(key[10:], 10, 64)
+		if err != nil {
+			return err
+		}
+
+		pbStats = append(pbStats, &pb.ServerDailyStat{
+			ServerId:  serverId,
+			RegionId:  sharedNodeConfig.RegionId,
+			Bytes:     bytes,
+			CreatedAt: timestamp,
+		})
+	}
+	if len(pbStats) == 0 {
+		return nil
+	}
+	_, err = client.ServerDailyStatRPC().UploadServerDailyStats(client.Context(), &pb.UploadServerDailyStatsRequest{Stats: pbStats})
+	return err
+}

internal/nodes/traffic_stat_manager_test.go

@@ -0,0 +1,35 @@
+package nodes
+
+import (
+	"runtime"
+	"testing"
+)
+
+func TestTrafficStatManager_Add(t *testing.T) {
+	manager := NewTrafficStatManager()
+	for i := 0; i < 100; i++ {
+		manager.Add(1, 10)
+	}
+	t.Log(manager.m)
+}
+
+func TestTrafficStatManager_Upload(t *testing.T) {
+	manager := NewTrafficStatManager()
+	for i := 0; i < 100; i++ {
+		manager.Add(1, 10)
+	}
+	err := manager.Upload()
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("ok")
+}
+
+func BenchmarkTrafficStatManager_Add(b *testing.B) {
+	runtime.GOMAXPROCS(1)
+
+	manager := NewTrafficStatManager()
+	for i := 0; i < b.N; i++ {
+		manager.Add(1, 1024)
+	}
+}

internal/rpc/rpc_client.go

@@ -93,6 +93,10 @@ func (this *RPCClient) ACMEAuthenticationRPC() pb.ACMEAuthenticationServiceClien
 	return pb.NewACMEAuthenticationServiceClient(this.pickConn())
 }
 
+func (this *RPCClient) ServerDailyStatRPC() pb.ServerDailyStatServiceClient {
+	return pb.NewServerDailyStatServiceClient(this.pickConn())
+}
+
 // 节点上下文信息
 func (this *RPCClient) Context() context.Context {
 	ctx := context.Background()