GoでHTTP/gRPCリクエストのレスポンスサイズを計測する方法
はじめに
この記事では HTTP のミドルウェアと gRPC のインターセプターでレスポンスサイズを取得する方法を紹介します。単にレスポンスサイズを取得するというよりは自前で計測する的な解決方法でしたのでこのような記事タイトルとなっています。
レスポンスサイズが必要になる場面はそんなに無いかもしれませんが、本記事で紹介する http.ResponseWriter や grpc.ServerStream のラッパーからデータを取得する方法はいろんな場面で応用できると思います。ちなみに私は Go のロガーライブラリの実装で活用しています。
HTTPリクエストのレスポンスサイズを計測する
HTTP リクエストの場合、リクエストのサイズはヘッダの Content-Length から簡単に取得できますがレスポンスサイズは Go 標準ライブラリの http.ResponseWriter にサイズを取得できるメソッドが存在しないため自前で計測する必要があります。まずは以下の http.ResponseWriter のラッパー構造体を用意します。
// http.ResponseWriter のラッパー構造体
type ResponseWriter struct {
http.ResponseWriter
size uint64
}
// http.ResponseWriterのWriteメソッドをラップする
func (rw *ResponseWriter) Write(buf []byte) (int, error) {
// 書き込み -> データサイズが戻ってくる
n, err := rw.ResponseWriter.Write(buf)
// 書き込んだサイズを足し合わせる
atomic.AddUint64(&rw.size, uint64(n))
return n, err
}
// レスポンスサイズを返すメソッドを新たに定義する
func (rw *ResponseWriter) Size() int {
return int(rw.size)
}
あとは以下のように、ミドルウェアで http.ResponseWriter のインスタンスを上記構造体でラップして ServeHTTP 関数に渡してあげます。これでレスポンスサイズの取得ができるようになります。
func FooMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// ラッパー構造体のインスタンスを生成する
rw := &ResponseWriter{
ResponseWriter: w,
}
defer func() {
// レスポンスサイズを取得する
fmt.Printf("response size: %d", rw.Size())
}()
// w ではなく rw を渡す
next.ServeHTTP(rw, r)
})
}
ちなみに Gin や Echo といったWebフレームワークを使っている場合は、レスポンスサイズは簡単に取得できます。
gRPCリクエストのレスポンスサイズを計測する
gRPC の場合、HTTP ヘッダの Content-Length 的なものが無いのでリクエストサイズの取得すらできません。なのでインターセプターでリクエストとレスポンスのオブジェクトを取得しサイズを計測します。まずは以下のようなオブジェクトのサイズ取得関数を用意します。サイズの取得には encoding/gob パッケージと encoding/binary パッケージを利用します。
import (
"bytes"
"encoding/binary"
"encoding/gob"
)
// オブジェクトのサイズ取得関数
func binarySize(val interface{}) int {
var buff bytes.Buffer
enc := gob.NewEncoder(&buff)
err := enc.Encode(val)
if err != nil {
return 0
}
return binary.Size(buff.Bytes())
}
この関数をインターセプターで使うことによりリクエストサイズとレスポンスサイズを計測することができます。Unary と Server streaming で実装方法が違うのでそれぞれ説明します。
Unaryの場合
Unary でリクエストサイズとレスポンスサイズを取得する場合はインターセプターの引数に渡ってくるリクエストオブジェクトとハンドラーを実行した結果からレスポンスを取得してそれぞれのオブジェクトのサイズを先ほどの binarySize 関数を使って取得します(Unaryの場合は計測というよりは取得と表現した方がしっくりきます)。
func UnaryServerInterceptor(
ctx context.Context,
req interface{},
info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler,
) (interface{}, error) {
//
var res interface{}
defer func() {
// リクエストサイズを取得する
reqSize := binarySize(req)
// レスポンスサイズを取得する
resSize := binarySize(res)
fmt.Printf("request size: %d, response size: %d", reqSize, resSize)
}()
res, err = handler(ctx, req)
return res, err
}
Server streamingの場合
Unary と違い Server streaming はリクエストオブジェクトやレスポンスオブジェクトを直接取得できないため以下のような grpc.ServerStream をラップする構造体を用意します。
// grpc.ServerStreamのラッパー構造体
type serverStream struct {
grpc.ServerStream
requestSize uint64
responseSize uint64
}
// クライアント側にデータを送信するメソッド
func (s *serverStream) SendMsg(m interface{}) error {
err := s.ServerStream.SendMsg(m)
if err == nil {
// 送信したデータのサイズを足し合わせる
atomic.AddUint64(&s.responseSize, uint64(binarySize(m)))
}
return err
}
// クライアントからのリクエストを受信するメソッド
func (s *serverStream) RecvMsg(m interface{}) error {
err := s.ServerStream.RecvMsg(m)
if err == nil {
// 受診したデータのサイズを足し合わせる
atomic.AddUint64(&s.requestSize, uint64(binarySize(m)))
}
return err
}
このラッパー構造体をインターセプターで使うことによりレスポンスサイズを取得することができます。
func StreamServerInterceptor(
srv interface{},
stream grpc.ServerStream,
info *grpc.StreamServerInfo,
handler grpc.StreamHandler,
) error {
// ラッパー構造体のインスタンスを生成する
wrapped := &serverStream{
ServerStream: stream,
}
defer func() {
// リクエストサイズを取得する
reqSize := wrapped.requestSize
// レスポンスサイズを取得する
resSize := wrapped.responseSize
fmt.Printf("request size: %d, response size: %d", reqSize, resSize)
}()
// stream ではなく wrapped を渡す
return handler(srv, wrapped)
}
まとめ
HTTP と gRPC のリクエストサイズとレスポンスサイズの取得方法を表にまとめると以下のようになります。
| リクエストサイズ | レスポンスサイズ | |
|---|---|---|
| HTTP | リクエストヘッダのContent-Length | http.ResponseWriterのラッパー |
| gRPC(Unary) | リクエストオブジェクトのサイズ | handlerの実行結果のサイズ |
| gRPC(Server streaming) | grpc.ServerStreamのラッパー | grpc.ServerStreamのラッパー |
開発中のライブラリのPR
この記事で紹介したコードは以下のライブラリで実際に使用しています。この記事を読んで興味をもたれた方はよかったらこちらのライブラリのコードをチェックしてみてください。
Discussion