Goで重い処理をtimeoutさせる

Goで重い処理を書いているとタイムアウトさせたいときがあると思います。

大抵のIOが発生するライブラリだとcontextを引数に加えると、context の終了通知が発生して終了してくれます。

例えば、以下のようにhttp requestであれば、contextにタイムアウト設定することでhttp requestがタイムアウトになるとリクエストをキャンセルしてくれます。

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

req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://www.google.com", bytes.NewReader([]byte("")))
resp, _ := r.client.Do(req)

しかし、ライブラリやその他IOによってはちゃんとcontextの終了が実装されていない場合があります。通常であれば、contextを渡すだけなはずが終了できない場合、自分でその機構を作ってしまうのも手です。

以下のようにキャンセルさせる関数を作ってみました。

func main(){
       // 重い関数
    heavyFunc := func() {
        time.Sleep(10 * time.Second)
    }
    if err := execWithTimeout(context.Background(), heavyFunc, 2*time.Second); err != nil {
        fmt.Println(err)
    }
}

// ctx 親context
// fn 実行させる重い関数
// timeout timeoutさせる時間
// この関数を通して実行すると、timeoutさせることができる。
// 注意するのは、goroutineで投げっぱなしになるので、一時的に重い処理ではなく常時重い処理の場合はどんどんリークしていくので使うべきではない
func execWithTimeout(ctx context.Context, fn func(), timeout time.Duration) error {
    timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
    defer cancel()


    go func() {
        fn()
        fmt.Println("finish heavy")
        cancel()
    }()
    return ticker(timeoutCtx)
}

// 入力されたcontextが終了するかどうかを監視する関数
func ticker(ctx context.Context) error {
    for {
        time.Sleep(10 * time.Millisecond)
        select {
        case <-ctx.Done():
            fmt.Println("Stop ticker")
            if ctx.Err() == context.Canceled {
                return nil
            }
            return ctx.Err()
        default:
        }
    }
}

以下にplay groundを貼っておきます

The Go Playground

Screaming Loud

日々是精進

Goで重い処理をtimeoutさせる