package redis import ( "context" "crypto/rand" "encoding/hex" "os" "strconv" "time" redis "github.com/redis/go-redis/v9" ) type Client struct { rdb *redis.Client } func NewFromEnv() (*Client, error) { addr := os.Getenv("REDIS_ADDR") if addr == "" { addr = "127.0.0.1:6379" } db := 0 if v := os.Getenv("REDIS_DB"); v != "" { n, err := strconv.Atoi(v) if err != nil { return nil, err } db = n } rdb := redis.NewClient(&redis.Options{ Addr: addr, Password: os.Getenv("REDIS_PASSWORD"), DB: db, }) return &Client{rdb: rdb}, nil } func (c *Client) Close() error { if c == nil || c.rdb == nil { return nil } return c.rdb.Close() } func (c *Client) Ping(ctx context.Context) error { if c == nil || c.rdb == nil { return nil } return c.rdb.Ping(ctx).Err() } func (c *Client) GetBytes(ctx context.Context, key string) ([]byte, error) { return c.rdb.Get(ctx, key).Bytes() } func (c *Client) SetBytes(ctx context.Context, key string, value []byte, ttl time.Duration) error { return c.rdb.Set(ctx, key, value, ttl).Err() } func (c *Client) Del(ctx context.Context, key string) error { return c.rdb.Del(ctx, key).Err() } func IsMiss(err error) bool { return err == redis.Nil } func randToken(n int) (string, error) { b := make([]byte, n) if _, err := rand.Read(b); err != nil { return "", err } return hex.EncodeToString(b), nil } func (c *Client) Lock(ctx context.Context, key string, ttl time.Duration) (token string, ok bool, err error) { if c == nil || c.rdb == nil { return "", false, nil } token, err = randToken(16) if err != nil { return "", false, err } ok, err = c.rdb.SetNX(ctx, key, token, ttl).Result() return token, ok, err } var unlockScript = redis.NewScript(` if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) else return 0 end `) func (c *Client) Unlock(ctx context.Context, key string, token string) error { if c == nil || c.rdb == nil { return nil } _, err := unlockScript.Run(ctx, c.rdb, []string{key}, token).Result() return err } func (c *Client) ZincrBy(ctx context.Context, key string, member string, score float64) error { if c == nil || c.rdb == nil { return nil } return c.rdb.ZIncrBy(ctx, key, score, member).Err() } func (c *Client) Expire(ctx context.Context, key string, ttl time.Duration) error { if c == nil || c.rdb == nil { return nil } return c.rdb.Expire(ctx, key, ttl).Err() } func (c *Client) ZUnionStore(ctx context.Context, dst string, keys []string, aggregate string) error { if c == nil || c.rdb == nil { return nil } return c.rdb.ZUnionStore(ctx, dst, &redis.ZStore{ Keys: keys, Aggregate: aggregate, }).Err() } func (c *Client) Exists(ctx context.Context, key string) (bool, error) { if c == nil || c.rdb == nil { return false, nil } n, err := c.rdb.Exists(ctx, key).Result() return n > 0, err } func (c *Client) ZRevRange(ctx context.Context, key string, start, stop int64) ([]string, error) { if c == nil || c.rdb == nil { return nil, nil } return c.rdb.ZRevRange(ctx, key, start, stop).Result() } func (c *Client) ZRevRangeByScore(ctx context.Context, key string, max, min string, offset, count int64) ([]string, error) { if c == nil || c.rdb == nil { return nil, nil } return c.rdb.ZRevRangeByScore(ctx, key, &redis.ZRangeBy{ Max: max, Min: min, Offset: offset, Count: count, }).Result() }