|
- package ratelimit
-
- import "net/http"
-
- // Throttle is a middleware that limits number of currently
- // processed requests at a time.
- func Throttle(limit int) func(http.Handler) http.Handler {
- if limit <= 0 {
- panic("Throttle expects limit > 0")
- }
-
- t := throttler{
- tokens: make(chan token, limit),
- }
- for i := 0; i < limit; i++ {
- t.tokens <- token{}
- }
-
- fn := func(h http.Handler) http.Handler {
- t.h = h
- return &t
- }
-
- return fn
- }
-
- // token represents a request that is being processed.
- type token struct{}
-
- // throttler limits number of currently processed requests at a time.
- type throttler struct {
- h http.Handler
- tokens chan token
- }
-
- // ServeHTTP implements http.Handler interface.
- func (t *throttler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
- select {
- case <-r.Context().Done():
- return
- case tok := <-t.tokens:
- defer func() {
- t.tokens <- tok
- }()
- t.h.ServeHTTP(w, r)
- }
- }
|