fusenapi/server/websocket/internal/logic/ws_statistics.go

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 data.Value <= 0 {
						continue
					}
					mapUserWsStat.Store(key, mapUserWsStatItem{
						CurWsConnectCount: data.Value,
					})
				}
				stat, ok := statData.(mapUserWsStatItem)
				if !ok {
					logx.Error("断言mapUserWsStatItem错误")
					continue
				}
				stat.CurWsConnectCount += data.Value
				if stat.CurWsConnectCount <= 0 { //小于等于0则移除
					mapUserWsStat.Delete(key)
				} else {
					mapUserWsStat.Store(key, stat)
				}
			case TYPE_COMBINE_IMAGE_REQUEST_COUNT: //请求算法合图计数
				statData, ok := mapUserWsStat.Load(key)
				//不存在
				if !ok {
					//数据是减少的
					if data.Value <= 0 {
						continue
					}
					mapUserWsStat.Store(key, mapUserWsStatItem{
						CurCombineCount: data.Value,
					})
				}
				//存在
				if stat, ok := statData.(mapUserWsStatItem); ok {
					stat.CurCombineCount += data.Value
					if stat.CurCombineCount < 0 {
						stat.CurCombineCount = 0
					}
					mapUserWsStat.Store(key, stat)
				} else {
					logx.Error("断言mapUserWsStatItem错误")
				}
			}
		}
	}
}