heritage-api/redis/redis.go

package redis

import (
	"encoding/json"
	"os"
	"strconv"
	"time"

	redigo "github.com/gomodule/redigo/redis"
)

/*
env REDIS_DSN=127.0.0.1:6379
env REDIS_PWD=
env REDIS_DB=0
*/
var pool *redigo.Pool

func init() {
	dbNumber := func() int {
		db := os.Getenv("REDIS_DB")
		database, err := strconv.Atoi(db)
		if err != nil {
			return 0
		}
		if database < 0 {
			database = 0
		}
		if database > 16 {
			database = 0
		}
		return database
	}()
	pool = &redigo.Pool{
		MaxActive:   100,
		MaxIdle:     1,
		IdleTimeout: time.Second * time.Duration(60),
		Dial: func() (redigo.Conn, error) {
			return redigo.Dial("tcp", os.Getenv("REDIS_DSN"),
				redigo.DialDatabase(dbNumber),
				redigo.DialPassword(os.Getenv("REDIS_PWD")),
			)
		},
		TestOnBorrow: func(conn redigo.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := conn.Do("PING")
			return err
		},
	}
}

// GetBytes 获取一个字节数组值
func getBytes(key *string) (*[]byte, error) {
	conn := pool.Get()
	defer conn.Close()

	data, err := redigo.Bytes(conn.Do("GET", *key))
	return &data, err
}

// Get 获取一个值
func get(key string) interface{} {
	conn := pool.Get()
	defer conn.Close()

	var data []byte
	var err error
	if data, err = redigo.Bytes(conn.Do("GET", key)); err != nil {
		return nil
	}
	var reply interface{}
	if err = json.Unmarshal(data, &reply); err != nil {
		return nil
	}

	return reply
}

// 集合Set增加元素
func setAdd(key string, data ...interface{}) error {
	conn := pool.Get()
	defer conn.Close()
	var err error
	for _, d := range data {
		_, e := conn.Do("SADD", key, d)
		if e != nil {
			err = e
			break
		}
	}
	return err
}

// 集合Set删除元素
func setRem(key string, data ...interface{}) error {
	conn := pool.Get()
	defer conn.Close()

	var err error
	for _, d := range data {
		_, e := conn.Do("SREM", key, d)
		if e != nil {
			err = e
			break
		}
	}

	return err
}

// 集合Set判断是否存在成员member，结果.(int64)==1表示存在
func setIsMember(key string, member interface{}) (interface{}, error) {
	conn := pool.Get()
	defer conn.Close()

	return conn.Do("SISMEMBER", key, member)
}

// 设置一个值
func set(key *string, val interface{}, timeout time.Duration) error {
	data, err := json.Marshal(val)
	if err != nil {
		return err
	}
	return setBytes(key, &data, timeout)
}

func setBytes(key *string, data *[]byte, timeout time.Duration) error {
	conn := pool.Get()
	defer conn.Close()

	_, err := conn.Do("SETEX", *key, int64(timeout/time.Second), *data)
	return err
}

// IsExist 判断key是否存在
func isExist(key string) bool {
	conn := pool.Get()
	defer conn.Close()

	a, _ := conn.Do("EXISTS", key)
	i := a.(int64)
	return i > 0
}

// Delete 删除
func delete(key string) error {
	conn := pool.Get()
	defer conn.Close()

	if _, err := conn.Do("DEL", key); err != nil {
		return err
	}

	return nil
}

// Expire 失效时间配置
func expire(key string, t int64) error {
	conn := pool.Get()
	defer conn.Close()

	if _, err := conn.Do("expire", key, t); err != nil {
		return err
	}
	return nil
}