crawlab/db/redis/client.go

package redis

import (
	"github.com/apex/log"
	"github.com/crawlab-team/crawlab/db"
	"github.com/crawlab-team/crawlab/db/errors"
	"github.com/crawlab-team/crawlab/db/utils"
	"github.com/crawlab-team/crawlab/trace"
	"github.com/gomodule/redigo/redis"
	"reflect"
	"strings"
	"time"
)

type Client struct {
	// settings
	backoffMaxInterval time.Duration
	timeout            int

	// internals
	pool *redis.Pool
}

func (client *Client) Ping() error {
	c := client.pool.Get()
	defer utils.Close(c)
	if _, err := redis.String(c.Do("PING")); err != nil {
		if err != redis.ErrNil {
			return trace.TraceError(err)
		}
		return err
	}
	return nil
}

func (client *Client) Keys(pattern string) (values []string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err = redis.Strings(c.Do("KEYS", pattern))
	if err != nil {
		return nil, trace.TraceError(err)
	}
	return values, nil
}

func (client *Client) AllKeys() (values []string, err error) {
	return client.Keys("*")
}

func (client *Client) Get(collection string) (value string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	value, err = redis.String(c.Do("GET", collection))
	if err != nil {
		return "", trace.TraceError(err)
	}
	return value, nil
}

func (client *Client) Set(collection string, value string) (err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	value, err = redis.String(c.Do("SET", collection, value))
	if err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) Del(collection string) error {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("DEL", collection); err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) RPush(collection string, value interface{}) error {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("RPUSH", collection, value); err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) LPush(collection string, value interface{}) error {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("LPUSH", collection, value); err != nil {
		if err != redis.ErrNil {
			return trace.TraceError(err)
		}
		return err
	}
	return nil
}

func (client *Client) LPop(collection string) (string, error) {
	c := client.pool.Get()
	defer utils.Close(c)

	value, err := redis.String(c.Do("LPOP", collection))
	if err != nil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	return value, nil
}

func (client *Client) RPop(collection string) (string, error) {
	c := client.pool.Get()
	defer utils.Close(c)

	value, err := redis.String(c.Do("RPOP", collection))
	if err != nil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	return value, nil
}

func (client *Client) LLen(collection string) (int, error) {
	c := client.pool.Get()
	defer utils.Close(c)

	value, err := redis.Int(c.Do("LLEN", collection))
	if err != nil {
		return 0, trace.TraceError(err)
	}
	return value, nil
}

func (client *Client) BRPop(collection string, timeout int) (value string, err error) {
	if timeout <= 0 {
		timeout = 60
	}
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := redis.Strings(c.Do("BRPOP", collection, timeout))
	if err != nil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	return values[1], nil
}

func (client *Client) BLPop(collection string, timeout int) (value string, err error) {
	if timeout <= 0 {
		timeout = 60
	}
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := redis.Strings(c.Do("BLPOP", collection, timeout))
	if err != nil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	return values[1], nil
}

func (client *Client) HSet(collection string, key string, value string) error {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("HSET", collection, key, value); err != nil {
		if err != redis.ErrNil {
			return trace.TraceError(err)
		}
		return err
	}
	return nil
}

func (client *Client) HGet(collection string, key string) (string, error) {
	c := client.pool.Get()
	defer utils.Close(c)
	value, err := redis.String(c.Do("HGET", collection, key))
	if err != nil && err != redis.ErrNil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	return value, nil
}

func (client *Client) HDel(collection string, key string) error {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("HDEL", collection, key); err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) HScan(collection string) (results map[string]string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	var (
		cursor int64
		items  []string
	)

	results = map[string]string{}

	for {
		values, err := redis.Values(c.Do("HSCAN", collection, cursor))
		if err != nil {
			if err != redis.ErrNil {
				return nil, trace.TraceError(err)
			}
			return nil, err
		}

		values, err = redis.Scan(values, &cursor, &items)
		if err != nil {
			if err != redis.ErrNil {
				return nil, trace.TraceError(err)
			}
			return nil, err
		}
		for i := 0; i < len(items); i += 2 {
			key := items[i]
			value := items[i+1]
			results[key] = value
		}
		if cursor == 0 {
			break
		}
	}
	return results, nil
}

func (client *Client) HKeys(collection string) (results []string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	results, err = redis.Strings(c.Do("HKEYS", collection))
	if err != nil {
		if err != redis.ErrNil {
			return results, trace.TraceError(err)
		}
		return results, err
	}
	return results, nil
}

func (client *Client) ZAdd(collection string, score float32, value interface{}) (err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	if _, err := c.Do("ZADD", collection, score, value); err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) ZCount(collection string, min string, max string) (count int, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	count, err = redis.Int(c.Do("ZCOUNT", collection, min, max))
	if err != nil {
		return 0, trace.TraceError(err)
	}
	return count, nil
}

func (client *Client) ZCountAll(collection string) (count int, err error) {
	return client.ZCount(collection, "-inf", "+inf")
}

func (client *Client) ZScan(collection string, pattern string, count int) (values []string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err = redis.Strings(c.Do("ZSCAN", collection, 0, pattern, count))
	if err != nil {
		if err != redis.ErrNil {
			return nil, trace.TraceError(err)
		}
		return nil, err
	}
	return values, nil
}

func (client *Client) ZPopMax(collection string, count int) (results []string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	results = []string{}

	values, err := redis.Strings(c.Do("ZPOPMAX", collection, count))
	if err != nil {
		if err != redis.ErrNil {
			return nil, trace.TraceError(err)
		}
		return nil, err
	}

	for i := 0; i < len(values); i += 2 {
		v := values[i]
		results = append(results, v)
	}

	return results, nil
}

