🧊

ソースコードを読んでcontextを理解する

2020/10/20に公開

tech

概要

contextパッケージは、生成したgoroutineの実行をキャンセルし、リソースを解放するための仕組みを提供しています。また、リクエストスコープの値を保持させることもできます。

ここでは、contextパッケージのソースコードから、どのようにgoroutineの実行がキャンセルされるかを見ていきます。

具体的には、基本的な以下のような使い方をした場合に何が行われているのかを確認していきます。

import (
	"context"
	"fmt"
	"sync"
)

func main() {
	var wg sync.WaitGroup
	wg.Add(1)

	ctx, cancel := context.WithCancel(context.Background())

	go func(ctx context.Context) {
		select {
		case <-ctx.Done():
			fmt.Println("----done----")
			wg.Done()
			return
		}
	}(ctx)

	cancel()
	wg.Wait()
}
// $ go run context.go 
// ----done----

これは、goroutineの生成側でcancelを実行し、ctx.Done()が返すchannelをcloseしています。そうすることで、実行中のgoroutineでそのchannelから受信することができ、selectを抜けます。なぜなら、閉じられたchannelからはゼロ値を受信することができるためです。

func main() {
	ch := make(chan struct{})
	close(ch)
	fmt.Println(<-ch)
}
// $ go run context.go 
// {}

contextパッケージを読む

派生したContextを作成する

まず、Contextはインターフェースです。そしてデフォルトのContextは以下の二つが用意されています。これらは、キャンセルすることもDeadlineを指定することもできません。主にmain関数から渡される最初のContextとして利用されます。

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

func TODO() Context {
	return todo
}

これをキャンセル可能にするには、context.WithCancelにContextを渡します。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	return &c, func() { c.cancel(true, Canceled) }
}

内部ではまず、newCancelCtxでcancelCtx構造体を作ります。これは元のContextが埋め込まれます。

func newCancelCtx(parent Context) cancelCtx {
	return cancelCtx{Context: parent}
}

このcancelCtx構造体は以下のように、childrenフィールドを持ち、キャンセル用のインターフェースを持つContext（canceler）をmapで保持しています。これは後で確認するように、派生したContextを表現するために用いられます。

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
}

そして次に、propagateCancelに元のContextとcancelCtxが渡されます。

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	...
	c := newCancelCtx(parent)
	propagateCancel(parent, &c)
	...
}

以下のように、propagateCancelは親のContextが*cancelCtx型であるか否かで処理が分かれます。

Contextが*cancelCtx型の場合は、そのchildrenに子を登録する
Contextが*cancelCtx型でないの場合は、終了すると子をキャンセルするgoroutineを起動する

どちらにせよ、親のContextが終了すると、子のContextが終了できるような準備をしています。

func propagateCancel(parent Context, child canceler) {
	...

	// Contextインターフェースのparentを*cancelCtx型にキャストする
	if p, ok := parentCancelCtx(parent); ok {
		p.mu.Lock()

		...
		// *cancelCtxの場合はchildrenにchildを登録する
		p.children[child] = struct{}{}
		...

		p.mu.Unlock()
	} else {
		// *cancelCtxでない場合はparentのchannelがcloseされるとchildをcancelするgoroutineを起動する
		atomic.AddInt32(&goroutines, +1)
		go func() {
			select {
			case <-parent.Done():
				child.cancel(false, parent.Err())
			case <-child.Done():
			}
		}()
	}
}

最後に、生成したContext（実際は*cancelCtx型）と、それをキャンセルする関数を返します。

// Canceled is the error returned by Context.Err when the context is canceled.
var Canceled = errors.New("context canceled")

func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
	...
	c := newCancelCtx(parent)
	...
	return &c, func() { c.cancel(true, Canceled) }
}

Contextをキャンセルする

次に先ほどの、cancelCtx.cancelが呼び出された場合を見ていきます。

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	...
	c.mu.Lock()
	...

	// 通信用channelを閉じる
	close(c.done)

	// 子のcancelCtxを全てキャンセルする
	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)
	}
}

内部ではまず、自身の終了を告げるchannelをcloseします。そうすることで、このContextに対してctx.Done()から受信することができるようになります。

また、このContextだけなくchildrenに格納されているContextを再起的にキャンセルしていきます。こうすることで、以下のような親のContext（c1）から何度も派生したContext（c2, c3）も、親が終了すると終了できるようになります。

import (
	"context"
	"fmt"
	"sync"
)

func main() {
	var wg sync.WaitGroup
	wg.Add(1)

	c1, can1 := context.WithCancel(context.Background())

	go func(ctx context.Context) {
		c2, _ := context.WithCancel(ctx)

		go func(ctx context.Context) {
			c3, _ := context.WithCancel(ctx)
			select {
			case <-c3.Done():
				fmt.Println("----done----")
				wg.Done()
				return
			}
		}(c2)
	}(c1)

	can1()
	wg.Wait()
}

そして最後にremoveFromParentがtrueの場合は、removeChildを呼び出し親のContextが*cancelCtx型の場合は、childrenから自身を除きます。

func (c *cancelCtx) cancel(removeFromParent bool, err error) {
	...

	if removeFromParent {
		// newCancelCtxでは、親のContextはcancelCtx.Contextフィールドに格納されている
		removeChild(c.Context, c)
	}
}

func removeChild(parent Context, child canceler) {
	p, ok := parentCancelCtx(parent)
	if !ok {
		return
	}
	p.mu.Lock()
	if p.children != nil {
		delete(p.children, child)
	}
	p.mu.Unlock()
}

context.WithDeadlineの場合

context.WithDeadlineも基本的にcontext.WithCancelと同じことをやっています。
違いとしては、time.AfterFuncを利用して指定時間を過ぎるとキャンセルするようにしている点くらいです。

func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
	...
	c := &timerCtx{
		cancelCtx: newCancelCtx(parent),
		deadline:  d,
	}
	...

	c.timer = time.AfterFunc(dur, func() {
		c.cancel(true, DeadlineExceeded)
	})

	...
	return c, func() { c.cancel(true, Canceled) }
}

分かったこと

親から子へ再起的にキャンセルされる

既に見てきたように、親のContextがキャンセルされると、派生した子であるContextは再起的にキャンセルされます。逆に、キャンセルされたContextの派生元である親や親から派生した他のContextはキャンセルされません。つまりA1がキャンセルされてもAはキャンセルされず、したがって他のA2, A3はキャンセルされません。キャンセルする場合にはどの階層のContextに対応したキャンセルなのかを意識する必要があります。

ctx A
 ├ ctx A1 <- cancel
 ├ ctx A2
 └ ctx A3

Context Leak

既に見たように、使われなくなったContextをキャンセルしないと、removeChildが呼ばれずに親のContextに残り続けることになります。したがって、余分にメモリを利用した状態になってしまいます。