Limit upload speed in Golang
11月 20, 2017 tech postWhen I am writing the wormhole (a tool to replicate docker images from one registry to another), I realize that it is dangerous without a rate limit.
So I found an implementation with juju/ratelimit which use the Token Bucket algorithm on StackOverflow:
package main
import (
"bytes"
"fmt"
"io"
"time"
"github.com/juju/ratelimit"
)
func main() {
// Source holding 1MB
src := bytes.NewReader(make([]byte, 1024*1024))
// Destination
dst := &bytes.Buffer{}
// Bucket adding 100KB every second, holding max 100KB
bucket := ratelimit.NewBucketWithRate(100*1024, 100*1024)
start := time.Now()
// Copy source to destination, but wrap our reader with rate limited one
io.Copy(dst, ratelimit.Reader(src, bucket))
fmt.Printf("Copied %d bytes in %s\n",dst.Len(),time.Since(start))
}
Yeah, It works for me:
// rl *float64
// if I want to limit speed with 5M , then rl = 5000 * 1024
b := ratelimit.NewBucketWithRate(*rl, int64((*rl)*1.2))
limitReader := ratelimit.Reader(reader, b)
err = t.Dst.UploadLayer(repository, digest, limitReader)