sync.Onceの内部実装から学ぶ実装アイデア

この記事は、2021年Go Advent Calendar9日目の記事です。

概要

Goで並行処理で利用されるsyncパッケージには、処理の一回のみの実行を担保するOnce型(構造体)が提供されていて、主に並行処理の中で重複処理や競合を避けるために使えます。
このOnceは、一見内部では複雑なことをしていそうですが、その内部実装は短く非常にシンプルになっています。
しかし、その短いコードの中では、処理の最適化、並行処理の安全性、パッケージの使いやすさなどのために、汎用的な実装上の工夫がいくつかされています。
そこで、この記事では、Onceの内部実装を追いながら、それらを見ていこうと思います。

sync.Onceとは

最初に、そもそもsync.Onceがなにかを説明します。すでにご存知の方はこのパートは飛ばしてください。

sync.Once型は、関数が一回のみ実行されることを担保する機能を提供しています。
実際に使い方を見てみましょう。

package main

import (
	"fmt"
	"sync"
)

func main() {
	var once sync.Once
	f := func() {
		fmt.Println("Done")
	}
	once.Do(f)
	once.Do(f)
	once.Do(f)
}

The Go Playground

sync.Onceの変数を初期化して、Doというメソッドに関数を渡して実行してます。

そして、複数回このDoを呼んでいますが、実行してみると、下記のように一度しか対象の関数が実行されません。

Done

このように、関数を一度だけ実行したい場合に有用な機能です。

そして、もちろん、並行処理でも一度しか実行されません。

package main

import (
	"fmt"
	"sync"
)

func main() {
	var once sync.Once
	f := func() {
		fmt.Println("Done")
	}
	done := make(chan bool)
	for i := 0; i < 10; i++ {
		go func() {
			once.Do(f)
			done <- true
		}()
	}
	for i := 0; i < 10; i++ {
		<-done
	}
}

