lingniu-vehicle-ingest/vehicle-data-platform/apps/api/internal/platform/redis_online.go

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
}