87.Go Redis实现可重入、自动续期分布式锁

在 86.分布式锁理论分析 中我们介绍了分布式锁的原理、"坑"以及解决办法。本文就给一下代码示例：

一、Redis实现分布式锁

package main

import (
  "fmt"
  "github.com/go-redis/redis"
  "time"
)

var client = redis.NewClient(&redis.Options{
  Addr:     "localhost:6379",
  Password: "", 
  DB:       0,  
})

func Lock(key string, value string, expiration time.Duration) (bool, error) {
  return client.SetNX(key, value, expiration).Result()
}

func Unlock(key string, value string) error {
  val, err := client.Get(key).Result()
  if err != nil {
    return err
  }
  // 避免释放它人的锁
  if val == value {
    return client.Del(key).Err()
  }
  return fmt.Errorf("key %s is locked by another value %s", key, val)
}

func main() {
  // my-value应该保证全局唯一，如uuid,雪花算法产生的ID等
  ok, err := Lock("my-key", "my-value", 10*time.Second)
  if err != nil {
    // handle error
  }
  if !ok {
    fmt.Println("Lock failed")
  } else {
    fmt.Println("Lock success")
    defer Unlock("my-key", "my-value")
    // do your job
   
  }
}

在上面的代码中，我们定义了两个函数：Lock和Unlock。Lock函数会尝试在Redis中设置一个键值对，并设置一个过期时间。如果这个键值对已经被其他地方设置了，那么SetNX函数会返回false，否则返回true。

Unlock函数则尝试删除这个键值对，但是在删除之前，会检查这个键值对的值是否符合输入的值，如果不符合，那么认为这个锁已经被其他地方获取，这时就不应该删除这个键值对。

二、可重入与自动续期

上面例子，锁缺失两个重要的性质：一个是可重入，一个是如何实现到期自动续期逻辑。

package main

import (
  "sync"
  "time"

  "github.com/go-redis/redis"
  "github.com/satori/go.uuid"
)

var client = redis.NewClient(&redis.Options{
  Addr:     "localhost:6379",
  Password: "", // no password set
  DB:       0,  // use default DB
})

type Lock struct {
  key        string
  value      string
  expiration time.Duration
  mu         sync.Mutex
  isLocked   bool
  count      int
}

func NewLock(key string, expiration time.Duration) *Lock {
  return &Lock{
    key:        key,
    value:      uuid.NewV4().String(), // 这里使用uuid作为value
    expiration: expiration,
  }
}

func (l *Lock) Lock() (bool, error) {
  l.mu.Lock()
  defer l.mu.Unlock()

  if l.isLocked {
    l.count++
    return true, nil
  }

  ok, err := client.SetNX(l.key, l.value, l.expiration).Result()
  if err != nil || !ok {
    return false, err
  }

  l.isLocked = true
  l.count++

  go l.renew()
  return true, nil
}

func (l *Lock) Unlock() error {
  l.mu.Lock()
  defer l.mu.Unlock()

  if !l.isLocked {
    return nil
  }

  l.count--

  if l.count > 0 {
    return nil
  }

  val, err := client.Get(l.key).Result()
  if err != nil {
    return err
  }
  // 保证释放的是自己的锁
  if val != l.value {
    return nil
  }

  l.isLocked = false
  return client.Del(l.key).Err()
}

func (l *Lock) renew() {
  ticker := time.NewTicker(l.expiration / 2)
  for range ticker.C {
    l.mu.Lock()

    if !l.isLocked {
      ticker.Stop()
      l.mu.Unlock()
      break
    }

    client.Expire(l.key, l.expiration)
    l.mu.Unlock()
  }
}

func main() {
  lock := NewLock("my-key", 10*time.Second)
  locked, err := lock.Lock()
  if err != nil {
    panic(err)
  }
  if !locked {
    return
  }
  defer lock.Unlock()
  // do something
}

这段代码中，Lock结构体中新加入了mu、isLocked和count字段，分别表示互斥锁、是否已经锁定还有重入次数。当再次获取锁的时候，已经锁定则重入次数增加，否则尝试获取锁。在unlock时，如果重入次数大于零，则直接减少重入次数而不释放锁。

同时加入了renew函数，这个函数会每过一段时间检查这个锁是否已经被释放，未被释放则续期，并在锁释放后停止续期。