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)
}