Goでioの帯域制限をする shapeio - fujiwara-ware advent calendar 2024 day 9

この記事は fujiwara-ware advent calendar 2024 の9日目です。

https://github.com/fujiwara/shapeio

shapeio とは

shapeio は、Goでio.Reader/io.Writerをラップして、帯域制限をかけるためのライブラリです。

なぜ作ったのか

昨日紹介した stretcher は Pull 型デプロイツールです。デプロイ先のサーバーが自分でデプロイメントを取得する方式のため、デプロイ元のネットワークストレージサーバーのIO帯域がボトルネックになることがあります。ファイルを取得する stretcher 側で帯域制限をかけることで、デプロイ元への負荷を調整できます。

shapeio はそのために作りました。

使い方

使い方は簡単です。io.Reader/io.Writerをラップして、SetRateLimit で帯域制限を設定します。

import "github.com/fujiwara/shapeio"

func ExampleReader() {
	// example for downloading http body with rate limit.
	resp, _ := http.Get("http://example.com")
	defer resp.Body.Close()

	reader := shapeio.NewReader(resp.Body)
	reader.SetRateLimit(1024 * 10) // 10KB/sec
	io.Copy(ioutil.Discard, reader)
}

func ExampleWriter() {
	// example for writing file with rate limit.
	src := bytes.NewReader(bytes.Repeat([]byte{0}, 32*1024)) // 32KB
	f, _ := os.Create("/tmp/foo")
	writer := shapeio.NewWriter(f)
	writer.SetRateLimit(1024 * 10) // 10KB/sec
	io.Copy(writer, src)
	f.Close()
}

実装

shapeio の実装には x/rate/limit パッケージを使っています。

x/rate/limit はトークンバケットアルゴリズムを実装している、流量調整のためのパッケージです。

どのように使用しているは、次の資料を参照してください。

https://speakerdeck.com/fujiwara3/rate

まとめ

shapeio は、io.Reader/io.Writerをラップして帯域制限をかけるためのライブラリです。

Go ではネットワークに限らずストリーミングでのIO処理は io パッケージ のインターフェースを使って実装するのが一般的です。shapeio はそのインターフェースを満たしているため、ネットワークに限らず汎用的な io 処理の帯域制限に使えます。

shapeio の最初のバージョンはもう8年も前に作ったものですが、その後も Go の互換性が保たれているため、今の最新の Go でも問題なく使えます。言語自体の互換性が高いのはありがたいですね。

それでは、明日もお楽しみに！