github.com/teng231/glock@v1.1.11/distributed_queue.go (about)

     1  package glock
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"errors"
     7  	"time"
     8  
     9  	"github.com/go-redis/redis/v8"
    10  )
    11  
    12  type IDistributedQueue interface {
    13  	// Push to the left
    14  	Unshift(key string, values ...interface{}) error
    15  	// Push to the right
    16  	Push(key string, values ...interface{}) error
    17  	// release 1 item right
    18  	Pop(key string, out interface{}) error
    19  	// release 1 item right
    20  	Shift(key string, out interface{}) error
    21  	// List item in list
    22  	List(key string, start, stop int64, f func([]string) error) error
    23  	// Size get current size of list
    24  	Size(key string) (int64, error)
    25  }
    26  
    27  type DistributedQueue struct {
    28  	client   *redis.Client
    29  	prefix   string
    30  	timelock time.Duration
    31  }
    32  
    33  func ConnectDistributedQueue(client *redis.Client, prefix string, timelock time.Duration) (*DistributedQueue, error) {
    34  	if client == nil {
    35  		return nil, errors.New("not found client")
    36  	}
    37  	if timelock < 0 {
    38  		return nil, errors.New("timelock is required")
    39  	}
    40  	ctx, cancel := context.WithTimeout(context.Background(), timelock)
    41  	defer cancel()
    42  	if err := client.Ping(ctx).Err(); err != nil {
    43  		return nil, err
    44  	}
    45  	return &DistributedQueue{
    46  		client:   client,
    47  		prefix:   prefix,
    48  		timelock: timelock,
    49  	}, nil
    50  }
    51  
    52  func (q *DistributedQueue) Set(key string, values ...interface{}) error {
    53  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
    54  	defer cancel()
    55  
    56  	if err := q.client.Del(ctx, q.prefix+key).Err(); err != nil {
    57  		return err
    58  	}
    59  	if len(values) > 0 {
    60  		inputs := []string{}
    61  		for _, val := range values {
    62  			bin, _ := json.Marshal(val)
    63  			inputs = append(inputs, string(bin))
    64  		}
    65  		if err := q.client.LPush(ctx, q.prefix+key, inputs).Err(); err != nil {
    66  			return err
    67  		}
    68  	}
    69  	return nil
    70  }
    71  
    72  func (q *DistributedQueue) Push(key string, values ...interface{}) error {
    73  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
    74  	defer cancel()
    75  	inputs := []string{}
    76  	for _, val := range values {
    77  		bin, _ := json.Marshal(val)
    78  		inputs = append(inputs, string(bin))
    79  	}
    80  	if err := q.client.RPush(ctx, q.prefix+key, inputs).Err(); err != nil {
    81  		return err
    82  	}
    83  	return nil
    84  }
    85  func (q *DistributedQueue) Pop(key string, out interface{}) error {
    86  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
    87  	defer cancel()
    88  	val, err := q.client.RPop(ctx, q.prefix+key).Result()
    89  	if err != nil && err == redis.Nil {
    90  		return errors.New(Empty)
    91  	}
    92  	if err != nil {
    93  		return err
    94  	}
    95  	if err := json.Unmarshal([]byte(val), out); err != nil {
    96  		return errors.New(CantParse)
    97  	}
    98  	return nil
    99  }
   100  
   101  func (q *DistributedQueue) Unshift(key string, values ...interface{}) error {
   102  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
   103  	defer cancel()
   104  	inputs := []string{}
   105  	for _, val := range values {
   106  		bin, _ := json.Marshal(val)
   107  		inputs = append(inputs, string(bin))
   108  	}
   109  	if err := q.client.LPush(ctx, q.prefix+key, inputs).Err(); err != nil {
   110  		return err
   111  	}
   112  	return nil
   113  }
   114  func (q *DistributedQueue) Shift(key string, out interface{}) error {
   115  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
   116  	defer cancel()
   117  	val, err := q.client.LPop(ctx, q.prefix+key).Result()
   118  	if err != nil && err == redis.Nil {
   119  		return errors.New(Empty)
   120  	}
   121  	if err != nil {
   122  		return err
   123  	}
   124  	if err := json.Unmarshal([]byte(val), out); err != nil {
   125  		return errors.New(CantParse)
   126  	}
   127  	return nil
   128  }
   129  
   130  // List start with page 0
   131  func (q *DistributedQueue) List(key string, start, stop int64, f func([]string) error) error {
   132  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
   133  	defer cancel()
   134  	vals, err := q.client.LRange(ctx, q.prefix+key, start, stop).Result()
   135  	if err != nil && err == redis.Nil {
   136  		return errors.New(Empty)
   137  	}
   138  	if err != nil {
   139  		return err
   140  	}
   141  	return f(vals)
   142  }
   143  
   144  func (q *DistributedQueue) Size(key string) (int64, error) {
   145  	ctx, cancel := context.WithTimeout(context.Background(), q.timelock)
   146  	defer cancel()
   147  	c, err := q.client.LLen(ctx, q.prefix+key).Result()
   148  	if err != nil {
   149  		return 0, err
   150  	}
   151  	return c, nil
   152  }