Go言語のContext調べてみた:キャンセル機能について

はじめに

最近 Go 言語を書いていて，よく Context に出くわしていました。

キャンセルや時間制限などを通知するための型ということはなんとなくわかっていたのですが、実践的な利用方法がいまいちわからないず、context。Background でいつも黙らせてきました。

ただ、このままでは何も学びがないかつ、まずいことをしているかもしれないと思い、Context について学んでみましたので、記事に残します。

思ったより長い記事になってしまったため、今回は Context の目的の一つである、実行される並行処理をキャンセルする API の提供に着目し説明します。

もう一つの目的であるデータを渡すためのデータの置き場所の提供という目的はまた今度別の記事を書こうかと思います。

Go でよく見る Context

並行処理を行う想定がされる関数(I/O 処理など)の第一引数によく定義されています。

Context は標準パッケージ context 内に定義されている interface です。

以下の例は AWS の SDK ですが、第一引数に context。Context interface を満たすデータを入れる必要があります。

func (c *Client) CreateCluster(ctx context.Context, params *CreateClusterInput, optFns func(*Options) (CreateClusterOutput,error) {...}

この Context ですが、何のためにあるのでしょうか？その理由を知るために並行処理で必要となるキャンセル処理についてまずは見ていきます。

並行処理のキャンセルについて

タイトルとは裏腹にまずは、逐次処理についてです。
逐次処理では上から下への一本道で処理が実行されるため、処理を実行するかしないかはその処理より上の結果に依存し、キャンセルという概念がそもそもないと思います。

// 逐次処理のケース
func sequential() error {
    for i := 0; i < 100; i++ {
        resp,err := registerData(fmt.Sprintf("hello %d",i));
        if err != nil {
            // もしエラーが発生したらその先の処理は実行しない
            return err
        }
    }
    return nil
}

一方、並行処理は複数の処理を少しずつ切り替えながら並行に実行するため、処理の順番依存性は少なく、各々の処理が独立して実行されることが多いはずです(逆にいうと独立性が低い処理を並行に実行することは難しいと思います)。

この時、並行に処理されている複数の処理をある条件によってキャンセルしたい場合があると思います。

例えば

• 並行に実行している処理の内どれか一つでも失敗したら、今後実行される予定の処理は全てキャンセルする
• 時間制限を設けて時間内に実行できなかった処理に関してはキャンセルする

などです。

上記には一本道の処理には無い難しさがあります。なぜなら独立して実行されている複数の処理に対して、キャンセルしたい！という情報を伝播させないといけないからです。

// 並行処理のケース
func parallel() error {
    for i := 0; i < 100; i++ {
        go func() {
            resp,err := registerData(fmt.Sprintf("hello %d",i));
            if err != nil {
                // もしエラーが発生したら他のgoroutineもキャンセルしたいけどどうしよう。。。

            }
        }()
    }
    return nil
}

done チャネルパターンによるキャンセル

並行処理のキャンセルを行うパターンとして done チャネルパターンというものがあるみたいです。

これは並行に実行予定の処理(下の例だと registerData)に done チャネルを渡し、その処理は実行する前に done チャネルが close されていないかを確認し、close されていたらそれ以降の処理をキャンセルするというものです。
並行処理を実行する親の関数が done チャネルを生成し、キャンセルしたくなった際に done チャネルを close します。

func main() {
    // doneチャネル
	done := make(chan interface{})
	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			resp, err := registerData(done,fmt.Sprintf("hello %d",i))
			if err != nil {
                // selectを使うことで二度doneがcloseされるのを防ぐ
				select {
                // doneがすでに閉じられているケース
				case <-done:
					println("cancel")
					return

                // doneが閉じらていないケース
				default:
                    // doneを閉じてキャンセルの旨を伝播
					close(done)
					println(err.Error())
					return
				}
			}
			println(resp)
		}()
	}
	wg.Wait()
}

func registerData(done <-chan interface{},data string)(string,error) {
    // doneチャネルが閉じられているかどうかを確認して処理を実行するか判断する
	select {
    // doneチャネルがcloseされたら何もしない(処理のキャンセル)
	case <-done:
		return "", nil
	default:
        // たまにエラーでこけるように実装
		if rand.Int()%7 == 0 {
			return "", fmt.Errorf("unlucky 7 error")
		}
		return "success", nil
	}
}

