日記マン

動画広告プロダクトしてます。Go, Kubernetesが好きです。

Goでプログラムを正常に終了する

CNCF系のOSSをコードリーディングしていると、(KubernetesとかPrometheusとかLokiとか)
Go言語の模範的な書き方を学ぶことができて有意義だったりする。
今回は読んでてためになったリファレンスな記述の1つとして、
OSシグナル(ctrl-C)などを受け取ったときに正常終了するためのメモ。(graceful shutdown)

以下の種類のプログラム別にまとめる。

  • デーモンプログラム
  • バッチジョブ

Daemon Jobs

デーモンプログラムの場合、基本的に for-loop や channel-blocking で処理の終了を防ぐ。

receive only な channel を引数に受けるメソッド設計のパターン、

// Client-Side
func main() {
    ctx, cancel := context.WithCancel(context.Background())
    hoge.Run(ctx.Done())
}

// Server-Side
func (h Hoge) Run(ch <-chan struct{}) {
    for {
        select {
        case <-ch:
            return
        default:
            // do something ...
        }
    }
}

もしくはContextを引数に受け取るパターン

// Client-Side
func main() {
    ctx, cancel := context.WithCancel(context.Background())
    hoge.Run(ctx)
}
func (h Hoge) Run(ctx context.Context) {
    for {
        select {
        case <-cxt.Done():
            return
        default:
            // do something ...
        }
    }
}

シグナルをハンドリングする goroutine を起動し、
シグナルを受信したことを通知するチャネルを返す signalHandler() を実装する。
念の為シグナルを二回受信した場合、直接強制終了する。

func signalHandler() <-chan struct{} {
    stop := make(chan struct{}, 0)
    go func() {
        quit := make(chan os.Signal, 2)
        signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
        <-quit
        close(stop)
        <-quit
        os.Exit(1)
    }()
    return stop
}

これにより、 Run() メソッドのブロッキングを解除する実装ができる。

func main() {
    stop := signalHandler()
    hoge.Run(stop) // Run(ch <-chan struct{})
}

もしくは

func main() {
    stop := signalHandler()
    ctx, cancel := context.WithCancel(context.Background())
    go func() {
        <-stop
        cancel()
    }()
    hoge.Run(ctx) // Run(ctx context.Context)
}

Batch Jobs

これを踏まえて、デーモンプログラムではなくとも、
長時間のワークロードが見込まれる Run() メソッドの終了もしくはキャンセルを扱う。
下の例では、1秒ごとにforループを一回実行、合計で100回実行するRunメソッドを扱う。
単純に100秒後for文を終了しRunメソッドを抜け出すか、
exit用のチャネルを受信しループの途中であっても抜け出すかという記述ができる。

func main() {
    stop := signalHandler()
    hoge.Run(stop) // blocks until finished or canceled
}

func signalHandler() <-chan struct{} {
    stop := make(chan struct{}, 0)
    go func() {
        quit := make(chan os.Signal, 2)
        signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
        <-quit
        close(stop)
        <-quit
        os.Exit(1)
    }()
    return stop
}

// Run blocks.
func (h Hoge) Run(ch <-chan struct{}) {
    t := time.NewTicker(1 * time.Second)
    defer t.Stop()

    // loop 100 times.
    for i := 0; i < 100; i++ {
        select {
        case <-t.C:
            // do something.
        case <-ch:
            // exit.
            return
        }
    }
}