Limit upload speed in Golang

11月 20, 2017 tech post

When 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)