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

     1  package glock
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"log"
     8  	"time"
     9  
    10  	"github.com/go-redis/redis/v8"
    11  )
    12  
    13  type ICountLock interface {
    14  	// Start is config start
    15  	Start(key interface{}, startWith int, expired time.Duration) error
    16  	// DecrBy decrease count
    17  	DecrBy(key interface{}, count int64) (int, error)
    18  	// Current get current count
    19  	Current(key interface{}) (int, error)
    20  	// Existed check key existed
    21  	Existed(key interface{}) (bool, error)
    22  	// IncrBy increase count
    23  	IncrBy(key interface{}, count int64) (int, error)
    24  	// IncrBy stop counter
    25  	StopCounter(key interface{}) error
    26  }
    27  type CountLock struct {
    28  	client   *redis.Client
    29  	timelock time.Duration
    30  	prefix   string
    31  }
    32  
    33  // StartCountLock
    34  func StartCountLock(cf *ConnectConfig) (*CountLock, error) {
    35  	client := redis.NewClient(&redis.Options{
    36  		Addr:            cf.RedisAddr,
    37  		Password:        cf.RedisPw,
    38  		MaxRetries:      10,
    39  		MinRetryBackoff: 15 * time.Millisecond,
    40  		MaxRetryBackoff: 1000 * time.Millisecond,
    41  		DialTimeout:     10 * time.Second,
    42  		DB:              cf.RedisDb, // use default DB
    43  	})
    44  	if cf.Timelock < 0 {
    45  		return nil, errors.New("timelock is required")
    46  	}
    47  	ctx, cancel := context.WithTimeout(context.Background(), cf.Timelock)
    48  	defer cancel()
    49  	if err := client.Ping(ctx).Err(); err != nil {
    50  		return nil, err
    51  	}
    52  	return &CountLock{
    53  		client:   client,
    54  		prefix:   cf.Prefix,
    55  		timelock: cf.Timelock,
    56  	}, nil
    57  }
    58  
    59  // CreateCountLock deprecated
    60  func CreateCountLock(redisAddr, redisPw, prefix string, timelock time.Duration) (*CountLock, error) {
    61  	client := redis.NewClient(&redis.Options{
    62  		Addr:            redisAddr,
    63  		Password:        redisPw,
    64  		MaxRetries:      10,
    65  		MinRetryBackoff: 15 * time.Millisecond,
    66  		MaxRetryBackoff: 1000 * time.Millisecond,
    67  		DialTimeout:     10 * time.Second,
    68  		DB:              1, // use default DB
    69  	})
    70  	if timelock < 0 {
    71  		return nil, errors.New("timelock is required")
    72  	}
    73  	ctx, cancel := context.WithTimeout(context.Background(), timelock)
    74  	defer cancel()
    75  	if err := client.Ping(ctx).Err(); err != nil {
    76  		return nil, err
    77  	}
    78  	return &CountLock{
    79  		client:   client,
    80  		prefix:   prefix,
    81  		timelock: timelock,
    82  	}, nil
    83  }
    84  
    85  // Start: khởi động 1 tiến trình bộ đếm. bắt đầu bằng startWith và thời hạn hiệu lực của bộ đếm là expired
    86  // Start có thể dùng để reset counter về 1 giá trị nào đó.
    87  func (cl *CountLock) Start(key interface{}, startWith int, expired time.Duration) error {
    88  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
    89  	defer cancel()
    90  
    91  	rdKey := fmt.Sprintf("%v_%v", cl.prefix, key)
    92  	if err := cl.client.Set(ctx, rdKey, startWith, expired).Err(); err != nil {
    93  		return err
    94  	}
    95  	if expired == -1 {
    96  		if err := cl.client.Persist(ctx, rdKey).Err(); err != nil {
    97  			log.Print(err)
    98  			return errors.New(StatusNotPersist)
    99  		}
   100  	}
   101  	return nil
   102  }
   103  
   104  func (cl *CountLock) DecrBy(key interface{}, count int64) (int, error) {
   105  	if count <= 0 {
   106  		return 0, errors.New(StatusInvalidCounter)
   107  	}
   108  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
   109  	defer cancel()
   110  
   111  	rdKey := fmt.Sprintf("%s_%v", cl.prefix, key)
   112  	curVal, err := cl.client.DecrBy(ctx, rdKey, count).Result()
   113  	if err != nil {
   114  		return 0, err
   115  	}
   116  	return int(curVal), nil
   117  }
   118  
   119  func (cl *CountLock) Current(key interface{}) (int, error) {
   120  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
   121  	defer cancel()
   122  
   123  	rdKey := fmt.Sprintf("%s_%v", cl.prefix, key)
   124  
   125  	current, err := cl.client.Get(ctx, rdKey).Int()
   126  	if err != nil {
   127  		return 0, err
   128  	}
   129  	return current, nil
   130  }
   131  
   132  func (cl *CountLock) IncrBy(key interface{}, count int64) (int, error) {
   133  	if count <= 0 {
   134  		return 0, errors.New(StatusInvalidCounter)
   135  	}
   136  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
   137  	defer cancel()
   138  
   139  	rdKey := fmt.Sprintf("%s_%v", cl.prefix, key)
   140  	curVal, err := cl.client.IncrBy(ctx, rdKey, count).Result()
   141  	if err != nil {
   142  		return 0, err
   143  	}
   144  	return int(curVal), nil
   145  }
   146  
   147  func (cl *CountLock) StopCounter(key interface{}) error {
   148  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
   149  	defer cancel()
   150  
   151  	rdKey := fmt.Sprintf("%s_%v", cl.prefix, key)
   152  	err := cl.client.Del(ctx, rdKey).Err()
   153  	if err != nil {
   154  		return err
   155  	}
   156  	return err
   157  }
   158  
   159  func (cl *CountLock) Existed(key interface{}) (bool, error) {
   160  	ctx, cancel := context.WithTimeout(context.Background(), cl.timelock)
   161  	defer cancel()
   162  
   163  	rdKey := fmt.Sprintf("%s_%v", cl.prefix, key)
   164  	val, err := cl.client.Exists(ctx, rdKey).Result()
   165  	if err != nil {
   166  		return false, err
   167  	}
   168  	if val == 0 {
   169  		return false, nil
   170  	}
   171  	return true, nil
   172  }