Go1.14のcontextは何が変わるのか

## 背景

Go1.14 で context パッケージが少し改善されるのは mattn さんの twitter を見て知った方も多いのではないでしょうか。このツイートされた時期と同じくらいにちょうど社内の勉強会で context パッケージをみんなで読んでおり、~~皆完全に context を理解したので~~ ある程度実装も把握していました。勉強会では GoDoc と最新 master ブランチのコードが結構違うね、みたいな話もありました。ということで、個人的にとても興味深いツイートでした。Go.1.14の[リリースノート](https://tip.golang.org/doc/go1.14#library)には記載されていないのがミソです(2020/02/23現在)。

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">Go 1.4 のリリースノートにまだ含まれてないけど context の WithCancel と WithTimeout の伝搬がこのコミットで速くなってる。<a href="https://t.co/gJiT81uVyj">https://t.co/gJiT81uVyj</a></p>&mdash; mattn (@mattn_jp) <a href="https://twitter.com/mattn_jp/status/1229785166341525505?ref_src=twsrc%5Etfw">February 18, 2020</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

そこで、本記事では Go1.14 でリリース予定の以下のコミットで何が変わったのか、そもそもどんな問題背景があったのか。ということを見てみたいと思います。

CL と proposal は以下です。正確な情報は下記をご覧ください。

- [196521: context: use fewer goroutines in WithCancel/WithTimeout](https://go-review.googlesource.com/c/go/+/196521/)
- [proposal: context: enable first class citizenship of third party implementations](https://github.com/golang/go/issues/28728)

## 結論

**Go1.14 の context はカスタムコンテキストを用いたときに、適切にコンテキストを埋め込むと、`WithCancel` や `WithTimeout` でゴルーチンが生成されなくなりました。**

## 基礎

上記の proposal を見る前に context について簡単に補足しておきます。以下のケースで考えてみます。(proposal のケースとは別です)

![context.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/dcd2ec10-0474-f1af-5b17-1114bd31a906.png)

起点になる場所で `Background` としてコンテキストを生成すると思いますが、これは `emptyCtx` としてコード上は表現されています。`WithValue` や `WithCancel` といったメソッドに応じて、子のコンテキストを生成すると思いますが、そのときに上記のような `valueCtx` とか `cancelCtx` 型のコンテキストが生成されていることになります。`emptyCtx` とか `valueCtx` とか `cancelCtx` はいずれも `Context` インターフェースを満たしている struct になります。

それぞれどのような struct になっているか以下に示します。どちらも `Context` インターフェースを埋め込んでいて、これは(実装上)親のコンテキストを参照するフィールドです。

```go:Go1.13のcontext/context.go
type valueCtx struct {
	Context
	key, val interface{}
}
```

```go:Go1.13のcontext/context.go
type cancelCtx struct {
	Context

	mu       sync.Mutex            // protects following fields
	done     chan struct{}         // created lazily, closed by first cancel call
	children map[canceler]struct{} // set to nil by the first cancel call
	err      error                 // set to non-nil by the first cancel call
}
```

キャンセルは子には伝播するが、親には伝播しないことはよく知られています。つまり以下のような状況です。

- 親のキャンセルが子に伝播する。親(この場合は `valueCtx` )には伝播しない

![context-cancal_parent.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/c467c9c6-9b68-b03b-6c54-af2ae197b288.png)

- 子のキャンセルは親には伝播しない

![context-cancel_child.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/bf22083e-63df-7727-3730-f11f1dc0b819.png)

実際、このようなキャンセル処理が実装上はどのようになっているのか示します。本質は `cancelCtx` の `children` フィールドである map と `propagateCancel` 関数です。

```go:Go1.13のcontext/context.go
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}
```

```go:Go1.13のcontext/context.go
func propagateCancel(parent Context, child canceler) {
	if parent.Done() == nil {
		return // parent is never canceled
	}
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()
		if p.err != nil {
			// parent has already been canceled
			child.cancel(false, p.err)
		} else {
			if p.children == nil {
				p.children = make(map[canceler]struct{})
			}
			p.children[child] = struct{}{}
		}
		p.mu.Unlock()
	} else {
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}
```

`propagateCancel` 関数を見るとわかるのですが `parentCancelCtx` 関数の結果に応じて、処理が変わっています。

`parentCancelCtx` 関数は何かというと、親のコンテキストがどのような型なのか、型アサーションをしながら調べて `*cencelCtx` (`timerCtx` も `cencelCtx` の仲間のようなものです)が出現するまでコンテキストのグラフを子コンテキストから親コンテキストの方向にたどる処理になっています。直感的に言うと、コンテキストから見たときに一番近いキャンセルのコンテキストを見つける、と言えます。

親のコンテキストはコンテキストに埋め込まれているので、子から親方向にたどることができます。

```go:Go1.13のcontext/context.go
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *timerCtx:
			return &c.cancelCtx, true
		case *valueCtx:
			parent = c.Context
		default:
			return nil, false
		}
	}
}
```

`propagateCancel` 関数に戻ると、あるコンテキストから逆向きにコンテキストのグラフたどって、キャンセルコンテキストがあれば、そのコンテキストを親として、子のコンテキストを親の `children map[canceler]struct{}` に格納しています。まだキャンセルのコンテキストが存在しない場合は新たにゴルーチンを 1 つ生成して、そのゴルーチンから今 `WithCancel` が呼ばれたときの親と生成する子コンテキストの両方を監視します。このゴルーチンは親がなんらかの理由で完了/キャンセル (`parent.done` が close される) すれば子もキャンセルするし、子がキャンセルされれば、ゴルーチンは監視する役目を終えます。

キャンセルの実装は(Go1.14 のプロポーザルには関係ないですが)以下のようになっていて、本質的にはチャネルの close と map に紐付いているコンテキストグラフ上のコンテキストのキャンセルの伝播します。(その他にグラフ上で親のコンテキストから子のコンテキストを切り離す場合もあります) 

```go
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	if err == nil {
		panic("context: internal error: missing cancel error")
	}
	c.mu.Lock()
	if c.err != nil {
		c.mu.Unlock()
		return // already canceled
	}
	c.err = err
	if c.done == nil {
		c.done = closedchan
	} else {
		close(c.done)
	}
	for child := range c.children {
		// NOTE: acquiring the child's lock while holding parent's lock.
		child.cancel(false, err)
	}
	c.children = nil
	c.mu.Unlock()

	if removeFromParent {
		removeChild(c.Context, c)
	}
}
```

---

少し細かい話になってしまいましたが、何が言いたかったかというと

- `context` パッケージの内部の struct のイメージ
- `WithCancel` の内部処理のイメージ
    - コンテキストのグラフを生成するときにある場合でゴルーチンが生成されている
    - 親のコンテキストから子のコンテキストに `map` を用いて有向グラフを生成している

がどのようなものになっているか、ということです。

## Proposal

それでは [Proposal](https://github.com/golang/go/issues/28728) を確認します。

> This proposal concerns a way to implement custom context.Context type, that is clever enough to cancel its children contexts during cancelation.

まずこの Proposal は `context.Context` を満たしたカスタムのコンテキストを用いて処理を実装しているケースの話です。

### 問題背景

プロポーザルをあげた @gobwas 氏のコメントをみると

> The point is that for the performance sensitive applications starting of a separate goroutine could be an issue.

と記載しています。

`WorkerContext` という `context.Context` インターフェースを満たしたコンテキストを定義していて、その場合のパフォーマンスに懸念がある。というのがプロポーザルの背景であると読み取れます。

https://github.com/golang/go/issues/28728#issue-379565484 よりサンプルの実装を引用します。

```go:Issueのサンプル実装
type WorkerContext struct {
    ID uint
    done chan struct{}
}

// WorkerContext implements context.Context.

func worker(id uint, tasks <-chan func(*WorkerContext)) {
    done := make(chan struct{})
    defer close(done)

    ctx := WorkerContext{
        ID: id,
        done: make(chan struct{}),
    }

    for task := range w.tasks {
        task(&ctx)
    }
}

go worker(1, tasks)
go worker(N, tasks)
```

```go:Issueのサンプル実装
tasks <- func(wctx *WorkerContext) {
    // Do some worker related things with wctx.ID.

    ctx, cancel := context.WithTimeout(wctx, time.Second)
    defer cancel()

    doSomeWork(ctx)
}
```

具体的にどのような問題か示します。

`WorkerContext` を親として Worker が処理するタスクの処理タイムアウトを `context.WithTimeout` を用いて設定したときの場合です。タスク数は本プロポーザルでは `N` と記載されてあります。十分大きい数と仮定して良いでしょう。以下の図のような状況です。

![context-proposal_isuue.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/29fe0dbb-c17c-9b83-d935-cbf9219f0209.png)

上記の[基礎](#基礎)でも見たように特定の条件下で `WithCancal` を呼び出したときの背景ではゴルーチンが起動していることになります。なのでカスタムコンテキストを用いて `N` つのタスクを起動したときにアプリケーションで `go worker(1, tasks)` として起動している `N` つのゴルーチンの他に、バッググラウンドでそれらを監視する `N` つのゴルーチンが起動していることになります。つまり以下のような状況です。

![context-proposal_isuue_2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/95ce416c-6450-0acc-2cd6-e3fe3c49ee4b.png)

いくつかの議論を重ね @rsc 氏が問題をまとめています。

https://github.com/golang/go/issues/28728#issuecomment-532793417

>To summarize the discussion so far, this issue concerns the implementation of context.WithCancel(parent) (or context.WithDeadline, but we can focus on WithCancel).

`context.WithCancel(parent)` の実装に関する問題と述べています。(実装を見ればわかりますが `context.WithDeadline` も期限付きの `context.WithCancel` のラッパーとみなせるので `context.WithCancel(parent)` に焦点を当てているのだと思います)

> The problem addressed by this issue is the creation of one goroutine per WithCancel call, which lasts until the created child is canceled. The goroutine is needed today because we need to watch for the parent to be done. The child is a known implementation; the parent is not.

**`WithCancel` の呼び出しごとにゴルーチンが 1 つ生成されることが課題である** 、と述べています。

### 対応

#### 修正内容

改善したときの commit はこちらです。(Go のコアチームの russ 氏がコミットしていますっ!!)

https://github.com/golang/go/commit/0ad368675bae1e3228c9146e092cd00cfb29ac27

対応する前提として `Simplicity` を保つということを強く感じられました。新しい API を追加するのは複雑になってしまうので可能であれば避ける、内部用の API も変更しない、ということです。上記のコメント以外にも Proposal の Isuue の中でコメントしています。

> But what if we make the mapping succeed even when parent is an unknown implementation? Sameer is trying to get at that with the optional Parent method, but again we don't really want to add new methods, even optional ones. New API, even optional new API, complicates usage for all users, and if we can avoid that, we usually prefer to avoid it.

#### 何が改善されたか？

> This CL changes the way we map a parent context back to the underlying data structure. Instead of walking up through known context implementations to reach the *cancelCtx, we look up parent.Value(&cancelCtxKey) to return the innermost *cancelCtx, which we use if it matches parent.Done().

CL のメッセージからわかるように、**子コンテキストを親コンテキストに紐付けるときの方法** を変更しています。

> This way, a custom context implementation wrapping a *cancelCtx but not changing Done-ness (and not refusing to return wrapped keys) will not require a goroutine anymore in WithCancel/WithTimeout.

この改善によって **カスタムコンテキストが条件を満たす場合(`*cancelCtx` をラップする、かつ `Done()` を想定外の変更をしない場合)には、カスタムコンテキストでもゴルーチンが生成されずに `WithCancel/WithTimeout` を用いることができます。**

#### 実装の詳細

Go1.13 までの実装は以下のようになっていました。上記の[基礎](#基礎)で記載したとおりですが、紐付ける親コンテキストを取得するために `for` ループで型アサーションをしながら `*cancelCtx` が見つかるまでコンテキストグラフを逆向きに探索します。このとき **独自のカスタムコンテキスト型から生成されたコンテキストが存在する場合 、以下の `default` の式に合致するため呼び出し元に `return nil, false` が返り、結果としてゴルーチンが生成されます**。

```go:Go1.13のcontext/context.go
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *timerCtx:
			return &c.cancelCtx, true
		case *valueCtx:
			parent = c.Context
		default:
			return nil, false
		}
	}
}
```

CL で変更した後の実装を見てみます。従来同様、一番近い `*cancelCtx` をコンテキストのグラフを子コンテキストから親コンテキストの方向に探索します。

```go:CLで修正された実装
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil {
		return nil, false
	}
	p, ok := parent.Value(&cancelCtxKey).(*cancelCtx)
	if !ok {
		return nil, false
	}
	p.mu.Lock()
	ok = p.done == done
	p.mu.Unlock()
	if !ok {
		return nil, false
	}
	return p, true
}
```

```go:CLで追加された実装
var cancelCtxKey int

// ..

func (c *cancelCtx) Value(key interface{}) interface{} {
	if key == &cancelCtxKey {
		return c
	}
	return c.Context.Value(key)
}
```

```go:Go1.13以前に存在しているメソッド
func (c *valueCtx) Value(key interface{}) interface{} {
	if c.key == key {
		return c.val
	}
	return c.Context.Value(key)
}
```

`Value` 関数はコンテキストから値を取得できる関数です。CL では `*cancelCtx` にも `Value` メソッドが実装されました。`*cancelCtx` でのみ使われる `int` 型の変数 `cancelCtxKey` を使い、変数のポインタアドレスを比較することで `cancelCtx` であることを判定しています。`cancelCtx` 型ではない場合は親コンテキスト方向にコンテキストの `Value` を探索します。親方向に値を探索する動作は従来の `Value` メソッドと同じです。

従来は型アサーションを用いてグラフを逆向きに探索していましたが、この場合は、context パッケージに含まれないコンテキスト、つまりカスタムコンテキストを用いた場合、ゴルーチンが生成されてしまいました。この修正によりカスタムコンテキストにおいても透過的に親のコンテキストを探索できるため、カスタムコンテキストよりも前に生成された `cancelCtx` を探索することができます。


またカスタムコンテキストが `Done()` メソッドを独自拡張している場合はゴルーチンを生成しない対象外です。これは探索して見つけた `*cancalCtx` と親コンテキストの `Done()` メソッドで取得できる done チャネルを比較しています。カスタムコンテキストで独自拡張した場合は `false` になります。

```go
ok = p.done == done
```

#### テスト

テストも明快です。むしろ最初にテストの内容を見て、どういう挙動になるか確認するほうが分かりやすかったかもしれません。以下のようなコンテキストのグラフでテストを実施していました。

![context-test.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/89391/f31460cc-e239-661e-30f9-4f00fe529f59.png)

以下のテストコードを確認するだけでも、キャンセルのコンテキストをフィールドにもつカスタムコンテキスト `ctx2` が `WithCancel` で子コンテキストを生成してもゴルーチンが生成されないことが分かります。同様に `Done()` メソッドを実装しているカスタムコンテキスト(Anonymous) が `WithCancel` で子コンテキストを生成した場合はゴルーチンが 1 つ生成されていることが分かります。

```go:context/context_test.go(Go1.14)
type myCtx struct {
	Context
}

type myDoneCtx struct {
	Context
}

func (d *myDoneCtx) Done() <-chan struct{} {
	c := make(chan struct{})
	return c
}

func XTestCustomContextGoroutines(t testingT) {
	g := atomic.LoadInt32(&goroutines)
	checkNoGoroutine := func() {
		t.Helper()
		now := atomic.LoadInt32(&goroutines)
		if now != g {
			t.Fatalf("%d goroutines created", now-g)
		}
	}
	checkCreatedGoroutine := func() {
		t.Helper()
		now := atomic.LoadInt32(&goroutines)
		if now != g+1 {
			t.Fatalf("%d goroutines created, want 1", now-g)
		}
		g = now
	}

	_, cancel0 := WithCancel(&myDoneCtx{Background()})
	cancel0()
	checkCreatedGoroutine()

	_, cancel0 = WithTimeout(&myDoneCtx{Background()}, 1*time.Hour)
	cancel0()
	checkCreatedGoroutine()

	checkNoGoroutine()
	defer checkNoGoroutine()

	ctx1, cancel1 := WithCancel(Background())
	defer cancel1()
	checkNoGoroutine()

	ctx2 := &myCtx{ctx1}
	ctx3, cancel3 := WithCancel(ctx2)
	defer cancel3()
	checkNoGoroutine()

	_, cancel3b := WithCancel(&myDoneCtx{ctx2})
	defer cancel3b()
	checkCreatedGoroutine() // ctx1 is not providing Done, must not be used

	ctx4, cancel4 := WithTimeout(ctx3, 1*time.Hour)
	defer cancel4()
	checkNoGoroutine()

	ctx5, cancel5 := WithCancel(ctx4)
	defer cancel5()
	checkNoGoroutine()

	cancel5()
	checkNoGoroutine()

	_, cancel6 := WithTimeout(ctx5, 1*time.Hour)
	defer cancel6()
	checkNoGoroutine()

	// Check applied to cancelled context.
	cancel6()
	cancel1()
	_, cancel7 := WithCancel(ctx5)
	defer cancel7()
	checkNoGoroutine()
}
```

## まとめ

Go1.14 の context はカスタムコンテキストを用いたときに、適切にコンテキストを埋め込むと、`WithCancel` や `WithTimeout` でゴルーチンが生成されなくなりました。