背景

以前Graceful shutdown機能の使い方を紹介しました。

christina04.hatenablog.com

これによって処理途中のリクエストが残っていても、きちんと完了するまで待ってからプロセスを終了してくれます。

しかしPubSubのPublish()のような非同期処理であればリクエストが完了しても内部でキューが残っており、単純にそのまま終了するとメッセージがロストしてしまいます。なのでプロセス終了までにStop()を呼び出す必要があります。
またミドルウェアとのコネクションクローズといった後処理も、SDKにメソッドが用意されているのであればプロセス終了までに呼び出すべきです。

このようなプロセスが終了するまでのやっておきたい後処理を、あらかじめ登録しておき最後にまとめて呼び出す仕組みをShutdown Hookと呼びます。

JavaではRuntimeに用意されていますが、goでは用意されていないので自前で実装する必要があります。

環境

Go 1.15.5

Shutdown Hook

イメージ

処理の流れとしては以下です。

終了シグナル
↓
http.Server.Shutdown()  // 新規リクエストを受け付けなくする
↓
Shutdown Hook  // 後処理
↓
プロセス終了

具体的なコード

以下のように

後処理の関数を保持する
保持した関数を並行に呼び出す

という２つの関数を用意します。

import (
    "context"
    "sync"
)

var (
        mu    sync.Mutex
        hooks []func(context.Context)
)

// Add appends a hook function for shutdown.
func Add(h func(ctx context.Context)) {
    mu.Lock()
    defer mu.Unlock()

    hooks = append(hooks, h)
}

// Invoke invokes shutdown hooks concurrently.
func Invoke(ctx context.Context) error {
    mu.Lock()
    defer mu.Unlock()

    wg := new(sync.WaitGroup)
    wg.Add(len(hooks))
    for i := range hooks {
        go func(idx int) {
            defer wg.Done()
            hooks[idx](ctx)
        }(i)
    }
    done := make(chan struct{})
    go func() {
        wg.Wait()
        close(done)
    }()

    select {
    case <-done:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

Shutdown Hookがずっとプロセスの終了をブロックしないよう、context.Contextでtimeoutを制御できるようにしています。

使い方

Add()で終了前にやっておきたい処理を登録しておき、最後にInvoke()でまとめて呼び出します。

func main()
    Add(func(ctx context.Context) {
        fmt.Println("start hook1")
        time.Sleep(3 * time.Second)
        fmt.Println("end hook1")
    })
    Add(func(ctx context.Context) {
        fmt.Println("start hook2")
        time.Sleep(3 * time.Second)
        fmt.Println("end hook2")
    })
    Add(func(ctx context.Context) {
        fmt.Println("start hook3")
        time.Sleep(10 * time.Second)
        fmt.Println("end hook3")
    })

    // サーバ起動

    // 終了シグナルのハンドリング
    // http.Server.Shutdown()

    // Shutdown Hookによる後処理
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    Invoke(ctx)
}

Invoke()の部分はdeferで前もって登録していても良いです。

実行

サーバ終了後に並行で登録した処理が実行されます。

start hook1
start hook2
start hook3
end hook1
end hook2
    hook_test.go:30: context deadline exceeded

サンプルコード

今回のサンプルコードはこちらです。

github.com

その他

RegisterOnShutdown()ではだめなのか

似たような役割と思われるRegisterOnShutdown()というメソッドが用意されています。

しかし現状だと↓のIssueにあるように途中で処理が中断されてしまうという話なので、解決するまでは利用を避けた方が良さそうです。