How would I limit upload and download speed from the server in golang?

How would I limit upload and download speed from the server in golang?

I'm writing a golang server to allow users to upload and download files. And file is big, about 1GB bytes. I want to limit the upload and download speed to (for instance) 1MB/s (configurable of course).

below is my upload code:

func uploadFile(w http.ResponseWriter, r *http.Request) {
	file, _, err := r.FormFile("file")

	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	defer file.Close()

	os.MkdirAll(`e:\test`, os.ModePerm)
	out, err := os.Create(`e:\test\test.mpg`)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	defer out.Close()

	_, err = io.Copy(out, file)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}

There's a token bucket algorithm that can be helpful to implement such the rate limit. I found an example implementation, which you can use: https://github.com/juju/ratelimit

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

After running it, the output is:

Copied 1048576 bytes in 9.239607694s

You can use different bucket implementations to provide desired behaviour. In your code, after setting up right token bucket, you would call:

_, err = io.Copy(out, ratelimit.Reader(file, bucket))

You could check out the implementation of PuerkitoBio/throttled, presented in this article:

> throttled, a Go package that implements various strategies to control access to HTTP handlers.
Out-of-the-box, it supports rate-limiting of requests, constant interval flow of requests and memory usage thresholds to grant or deny access, but it also provides mechanisms to extend its functionality.

The rate limit isn't exactly what you need, but can give a good idea for implementing a similar feature.

You can use https://github.com/ConduitIO/bwlimit to limit the bandwidth of requests on the server and the client. It differs from other libraries, because it respects read/write deadlines (timeouts) and limits the bandwidth of the whole request including headers, not only the request body.

If you are interested in only limiting the upload and download speed of the file for a single HTTP handler, you can use the Reader and Writer objects provided by the library.

package example

import (
	"io"
	"net/http"

	"github.com/conduitio/bwlimit"
)

const (
	writeLimit = 1 * bwlimit.Mebibyte // write limit is 1048576 B/s
	readLimit  = 4 * bwlimit.KB       // read limit is 4000 B/s
)

func uploadFile(w http.ResponseWriter, r *http.Request) {
	file, _, _ := r.FormFile("file")

	// apply bandwidth limit
	fileReader := bwlimit.NewReader(file, readLimit)

	// prepare out ...

	// copy using the bandwidth limit
	_, _ = io.Copy(out, fileReader)
}

func downloadFile(w http.ResponseWriter, r *http.Request) {
	// prepare file ...
	in, _ := os.Open("file")

	// apply bandwidth limit
	responseWriter := bwlimit.NewWriter(w, writeLimit)

	// write body with bandwidth limit
	io.Copy(responseWriter, in)
}