github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/mq/redis/redis_backend.go (about)

     1  package redis
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/gomodule/redigo/redis"
    10  
    11  	"github.com/johnnyeven/libtools/mq"
    12  )
    13  
    14  func NewRedisBackend(pool *redis.Pool, prefix string) *RedisBackend {
    15  	return &RedisBackend{
    16  		feedback:     mq.Prefix("feedback", prefix),
    17  		cancellation: mq.Prefix("cancellation", prefix),
    18  		pool:         pool,
    19  	}
    20  }
    21  
    22  type RedisBackend struct {
    23  	feedback     string
    24  	cancellation string
    25  	pool         *redis.Pool
    26  }
    27  
    28  var _ mq.Backend = (*RedisBackend)(nil)
    29  
    30  var ChannelNamePrefix = "mq-channel-"
    31  
    32  func (backend *RedisBackend) ListChannel() ([]string, error) {
    33  	list, err := redis.Strings(ConnRedis(backend.pool).Do("KEYS", ChannelNamePrefix+"*"))
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	channelList := make([]string, len(list))
    39  
    40  	for i, k := range list {
    41  		channelList[i] = strings.TrimPrefix(k, ChannelNamePrefix)
    42  	}
    43  
    44  	return channelList, nil
    45  }
    46  
    47  func (backend *RedisBackend) ListSubject(channel string) ([]string, error) {
    48  	return redis.Strings(ConnRedis(backend.pool).Do("SMEMBERS", ChannelNamePrefix+channel))
    49  }
    50  
    51  func (backend *RedisBackend) RegisterChannel(channel string, subjects []string) error {
    52  	channel = ChannelNamePrefix + channel
    53  
    54  	args := []interface{}{channel}
    55  	for _, subject := range subjects {
    56  		args = append(args, subject)
    57  	}
    58  
    59  	pingRedis := func() error {
    60  		conn := ConnRedis(backend.pool)
    61  
    62  		conn.Send("DEL", channel)
    63  		conn.Send("SADD", args...)
    64  		conn.Send("EXPIRE", channel, 70)
    65  
    66  		_, err := conn.Exec()
    67  		return err
    68  	}
    69  
    70  	if err := pingRedis(); err != nil {
    71  		return err
    72  	}
    73  
    74  	go func() {
    75  		for {
    76  			time.Sleep(1 * time.Minute)
    77  			pingRedis()
    78  		}
    79  	}()
    80  
    81  	return nil
    82  }
    83  
    84  func (backend *RedisBackend) IsCancelled(id string) (bool, error) {
    85  	i, err := redis.Int(ConnRedis(backend.pool).Do("HEXISTS", backend.cancellation, id))
    86  	return i == 1, err
    87  }
    88  
    89  func (backend *RedisBackend) Cancel(id string) error {
    90  	_, err := ConnRedis(backend.pool).Do("HSET", backend.cancellation, id, "1")
    91  	return err
    92  }
    93  
    94  func (backend *RedisBackend) ClearCancellation(id string) error {
    95  	_, err := ConnRedis(backend.pool).Do("HDEL", backend.cancellation, id)
    96  	return err
    97  }
    98  
    99  func (backend *RedisBackend) FeedBack(taskStatus *mq.TaskStatus) error {
   100  	data, err := json.Marshal(taskStatus)
   101  	if err != nil {
   102  		return err
   103  	}
   104  	_, err = ConnRedis(backend.pool).Do("RPUSH", backend.feedback, data)
   105  	return err
   106  }
   107  
   108  func (backend *RedisBackend) GetFeedback() (*mq.TaskStatus, error) {
   109  	ret, err := ConnRedis(backend.pool).Do("BLPOP", backend.feedback, "1")
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	if ret == nil {
   115  		return nil, fmt.Errorf("null message received from redis")
   116  	}
   117  
   118  	msgPair := ret.([]interface{})
   119  
   120  	if string(msgPair[0].([]byte)) != backend.feedback {
   121  		return nil, fmt.Errorf("not a backend message: %v", msgPair[0])
   122  	}
   123  
   124  	task := (*mq.TaskStatus)(nil)
   125  	if err := json.Unmarshal(msgPair[1].([]byte), &task); err != nil {
   126  		return nil, err
   127  	}
   128  	return task, nil
   129  }