diff --git a/internal/nodes/http_request.go b/internal/nodes/http_request.go index 4d02981..cc4886b 100644 --- a/internal/nodes/http_request.go +++ b/internal/nodes/http_request.go @@ -403,7 +403,7 @@ func (this *HTTPRequest) doEnd() { attackBytes = this.CalculateSize() } - stats.SharedTrafficStatManager.Add(this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes()+this.writer.SentHeaderBytes(), cachedBytes, 1, countCached, countAttacks, attackBytes, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId()) + stats.SharedTrafficStatManager.Add(this.ReqServer.UserId, this.ReqServer.Id, this.ReqHost, this.writer.SentBodyBytes()+this.writer.SentHeaderBytes(), cachedBytes, 1, countCached, countAttacks, attackBytes, this.ReqServer.ShouldCheckTrafficLimit(), this.ReqServer.PlanId()) // 指标 if metrics.SharedManager.HasHTTPMetrics() { diff --git a/internal/nodes/listener_tcp.go b/internal/nodes/listener_tcp.go index afe7b7c..a41ed74 100644 --- a/internal/nodes/listener_tcp.go +++ b/internal/nodes/listener_tcp.go @@ -114,14 +114,14 @@ func (this *TCPListener) handleConn(conn net.Conn) error { serverName = tlsConn.ConnectionState().ServerName if len(serverName) > 0 { // 统计 - stats.SharedTrafficStatManager.Add(server.Id, serverName, 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) + stats.SharedTrafficStatManager.Add(server.UserId, server.Id, serverName, 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) recordStat = true } } // 统计 if !recordStat { - stats.SharedTrafficStatManager.Add(server.Id, "", 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) + stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) } originConn, err := this.connectOrigin(server.Id, serverName, server.ReverseProxy, conn.RemoteAddr().String()) @@ -176,7 +176,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error { // 记录流量 if server != nil { - stats.SharedTrafficStatManager.Add(server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) + stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) } } if err != nil { diff --git a/internal/nodes/listener_udp.go b/internal/nodes/listener_udp.go index 1760515..3a8a408 100644 --- a/internal/nodes/listener_udp.go +++ b/internal/nodes/listener_udp.go @@ -370,7 +370,7 @@ func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyListener // 统计 if server != nil { - stats.SharedTrafficStatManager.Add(server.Id, "", 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) + stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", 0, 0, 1, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) } // 处理ControlMessage @@ -401,7 +401,7 @@ func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyListener // 记录流量和带宽 if server != nil { // 流量 - stats.SharedTrafficStatManager.Add(server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) + stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId()) // 带宽 stats.SharedBandwidthStatManager.AddBandwidth(server.UserId, server.Id, int64(n), int64(n)) diff --git a/internal/stats/traffic_stat_manager.go b/internal/stats/traffic_stat_manager.go index eb44004..82eb114 100644 --- a/internal/stats/traffic_stat_manager.go +++ b/internal/stats/traffic_stat_manager.go @@ -22,6 +22,7 @@ import ( var SharedTrafficStatManager = NewTrafficStatManager() type TrafficItem struct { + UserId int64 Bytes int64 CachedBytes int64 CountRequests int64 @@ -107,7 +108,7 @@ func (this *TrafficStatManager) Start() { } // Add 添加流量 -func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64, cachedBytes int64, countRequests int64, countCachedRequests int64, countAttacks int64, attackBytes int64, checkingTrafficLimit bool, planId int64) { +func (this *TrafficStatManager) Add(userId int64, serverId int64, domain string, bytes int64, cachedBytes int64, countRequests int64, countCachedRequests int64, countAttacks int64, attackBytes int64, checkingTrafficLimit bool, planId int64) { if serverId == 0 { return } @@ -128,7 +129,9 @@ func (this *TrafficStatManager) Add(serverId int64, domain string, bytes int64, // 总的流量 item, ok := this.itemMap[key] if !ok { - item = &TrafficItem{} + item = &TrafficItem{ + UserId: userId, + } this.itemMap[key] = item } item.Bytes += bytes @@ -200,6 +203,7 @@ func (this *TrafficStatManager) Upload() error { } pbServerStats = append(pbServerStats, &pb.ServerDailyStat{ + UserId: item.UserId, ServerId: serverId, NodeRegionId: regionId, Bytes: item.Bytes, diff --git a/internal/stats/traffic_stat_manager_test.go b/internal/stats/traffic_stat_manager_test.go index 4d632ca..fb3db47 100644 --- a/internal/stats/traffic_stat_manager_test.go +++ b/internal/stats/traffic_stat_manager_test.go @@ -10,7 +10,7 @@ import ( func TestTrafficStatManager_Add(t *testing.T) { manager := NewTrafficStatManager() for i := 0; i < 100; i++ { - manager.Add(1, "goedge.cn", 1, 0, 0, 0, 0, 0, false, 0) + manager.Add(1, 1, "goedge.cn", 1, 0, 0, 0, 0, 0, false, 0) } t.Log(manager.itemMap) } @@ -18,7 +18,7 @@ func TestTrafficStatManager_Add(t *testing.T) { func TestTrafficStatManager_Upload(t *testing.T) { manager := NewTrafficStatManager() for i := 0; i < 100; i++ { - manager.Add(1, "goedge.cn"+types.String(rands.Int(0, 10)), 1, 0, 1, 0, 0, 0, false, 0) + manager.Add(1, 1, "goedge.cn"+types.String(rands.Int(0, 10)), 1, 0, 1, 0, 0, 0, false, 0) } err := manager.Upload() if err != nil { @@ -32,6 +32,6 @@ func BenchmarkTrafficStatManager_Add(b *testing.B) { manager := NewTrafficStatManager() for i := 0; i < b.N; i++ { - manager.Add(1, "goedge.cn", 1024, 1, 0, 0, 0, 0, false, 0) + manager.Add(1, 1, "goedge.cn", 1024, 1, 0, 0, 0, 0, false, 0) } }