Last active
June 30, 2021 09:34
-
-
Save linnv/1919c6b10f26e155590703508e711006 to your computer and use it in GitHub Desktop.
HTTP QPS rate limiting in Golang
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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