# 出力結果例
$ go run main.go
success
success
unlucky 7 error
success
success
cancel
cancel
cancel
...

上の出力結果例ではエラーの後に数件 success してしまっています。
これは恐らく select 文で done が close されたことを検知する前に実行されてしまったものと考えられます。

このように done チャネルを利用することで並行に実行されている処理に対してある程度キャンセルを行うことができます。

キャンセルの要件と Context の機能

キャンセルは上の例で説明したエラー発生時以外にも

• 一定時間経過したら処理をキャンセルするといったタイムアウト設定

• ある時刻までに実行できなかったらキャンセルするといった時間制限

などを実装したい機会は多いと思います。

done チャネルを利用して自分で実装しても良いのですが、Go 言語ではこれらの要件を簡単に達成できるように Context が実装されています。

Context は、done チャネル同様にキャンセルされたかどうかを Done メソッドの返り値のチャネルから検知可能です。Done メソッドの返り値が close されていればキャンセルされたという実装になっているので以下のプログラムで検査可能です。

select {
    case <-ctx.Done():
        // キャンセルされた...
}

そして、Context はタイムアウトや時間制限を簡単に実装できます。

ある時刻になったらキャンセルしたい場合は以下のメソッドで Context を生成し、並行処理に渡せば良いです。

// 今から1秒後を表すデータ
deadline := time.Now().Add(1*time.Second)
// 今から1秒後に自動でキャンセルされるContextの生成
ctx,cancel := context.WithDeadline(ctx,deadline)

タイムアウトを実装する場合は以下です。

// タイムアウトを1秒に設定
timeout := 1 * time.Second
// 今から1秒後に自動でキャンセルされるContextの生成
ctx,cancel := context.WithTimeout(ctx,timeout)

上記 2 つに共通する第 2 返り値である cancel は手動でキャンセルを通知するための関数です。以下のように呼び出すだけで、紐づいている Context をキャンセルすることができます。

// タイムアウトを1秒に設定
timeout := 1 * time.Second
// 今から1秒後に自動でキャンセルされる
ctx,cancel := context.WithTimeout(ctx,timeout)

for {
    err := maybeCauseError(ctx)
    if err != nil {
        // 並行で実行されているmaybeCauseErrorをキャンセルする
        // timeoutせずともキャンセルすることができる
        cancel()
        return
    }
}

キャンセルのみを取得可能な WithCancel 関数もあります。cancel は特定の条件か、defer を用いて使うことが一般的です。
以下はどれか一つの並行処理が失敗すれば未実行の処理はキャンセルされます。

var wg sync.WaitGroup
ctx, cancel := context.WithCancel(ctx)

defer cancel()

wg.Add(3)
go func() {
    defer wg.Done()
    err := f1(ctx)
    if err != nil {
        // f1が失敗したら他をキャンセル
        cancel()
    }
}()

go func(){
    defer wg.Done()
    err := f2(ctx)
    if err != nil {
        // f2が失敗したら他をキャンセル
        cancel()
    }
}()

go func(){
    defer wg.Done()
    err := f3(ctx)
    if err != nil {
        // f3が失敗したら他をキャンセル
        cancel()
    }
}()

wg.Wait()

WithXXX 関数は親の Context から新しい Context を返す

先ほど説明した WithXXX 関数はどれも第一引数に Context を取り、新しい Context を返していました。

第一引数の Context は親の関数から渡された Context が想定されています。(実装されている関数のシグネチャは func WithCancel(parent Context)... となっています。)

つまり、親の Context に対して新たにタイムアウトなどの設定を足すことができることを意味しています。

これは便利な半分、親が設定したタイムアウト値を変更できてしまうのでは無いかと僕は思いました。

この変更が、厳しい方向(1 分->1 秒)になるのであれば問題ないですが、緩い方向(1 秒->1 分)になっては、1 秒で処理が終わると想定していた親はびっくりです。

どのようになるのか以下のコードで検証してみました。