func (client *Client) ZPopMin(collection string, count int) (results []string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	results = []string{}

	values, err := redis.Strings(c.Do("ZPOPMIN", collection, count))
	if err != nil {
		if err != redis.ErrNil {
			return nil, trace.TraceError(err)
		}
		return nil, err
	}

	for i := 0; i < len(values); i += 2 {
		v := values[i]
		results = append(results, v)
	}

	return results, nil
}

func (client *Client) ZPopMaxOne(collection string) (value string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := client.ZPopMax(collection, 1)
	if err != nil {
		return "", err
	}
	if values == nil || len(values) == 0 {
		return "", nil
	}
	return values[0], nil
}

func (client *Client) ZPopMinOne(collection string) (value string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := client.ZPopMin(collection, 1)
	if err != nil {
		return "", err
	}
	if values == nil || len(values) == 0 {
		return "", nil
	}
	return values[0], nil
}

func (client *Client) BZPopMax(collection string, timeout int) (value string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := redis.Strings(c.Do("BZPOPMAX", collection, timeout))
	if err != nil {
		if err != redis.ErrNil {
			return "", trace.TraceError(err)
		}
		return "", err
	}
	if len(values) < 3 {
		return "", trace.TraceError(errors.ErrorRedisInvalidType)
	}
	return values[1], nil
}

func (client *Client) BZPopMin(collection string, timeout int) (value string, err error) {
	c := client.pool.Get()
	defer utils.Close(c)

	values, err := redis.Strings(c.Do("BZPOPMIN", collection, timeout))
	if err != nil {
		if err != redis.ErrNil {
			return "", trace.TraceError(err)
		}
		return "", err
	}
	if len(values) < 3 {
		return "", trace.TraceError(errors.ErrorRedisInvalidType)
	}
	return values[1], nil
}

func (client *Client) Lock(lockKey string) (value int64, err error) {
	c := client.pool.Get()
	defer utils.Close(c)
	lockKey = client.getLockKey(lockKey)

	ts := time.Now().Unix()
	ok, err := c.Do("SET", lockKey, ts, "NX", "PX", 30000)
	if err != nil {
		if err != redis.ErrNil {
			return value, trace.TraceError(err)
		}
		return value, err
	}
	if ok == nil {
		return 0, trace.TraceError(errors.ErrorRedisLocked)
	}
	return ts, nil
}

func (client *Client) UnLock(lockKey string, value int64) {
	c := client.pool.Get()
	defer utils.Close(c)
	lockKey = client.getLockKey(lockKey)

	getValue, err := redis.Int64(c.Do("GET", lockKey))
	if err != nil {
		log.Errorf("get lockKey error: %s", err.Error())
		return
	}

	if getValue != value {
		log.Errorf("the lockKey value diff: %d, %d", value, getValue)
		return
	}

	v, err := redis.Int64(c.Do("DEL", lockKey))
	if err != nil {
		log.Errorf("unlock failed, error: %s", err.Error())
		return
	}

	if v == 0 {
		log.Errorf("unlock failed: key=%s", lockKey)
		return
	}
}

func (client *Client) MemoryStats() (stats map[string]int64, err error) {
	stats = map[string]int64{}
	c := client.pool.Get()
	defer utils.Close(c)
	values, err := redis.Values(c.Do("MEMORY", "STATS"))
	for i, v := range values {
		t := reflect.TypeOf(v)
		if t.Kind() == reflect.Slice {
			vc, _ := redis.String(v, err)
			if utils.ContainsString(MemoryStatsMetrics, vc) {
				stats[vc], _ = redis.Int64(values[i+1], err)
			}
		}
	}
	if err != nil {
		if err != redis.ErrNil {
			return stats, trace.TraceError(err)
		}
		return stats, err
	}
	return stats, nil
}

func (client *Client) SetBackoffMaxInterval(interval time.Duration) {
	client.backoffMaxInterval = interval
}

func (client *Client) SetTimeout(timeout int) {
	client.timeout = timeout
}

func (client *Client) init() (err error) {
	b := backoff.NewExponentialBackOff()
	b.MaxInterval = client.backoffMaxInterval
	if err := backoff.Retry(func() error {
		err := client.Ping()
		if err != nil {
			log.WithError(err).Warnf("waiting for redis pool active connection. will after %f seconds try again.", b.NextBackOff().Seconds())
		}
		return nil
	}, b); err != nil {
		return trace.TraceError(err)
	}
	return nil
}

func (client *Client) getLockKey(lockKey string) string {
	lockKey = strings.ReplaceAll(lockKey, ":", "-")
	return "nodes:lock:" + lockKey
}

func (client *Client) getTimeout(timeout int) (res int) {
	if timeout == 0 {
		return client.timeout
	}
	return timeout
}

var client db.RedisClient

func NewRedisClient(opts ...Option) (client *Client, err error) {
	// client
	client = &Client{
		backoffMaxInterval: 20 * time.Second,
		pool:               NewRedisPool(),
	}

	// apply options
	for _, opt := range opts {
		opt(client)
	}

	// init
	if err := client.init(); err != nil {
		return nil, err
	}

	return client, nil
}

func GetRedisClient() (c db.RedisClient, err error) {
	if client != nil {
		return client, nil
	}
	c, err = NewRedisClient()
	if err != nil {
		return nil, err
	}

	return c, nil
}