首頁 > 軟體

Golang實現可重入鎖的範例程式碼

2022-05-25 14:01:25

專案中遇到了可重入鎖的需求和實現,具體記錄下。

什麼是可重入鎖

我們平時說的分散式鎖,一般指的是在不同伺服器上的多個執行緒中,只有一個執行緒能搶到一個鎖,從而執行一個任務。而我們使用鎖就是保證一個任務只能由一個執行緒來完成。所以我們一般是使用這樣的三段式邏輯:

Lock();
DoJob();
Unlock();

但是由於我們的系統都是分散式的,這個鎖一般不會只放在某個程序中,我們會借用第三方儲存,比如 Redis 來做這種分散式鎖。但是一旦藉助了第三方儲存,我們就必須面對這個問題:Unlock是否能保證一定執行呢?

這個問題,我們面對的除了程式的bug之外,還有網路的不穩定,程序被殺死,伺服器被down機等。我們是無法保證Unlock一定被執行的。

那麼我們就一般在Lock的時候為這個鎖加一個超時時間作為兜底。

LockByExpire(duration);
DoJob();
Unlock();

這個超時時間是為了一旦出現異常情況導致Unlock沒有被執行,這個鎖在duration時間內也會被自動釋放。這個在redis中我們一般就是使用set ex 來進行鎖超時的設定。

但是有這個超時時間我們又遇上了問題,超時時間設定多久合適呢?當然要設定的比 DoJob 消耗的時間更長,否則的話,在任務還沒結束的時候,鎖就被釋放了,還是有可能導致並行任務的存在。

但是實際上,同樣由於網路超時問題,系統執行狀況問題等,我們是無法準確知道DoJob這個函數要執行多久的。那麼這時候怎麼辦呢?

有兩個辦法:

第一個方法,我們可以對DoJob做一個超時設定。讓DoJob最多隻能執行n秒,那麼我的分散式鎖的超時時長設定比n秒長就可以了。為一個任務設定超時時間在很多語言是可以做到的。比如golang 中的 TimeoutContext。

而第二種方法,就是我們先為鎖設定一個比較小的超時時長,然後不斷續期這個鎖。對一個鎖的不斷需求,也可以理解為重新開始加鎖,這種可以不斷續期的鎖,就叫做可重入鎖。

除了主執行緒之外,可重入鎖必然有一個另外的執行緒(或者攜程)可以對這個鎖進行續期,我們叫這個額外的程式叫做watchDog(看門狗)。

具體實現

在Golang中,語言級別天生支援協程,所以這種可重入鎖就非常容易實現:

// DistributeLockRedis 基於redis的分散式可重入鎖,自動續租
type DistributeLockRedis struct {
	key       string             // 鎖的key
	expire    int64              // 鎖超時時間
	status    bool               // 上鎖成功標識
	cancelFun context.CancelFunc // 用於取消自動續租攜程
	redis     redis.Client       // redis控制程式碼
}

// 建立可
func NewDistributeLockRedis(key string, expire int64) *DistributeLockRedis {
	return &DistributeLockRedis{
		 key : key,
		 expire : expire,
	}
}

// TryLock 上鎖
func (dl *DistributeLockRedis) TryLock() (err error) {
	if err = dl.lock(); err != nil {
		return err
	}
	ctx, cancelFun := context.WithCancel(context.Background())
	dl.cancelFun = cancelFun
	dl.startWatchDog(ctx) // 建立守護協程,自動對鎖進行續期
	dl.status = true
	return nil
}

// competition 競爭鎖
func (dl *DistributeLockRedis) lock() error {
	if res, err := redis.String(dl.redis.Do(context.Background(), "SET", dl.key, 1, "NX", "EX", dl.expire)); err != nil {
		return err
	} 
	return nil
}


// guard 建立守護協程,自動續期
func (dl *DistributeLockRedis) startWatchDog(ctx context.Context) {
	safeGo(func() error {
		for {
			select {
			// Unlock通知結束
			case <-ctx.Done():
				return nil
			default:
				// 否則只要開始了,就自動重入(續租鎖)
				if dl.status {
					if res, err := redis.Int(dl.redis.Do(context.Background(), "EXPIRE", dl.key, dl.expire)); err != nil {
						return nil
					} 
					// 續租時間為 expire/2 秒
					time.Sleep(time.Duration(dl.expire/2) * time.Second)
				}
			}
		}
	})
}

// Unlock 釋放鎖
func (dl *DistributeLockRedis) Unlock() (err error) {
	// 這個重入鎖必須取消,放在第一個地方執行
	if dl.cancelFun != nil {
		dl.cancelFun() // 釋放成功,取消重入鎖
	}
	var res int
	if dl.status {
		if res, err = redis.Int(dl.redis.Do(context.Background(), "Del", dl.key)); err != nil {
			return fmt.Errorf("釋放鎖失敗")
		}
		if res == 1 {
			dl.status = false
			return nil
		}
	}
	return fmt.Errorf("釋放鎖失敗")
}

這段程式碼的邏輯基本上都以註釋的形式來寫了。其中主要就在startWatchDog,對鎖進行重新續期

ctx, cancelFun := context.WithCancel(context.Background())
dl.cancelFun = cancelFun
dl.startWatchDog(ctx) // 建立守護協程,自動對鎖進行續期
dl.status = true

首先建立一個cancelContext,它的context函數cancelFunc是給Unlock進行呼叫的。然後啟動一個goroutine程序來回圈續期。

這個新啟動的goroutine在主goroutine處理結束,呼叫Unlock的時候,才會結束,否則會在 過期時間/2 的時候,呼叫一次redis的expire命令來進行續期。

至於外部,在使用的時候如下

func Foo() error {
  key := foo
  
  // 建立可重入的分散式鎖
	dl := NewDistributeLockRedis(key, 10)
	// 爭搶鎖
	err := dl.TryLock()
	if err != nil {
		// 沒有搶到鎖
		return err
	}
	
	// 搶到鎖的記得釋放鎖
	defer func() {
		dl.Unlock()
	}
	
	// 做真正的任務
	DoJob()
}

以上就是Golang實現可重入鎖的範例程式碼的詳細內容,更多關於Golang可重入鎖的資料請關注it145.com其它相關文章!


IT145.com E-mail:sddin#qq.com