func child(ctx context.Context) (string, error) {
	// 親のContextのtimeoutを1分に変更
	ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
	defer cancel()
	userName, err := fetchUserName(ctx)
	if err != nil {
		return "", err
	}
	return userName, nil
}
func fetchUserName(ctx context.Context) (string, error) {
	// ここでcontextを渡す
	select {
	case <-ctx.Done():
		return "", ctx.Err()
	// 3秒後にkaiを返す
	case <-time.After(3 * time.Second):
		return "kai", nil
	}
}
func main() {
    // timeoutを1秒に設定
	timeout := 1 * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	// child内で実行されるfetchUserNameは実行に3秒かかるため、タイムアウトするはずだが...
	// child内で新しいContextを作成しているので、タイムアウトしないかも...
	userName, err := child(ctx)
	if err != nil {
		panic(err)
	}
	fmt.Println(userName)
}

コメントにもある通り、main ではタイムアウト値を 1 秒に設定していますが、child 関数内ではその Context を引数にタイムアウトを 1 分後に設定して再度 Context を生成しています。

child によって生成された Context を利用して実行される fetchUserName 関数は実行までに 3 秒かかるので、main のタイムアウト値が優先されれば実行はキャンセルされ、child のタイムアウト値が優先されればキャンセルされずに実行されます。

安全性の面から言うと親のタイムアウト値を延長することは子供にさせたくないため、キャンセルされれば良いです。

結果は以下です.

$ go run main.go
panic: context deadline exceeded

おお、child が Context の timeout を 1 分に変更したのにすぐにキャンセルされているぞ！

ということで、親のタイムアウト値が優先されていることがわかりました。

これは実装を見るしかなさそうです。
WithDeadline を覗いてみると以下の記述があります。

    ...
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	propagateCancel(parent, c)

d は WithDeadline で指定された時間で、parent は親の Context です。
c で新しい Context を作成していますが、その次の propagateCancel で parent と新しい Context を入れています。その内部が以下です。

func propagateCancel(parent Context, child canceler) {
	done := parent.Done()
	if done == nil {
		return // parent is never canceled
	}

	select {
	case <-done:
		// parent is already canceled
		child.cancel(false, parent.Err(), Cause(parent))
        ...

このように propagateCancel 関数内で parent の Context が done しているかどうかをチェックしています。つまり、子供がどのような値を Timeout などに設定してもまずは親の Context がキャンセルをしていないか確認するみたいです。
この仕組みにより安全に Context の生成ができるのだと思います。

また、 Context イミュータブルで変更不可のようになっています。
そのため、第一引数で受け取り、用意されている関数を通してのみ新しい Context を生成できるのです。
可変オブジェクトを引き回しているわけではないので安全に Context を生成、利用できるということでしょう。

ただし、先ほどの child 関数の第一引数を親の Context ではなく、context.TODO や context.Background 関数によって新たに生成した context を渡すとどうなるでしょうか？

先ほどのコードの child 関数を以下のように変更します。

func child(ctx context.Context) (string, error) {
	// timeoutを1分に変更
    // 親のctxを無視して新しいContextを生成
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
	defer cancel()
	userName, err := fetchUserName(ctx)
	if err != nil {
		return "", err
	}
	return userName, nil
}

結果は以下です。

$ go run main.go
kai

当たり前ですが、新しく Context を生成してしまうと、このように親の Context は無視されてしまいます。

どのように利用するのが良いのか？

これまでの話から、Context をキャンセルを行う目的で利用する場合は以下のようにすると良いと思います。

• main 関数などの最上位の関数で初めて Context を作成するときは context。Background、またはタイムアウトなどを設けたい場合は context.WithTimeout(context.Background(),timeout)のようにして生成する。

• 子供の関数に Context を渡す場合は基本親から渡された Context を渡すか、親から渡された Context を利用して生成した Context を渡す

終わりに

今回は Context のキャンセル機能について調査し、まとめました。

恥ずかしながら今まで Context が必要なところでは、どこでも context.Background で黙らせていたので、今回の学びを機に変えていこうと思います。

また、Context のもう一つの目的でもある情報の伝達も面白く、採用するかどうかは難しい機能と感じたためまた記事にまとめていこうと思います。
ここまで読んでくださりありがとうございました。もし、間違いなどがありましたらご指摘いただけると幸いです。