Zenn
🔐

[Go] Context-awareなLockをつくってみた

2023/12/28に公開
Go
tech

0. TL;DR

ロック取得において context.Context を考慮する排他ロック・共有ロックをGoで実装して公開しました。(GitHubに公開するだけで、 pkg.go.dev にも表示されるの楽でいいですね。)

https://github.com/ymd-h/go/tree/master/ctxlock
https://pkg.go.dev/github.com/ymd-h/go/ctxlock

1. はじめに

少し前からちょこちょこGoでプログラムを書く機会があり、goroutineやchannel、contextをベースとした並行プログラミングが、独特だけれども非常に強力で素晴らしいなと感じました。

初めは少し戸惑いもありましたが、Zennにもわかりやすくて良い記事・本がたくさんあり、とても助かりました。全部は挙げきれないので、ここでは一部のみ紹介します。
https://zenn.dev/nobonobo/articles/53b1539ae1dd09
https://zenn.dev/hsaki/books/golang-context

今回の記事は、個人的な習作として色々書いている中で、作ってみたものについての記事です。

2. 標準 sync.Mutex / sync.RWMutex

並行（ないし並列）プログラミングをしていると不整合を起こさないようにデータを取り扱う事が重要です。Goにおいては（おそらく）多くの場合、channelを利用することで解決できる（しそうするべき）と思いますが、それでもユースケースによっては他のプログラミング言語のようにロックを使うこともあると考えています。

Goでは標準で排他ロックのための sync.Mutex(ドキュメント)、共有ロックのための sync.RWMutex(ドキュメント) が提供されています。

https://go.dev/tour/concurrency/9

これらは、 Lock() メソッド、 Unlock() メソッドによってロックを取得したり解放したりします。Lock() はロックが取得できないと待たされますが、知る限りでは途中でキャンセルする方法は存在せず、ロックを取得した後に Unlock() するか、そもそも全く待機をしない TryLock() を代わりに使うことになります。[1]

「30秒待って駄目だったら諦めたい」「待っている間に別のエラーが発生したら取りやめたい」といったGoの並行処理で頻出の context.Context が担うようなニーズと併せたような使い方はできなさそうというのが、今回の開発のモチベーションになっています。

3. ctxlock モジュール

今回開発した ctxlock モジュールは以下のように、排他ロック用の ctxlock.Lock、共有ロック用の ctxlock.SharableLock を提供しています。

  • Lock (排他ロック)
    • func NewLock() *Lock
    • func (*Lock) Lock(context.Context) (UnlockFunc, error)
  • SharableLock (共有ロック)
    • NewSharableLock() *SharableLock
    • func (*SharableLock) ExclusiveLock(context.Context) (UnlockFunc, error)
    • func (*SharableLock) SharedLock(context.Context) (UnlockFunc, error)
  • UnlockFunc (aka. func())
    • func (UnlockFunc) UnlockOnCancel(context.Context)
排他ロック
L := ctxlock.NewLock()

// Lock() method tries to lock and returns unlock function.
unlock, err := L.Lock(context.Background())
共有ロック
L := ctxlock.NewSharableLock()

// Reader Lock
unlock, err := L.SharedLock(context.Background())

// Writer Lock
unlock, err := L.ExclusiveLock(context.Background())

標準の sync.(RW)Mutex との違いは、「ロック取得で context.Context を引数にとる事」「ロック解放はメソッドではなく、ロック取得時に構築される関数(ctxlock.UnlockFunc)によって行う事[2]」です。
ctxlock.UnlockFunc はデフォルトでは手動で呼び出す設計ですが、UnlockOnCancel(context.Context) メソッドで、context.Context のキャンセルに紐付ける事ができます。

4. 内部実装

4.1 排他ロック (ctxlock.Lock)

排他ロックは、サイズ 1 のバッファを持つ channel によって実現しています。
すごい簡略化して書くと次のようになります。

lck := make(chan struct{}, 1)

// Lock
select {
case lck <- struct{}{}:
	// バッファがいっぱいだと書けなくて待機するので、排他ロックになる
	// ...
case <- ctx.Done():
	// ...
}

// Unlock
<- lck

4.2 共有ロック (ctxlock.SharableLock)

排他ロックと比べて、共有ロックは設計にかなり苦労しました。

  1. 複数の reader が同時にロックを取得できる
  2. writer は、reader や他の writer と排他でロックを取得する

上記の仕様を満たしつつ、context.Context を考慮するために select 構文で待機をする必要があります。
最終的に、reader と writer で排他ロックを取り合い、reader は他の reader からのロック取得要望も受け付けられるように、別 goroutine で排他ロックを保持する設計としました。[3]
ポイントとしては後からどんどん新しい reader が入ってきて writer が無限待機にならないように、writer がロックを取得したいと希望したら、新しい reader はロックに入ってこないようにする事です。[4]

