From 485319d4f8df59155cca83fd642d46be79f22b18 Mon Sep 17 00:00:00 2001
From: laodaming <11058467+laudamine@user.noreply.gitee.com>
Date: Tue, 31 Oct 2023 15:03:15 +0800
Subject: [PATCH] fix

---
 .../websocket/internal/logic/getstatlogic.go  | 15 +++++-
 .../websocket/internal/logic/ws_statistics.go | 54 ++++++++++++-------
 2 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/server/websocket/internal/logic/getstatlogic.go b/server/websocket/internal/logic/getstatlogic.go
index 8f56fcdb..5bbf0d21 100644
--- a/server/websocket/internal/logic/getstatlogic.go
+++ b/server/websocket/internal/logic/getstatlogic.go
@@ -37,9 +37,20 @@ func (l *GetStatLogic) GetStat(req *types.GetStatReq, userinfo *auth.UserInfo) (
 	}
 	GetStatMutex.Lock()
 	defer GetStatMutex.Unlock()
-	userStat := make(map[string]int)
+	userStat := make(map[string]interface{})
+	currentWebsocketConnectCount := 0
+	currentRequestCombineApiCount := 0
 	mapUserWsStat.Range(func(key, value any) bool {
-		userStat[key.(string)] = value.(int)
+		statData, ok := mapUserWsStat.Load(key)
+		if ok { //存在就累加
+			if stat, ok := statData.(mapUserWsStatItem); ok {
+				currentWebsocketConnectCount += stat.CurWsConnectCount
+				currentRequestCombineApiCount += stat.CurCombineCount
+			} else {
+				logx.Error("断言mapUserWsStatItem错误")
+			}
+		}
+		userStat[key.(string)] = value
 		return true
 	})
 	return resp.SetStatusWithMessage(basic.CodeOK, "success", types.GetStatRsp{
diff --git a/server/websocket/internal/logic/ws_statistics.go b/server/websocket/internal/logic/ws_statistics.go
index 176ee2a5..98d7bbc1 100644
--- a/server/websocket/internal/logic/ws_statistics.go
+++ b/server/websocket/internal/logic/ws_statistics.go
@@ -7,18 +7,6 @@ import (
 	"sync"
 )
 
-// 统计信息
-var (
-	//当前ws连接数
-	currentWebsocketConnectCount = 0
-	//当前合图进行的请求总数
-	currentRequestCombineApiCount = 0
-	//用户连接统计
-	mapUserWsStat = sync.Map{}
-	//添加or减少连接的控制chan
-	websocketStat = make(chan websocketStatItem, 20)
-)
-
 type websocketStatType string
 
 const (
@@ -32,6 +20,18 @@ type websocketStatItem struct {
 	Type    websocketStatType `json:"type"`  //类型
 	Value   int               `json:"value"` //数值
 }
+type mapUserWsStatItem struct {
+	CurCombineCount   int `json:"cur_combine_count"`    //当前合图数
+	CurWsConnectCount int `json:"cur_ws_connect_count"` //当前连接数
+}
+
+// 统计信息
+var (
+	//用户连接统计
+	mapUserWsStat = sync.Map{}
+	//添加or减少连接的控制chan
+	websocketStat = make(chan websocketStatItem, 20)
+)
 
 // 累增ws连接数计数
 func increaseWebsocketConnectCount(userId, guestId int64) {
@@ -89,18 +89,36 @@ func ConsumeWebsocketStatData(ctx context.Context) {
 	for {
 		select {
 		case data := <-websocketStat:
+			key := fmt.Sprintf("%d_%d", data.UserId, data.GuestId)
 			switch data.Type {
 			case TYPE_CONNECT_COUNT: //ws连接计数
-				key := fmt.Sprintf("%d_%d", data.UserId, data.GuestId)
-				statCount, ok := mapUserWsStat.Load(key)
+				statData, ok := mapUserWsStat.Load(key)
 				if ok { //存在就累加
-					mapUserWsStat.Store(key, statCount.(int)+data.Value)
+					if stat, ok := statData.(mapUserWsStatItem); ok {
+						stat.CurWsConnectCount += data.Value
+						mapUserWsStat.Store(key, stat)
+					} else {
+						logx.Error("断言mapUserWsStatItem错误")
+					}
 				} else { //不存在就赋值
-					mapUserWsStat.Store(key, data.Value)
+					mapUserWsStat.Store(key, mapUserWsStatItem{
+						CurWsConnectCount: data.Value,
+					})
 				}
-				currentWebsocketConnectCount += data.Value
 			case TYPE_COMBINE_IMAGE_REQUEST_COUNT: //请求算法合图计数
-				currentRequestCombineApiCount += data.Value
+				statData, ok := mapUserWsStat.Load(key)
+				if ok { //存在就累加
+					if stat, ok := statData.(mapUserWsStatItem); ok {
+						stat.CurCombineCount += data.Value
+						mapUserWsStat.Store(key, stat)
+					} else {
+						logx.Error("断言mapUserWsStatItem错误")
+					}
+				} else { //不存在就赋值
+					mapUserWsStat.Store(key, mapUserWsStatItem{
+						CurCombineCount: data.Value,
+					})
+				}
 			}
 		}
 	}