项目实训day12-13
day12，day13
- 数据库方面的接口
- redis部分功能
- 有关节点通信的pubsub
数据库方面的接口
redis部分功能
与节点有关的部分就是在redis的node集合中对节点信息进行操作，然后涉及到节点通信的pub和sub。
package database
var RedisClient *Redis
type Redis struct {
pool *redis.Pool
}
type Mutex struct {
Name string
expiry time.Duration
tries int
delay time.Duration
value string
}
func NewRedisClient() *Redis {
return &Redis{pool: NewRedisPool()}
}
func (r *Redis) RPush(collection string, value interface{}) error {//对应list尾部添加字符串元素
c := r.pool.Get()
defer utils.Close(c)
if _, err := c.Do("RPUSH", collection, value); err != nil {
log.Error(err.Error())
debug.PrintStack()
return err
}
return nil
}
func (r *Redis) LPush(collection string, value interface{}) error {//对应list头部添加字符串元素
c := r.pool.Get()
defer utils.Close(c)
if _, err := c.Do("RPUSH", collection, value); err != nil {
log.Error(err.Error())
debug.PrintStack()
return err
}
return nil
}
func (r *Redis) LPop(collection string) (string, error) {
c := r.pool.Get()
defer utils.Close(c)
value, err2 := redis.String(c.Do("LPOP", collection))
if err2 != nil {
return value, err2
}
return value, nil
}
func (r *Redis) HSet(collection string, key string, value string) error {//相当于hashmap的set操作
c := r.pool.Get()
defer utils.Close(c)
if _, err := c.Do("HSET", collection, key, value); err != nil {
log.Error(err.Error())
debug.PrintStack()
return err
}
return nil
}
func (r *Redis) HGet(collection string, key string) (string, error) {//相当于hashmap的get操作
c := r.pool.Get()
defer utils.Close(c)
value, err2 := redis.String(c.Do("HGET", collection, key))
if err2 != nil && err2 != redis.ErrNil {
log.Error(err2.Error())
debug.PrintStack()
return value, err2
}
return value, nil
}
func (r *Redis) HDel(collection string, key string) error {
c := r.pool.Get()
defer utils.Close(c)
if _, err := c.Do("HDEL", collection, key); err != nil {
log.Error(err.Error())
debug.PrintStack()
return err
}
return nil
}
func (r *Redis) HKeys(collection string) ([]string, error) {
c := r.pool.Get()
defer utils.Close(c)
value, err2 := redis.Strings(c.Do("HKEYS", collection))
if err2 != nil {
log.Error(err2.Error())
debug.PrintStack()
return []string{}, err2
}
return value, nil
}
func (r *Redis) BRPop(collection string, timeout int) (string, error) {
if timeout <= 0 {
timeout = 60
}
c := r.pool.Get()
defer utils.Close(c)
values, err := redis.Strings(c.Do("BRPOP", collection, timeout))
if err != nil {
return "", err
}
return values[1], nil
}
func NewRedisPool() *redis.Pool {
var address = viper.GetString("redis.address")
var port = viper.GetString("redis.port")
var database = viper.GetString("redis.database")
var password = viper.GetString("redis.password")
var url string
if password == "" {
url = "redis://" + address + ":" + port + "/" + database
} else {
url = "redis://x:" + password + "@" + address + ":" + port + "/" + database
}
return &redis.Pool{
Dial: func() (conn redis.Conn, e error) {
return redis.DialURL(url,
redis.DialConnectTimeout(time.Second*10),
redis.DialReadTimeout(time.Second*600),
redis.DialWriteTimeout(time.Second*10),
)
},
TestOnBorrow: func(c redis.Conn, t time.Time) error {
if time.Since(t) < time.Minute {
return nil
}
_, err := c.Do("PING")
return err
},
MaxIdle: 10,
MaxActive: 0,
IdleTimeout: 300 * time.Second,
Wait: false,
MaxConnLifetime: 0,
}
}
func InitRedis() error {
RedisClient = NewRedisClient()
return nil
}
func Pub(channel string, msg entity.NodeMessage) error {
if _, err := RedisClient.Publish(channel, utils.GetJson(msg)); err != nil {
log.Errorf("publish redis error: %s", err.Error())
debug.PrintStack()
return err
}
return nil
}
func Sub(channel string, consume ConsumeFunc) error {
ctx := context.Background()
if err := RedisClient.Subscribe(ctx, consume, channel); err != nil {
log.Errorf("subscribe redis error: %s", err.Error())
debug.PrintStack()
return err
}
return nil
}
// 构建同步锁key
func (r *Redis) getLockKey(lockKey string) string {
lockKey = strings.ReplaceAll(lockKey, ":", "-")
return "nodes:lock:" + lockKey
}
// 获得锁
func (r *Redis) Lock(lockKey string) (int64, error) {
c := r.pool.Get()
defer utils.Close(c)
lockKey = r.getLockKey(lockKey)
ts := time.Now().Unix()
ok, err := c.Do("SET", lockKey, ts, "NX", "PX", 30000)
if err != nil {
log.Errorf("get lock fail with error: %s", err.Error())
debug.PrintStack()
return 0, err
}
if ok == nil {
log.Errorf("the lockKey is locked: key=%s", lockKey)
return 0, errors.New("the lockKey is locked")
}
return ts, nil
}
func (r *Redis) UnLock(lockKey string, value int64) {
c := r.pool.Get()
defer utils.Close(c)
lockKey = r.getLockKey(lockKey)
getValue, err := redis.Int64(c.Do("GET", lockKey))
if err != nil {
log.Errorf("get lockKey error: %s", err.Error())
debug.PrintStack()
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())
debug.PrintStack()
return
}
if v == 0 {
log.Errorf("unlock failed: key=%s", lockKey)
return
}
}
有关节点通信的pubsub
package database
type ConsumeFunc func(message redis.Message) error
func (r *Redis) Close() {
err := r.pool.Close()
if err != nil {
log.Errorf("redis close error.")
}
}
func (r *Redis) subscribe(ctx context.Context, consume ConsumeFunc, channel ...string) error {
psc := redis.PubSubConn{Conn: r.pool.Get()}
if err := psc.Subscribe(redis.Args{}.AddFlat(channel)...); err != nil {
return err
}
done := make(chan error, 1)
tick := time.NewTicker(time.Second * 3)
defer tick.Stop()
go func() {
defer utils.Close(psc)
for {
switch msg := psc.Receive().(type) {
case error:
done <- fmt.Errorf("redis pubsub receive err: %v", msg)
return
case redis.Message:
if err := consume(msg); err != nil {
fmt.Printf("redis pubsub consume message err: %v", err)
continue
}
case redis.Subscription:
fmt.Println(msg)
if msg.Count == 0 {
// all channels are unsubscribed
return
}
}
}
}()
// start a new goroutine to receive message
for {
select {
case <-ctx.Done():
if err := psc.Unsubscribe(); err != nil {
fmt.Printf("redis pubsub unsubscribe err: %v \n", err)
}
done <- nil
case <-tick.C:
if err := psc.Ping(""); err != nil {
fmt.Printf("ping message error: %s \n", err)
//done <- err
}
case err := <-done:
close(done)
return err
}
}
}
func (r *Redis) Subscribe(ctx context.Context, consume ConsumeFunc, channel ...string) error {
index := 0
go func() {
for {
err := r.subscribe(ctx, consume, channel...)
fmt.Println(err)
if err == nil {
break
}
time.Sleep(5 * time.Second)
index += 1
fmt.Printf("try reconnect %d times \n", index)
}
}()
return nil
}
func (r *Redis) Publish(channel, message string) (n int, err error) {
conn := r.pool.Get()
defer utils.Close(conn)
n, err = redis.Int(conn.Do("PUBLISH", channel, message))
if err != nil {
return 0, errors2.Wrapf(err, "redis publish %s %s", channel, message)
}
return
}