Go: selectでctx.Done()を受信するときの注意点

go

Goで非同期的な処理の中断を検知したい場合、次のように select と ctx.Done() を使って書くことが多い。 とても便利なパターンなのだが、いくつか使うときの注意点がある。

select {
case <-ctx.Done():
    // done
case <-ch:
    // ...
}

selectによる受信のランダム性

channelに届いたタスクを順番に処理しつつ、キャンセル要求が来たら終了する簡単なプログラムを例にする。 コードにある通り次のような順序で実行される。

  • task1を送る
  • キャンセル要求を送る
  • task2を送る
package main

import (
    "context"
    "log"
    "time"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    tasks := make(chan func())

    go func() {
        tasks <- func() {
            log.Printf("start task1")
            time.Sleep(1 * time.Second)
        }

        time.Sleep(100 * time.Millisecond)
        cancel()

        tasks <- func() {
            log.Printf("start task2")
            time.Sleep(1 * time.Second)
        }
    }()

    for {
        select {
        case <-ctx.Done():
            log.Printf("done")
            return
        case task := <-tasks:
            task()
        }
    }
}

キャンセル要求を送った後にtask2を送っているので、task1だけ実行されて終了するように見える。

しかし、このプログラムの出力結果は不定である。何度も実行するとtask2が実行されたりされなかったりする。

2021/07/29 15:28:36 start task1
2021/07/29 15:28:37 start task2
2021/07/29 15:28:38 done
2021/07/29 15:28:47 start task1
2021/07/29 15:28:48 done

なぜこうなるのか？ その理由は task1を実行している間にキャンセル要求とtask2が両方channelに入り、select がどちらを選ぶかは不定（ランダム）だから 1 である。

では、常にキャンセル要求を優先したい場合はどうすればいいか？ ctx.Done()を2重でチェックするという小技が使える。次のように ctx.Done() とは違うcaseに入ってしまった場合に ctx.Done() を再度チェックするとよい。

       select {
        case <-ctx.Done():
            log.Printf("done")
            return
        case task := <-tasks:
            // Prefer the context's cancel.
            select {
            default:
            case <-ctx.Done():
                log.Printf("done")
                return
            }
            task()
        }

Go本体のテストでもこの select のランダム性によってテストの実行結果が不安定となっていて、ctx.Done() の二重チェックを入れて修正したようなコミットがあった。

http2: make Transport prefer HTTP response header recv before body wr… · c3mb0/net@e325faf · GitHub

複数channel受信時のgoroutine leak

時間がかかる処理 (longRunningTask) の完了を待ちつつ、途中でキャンセル可能にする次のプログラムを例にする。 時間がかかる処理が終わる前に意図的にキャンセルしている。

時間がかかる処理がエラーを返した場合は、エラーを受け取って表示する仕組みも入れている。 一見よくありそうな書き方に見えるが、このままだとgoroutine leakが発生する。

package main

import (
    "context"
    "errors"
    "log"
    "time"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    go process(ctx)
 
    time.Sleep(100 * time.Millisecond)
    cancel()
 
    time.Sleep(3 * time.Second)
}

func process(ctx context.Context) {
    errCh := make(chan error)
    go func() {
        if err := longRunningTask(); err != nil {
            errCh <- err
        }
    }()

    select {
    case <-ctx.Done():
        log.Printf("done")
    case err := <-errCh:
        log.Printf("%+v", err)
    }
}

func longRunningTask() error {
    time.Sleep(1 * time.Second)
    return errors.New("error")
}

これを実行すると次のような順序で処理される。

  • longRunningTask 開始
  • 100ms待機
  • キャンセル要求
  • 900ms待機
  • longRunningTask 完了

キャンセル要求を送った時点で ctx.Done() を受け取るので select は終了する。 そのあとに longRunningTask が完了してエラーがあれば errCh にエラーを送ろうとする。

しかし、受信する select が先に消えてしまったので channel に値を送信できず、longRunningTaskのgoroutineがずっと残り続ける。

解決策としては errCh に送るときに ctx.Done() も一緒にチェックするなどだろうか。 これだと select で受信するchannelが3つ以上になったときに面倒かもしれない。

func process(ctx context.Context) {
    errCh := make(chan error)
    go func() {
        if err := longRunningTask(); err != nil {
            select {
            case errCh <- err:
            case <-ctx.Done():
            }
        }
    }()

    select {
    case <-ctx.Done():
        log.Printf("done")
    case err := <-errCh:
        log.Printf("%+v", err)
    }
}

  1. Twitterでnanasi880さんから教えてもらいました 感謝！ https://twitter.com/_NANASI880/status/1222740812238180353↩