github.com/binbinly/pkg@v0.0.11-0.20240321014439-f4fbf666eb0f/storage/redis/idalloc.go (about)

     1  package redis
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"strings"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/redis/go-redis/v9"
    10  )
    11  
    12  // IDAlloc id生成器
    13  // key 为业务key, 由业务前缀+功能前缀+具体场景id组成
    14  // 比如生成用户id, 可以传入user_id， 完整示例: eagle:idalloc:user_id
    15  type IDAlloc struct {
    16  	// redis 实例，最好使用和业务独立的实例，最好可以部署集群，让 id alloc做到高可用
    17  	client *redis.Client
    18  }
    19  
    20  // NewIDAlloc create a id alloc instance
    21  func NewIDAlloc(conn *redis.Client) *IDAlloc {
    22  	return &IDAlloc{
    23  		client: conn,
    24  	}
    25  }
    26  
    27  // GetNewID 生成id
    28  func (ia *IDAlloc) GetNewID(key string, step int64) (int64, error) {
    29  	key = ia.GetKey(key)
    30  	id, err := ia.client.IncrBy(context.Background(), key, step).Result()
    31  	if err != nil {
    32  		return 0, errors.Wrapf(err, "redis incr err, key: %s", key)
    33  	}
    34  
    35  	if id == 0 {
    36  		return 0, errors.Wrapf(err, "[redis.idalloc] %s GetNewID failed", key)
    37  	}
    38  	return id, nil
    39  }
    40  
    41  // GetCurrentID 获取当前id
    42  func (ia *IDAlloc) GetCurrentID(key string) (int64, error) {
    43  	key = ia.GetKey(key)
    44  	ret, err := ia.client.Get(context.Background(), key).Result()
    45  	if err != nil {
    46  		return 0, errors.Wrapf(err, "redis get err, key: %s", key)
    47  	}
    48  	id, err := strconv.Atoi(ret)
    49  	if err != nil {
    50  		return 0, errors.Wrap(err, "str convert err")
    51  	}
    52  	return int64(id), nil
    53  }
    54  
    55  // GetKey 获取key
    56  func (ia *IDAlloc) GetKey(key string) string {
    57  	lockKey := "idalloc"
    58  	return strings.Join([]string{lockKey, key}, ":")
    59  }