Goのサーバーでcontext canceledをハンドリングする

はじめに

Goでechoを使って以下のようなサーバーを作っていました。

var db *sql.DB

func main() {
	e := echo.New()

	e.POST("/", func(c echo.Context) error {
		ctx := c.Request().Context()

		// 自身の管理するDBに保存する
		db.ExecContext(ctx, "INSERT INTO ...")

		// 別のサーバーで管理するDBに保存する
		req, err := http.NewRequestWithContext(
			ctx,
			http.MethodPost,
			"http://localhost:8081/",
			nil,
		)
		if err != nil {
			return err
		}
		res, err := http.DefaultClient.Do(req)
		...

		return c.NoContent(http.StatusOK)
	})

	log.Fatal(e.Start(":8080"))
}

省略していますが、リクエストされたデータを自身の管理するDBに保存し、関連するデータを別のサーバーで管理するDBに保存するような処理です。

このとき、自身の管理するDBへの保存が完了し、別のサーバーで管理するDBへの保存が完了する前にクライアントがリクエストを切断した場合、別のサーバーでの保存が失敗する可能性があります。その場合、2つのDBの間でデータの不整合が起きます。

		// 別のサーバーで管理するDBに保存する
		req, err := http.NewRequestWithContext(
			ctx,
			http.MethodPost,
			"http://localhost:8081/",
			nil,
		)
		if err != nil {
			return err
		}
        
		// ここで失敗した場合、データの不整合が起きる
		res, err := http.DefaultClient.Do(req)
		...

そのため、クライアントが切断した際に不整合が起きないような実装を入れました。

原因

クライアントが通信を切断すると、 ctx := c.Request().Context() のcontextはキャンセルされ、 ctx.Err() は context.Canceled のエラーを返すようになります。

そのため、例えばDBを操作するライブラリ内で ctx.Done() を適切にハンドリングしている場合、クライアントからの通信が切断された時点でDB操作をキャンセルし、エラーを返します。

この挙動は、DBに保存する処理のないGETリクエスト等の場合は有り難いものです。クライアント側が切断した時点でDBからデータを取得する必要がなくなるので、無駄な取得処理をキャンセルできます。また、DBに保存する処理のあるPOSTリクエスト等の場合でも、適切にトランザクションを貼っていれば問題にはなりません。

しかし、今回のケースではそれができず、途中で切断されると不整合が起きる可能性がありました。そのため、クライアントが通信を切断しても不整合が起きないような実装を入れました。

修正方法

以下のようなミドルウェアを追加します。


	e := echo.New()
	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			ctx := c.Request().Context()

			// タイムアウトを指定したContextを生成し、差し替える
			timeoutCtx, cancel := context.WithTimeout(
				context.Background(),
				time.Minute,
			)
			defer cancel()
			c.SetRequest(c.Request().WithContext(timeoutCtx))

			resultCh := make(chan error)
			// ハンドラーの処理を非同期で実行する
			go func() {
				resultCh <- next(c)
			}()
			var err error
			select {
			case err = <-resultCh:
			case <-ctx.Done():
				// clientが通信を切断した場合、context canceledのエラーが発生する
				if err := ctx.Err(); err != nil {
					fmt.Printf("echo context done: %v\n", err)
				}
				err = <-resultCh
			}
			return err
		}
	})

新しいcontextを作成し、c.Request() が持っているcontextを差し替えることで、クライアントが切断してもcontext canceledが伝播されなくなります。念の為、タイムアウトを設定しています。

実際のハンドラーの処理 (next 関数) を非同期で実行し、結果をチャネル (resultCh) に格納します。

select で resultCh の結果を受信するケースと ctx.Done() チャネルからの通知を受信するケースを待機します。通常の処理フローでは、resultCh から受信します。しかし、クライアントが通信を切断した場合はctx.Done()から受信します。その場合、ログにエラーメッセージを出力し、ハンドラーの処理が完了してresultCh から受信するまで待機します。

これでcontext canceledが起きても、即座にエラーにはならず、ハンドラーの処理は実行され続けます。

非同期でnextを実行する箇所は recover() しておいた方が安全です。

go func() {
	defer func() {
		if r := recover(); r != nil {
			resultCh <- fmt.Errorf("%v", r)
		}
	}()
	resultCh <- next(c)
}()

動作確認

テストとして、ハンドラーを以下のように書き換えて試してみます。

	e.POST("/", func(c echo.Context) error {
		ctx := c.Request().Context()

		time.Sleep(10 * time.Second)
		if err := ctx.Err(); err != nil {
			fmt.Printf("ctx error: %v\n", err)
			return err
		}
		fmt.Println("created")

		return c.NoContent(http.StatusOK)
	})

10秒待ってからcontextのエラーを確認し、エラーであればその時点でreturnします。エラーでなければ、ログを出して200を返します。

この状態でサーバーを起動し、リクエストを投げて10秒以内に切断します。

$ curl -XPOST http://localhost:8080
^C

上で示したミドルウェアを追加しなかった場合、以下のようなログが出てエラーになります。contextの終了を検知し、それ以降の処理は実行されません。

ctx error: context canceled

ミドルウェアを追加した場合、以下のようなログが出ます。ミドルウェア内でcontextの終了を検知してログが出ていますが、ハンドラーの処理は実行されていることがわかります。

echo context done: context canceled
created

おわりに

不整合が起きる可能性を減らすことができました。今回はechoを使ったサーバーでしたが、他のライブラリでも同じような方法で対応できると思います。