57 lines
1012 B
Go
57 lines
1012 B
Go
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
|
|
}
|
|
}
|
|
}
|