gin-gonicのタイムアウト用ハンドラを自前実装する

【㈱フクロウラボ】エンジニアブログ

経緯


こんにちは。バックエンドエンジニアの田村です。Go言語を主に書きます。

gin-gonicで書かれたサーバで、タイムアウト処理を書く必要ができたので、何かいい方法がないか探していたところ、こちらの記事/動画に出会いました。
https://dev.to/jacobsngoodwin/13-gin-handler-timeout-middleware-4bhg

2020年の記事なので少し古いですが、面白かった＆大いに参考にさせていただいたので、導入部分を軽くまとめると、

  • 記事の筆者がやりたいこと

  • 記事の中で候補に挙がったライブラリのmiddlewareと、それぞれの問題点

    • http.TimeoutHandler
      問題点：
      ginのRouterやGroupに直接すぐに設定することができない。
      田村注） gin.WrapF()で使えないことはないかもしれないです。
      デフォルトでHTMLを送るようになっており、他のcontent-typeを送ることができない。

    • gin-contrib/timeout
      問題点：
      タイムアウト後にレスポンスヘッダを上書きしようとすると、panicになる
      田村注）これは2023年10月時点では見られない挙動なので、修正された可能性が高い。全然使えそう。

    • vearne/gin-timeout
      問題点：
      エラー時のレスポンスをカスタマイズすることができない。
      田村注）こちらも2023年10月時点では設定できる。便利そう。

これらの理由で、この方は自分でタイムアウト用のmiddlewareを実装することにしたそうです。そしてソースはこちらに公開されている通りです。

私の実装と解説

私は、実装前に自分で使用感を確認せずに、これらのライブラリが自分のニーズに合わないと勝手に思い込んでいたのでこのソースコードをさらにカスタマイズして実装しようと思い立ちました。私が変更しようと思った点は以下です。

そうしてできたソースがこちらです。コメントアウトに日本語で解説を書きましたのでご参考までに。

func NewTimeoutHandler(timeout time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		// ginのwriterをカスタムのwriterに置き換える
		tw := &timeoutWriter{ResponseWriter: c.Writer, h: make(http.Header)}
		c.Writer = tw

		// request contextをタイムアウト付きのcontextでラップする
		ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
		defer cancel()

		// ラップしたcontextをginのrequest contextに再度セットする
		c.Request = c.Request.WithContext(ctx)

		finished := make(chan struct{})

		// 別スレッドで後段のmiddlewareとhandlerを実行する
		go func() {
			c.Next()
			finished <- struct{}{}
		}()

		select {
		case <-finished:
			// 通常の終了時は、他のmiddleware/handlerが書き込んだheaderとbodyをそのまま書き込む
			tw.mu.Lock()
			defer tw.mu.Unlock()

			// 後段のmiddleware/handlerが tw.h に書き込んだヘッダを tw.ResponseWriter （c.Writer） にコピーする
			dst := tw.ResponseWriter.Header()
			for k, vv := range tw.Header() {
				dst[k] = vv
			}
			// 後段のmiddleware/handlerが　tw.WriteHeader()で書き込んだstatus codeをtw.ResponseWriterにコピーする
			tw.ResponseWriter.WriteHeader(tw.code)

			// 後段のmiddleware/handlerが　tw.Write() で書き込んだ tw.wbuf を tw.ResponseWriter にコピーする
			tw.ResponseWriter.Write(tw.wbuf.Bytes())

		case <-ctx.Done():
			// タイムアウト時は、それがわかるエラーを返す
			tw.mu.Lock()
			defer tw.mu.Unlock()
			tw.ResponseWriter.Header().Set("Content-Type", "text/plain")
			tw.ResponseWriter.WriteHeader(http.StatusGatewayTimeout)
			tw.ResponseWriter.Write([]byte(`server timed out`))
			tw.SetTimedOut()

			// 自前のslack通知の関数を呼び出す
			errMsg := "timeout!"
			if err := slack.NewNotifier().Notify(errMsg); err != nil {
				fmt.Println("slack error")
			}
			
		}

	}
}

// http.Writer interfaceを実装する
// 通常の処理に加えて、タイムアウトが発生したか、すでにheaderが書き込まれているかを監視する
// マルチスレッド（このmiddlewareと次に呼び出されるhandler）でフィールドにアクセスするため、競合を防ぐために、mutexを使用する
// gin.ResponseWriterをラップして、レスポンスの書き込みをコントロールする
type timeoutWriter struct {
	gin.ResponseWriter
	h    http.Header
	wbuf bytes.Buffer

	mu          sync.Mutex
	timedOut    bool
	wroteHeader bool
	code        int
}

// レスポンスをwriteするが、最初にタイムアウトが発生していないか確認する
// （別ののmiddlewareとhandlerも使用するため、上書きを避ける）
// http.ResponseWriter interfaceを実装する
func (tw *timeoutWriter) Write(b []byte) (int, error) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
		return 0, nil
	}

	return tw.wbuf.Write(b)
}

// http.ResponseWriter interfaceを実装する
func (tw *timeoutWriter) WriteHeader(code int) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	// timeoutの場合やすでにheaderが存在している場合は何もしない
	if tw.timedOut || tw.wroteHeader {
		return
	}
	tw.writeHeader(code)
}

// ヘッダが書き込まれたかどうかを示すフラグをセットする
func (tw *timeoutWriter) writeHeader(code int) {
	tw.wroteHeader = true
	tw.code = code
}

// http.ResponseWriter interfaceを実装する
func (tw *timeoutWriter) Header() http.Header {
	return tw.h
}

// すでにタイムアウトの処理を行ったことを示すフラグをセットする
func (tw *timeoutWriter) SetTimedOut() {
	tw.timedOut = true
}

さいごに

ここまで実装して、ではhttp.TimeoutHandlerのソースコードはどうなっているんだろうと見ていたら、仕組みがかなり似ていることに気づきました。
Goの標準ライブラリの実装をよく読むことで面白いコードがもっと書けるようになりそうですね。


この記事が気に入ったらサポートをしてみませんか？