lck := make(chan struct{}, 1)
add := make(chan struct{})  // reader 追加
done := make(chan struct{}) // reader 完了
var want atomic.Int32 // writer 待機数

// reader 用に排他ロックを保持する補助 goroutine 関数
func readThread() {
	defer func(){ <- lck }() // 終わったら排他ロックを解放

	i := 1 // reader の数
	for i > 0 {
		if want.Load() > 0 {
			<- done
			i -= 1
		} else {
			select {
			case <- add:
				i += 1
			case <- done:
				i -=1
			}
		}
	}
}

// SharedLock
select {
case lck <- struct{}{}:
	// 排他ロックを取得できたので、新規に reader 用の補助 goroutine を建てる
	go readThread()
	// ...
case add <- struct{}{}:
	// 既に reader 用の補助 goroutine が建っている
	// ...
case <- ctx.Done():
	// ...
}

// ExclusiveLock
want.Add(1) // writer が待機している事を知らせる
defer want.Add(-1)
select {
case lck <- struct{}{}:
	// ...
case <- ctx.Done():
	// ...
}

4.3 ロック解除 (ctxlock.UnlockFunc)

ctxlock.Unlock は安全性のため、1回しか実行されない (複数回呼んでも、2回目以降何も起きない) ようにしています。

このニーズでは、標準の sync.Once (ドキュメント) や、それをラップした sync.OnceFunc(func()) func() (ドキュメント) が思い浮かびますが、採用せず独自に簡易版を実装して使っています。

こちらの記事で紹介されていますが、sync.Once は「1回しか実行されない事」以外に「呼び出し後には実行が完了している事」も保証しています。
https://zenn.dev/sryoya/articles/b0e8e8d83032b0

ctxlock におけるロック解放は、実行完了を厳密に待つ必要はなく、sync.Onceにおいて「間違った実装」として紹介されている簡易版に差し替える事でパフォーマンス向上が確認できたため、そちらを採用しています。

func onceFunc(f func()) func() {
	var done atomic.Bool

	return func(){
		if done.CompareAndSwap(false, true) {
			f()
		}
	}
}

5. パフォーマンス

さてパフォーマンスですが、 context.Context を考慮するため、 sync.(RW)Lock よりも遅いです。特に共有ロックでその差は顕著で、残念ながら手元でベンチマークを取った限りでは100倍ぐらい遅いです。

ベンチマーク
ベンチマークコード
type (
	NaiveLock struct {
		mu sync.Mutex
	}
)

func (n *NaiveLock) Lock(ctx context.Context) (UnlockFunc, error) {
	select {
	case <- ctx.Done():
		return nil, context.Cause(ctx)
	default:
	}

	ch := make(chan struct{})
	go func(){
		n.mu.Lock()
		close(ch)
	}()

	select {
	case <- ch:
		return sync.OnceFunc(func(){ n.mu.Unlock() }), nil
	case <- ctx.Done():
		go func(){
			<- ch
			n.mu.Unlock()
		}()
		return nil, context.Cause(ctx)
	}
}


func BenchmarkMutex(b *testing.B){
	var mu sync.Mutex

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		mu.Lock()
		mu.Unlock()
	}
}

func BenchmarkCtxLock(b *testing.B){
	L := NewLock()
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		unlock, _ := L.Lock(ctx)
		unlock()
	}
}

func BenchmarkNaiveLock(b *testing.B){
	L := &NaiveLock{}
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		unlock, _ := L.Lock(ctx)
		unlock()
	}
}

func BenchmarkMutexContextSwitch(b *testing.B){
	var mu sync.Mutex
	var wg sync.WaitGroup

	mu.Lock()

	for i := 0; i < b.N; i++ {
		wg.Add(1)
		go func(){
			mu.Lock()
			mu.Unlock()
			wg.Done()
		}()
	}

	b.ResetTimer()
	mu.Unlock()
	wg.Wait()
}

func BenchmarkCtxLockContextSwitch(b *testing.B){
	L := NewLock()
	ctx := context.Background()
	var wg sync.WaitGroup

	unlock, _ := L.Lock(ctx)

	for i := 0; i < b.N; i++ {
		wg.Add(1)
		go func(){
			unlock, _ := L.Lock(ctx)
			unlock()
			wg.Done()
		}()
	}

	b.ResetTimer()
	unlock()
	wg.Wait()
}

func BenchmarkRWMutexExclusive(b *testing.B){
	var mu sync.RWMutex

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		mu.Lock()
		mu.Unlock()
	}
}

func BenchmarkCtxSharableLockExclusive(b *testing.B){
	L := NewSharableLock()
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		unlock, _ := L.ExclusiveLock(ctx)
		unlock()
	}
}

