lock.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. package redislock
  2. import (
  3. "context"
  4. "github.com/pkg/errors"
  5. "github.com/redis/go-redis/v9"
  6. "time"
  7. )
  8. const (
  9. tryLockTimeoutSec = 5
  10. )
  11. type Lock struct {
  12. redisClient *redis.Client
  13. ownerID string
  14. }
  15. func NewLock(address string, password string, db int) (*Lock, error) {
  16. if address == "" {
  17. return nil, errors.New("redis address不能为空")
  18. }
  19. return &Lock{redisClient: redis.NewClient(&redis.Options{
  20. Addr: address,
  21. Password: password,
  22. DB: db,
  23. })}, nil
  24. }
  25. func DestroyLock(lock *Lock) error {
  26. if lock == nil || lock.redisClient == nil {
  27. return nil
  28. }
  29. err := lock.redisClient.Close()
  30. if err != nil {
  31. return err
  32. }
  33. lock.redisClient = nil
  34. return nil
  35. }
  36. func (lock *Lock) Lock(ownerID string, key string, expireSec int) error {
  37. if ownerID == "" {
  38. return errors.New("ownerID不能为空")
  39. }
  40. if key == "" {
  41. return errors.New("key不能为空")
  42. }
  43. if expireSec == 0 {
  44. return errors.New("expireSec不能为0")
  45. }
  46. for {
  47. locked, err := lock.tryLock(ownerID, key, expireSec)
  48. if err != nil {
  49. return err
  50. }
  51. if locked {
  52. return nil
  53. }
  54. time.Sleep(200 * time.Millisecond)
  55. }
  56. }
  57. func (lock *Lock) TryLock(ownerID string, key string, expireSec int) (bool, error) {
  58. if ownerID == "" {
  59. return false, errors.New("ownerID不能为空")
  60. }
  61. if key == "" {
  62. return false, errors.New("key不能为空")
  63. }
  64. if expireSec == 0 {
  65. return false, errors.New("expireSec不能为0")
  66. }
  67. return lock.tryLock(ownerID, key, expireSec)
  68. }
  69. func (lock *Lock) Unlock(ownerID string, key string) error {
  70. if ownerID == "" {
  71. return errors.New("ownerID不能为空")
  72. }
  73. if key == "" {
  74. return errors.New("key不能为空")
  75. }
  76. if ownerID != lock.ownerID {
  77. return nil
  78. }
  79. return lock.deleteLockTime(key)
  80. }
  81. func (lock *Lock) tryLock(ownerID string, key string, expireSec int) (bool, error) {
  82. expire := time.Duration(expireSec) * time.Second
  83. nowTime := time.Now()
  84. nowNano := nowTime.UnixNano()
  85. currentLockNano := nowTime.Add(expire).UnixNano()
  86. // 尝试以当前lockTime上锁,如果返回true,则上锁成功,否则锁正被占用
  87. locked, err := lock.setLockTime(key, currentLockNano, expire)
  88. if err != nil {
  89. return false, err
  90. }
  91. if locked {
  92. lock.ownerID = ownerID
  93. return true, nil
  94. }
  95. // 获取保存的lockTime
  96. savedLockTimeNano, err := lock.getLockTime(key)
  97. if err != nil {
  98. return false, err
  99. }
  100. // 锁已经过期
  101. if savedLockTimeNano != 0 && savedLockTimeNano < nowNano {
  102. // 尝试以当前lockTime上锁,返回之前设置的值
  103. oldLockTime, err := lock.getAndSetLockTime(key, currentLockNano)
  104. if err != nil {
  105. return false, err
  106. }
  107. // 排除上一步操作存在的并发问题
  108. if oldLockTime != 0 && oldLockTime == savedLockTimeNano {
  109. lock.ownerID = ownerID
  110. return true, nil
  111. }
  112. }
  113. return false, nil
  114. }
  115. func (lock *Lock) setLockTime(key string, lockTime int64, expireSec time.Duration) (bool, error) {
  116. ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
  117. defer cancel()
  118. cmd := lock.redisClient.SetNX(ctx, key, lockTime, expireSec)
  119. if cmd.Err() != nil {
  120. return false, cmd.Err()
  121. }
  122. return cmd.Val(), nil
  123. }
  124. func (lock *Lock) getLockTime(key string) (int64, error) {
  125. ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
  126. defer cancel()
  127. cmd := lock.redisClient.Get(ctx, key)
  128. if cmd.Err() != nil {
  129. if cmd.Err().Error() == redis.Nil.Error() {
  130. return 0, nil
  131. }
  132. return 0, cmd.Err()
  133. }
  134. lockTime, err := cmd.Int64()
  135. if err != nil {
  136. return 0, err
  137. }
  138. return lockTime, nil
  139. }
  140. func (lock *Lock) getAndSetLockTime(key string, newLockTime int64) (int64, error) {
  141. ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
  142. defer cancel()
  143. cmd := lock.redisClient.GetSet(ctx, key, newLockTime)
  144. if cmd.Err() != nil {
  145. if cmd.Err().Error() == redis.Nil.Error() {
  146. return 0, nil
  147. }
  148. return 0, cmd.Err()
  149. }
  150. lockTime, err := cmd.Int64()
  151. if err != nil {
  152. return 0, err
  153. }
  154. return lockTime, nil
  155. }
  156. func (lock *Lock) deleteLockTime(key string) error {
  157. ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
  158. defer cancel()
  159. cmd := lock.redisClient.Del(ctx, key)
  160. if cmd.Err() != nil {
  161. return cmd.Err()
  162. }
  163. return nil
  164. }