[The Go Playground](https://go.dev/play/p/b6QDJnP0kpE

Done

sync.Onceの内部実装

それでは、sync.Onceの内部実装を見ていきましょう。

全体像

まず、コード全体を見てみましょう。コメントを取り除くと、25行程度の非常に短い実装になっています。
しかも、Goの組み込みの機能と標準パッケージのpublicな機能しか使っておらず、言語内部専用の機能は利用していません。

package sync

import (
	"sync/atomic"
)

type Once struct {
	done uint32
	m    Mutex
}

func (o *Once) Do(f func()) {
	if atomic.LoadUint32(&o.done) == 0 {
		o.doSlow(f)
	}
}

func (o *Once) doSlow(f func()) {
	o.m.Lock()
	defer o.m.Unlock()
	if o.done == 0 {
		defer atomic.StoreUint32(&o.done, 1)
		f()
	}
}

それでは、順番に見ていきましょう。

sync.Onceの構造

sync.Onceは下記のような構造体となっています。

// Once is an object that will perform exactly one action.
//
// A Once must not be copied after first use.
type Once struct {
	// done indicates whether the action has been performed.
	// It is first in the struct because it is used in the hot path.
	// The hot path is inlined at every call site.
	// Placing done first allows more compact instructions on some architectures (amd64/386),
	// and fewer instructions (to calculate offset) on other architectures.
	done uint32
	m    Mutex
}

フィールドが２つだけのシンプルな構造になっています。

• done
  
  フィールドが実効済みかを管理するフラグの役割をしています。
  uint32になっているのは、sync/atomicパッケージの機能を利用して、排他操作をするためです。
  (sync/atomicについてはこちらをご覧ください。)

• m
  
  排他制御のためにsync.Mutexを持っています。(sync.Mutexについてはこちらをご覧ください。)

このOnceの構造体で、注目すべき点が２つほどあるので、さらに見てみましょう。

ゼロ値で使える構造体

構造体をゼロ値でそのまま使えるという点があります。
下記のように、初期化用の関数や値の設定等が不要で、ただ変数に初期化するだけでそのまま使えるので、使いやすい作りになっています。

var once sync.Once

もちろん、このような構造体にできるケースは限られますが、できるときにはやった方が良いデザインだと思います。

Hot Path

さらに、コメントも一部見てみましょう。
doneフィールドについて、下記のコメントがあります。

// It is first in the struct because it is used in the hot path.
// The hot path is inlined at every call site.
// Placing done first allows more compact instructions on some architectures (amd64/386),
// and fewer instructions (to calculate offset) on other architectures.

このコメントでは、このdoneがstructの最初のフィールドである理由を説明しています。
この中で、Hot Pathという単語で出てきますが、これは頻繁に実行されるパスという意味です。
つまり、doneが一番頻繁に参照されるフィールドなので、structの最初のフィールドになっているということです。
なぜ、頻繁に参照されるフィールドが最初に来るべきかというと、structの最初のフィールドのメモリアドレスはstruct自体のメモリアドレスと同じなので、フィールドにアクセスするための処理が少なくて済むからです。
逆に、二番目以降のフィールドにアクセスする場合は、そのメモリアドレスを計算して取得するための処理が追加で必要になります。コメントでoffsetという言葉が出てきますが、offsetはstructなどのデータ構造の先頭から特定の要素の位置までの距離のことなので、ここではもしdoneが二番目以降のフィールドだったらそれを計算する必要があるということを言っています。
このように、CPUの処理を少しでも最適化するために、構造体の最初のフィールドを一番参照されるdoneにしているのでした。

以上がOnceの構造でした。

Doの実装

次に、Doメソッドの内部実装も見ていきましょう。

func (o *Once) Do(f func()) {
	// Note: Here is an incorrect implementation of Do:
	//
	//	if atomic.CompareAndSwapUint32(&o.done, 0, 1) {
	//		f()
	//	}
	//
	// Do guarantees that when it returns, f has finished.
	// This implementation would not implement that guarantee:
	// given two simultaneous calls, the winner of the cas would
	// call f, and the second would return immediately, without
	// waiting for the first's call to f to complete.
	// This is why the slow path falls back to a mutex, and why
	// the atomic.StoreUint32 must be delayed until after f returns.

	if atomic.LoadUint32(&o.done) == 0 {
		// Outlined slow-path to allow inlining of the fast-path.
		o.doSlow(f)
	}
}

Doメソッドは、Onceの中のdoneフィールドをチェックして、まだ未完了の状態であればdoSlowという別の関数を呼び出しています。
すでに関数が実行済みである場合は、ここでチェックだけしてなにもされずに終わることになります。
doneのチェックは、sync/atomicパッケージのLoadUint32を使って、他の処理の割り込みを防ぎながら実行しています。

doSlowも見てみましょう。

func (o *Once) doSlow(f func()) {
	o.m.Lock()
	defer o.m.Unlock()
	if o.done == 0 {
		defer atomic.StoreUint32(&o.done, 1)
		f()
	}
}

このメソッドで実際に関数が実行されるわけですが、最初にsync.Mutexでロックをしています。
その上で、再度実行済みじゃないことを確認した上で、対象の関数f()を実行しています。
そして、f()の実行が終わった後に、doneフィールドを更新しています。

ここで、Mutexロックの解除とdoneの更新は、deferで実行しています。
そのため、fの中でpanicが発生しても、実行済みとなります。
また、deferはLIFO(last-in-first-out)なので、後にセットされたdoneの更新処理が、Mutexロックの解除より先に実行されます。そのおかげで、fが二回実行されることが絶対起きないようにしています。

後続処理の安全性確保: 関数の実行が終わるまですべての呼び出しがブロックされる

doneの更新とMutexのロック解除は、関数f()の実行が終わった後に行われています。つまり、その間、他のDoを呼び出している処理はブロックされてます。

そこについては、こちらのコメントで説明されています。

// Note: Here is an incorrect implementation of Do:
//
//	if atomic.CompareAndSwapUint32(&o.done, 0, 1) {
//		f()
//	}
//
// Do guarantees that when it returns, f has finished.
// This implementation would not implement that guarantee:
// given two simultaneous calls, the winner of the cas would
// call f, and the second would return immediately, without
// waiting for the first's call to f to complete.
// This is why the slow path falls back to a mutex, and why
// the atomic.StoreUint32 must be delayed until after f returns.

そもそも、とにかく二回関数が呼び出されてないことだけを担保するなら、コメントにあるように、下記のコードだけで済みます。

if atomic.CompareAndSwapUint32(&o.done, 0, 1) {
	f()
}

ただ、意図的にこの実装にしないで、対象の関数f()の実行が完了するまで、すべてのDoの実行がブロックされるようにしています。
これは、すべての後続処理が安全に実行されるようにするためのようです。特に、Doもf()も返り値を持たないので、後続処理はf()が実行されたかどうかをDoの返り値などから直接確かめる方法がありません。そこで、他のプロセスでも同時にDoが呼ばれていたとしても、f()が完了するまでブロックされることで、安全かつ簡潔に後続処理が書けます。
この仕様は、例えば、初期化処理にOnceを利用した際などに有用となります。

Fast Path / Slow Path

Doが呼び出しているメソッドは、doSlowという名前でした。
なぜslowなのかを見てみましょう。
ここで、Fast Pathという概念を紹介します。fast pathというのは、名前の通り、通常よりも速い(処理の少ない)ショートカットのパスということになります。
そして、そうじゃない遅いパスがSlow Pathとなります。
このDoにはこの２つのパスが存在しています。

まず、上記の通り、「f()の実行が完了するまで、すべてのDoが完了しない」という仕様がありました。そして、そのためにdoSlowメソッドの中では、都度Mutexロックをかけてdoneの値をチェックする(そして最後にロックを解除する)という処理をしていました。これは、f()がすでに実行済みの状態でDoが呼ばれたときのことを考えると、やや冗長ではあります。
そこで、Doは、最初にatomic.LoadUint32(&o.done)でその時点でもう実行済みかというのを確認しています。その上で実行済みじゃなかったら、doSlowメソッドが呼ばれています。

つまり、ここでは、

• f()がすでに実行済みの場合は、すぐにreturnする -> Fast Path
• f()が未実行(実行中含む)の場合は、doSlowでブロック(&実行)する -> Slow Path

というフローになっています。

以上がDoの中身でした。

まとめ

これで、sync.Onceの内部実装を全部確認しました。想定よりも短く読みやすいコードだったと思います。
しかし、その短いコードは、パッケージとしての使いやすさ、並行処理の安全な実装、処理の最適化などのために、いくつかの工夫が見受けられるのでした。