はじめに

特定のバッチの特定の区間で行われる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,
}