goでhttp clientを利用したrequest数をカウント(transportとcontextを使ってごまかす)

はじめに

特定のバッチの特定の区間で行われるreuest数をカウントしたいということがあった。
内部的には、paginationされた処理を繰り返しrequestし辿る部分などもあり、明示的にカウントする処理を入れていくのは大変そうだった。

方針として以下２つなどが考えられた

• proxyを別途立てて、何らかのmarkerと一緒にrequestする(proxy側でカウント)
• contextにrequest数のカウンターを詰めて、カウントする処理を行うtransportを挟む

今回は後者の話。

transport

httpClientのTransportに http.RoundTripper のinterfaceを実装した値を渡すとmiddlewareのようにして使える。例えば以下の様なstructの定義を行うことでrequest数のカウントが実現できる。

// Transport :
type Transport struct {
    Base http.RoundTripper
}

// RoundTrip :
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
    if sc := GetCollectorFromContext(req.Context()); sc != nil {
        sc.Collect(req)
    }
    return t.Base.RoundTrip(req)
}

// GetCollectorFromContext :
func GetCollectorFromContext(ctx context.Context) *StatCollector {
    switch sc := ctx.Value(K).(type) {
    case *StatCollector:
        return sc
    default:
        return nil
    }
}

// K :
var K struct{}

名前はてきとう。requestのcontextからcollector(request数カウンター)を取り出してカウントしている。

collector

collectorの定義は以下のようなもの。柔軟にしたければinterfaceにするべきかもしれないけれど。特に今回はそこまでする必要がなかったのでただのstructの定義。

// StatCollector :
type StatCollector struct {
    Stat map[string]int
}

// NewStatCollector :
func NewStatCollector() *StatCollector {
    return &StatCollector{Stat: map[string]int{}}
}

// Collect :
func (sc *StatCollector) Collect(req *http.Request) {
    sc.Stat[req.Host]++
}

真面目にやるならlockなどを追加する必要もあるかもしれない。

client

http.ClientにはTransportという引数を取る事ができる。ここに作ったTransportを渡す。

func tick(ctx context.Context, client *http.Client, url string) error {
    req, err := http.NewRequest("GET", url, nil)
    if err != nil {
        return err
    }
    req = req.WithContext(ctx)
    response, err := client.Do(req)
    if err != nil {
        return err
    }
    defer response.Body.Close()
    io.Copy(os.Stdout, response.Body)
    return nil
}

func main() {
    client := &http.Client{Transport: &Transport{Base: http.DefaultTransport}}
    sc := NewStatCollector()
    ctx := context.Background()
    ctx = context.WithValue(ctx, K, sc)
    if err := tick(ctx, client, "http://localhost:54321/foo"); err != nil {
        log.Fatal(err)
    }
    if err := tick(ctx, client, "http://localhost:54321/foo"); err != nil {
        log.Fatal(err)
    }
    if err := tick(ctx, client, "http://localhost:54321/bar"); err != nil {
        log.Fatal(err)
    }
    pp.Print(sc.Stat)
}

リクエスト数のカウントができた。

{
  "localhost:54321": 3,
}