frostfs-s3-gw/api/max-clients.go

57 lines
1,012 B
Go
Raw Normal View History

package api
import (
"net/http"
"time"
)
type (
MaxClients interface {
Handle(http.HandlerFunc) http.HandlerFunc
}
maxClients struct {
pool chan struct{}
timeout time.Duration
}
)
const defaultRequestDeadline = time.Second * 30
func NewMaxClientsMiddleware(count int, timeout time.Duration) MaxClients {
if timeout <= 0 {
timeout = defaultRequestDeadline
}
return &maxClients{
pool: make(chan struct{}, count),
timeout: timeout,
}
}
func (m *maxClients) Handle(f http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if m.pool == nil {
f.ServeHTTP(w, r)
return
}
deadline := time.NewTimer(m.timeout)
defer deadline.Stop()
select {
case m.pool <- struct{}{}:
defer func() { <-m.pool }()
f.ServeHTTP(w, r)
case <-deadline.C:
// Send a http timeout message
WriteErrorResponse(r.Context(), w,
errorCodes.ToAPIErr(ErrOperationMaxedOut),
r.URL)
return
case <-r.Context().Done():
return
}
}
}