func BenchmarkRWMutexExclusiveContextSwitch(b *testing.B){
	var mu sync.RWMutex
	var wg sync.WaitGroup

	mu.Lock()
	for i := 0; i < b.N; i++ {
		wg.Add(1)
		go func(){
			mu.Lock()
			mu.Unlock()
			wg.Done()
		}()
	}

	b.ResetTimer()
	mu.Unlock()
	wg.Wait()
}

func BenchmarkCtxSharableLockExclusiveContextSwitch(b *testing.B){
	L := NewSharableLock()
	ctx := context.Background()
	var wg sync.WaitGroup

	unlock, _ := L.ExclusiveLock(ctx)
	for i := 0; i < b.N; i++ {
		wg.Add(1)
		go func(){
			unlock, _ := L.ExclusiveLock(ctx)
			unlock()
			wg.Done()
		}()
	}

	b.ResetTimer()
	unlock()
	wg.Wait()
}

func BenchmarkRWMutexSharedLockOnly(b *testing.B){
	var mu sync.RWMutex

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		mu.RLock()
	}
}

func BenchmarkCtxSharableLockSharedLockOnly(b *testing.B){
	L := NewSharableLock()
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		L.SharedLock(ctx)
	}
}

func BenchmarkRWMutexSharedLockUnlock(b *testing.B){
	var mu sync.RWMutex

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		mu.RLock()
		mu.RUnlock()
	}
}

func BenchmarkCtxSharableLockSharedLockUnlock(b *testing.B){
	L := NewSharableLock()
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		unlock, _ := L.SharedLock(ctx)
		unlock()
	}
}
ベンチマーク結果
go test -bench .
goos: darwin
goarch: amd64
pkg: github.com/ymd-h/go/ctxlock
cpu: Intel(R) Core(TM) i7-4770HQ CPU @ 2.20GHz
BenchmarkMutex-8                                   	95170290	        11.89 ns/op
BenchmarkCtxLock-8                                 	 6370436	       192.5 ns/op
BenchmarkNaiveLock-8                               	  943366	      1157 ns/op
BenchmarkMutexContextSwitch-8                      	 1000000	      2215 ns/op
BenchmarkCtxLockContextSwitch-8                    	 1000000	      1056 ns/op
BenchmarkRWMutexExclusive-8                        	41585079	        29.02 ns/op
BenchmarkCtxSharableLockExclusive-8                	 5418589	       205.7 ns/op
BenchmarkRWMutexExclusiveContextSwitch-8           	 1412592	      1081 ns/op
BenchmarkCtxSharableLockExclusiveContextSwitch-8   	 1000000	      1155 ns/op
BenchmarkRWMutexSharedLockOnly-8                   	198582211	         6.177 ns/op
BenchmarkCtxSharableLockSharedLockOnly-8           	 1788302	       653.6 ns/op
BenchmarkRWMutexSharedLockUnlock-8                 	93392736	        12.05 ns/op
BenchmarkCtxSharableLockSharedLockUnlock-8         	  962224	      1234 ns/op

とはいっても、操作に応じてナノ秒〜マイクロ秒程度の所要時間なので、context.Context のキャンセルを考慮するようなユースケースでは十分だと思っています。
これが問題になるぐらいロック取得性能がクリティカルなユースケースであるならば、おそらく context.Context のキャンセルはもっと上位の大きなループで実行する方が適切なのではないかと思います。

6. まとめ

goの習作として、context.Contex を考慮する排他・共有ロックを提供するモジュール ctxlock を実装して公開しました。
もし興味を持ってもらえたらなら、使ってみてもらえると嬉しいです。

https://github.com/ymd-h/go/tree/master/ctxlock
https://pkg.go.dev/github.com/ymd-h/go/ctxlock

2月末から長らく記事を書けていませんでしたが、やっと書けました。来年はもっと沢山記事を書いていけるようにしたいです。

脚注

  1. TryLock() を使うような状況はそもそも設計に問題があると言うのがGo公式の考えのようです。 ↩︎

  2. sync.(RW)Lock は他の場所からもロック解放ができてしまうし、ロック未取得で Unlock() を呼ぶとエラーになる挙動が安全ではないと感じて、このように設計しています。 ↩︎

  3. 当初はコンストラクタで、reader 管理用の goroutine を建てていましたが、それだとオブジェクトを構築するだけ goroutine が増えていくので、紆余曲折を経て今のように reader がロックを取得している間のみ goroutine を建てる形になりました。 ↩︎

  4. 設計の都合上、厳密には writer がロックを希望した瞬間ではなく、1回 reader が増減して次の待機時に、writer の希望を考慮します。reader-writerの優先順位は厳密な要求仕様ではないので、このぐらいは許容範囲だと考えています。 ↩︎

Discussion