142 lines
3.7 KiB
Go
142 lines
3.7 KiB
Go
package centrifugo
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"strconv"
|
||
"time"
|
||
|
||
"github.com/zeromicro/go-zero/core/stores/redis"
|
||
)
|
||
|
||
// RedisOnlineStore 使用 Redis 實作的在線狀態存儲
|
||
type RedisOnlineStore struct {
|
||
client *redis.Redis
|
||
keyPrefix string
|
||
}
|
||
|
||
// NewRedisOnlineStore 創建 Redis 在線狀態存儲
|
||
func NewRedisOnlineStore(client *redis.Redis) *RedisOnlineStore {
|
||
return &RedisOnlineStore{
|
||
client: client,
|
||
keyPrefix: "online:",
|
||
}
|
||
}
|
||
|
||
// NewRedisOnlineStoreWithPrefix 創建帶自定義前綴的 Redis 在線狀態存儲
|
||
func NewRedisOnlineStoreWithPrefix(client *redis.Redis, prefix string) *RedisOnlineStore {
|
||
return &RedisOnlineStore{
|
||
client: client,
|
||
keyPrefix: prefix,
|
||
}
|
||
}
|
||
|
||
// key 生成 Redis key
|
||
func (s *RedisOnlineStore) key(userID string) string {
|
||
return s.keyPrefix + userID
|
||
}
|
||
|
||
// clientCountKey 生成連線數 key
|
||
func (s *RedisOnlineStore) clientCountKey(userID string) string {
|
||
return s.keyPrefix + "clients:" + userID
|
||
}
|
||
|
||
// SetOnline 設置用戶在線
|
||
func (s *RedisOnlineStore) SetOnline(ctx context.Context, userID string, ttl time.Duration) error {
|
||
return s.client.SetexCtx(ctx, s.key(userID), fmt.Sprintf("%d", time.Now().Unix()), int(ttl.Seconds()))
|
||
}
|
||
|
||
// SetOffline 設置用戶離線
|
||
func (s *RedisOnlineStore) SetOffline(ctx context.Context, userID string) error {
|
||
// 刪除在線狀態
|
||
_, err := s.client.DelCtx(ctx, s.key(userID))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
// 刪除連線數
|
||
_, err = s.client.DelCtx(ctx, s.clientCountKey(userID))
|
||
return err
|
||
}
|
||
|
||
// IsOnline 檢查用戶是否在線
|
||
func (s *RedisOnlineStore) IsOnline(ctx context.Context, userID string) (bool, error) {
|
||
return s.client.ExistsCtx(ctx, s.key(userID))
|
||
}
|
||
|
||
// GetOnlineUsers 批量獲取在線用戶狀態
|
||
func (s *RedisOnlineStore) GetOnlineUsers(ctx context.Context, userIDs []string) (map[string]bool, error) {
|
||
if len(userIDs) == 0 {
|
||
return make(map[string]bool), nil
|
||
}
|
||
|
||
result := make(map[string]bool, len(userIDs))
|
||
for _, userID := range userIDs {
|
||
exists, err := s.client.ExistsCtx(ctx, s.key(userID))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to check online status for %s: %w", userID, err)
|
||
}
|
||
result[userID] = exists
|
||
}
|
||
|
||
return result, nil
|
||
}
|
||
|
||
// IncrClient 增加用戶連線數
|
||
func (s *RedisOnlineStore) IncrClient(ctx context.Context, userID string) (int64, error) {
|
||
count, err := s.client.IncrCtx(ctx, s.clientCountKey(userID))
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
return int64(count), nil
|
||
}
|
||
|
||
// DecrClient 減少用戶連線數
|
||
func (s *RedisOnlineStore) DecrClient(ctx context.Context, userID string) (int64, error) {
|
||
count, err := s.client.DecrCtx(ctx, s.clientCountKey(userID))
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
|
||
// 確保不會變成負數
|
||
if count < 0 {
|
||
_ = s.client.SetCtx(ctx, s.clientCountKey(userID), "0")
|
||
return 0, nil
|
||
}
|
||
|
||
return int64(count), nil
|
||
}
|
||
|
||
// GetClientCount 獲取用戶連線數
|
||
func (s *RedisOnlineStore) GetClientCount(ctx context.Context, userID string) (int64, error) {
|
||
val, err := s.client.GetCtx(ctx, s.clientCountKey(userID))
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
if val == "" {
|
||
return 0, nil
|
||
}
|
||
return strconv.ParseInt(val, 10, 64)
|
||
}
|
||
|
||
// GetAllOnlineUserIDs 獲取所有在線用戶 ID
|
||
// 注意：此方法使用 KEYS 命令，在大規模生產環境中可能有性能問題
|
||
// 建議在需要時使用 Centrifugo Presence API 替代
|
||
func (s *RedisOnlineStore) GetAllOnlineUserIDs(ctx context.Context) ([]string, error) {
|
||
// 使用 KEYS 查找所有在線用戶（排除 clients: 開頭的 key）
|
||
pattern := s.keyPrefix + "[^c]*"
|
||
keys, err := s.client.KeysCtx(ctx, pattern)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
userIDs := make([]string, 0, len(keys))
|
||
prefixLen := len(s.keyPrefix)
|
||
for _, key := range keys {
|
||
if len(key) > prefixLen {
|
||
userIDs = append(userIDs, key[prefixLen:])
|
||
}
|
||
}
|
||
|
||
return userIDs, nil
|
||
}
|