package httpserver import ( "net/http" ) // RequestCounter is an http.Handler that tracks the number of // requests in progress. type RequestCounter interface { http.Handler // Current() returns the number of requests in progress. Current() int // Max() returns the maximum number of concurrent requests // that will be accepted. Max() int } type limiterHandler struct { requests chan struct{} handler http.Handler } // NewRequestLimiter returns a RequestCounter that delegates up to // maxRequests at a time to the given handler, and responds 503 to all // incoming requests beyond that limit. func NewRequestLimiter(maxRequests int, handler http.Handler) RequestCounter { return &limiterHandler{ requests: make(chan struct{}, maxRequests), handler: handler, } } func (h *limiterHandler) Current() int { return len(h.requests) } func (h *limiterHandler) Max() int { return cap(h.requests) } func (h *limiterHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) { select { case h.requests <- struct{}{}: default: // reached max requests resp.WriteHeader(http.StatusServiceUnavailable) return } h.handler.ServeHTTP(resp, req) <-h.requests }