Skip to content

Instantly share code, notes, and snippets.

@linnv
Last active June 30, 2021 09:34
Show Gist options
  • Save linnv/1919c6b10f26e155590703508e711006 to your computer and use it in GitHub Desktop.
Save linnv/1919c6b10f26e155590703508e711006 to your computer and use it in GitHub Desktop.
HTTP QPS rate limiting in Golang
package main
import (
"flag"
"fmt"
"log"
"net/http"
"os"
"os/signal"
"time"
)
type midHTTP struct {
}
const burstLimit = 100
var throttle = make(chan struct{}, burstLimit)
func (m *midHTTP) ServeHTTP(w http.ResponseWriter, r *http.Request) {
select {
case <-throttle:
http.DefaultServeMux.ServeHTTP(w, r)
default:
w.WriteHeader(http.StatusTooManyRequests)
return
}
}
func main() {
// 100 qps
rate := time.Second / 100
tick := time.NewTicker(rate)
defer tick.Stop()
go func() {
for _ = range tick.C {
select {
case throttle <- struct{}{}:
default:
}
} // does not exit after tick.Stop()
}()
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
fmt.Fprintln(w, "good")
})
flagPort := flag.String("p", "9091", "listen port")
flag.Parse()
mid := &midHTTP{}
port := ":" + *flagPort
server := http.Server{
Addr: port,
Handler: mid,
}
go func() {
if err := server.ListenAndServe(); err != nil {
if err == http.ErrServerClosed {
} else {
panic(err)
}
}
}()
sigChan := make(chan os.Signal, 2)
signal.Notify(sigChan, os.Interrupt, os.Kill)
log.Print("http server listen on port ", port)
log.Print("use c-c to exit: \n")
<-sigChan
server.Shutdown(nil)
}
// curl http://127.0.0.1:9091
// wrk -c 10 -d 10 -t 2 http://127.0.0.1:9091
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment