open-im-server-deploy/pkg/common/storage/cache/redis/online_count.go

package redis

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"git.imall.cloud/openim/open-im-server-deploy/pkg/common/storage/cache/cachekey"
	"git.imall.cloud/openim/protocol/constant"
	"github.com/openimsdk/tools/errs"
	"github.com/redis/go-redis/v9"
)

const onlineUserCountHistorySeparator = ":"

// OnlineUserCountSample 在线人数历史采样点
type OnlineUserCountSample struct {
	// Timestamp 采样时间（毫秒时间戳）
	Timestamp int64
	// Count 采样在线人数
	Count int64
}

// GetOnlineUserCount 读取在线人数缓存
func GetOnlineUserCount(ctx context.Context, rdb redis.UniversalClient) (int64, error) {
	if rdb == nil {
		return 0, errs.ErrInternalServer.WrapMsg("redis client is nil")
	}
	val, err := rdb.Get(ctx, cachekey.OnlineUserCountKey).Result()
	if err != nil {
		if errors.Is(err, redis.Nil) {
			return 0, err
		}
		return 0, errs.Wrap(err)
	}
	count, err := strconv.ParseInt(val, 10, 64)
	if err != nil {
		return 0, errs.WrapMsg(err, "parse online user count failed")
	}
	return count, nil
}

// RefreshOnlineUserCount 刷新在线人数缓存
func RefreshOnlineUserCount(ctx context.Context, rdb redis.UniversalClient) (int64, error) {
	if rdb == nil {
		return 0, errs.ErrInternalServer.WrapMsg("redis client is nil")
	}
	var (
		cursor uint64
		total  int64
	)
	now := strconv.FormatInt(time.Now().Unix(), 10)
	for {
		keys, nextCursor, err := rdb.Scan(ctx, cursor, fmt.Sprintf("%s*", cachekey.OnlineKey), constant.ParamMaxLength).Result()
		if err != nil {
			return 0, errs.Wrap(err)
		}
		for _, key := range keys {
			count, err := rdb.ZCount(ctx, key, now, "+inf").Result()
			if err != nil {
				return 0, errs.Wrap(err)
			}
			if count > 0 {
				total++
			}
		}
		cursor = nextCursor
		if cursor == 0 {
			break
		}
	}
	if err := rdb.Set(ctx, cachekey.OnlineUserCountKey, total, 0).Err(); err != nil {
		return 0, errs.Wrap(err)
	}
	return total, nil
}

// AppendOnlineUserCountHistory 写入在线人数历史采样
func AppendOnlineUserCountHistory(ctx context.Context, rdb redis.UniversalClient, timestamp int64, count int64) error {
	if rdb == nil {
		return errs.ErrInternalServer.WrapMsg("redis client is nil")
	}
	if timestamp <= 0 {
		return errs.ErrArgs.WrapMsg("invalid timestamp")
	}
	member := fmt.Sprintf("%d%s%d", timestamp, onlineUserCountHistorySeparator, count)
	if err := rdb.ZAdd(ctx, cachekey.OnlineUserCountHistoryKey, redis.Z{
		Score:  float64(timestamp),
		Member: member,
	}).Err(); err != nil {
		return errs.Wrap(err)
	}
	// 清理历史数据，避免无界增长
	retentionMs := int64(cachekey.OnlineUserCountHistoryRetention / time.Millisecond)
	cutoff := timestamp - retentionMs
	if cutoff > 0 {
		if err := rdb.ZRemRangeByScore(ctx, cachekey.OnlineUserCountHistoryKey, "0", strconv.FormatInt(cutoff, 10)).Err(); err != nil {
			return errs.Wrap(err)
		}
	}
	return nil
}

// GetOnlineUserCountHistory 读取在线人数历史采样
func GetOnlineUserCountHistory(ctx context.Context, rdb redis.UniversalClient, startTime int64, endTime int64) ([]OnlineUserCountSample, error) {
	if rdb == nil {
		return nil, errs.ErrInternalServer.WrapMsg("redis client is nil")
	}
	if startTime <= 0 || endTime <= 0 || endTime <= startTime {
		return nil, nil
	}
	// 包含endTime的数据，使用endTime作为最大值
	values, err := rdb.ZRangeByScore(ctx, cachekey.OnlineUserCountHistoryKey, &redis.ZRangeBy{
		Min: strconv.FormatInt(startTime, 10),
		Max: strconv.FormatInt(endTime, 10),
	}).Result()
	if err != nil {
		if errors.Is(err, redis.Nil) {
			return nil, nil
		}
		return nil, errs.Wrap(err)
	}
	if len(values) == 0 {
		return nil, nil
	}
	samples := make([]OnlineUserCountSample, 0, len(values))
	for _, val := range values {
		parts := strings.SplitN(val, onlineUserCountHistorySeparator, 2)
		if len(parts) != 2 {
			continue
		}
		ts, err := strconv.ParseInt(parts[0], 10, 64)
		if err != nil {
			continue
		}
		cnt, err := strconv.ParseInt(parts[1], 10, 64)
		if err != nil {
			continue
		}
		samples = append(samples, OnlineUserCountSample{
			Timestamp: ts,
			Count:     cnt,
		})
	}
	return samples, nil
}