package logic import ( "context" "fmt" "github.com/zeromicro/go-zero/core/logx" "sync" ) type websocketStatType string const ( TYPE_CONNECT_COUNT websocketStatType = "connect_count" //ws连接数 TYPE_COMBINE_IMAGE_REQUEST_COUNT websocketStatType = "combine_count" ) type websocketStatItem struct { UserId int64 `json:"user_id"` GuestId int64 `json:"guest_id"` 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) { websocketStat <- websocketStatItem{ UserId: userId, GuestId: guestId, Type: TYPE_CONNECT_COUNT, Value: 1, } } // 减少ws连接数计数 func decreaseWebsocketConnectCount(userId, guestId int64) { websocketStat <- websocketStatItem{ UserId: userId, GuestId: guestId, Type: TYPE_CONNECT_COUNT, Value: -1, } } // 累增合图请求数计数 func increaseCombineRequestCount(userId, guestId int64) { websocketStat <- websocketStatItem{ UserId: userId, GuestId: guestId, Type: TYPE_COMBINE_IMAGE_REQUEST_COUNT, Value: 1, } } // 减少合图请求数计数 func decreaseCombineRequestCount(userId, guestId int64) { websocketStat <- websocketStatItem{ UserId: userId, GuestId: guestId, Type: TYPE_COMBINE_IMAGE_REQUEST_COUNT, Value: -1, } } // 消费数据 func ConsumeWebsocketStatData(ctx context.Context) { defer func() { if err := recover(); err != nil { logx.Error("ConsumeWebsocketStatData panic:", err) } }() go func() { select { case <-ctx.Done(): panic("ConsumeWebsocketStatData ctx deadline") } }() for { select { case data := <-websocketStat: key := fmt.Sprintf("%d_%d", data.UserId, data.GuestId) switch data.Type { case TYPE_CONNECT_COUNT: //ws连接计数 statData, ok := mapUserWsStat.Load(key) if ok { //存在就累加 if stat, ok := statData.(mapUserWsStatItem); ok { stat.CurWsConnectCount += data.Value if stat.CurWsConnectCount == 0 { //为0则移除 mapUserWsStat.Delete(key) } else { mapUserWsStat.Store(key, stat) } } else { logx.Error("断言mapUserWsStatItem错误") } } else { //不存在就赋值 mapUserWsStat.Store(key, mapUserWsStatItem{ CurWsConnectCount: data.Value, }) } case TYPE_COMBINE_IMAGE_REQUEST_COUNT: //请求算法合图计数 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, }) } } } } }