Go の context を使ったキャンセル、タイムアウト

context でキャンセルやタイムアウトを実装してみます。

目次

キャンセル

context.WithCancelでコンテキストから子コンテキストとキャンセル関数を生成し、子コンテキストをfunに渡して サブの goroutine で実行します。
関数funは基本的には、1秒毎に"working"と出力し続けます。
main の goroutine では、5秒待機してからcancel()します。

package main

import (
    "context"
    "fmt"
    "time"
)

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

    go fun(ctx)

    time.Sleep(5 * time.Second)
    cancel()
    time.Sleep(5 * time.Second)
}

func fun(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("canceled")
            return
        default:
            fmt.Println("working")
            time.Sleep(1 * time.Second)
        }
    }
}
$ go run main.go
working
working
working
working
working
canceled

cancel()が実行されると、ctx.Done()のチャネルに送信されます。
funのselectでは、ctx.Done()のチャネルから受信されるため、case <-ctx.Done():のcase に落ちます。
子コンテキストを渡していたサブ goroutine の処理を終わらせることができました。

タイムアウト

上のキャンセルでは、cancel()を実行したタイミングでコンテキストをキャンセルさせていました。
タイムアウトでは、どれくらいの時間経過したらキャンセルするかを設定しておきます。

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()

    go fun(ctx)

    time.Sleep(5 * time.Second)
}

func fun(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("canceled")
            return
        default:
            fmt.Println("working")
            time.Sleep(1 * time.Second)
        }
    }
}

$ go run main.go
working
working
canceled

2秒経過するとタイムアウトになり、ctx.Done()のチャネルに送信されます。
タイムアウトした場合も、defer cancel()でコンテキストをクローズする必要があるっぽいです。

デッドライン

デッドラインは、キャンセルする時刻を設定します。

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Second))
    defer cancel()

    go fun(ctx)

    time.Sleep(5 * time.Second)
}

func fun(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            fmt.Println("canceled")
            return
        default:
            fmt.Println("working")
            time.Sleep(1 * time.Second)
        }
    }
}

$ go run main.go
working
working
canceled

現在時刻の2秒後の時刻をデッドラインに指定しているため、タイムアウトと同じ挙動になります。

おわりに

コンテキストのキャンセルを試してみました。
無駄な処理や、goroutineリークを防げるようになりました〜。