THINKING MEGANE

今日のコードリーティング: tsenart/vegeta

HTTP load testing tool and library. It’s over 9000! - tsenart/vegeta

概要

前回読んだ rakyll/boomと同じ負荷測定ツール。

$ echo "GET http://google.com/" | vegeta attack -duration=2s | vegeta report

バージョン

2015/4/30時点のmaster (3739d8a2090cd61f4fe018e00adde281d824d3c6)

コード

vegeta attack 処理を中心に必要なところだけを抜粋

サブコマンドと引数

サブコマンド(attack, report…)ごとのコマンド構造体を定義。

type command struct {
	fs *flag.FlagSet	     // コマンドごとのオプション
	fn func(args []string) error // コマンドごとの起点となるfunction
}

標準のflag.Parseではなくてflag.FlagSet.Parseを使うことでParse対象に任意の引数を渡すことができるのでサブコマンドに必要な引数だけを渡している。便利。

(*Attacker) Attack

並行数分のworkerの立ち上げと時間あたりのリクエストの配分を制御しているところ。 処理に時間がかかるので戻り値として結果のchannelを返す。

func (a *Attacker) Attack(tr Targeter, rate uint64, du time.Duration) <-chan *Result {
	workers := &sync.WaitGroup{}
	results := make(chan *Result)
	ticks := make(chan time.Time)
	for i := uint64(0); i < a.workers; i++ {
		go a.attack(tr, workers, ticks, results)
	}

	go func() {
		for began, done := time.Now(), uint64(0); done < hits; done++ {
			select {
			case ticks <- max(next, now):
				return
			}
		}
	}()

	return results
}

workerとして実際にリクエストを行う処理。time.Tick相当のタイマーでリクエストを行う。

func (a *Attacker) attack(tr Targeter, workers *sync.WaitGroup, ticks <-chan time.Time, results chan<- *Result) {
	workers.Add(1)
	defer workers.Done()
	for tm := range ticks {
		results <- a.hit(tr, tm)
	}
}

このあたりもう少し改善の余地あると思った。

workerはタイマーを意識する必要はなくて単純にリクエストのjobキューを順番に処理していくだけにして、jobキューに投入する量をタイマーで制御するようにしたほうが見通しがよくなりそう。

attack

上記Attackの結果を受信して、gob.NewEncoderにシリアライズして出力する。

// gobパッケージを使ったエンコーダー
enc := gob.NewEncoder(out)

// 結果を受信
for {
	select {
	// Attackの結果を受信
	case r, ok := <-res:
		// 結果をシリアライズして出力
		if err = enc.Encode(r); err != nil {
			return err
		}
	}
}

gobはGoで使えるシリアライザ。送信側と受信側でフィールド名と型があっていればデシリアライズしてくれる。vegeta attack | vegeta report 間でシリアライズしたデータをやりとりしている。

雑感

  • flag.FlagSetを使ったサブコマンドのオプション実装は参考になった
  • 実装はrakyll/boomのほうが洗練されてると思う
  • 負荷測定は長くかかる処理なのでsignal.Notify(sig, os.Interrupt)を使って明示的に終了処理を呼び出すようにしているのはよいと思った
  • gob知らなかったのでRPC的な処理をやるときの候補に入れておく
このエントリーをはてなブックマークに追加