package platform import ( "bufio" "context" "errors" "fmt" "io" "net" "strconv" "strings" "time" ) type RedisOnlineKeyCounter struct { addr string username string password string db int } func NewRedisOnlineKeyCounter(addr, username, password string, db int) *RedisOnlineKeyCounter { return &RedisOnlineKeyCounter{addr: strings.TrimSpace(addr), username: username, password: password, db: db} } func (c *RedisOnlineKeyCounter) CountOnlineKeys(ctx context.Context) (int, error) { if c.addr == "" { return 0, errors.New("redis addr is empty") } dialer := net.Dialer{Timeout: 2 * time.Second} conn, err := dialer.DialContext(ctx, "tcp", c.addr) if err != nil { return 0, err } defer conn.Close() if deadline, ok := ctx.Deadline(); ok { _ = conn.SetDeadline(deadline) } else { _ = conn.SetDeadline(time.Now().Add(3 * time.Second)) } reader := bufio.NewReader(conn) if c.password != "" { args := []string{"AUTH"} if c.username != "" { args = append(args, c.username) } args = append(args, c.password) if _, err := redisCommand(conn, reader, args...); err != nil { return 0, err } } if _, err := redisCommand(conn, reader, "PING"); err != nil { return 0, err } if c.db >= 0 { if _, err := redisCommand(conn, reader, "SELECT", strconv.Itoa(c.db)); err != nil { return 0, err } } total := 0 cursor := "0" for { value, err := redisCommand(conn, reader, "SCAN", cursor, "MATCH", "vehicle:online:*", "COUNT", "1000") if err != nil { return 0, err } reply, ok := value.([]any) if !ok || len(reply) != 2 { return 0, fmt.Errorf("unexpected redis scan reply: %#v", value) } nextCursor, ok := reply[0].(string) if !ok { return 0, fmt.Errorf("unexpected redis scan cursor: %#v", reply[0]) } keys, ok := reply[1].([]any) if !ok { return 0, fmt.Errorf("unexpected redis scan keys: %#v", reply[1]) } total += len(keys) if nextCursor == "0" { return total, nil } cursor = nextCursor } } func redisCommand(conn net.Conn, reader *bufio.Reader, args ...string) (any, error) { if _, err := fmt.Fprintf(conn, "*%d\r\n", len(args)); err != nil { return nil, err } for _, arg := range args { if _, err := fmt.Fprintf(conn, "$%d\r\n%s\r\n", len(arg), arg); err != nil { return nil, err } } return readRedisReply(reader) } func readRedisReply(reader *bufio.Reader) (any, error) { prefix, err := reader.ReadByte() if err != nil { return nil, err } switch prefix { case '+': return readRedisLine(reader) case '-': line, err := readRedisLine(reader) if err != nil { return nil, err } return nil, errors.New(line) case ':': line, err := readRedisLine(reader) if err != nil { return nil, err } return strconv.Atoi(line) case '$': line, err := readRedisLine(reader) if err != nil { return nil, err } size, err := strconv.Atoi(line) if err != nil { return nil, err } if size < 0 { return "", nil } buf := make([]byte, size+2) if _, err := io.ReadFull(reader, buf); err != nil { return nil, err } return string(buf[:size]), nil case '*': line, err := readRedisLine(reader) if err != nil { return nil, err } count, err := strconv.Atoi(line) if err != nil { return nil, err } values := make([]any, 0, count) for i := 0; i < count; i++ { value, err := readRedisReply(reader) if err != nil { return nil, err } values = append(values, value) } return values, nil default: return nil, fmt.Errorf("unexpected redis reply prefix %q", prefix) } } func readRedisLine(reader *bufio.Reader) (string, error) { line, err := reader.ReadString('\n') if err != nil { return "", err } return strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r"), nil }