添加redis_lock

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/utils.iml" filepath="$PROJECT_DIR$/.idea/utils.iml" />
+    </modules>
+  </component>
+</project>

.idea/utils.iml

@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="WEB_MODULE" version="4">
+  <component name="Go" enabled="true" />
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

go.mod

@@ -0,0 +1,10 @@
+module git.sxidc.com/go-tools/utils
+
+go 1.21.3
+
+require github.com/redis/go-redis/v9 v9.4.0
+
+require (
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+)

go.sum

@@ -0,0 +1,10 @@
+github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
+github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
+github.com/redis/go-redis/v9 v9.4.0 h1:Yzoz33UZw9I/mFhx4MNrB6Fk+XHO1VukNcCa1+lwyKk=
+github.com/redis/go-redis/v9 v9.4.0/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=

redis-lock/lock.go

@@ -0,0 +1,210 @@
+package redis_lock
+
+import (
+	"context"
+	"errors"
+	"github.com/redis/go-redis/v9"
+	"time"
+)
+
+const (
+	tryLockTimeoutSec = 5
+)
+
+type Lock struct {
+	redisClient *redis.Client
+	ownerID     string
+}
+
+func NewLock(address string, password string, db int) (*Lock, error) {
+	if address == "" {
+		return nil, errors.New("redis address不能为空")
+	}
+
+	return &Lock{redisClient: redis.NewClient(&redis.Options{
+		Addr:     address,
+		Password: password,
+		DB:       db,
+	})}, nil
+}
+
+func DestroyLock(lock *Lock) error {
+	if lock == nil || lock.redisClient == nil {
+		return nil
+	}
+
+	err := lock.redisClient.Close()
+	if err != nil {
+		return err
+	}
+
+	lock.redisClient = nil
+
+	return nil
+}
+
+func (lock *Lock) Lock(ownerID string, key string, expireSec int) error {
+	if ownerID == "" {
+		return errors.New("ownerID不能为空")
+	}
+
+	if key == "" {
+		return errors.New("key不能为空")
+	}
+
+	if expireSec == 0 {
+		return errors.New("expireSec不能为0")
+	}
+
+	for {
+		locked, err := lock.tryLock(ownerID, key, expireSec)
+		if err != nil {
+			return err
+		}
+
+		if locked {
+			return nil
+		}
+
+		time.Sleep(200 * time.Millisecond)
+	}
+}
+
+func (lock *Lock) TryLock(ownerID string, key string, expireSec int) (bool, error) {
+	if ownerID == "" {
+		return false, errors.New("ownerID不能为空")
+	}
+
+	if key == "" {
+		return false, errors.New("key不能为空")
+	}
+
+	if expireSec == 0 {
+		return false, errors.New("expireSec不能为0")
+	}
+
+	return lock.tryLock(ownerID, key, expireSec)
+}
+
+func (lock *Lock) Unlock(ownerID string, key string) error {
+	if ownerID == "" {
+		return errors.New("ownerID不能为空")
+	}
+
+	if key == "" {
+		return errors.New("key不能为空")
+	}
+
+	if ownerID != lock.ownerID {
+		return nil
+	}
+
+	return lock.deleteLockTime(key)
+}
+
+func (lock *Lock) tryLock(ownerID string, key string, expireSec int) (bool, error) {
+	expire := time.Duration(expireSec) * time.Second
+	nowTime := time.Now()
+	nowNano := nowTime.UnixNano()
+	currentLockNano := nowTime.Add(expire).UnixNano()
+
+	// 尝试以当前lockTime上锁，如果返回true，则上锁成功，否则锁正被占用
+	locked, err := lock.setLockTime(key, currentLockNano, expire)
+	if err != nil {
+		return false, err
+	}
+
+	if locked {
+		lock.ownerID = ownerID
+		return true, nil
+	}
+
+	// 获取保存的lockTime
+	savedLockTimeNano, err := lock.getLockTime(key)
+	if err != nil {
+		return false, err
+	}
+
+	// 锁已经过期
+	if savedLockTimeNano != 0 && savedLockTimeNano < nowNano {
+		// 尝试以当前lockTime上锁，返回之前设置的值
+		oldLockTime, err := lock.getAndSetLockTime(key, currentLockNano)
+		if err != nil {
+			return false, err
+		}
+
+		// 排除上一步操作存在的并发问题
+		if oldLockTime != 0 && oldLockTime == savedLockTimeNano {
+			lock.ownerID = ownerID
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func (lock *Lock) setLockTime(key string, lockTime int64, expireSec time.Duration) (bool, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
+	defer cancel()
+
+	cmd := lock.redisClient.SetNX(ctx, key, lockTime, expireSec)
+	if cmd.Err() != nil {
+		return false, cmd.Err()
+	}
+
+	return cmd.Val(), nil
+}
+
+func (lock *Lock) getLockTime(key string) (int64, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
+	defer cancel()
+
+	cmd := lock.redisClient.Get(ctx, key)
+	if cmd.Err() != nil {
+		if cmd.Err().Error() == redis.Nil.Error() {
+			return 0, nil
+		}
+
+		return 0, cmd.Err()
+	}
+
+	lockTime, err := cmd.Int64()
+	if err != nil {
+		return 0, err
+	}
+
+	return lockTime, nil
+}
+
+func (lock *Lock) getAndSetLockTime(key string, newLockTime int64) (int64, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
+	defer cancel()
+
+	cmd := lock.redisClient.GetSet(ctx, key, newLockTime)
+	if cmd.Err() != nil {
+		if cmd.Err().Error() == redis.Nil.Error() {
+			return 0, nil
+		}
+
+		return 0, cmd.Err()
+	}
+
+	lockTime, err := cmd.Int64()
+	if err != nil {
+		return 0, err
+	}
+
+	return lockTime, nil
+}
+
+func (lock *Lock) deleteLockTime(key string) error {
+	ctx, cancel := context.WithTimeout(context.Background(), tryLockTimeoutSec*time.Second)
+	defer cancel()
+
+	cmd := lock.redisClient.Del(ctx, key)
+	if cmd.Err() != nil {
+		return cmd.Err()
+	}
+
+	return nil
+}

redis-lock/lock_test.go

@@ -0,0 +1,118 @@
+package redis_lock
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestLock(t *testing.T) {
+	lock, err := NewLock("localhost:30379", "mtyzxhc", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func(lock *Lock) {
+		err := DestroyLock(lock)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}(lock)
+
+	wg := sync.WaitGroup{}
+	wg.Add(5)
+
+	for i := 0; i < 5; i++ {
+		go func(index int) {
+			err := lock.Lock("test", "test", 60)
+			if err != nil {
+				panic(err)
+			}
+
+			fmt.Println("Lock: ", index, "Time: ", time.Now())
+
+			time.Sleep(1 * time.Second)
+
+			err = lock.Unlock("test", "test")
+			if err != nil {
+				panic(err)
+			}
+
+			fmt.Println("Unlock: ", index, "Time: ", time.Now())
+
+			wg.Done()
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestTryLock(t *testing.T) {
+	lock, err := NewLock("localhost:30379", "mtyzxhc", 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func(lock *Lock) {
+		err := DestroyLock(lock)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}(lock)
+
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+
+	go func() {
+		err := lock.Lock("test", "test", 60)
+		if err != nil {
+			panic(err)
+		}
+
+		fmt.Println("Lock: ", 1, "Time: ", time.Now())
+
+		time.Sleep(6 * time.Second)
+
+		err = lock.Unlock("test", "test")
+		if err != nil {
+			panic(err)
+		}
+
+		fmt.Println("Unlock: ", 1, "Time: ", time.Now())
+
+		wg.Done()
+	}()
+
+	time.Sleep(1 * time.Second)
+
+	go func() {
+		for {
+			locked, err := lock.TryLock("test", "test", 60)
+			if err != nil {
+				panic(err)
+			}
+
+			if !locked {
+				fmt.Println("TryLock Failed: ", 2, "Time: ", time.Now())
+				time.Sleep(1 * time.Second)
+				continue
+			}
+
+			fmt.Println("Lock: ", 2, "Time: ", time.Now())
+
+			err = lock.Unlock("test", "test")
+			if err != nil {
+				panic(err)
+			}
+
+			fmt.Println("Unlock: ", 2, "Time: ", time.Now())
+
+			break
+		}
+
+		wg.Done()
+	}()
+
+	wg.Wait()
+}