Golang基于Redis实现分布式读写锁

使用Redis实现的读写锁，在使用时，只有2个返回值:
0,nil 可正常往下执行了
1,nil 超出了锁定最大重试次数

package redistool

import (
	"github.com/fwhezfwhez/errorx"
	"github.com/garyburd/redigo/redis"
	"time"
)

type RedisRWMutex struct {
	key           string
	maxlocksecond int           // 锁定状态标记的最大时间
	maxretryTimes int           // 最大阻塞重试次数
	retryInteval  time.Duration // 重试的间隔
}

func NewRedisRWMutex(key string, maxlocksecond int, maxretrytimes int, retryInteval time.Duration) *RedisRWMutex {
	return &RedisRWMutex{
		key:           key,
		maxlocksecond: maxlocksecond,
		maxretryTimes: maxretrytimes,
		retryInteval:  retryInteval,
	}
}

func (m *RedisRWMutex) RLock(conn redis.Conn) (int, error) {
	maxRetry := m.maxretryTimes

	rs, e := m.rlockloop(&maxRetry, conn)

	if e != nil {
		return 0, errorx.Wrap(e)
	}

	return rs, nil
}

func (m *RedisRWMutex) rlockloop(retryTimes *int, conn redis.Conn) (int, error) {
	// 写锁定时, 锁状态置为2, 阻塞其他读写
	// 读锁定时,锁状态置为1,不做任何阻塞
	// 无锁时，锁状态为0，或者不存在该key

	// 返回3表示需要阻塞
	// 返回2表示可执行
	var script = `
        local stat = redis.call('GET',KEYS[1]);
        
        -- 不存在,无锁时,返回可执行，并标记为读锁中
        if not stat then
            redis.call('SETEX', KEYS[1],ARGV[1],1)
            return 2;
        end
        
        -- 存在，但是出于无锁状态,返回可执行，标记为读锁中
        if tonumber(stat) == 0 then
            redis.call('SETEX', KEYS[1],ARGV[1],1)
            return 2;
        end

        -- 写锁定时，返回阻塞
         if tonumber(stat) == 2 then
            return 3;
         end

         -- 读锁定时，返回放行
         if tonumber(stat) == 1 then
            return 2;
         end
 
         -- 预期之外的结果
         return 4;
`
	vint, e := redis.Int(conn.Do("eval", script, 1, m.key, m.maxlocksecond))
	if e != nil {
		return 0, errorx.Wrap(e)
	}

	// 可执行
	if vint == 2 {
		return 0, nil
	}

	if vint == 3 {

		*retryTimes --
		if *retryTimes == 0 {
			return 1, nil
		}

		time.Sleep(m.retryInteval)
		return m.rlockloop(retryTimes, conn)
	}

	return 0, errorx.NewFromStringf("unexpected lock stat return %d", vint)
}

func (m *RedisRWMutex) RUnLock(
	conn redis.Conn) {

	conn.Do("del", m.key)

}

// 0 可执行
// 1 超出了最大重试次数了
func (m RedisRWMutex) Lock(conn redis.Conn) (int, error) {
	var max = m.maxretryTimes

	rs, e := m.loopLock(&max, conn)
	if e != nil {
		return 0, errorx.Wrap(e)
	}

	return rs, nil
}

func (m *RedisRWMutex) loopLock(
	retryTimes *int,
	conn redis.Conn) (int, error) {

	// 写锁定时, 锁状态置为2, 阻塞其他读写
	// 读锁定时,锁状态置为1,不做任何阻塞
	// 无锁时，锁状态为0，或者不存在该key

	// 返回3表示需要阻塞
	// 返回2表示可执行
	var script = `
        local stat = redis.call('GET',KEYS[1]);
        
        -- 无锁时,返回可执行，并标记为写锁中
        if not stat then
            redis.call('SETEX', KEYS[1],ARGV[1],2)
            return 2;
        end
        
        -- 无锁,返回可执行，标记为写锁中
        if math.abs(tonumber(stat)) < 0.1 then
            redis.call('SETEX', KEYS[1],ARGV[1],2)
            return 2;
        end

        -- 写锁定时，返回阻塞
         if math.abs(tonumber(stat)-2) < 0.1 then
            return 3;
         end

         -- 读锁定时,返回阻塞
          if math.abs(tonumber(stat)-1) < 0.1 then
            return 3;
         end

         -- 预期之外的结果
         return 4;
`
	vint, e := redis.Int(conn.Do("eval", script, 1, m.key, m.maxlocksecond))
	if e != nil {
		return 0, errorx.Wrap(e)
	}

	// 可执行
	if vint == 2 {
		return 0, nil
	}

	if vint == 3 {

		*retryTimes --
		if *retryTimes == 0 {
			return 1, nil
		}
		time.Sleep(m.retryInteval)

		return m.loopLock(retryTimes, conn)
	}

	return 0, errorx.NewFromStringf("unexpected lock stat return %d", vint)
}

func (m *RedisRWMutex) UnLock(conn redis.Conn) {
	conn.Do("del", m.key)
}

测试用例

执行测试用例时，注意并发度太高超过了连接池连接数设定。

func TestRWMutextRLock(t *testing.T) {
	var num = 26

	wg := sync.WaitGroup{}

	wg.Add(num)

	var a = 0
	var l = NewRedisRWMutex("locka", 15, 10, 10*time.Millisecond, )

	for i := 0; i < num; i++ {
		go func(i int ) {
			defer wg.Done()

			conn := RedisPool.Get()
			defer conn.Close()

			switch i % 2 {
			case 0:
				l.RLock(conn)
				defer l.RUnLock(conn)

				fmt.Println(a)
			case 1:
				l.Lock(conn)
				defer l.UnLock(conn)

				a++
				// fmt.Println(a)
			}
		}(i)
	}

	wg.Wait()

	fmt.Println("最后结果:", a)
}