modelRT/distributedlock/rwlock_test.go

package distributedlock

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/stretchr/testify/assert"
	"go.uber.org/zap"
)

var log *zap.Logger

func init() {
	log = zap.Must(zap.NewDevelopment())
}

func TestRWLockRLockAndUnRLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker.logger = log

	duration := 10 * time.Second
	// 第一次加读锁
	err := rwLocker.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey := strings.Join([]string{rwLocker.rwTokenTimeoutPrefix, rwLocker.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	err = rwLocker.UnRLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, redis.Nil, err)
	assert.Equal(t, 0, num)
	t.Log("rwLock rlock and unrlock test success")
	return
}

func TestRWLockReentrantRLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker.logger = log

	duration := 10 * time.Second
	// 第一次加读锁
	err := rwLocker.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey := strings.Join([]string{rwLocker.rwTokenTimeoutPrefix, rwLocker.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// 第二次加读锁
	err = rwLocker.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 2, num)

	// 第一次解读锁
	err = rwLocker.UnRLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// 第二次解读锁
	err = rwLocker.UnRLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, redis.Nil, err)
	assert.Equal(t, 0, num)
	t.Log("rwLock reentrant lock test success")
	return
}

func TestRWLockRefreshRLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 10,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker.logger = log

	duration := 10 * time.Second
	// 第一次加读锁
	err := rwLocker.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey := strings.Join([]string{rwLocker.rwTokenTimeoutPrefix, rwLocker.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	time.Sleep(10 * time.Second)
	script := `return redis.call('httl', KEYS[1], 'fields', '1', ARGV[1]);`
	result, err := rdb.Eval(ctx, script, []string{rwLocker.key}, tokenKey).Result()
	assert.Equal(t, nil, err)
	ttls, ok := result.([]interface{})
	assert.Equal(t, true, ok)
	ttl, ok := ttls[0].(int64)
	assert.Equal(t, true, ok)
	compareValue := int64(8)
	assert.Greater(t, ttl, compareValue)

	err = rwLocker.UnRLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, redis.Nil, err)
	assert.Equal(t, 0, num)
	t.Log("rwLock refresh lock test success")
	return
}

// TODO 设计两个客户端分别加读锁,测试是否可以加锁成功
func TestRWLock2ClientRLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker1 := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker1.logger = log

	rwLocker2 := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc5577",
	})
	rwLocker2.logger = log

	duration := 10 * time.Second
	// locker1加读锁
	err := rwLocker1.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey1 := strings.Join([]string{rwLocker1.rwTokenTimeoutPrefix, rwLocker1.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker1.key, tokenKey1).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// locker2加读锁
	err = rwLocker2.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey2 := strings.Join([]string{rwLocker2.rwTokenTimeoutPrefix, rwLocker2.token}, ":")
	num, err = rdb.HGet(ctx, rwLocker2.key, tokenKey2).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	err = rdb.HLen(ctx, rwLocker1.key).Err()
	assert.Equal(t, nil, err)
	hLen := rdb.HLen(ctx, rwLocker1.key).Val()
	assert.Equal(t, 3, hLen)

	// locker1解读锁
	err = rwLocker1.UnRLock(ctx)
	assert.Equal(t, nil, err)

	// locker1解读锁
	err = rwLocker2.UnRLock(ctx)
	assert.Equal(t, nil, err)

	err = rdb.Exists(ctx, rwLocker1.key).Err()
	assert.Equal(t, redis.Nil, err)
	existNum := rdb.Exists(ctx, rwLocker1.key).Val()
	assert.Equal(t, 0, existNum)
	t.Log("rwLock 2 client lock test success")
	return
}

// TODO 设计两个客户端分别加时间不同的读锁,测试ttl时间在有一个key删除后是否可以变换成功
func TestRWLock2CWith2DifTimeRLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker1 := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker1.logger = log

	rwLocker2 := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 30,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc5577",
	})
	rwLocker2.logger = log

	duration := 10 * time.Second
	// locker1加读锁
	err := rwLocker1.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey1 := strings.Join([]string{rwLocker1.rwTokenTimeoutPrefix, rwLocker1.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker1.key, tokenKey1).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// locker2加读锁
	err = rwLocker2.RLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey2 := strings.Join([]string{rwLocker2.rwTokenTimeoutPrefix, rwLocker2.token}, ":")
	num, err = rdb.HGet(ctx, rwLocker2.key, tokenKey2).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	err = rdb.HLen(ctx, rwLocker1.key).Err()
	assert.Equal(t, nil, err)
	hLen := rdb.HLen(ctx, rwLocker1.key).Val()
	assert.Equal(t, 3, hLen)

	// locker1解读锁
	err = rwLocker1.UnRLock(ctx)
	assert.Equal(t, nil, err)

	// locker1解读锁
	err = rwLocker2.UnRLock(ctx)
	assert.Equal(t, nil, err)

	err = rdb.Exists(ctx, rwLocker1.key).Err()
	assert.Equal(t, redis.Nil, err)
	existNum := rdb.Exists(ctx, rwLocker1.key).Val()
	assert.Equal(t, 0, existNum)
	t.Log("rwLock 2 client lock test success")
	return
}

func TestRWLockWLockAndUnWLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker.logger = log

	duration := 10 * time.Second
	// 第一次加读锁
	err := rwLocker.WLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey := strings.Join([]string{rwLocker.rwTokenTimeoutPrefix, rwLocker.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	err = rwLocker.UnWLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, redis.Nil, err)
	assert.Equal(t, 0, num)
	t.Log("rwLock rlock and unrlock test success")
	return
}

func TestRWLockReentrantWLock(t *testing.T) {
	ctx := context.TODO()
	rdb := redis.NewClient(&redis.Options{
		Network:     "tcp",
		Addr:        "192.168.2.104:6379",
		Password:    "cnstar",
		PoolSize:    50,
		DialTimeout: 10 * time.Second,
	})

	rwLocker := GetRWLocker(rdb, &RedissionLockConfig{
		LockLeaseTime: 120,
		NeedRefresh:   true,
		Key:           "component",
		Token:         "fd348a84-e07c-4a61-8c19-f753e6bc556a",
	})
	rwLocker.logger = log

	duration := 10 * time.Second
	// 第一次加写锁
	err := rwLocker.WLock(ctx, duration)
	assert.Equal(t, nil, err)

	tokenKey := strings.Join([]string{rwLocker.rwTokenTimeoutPrefix, rwLocker.token}, ":")
	num, err := rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// 第二次加写锁
	err = rwLocker.WLock(ctx, duration)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 2, num)

	// 第一次解写锁
	err = rwLocker.UnWLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, nil, err)
	assert.Equal(t, 1, num)

	// 第二次解写锁
	err = rwLocker.UnWLock(ctx)
	assert.Equal(t, nil, err)

	num, err = rdb.HGet(ctx, rwLocker.key, tokenKey).Int()
	assert.Equal(t, redis.Nil, err)
	assert.Equal(t, 0, num)
	t.Log("rwLock reentrant lock test success")
	return
}