sdk/go/httpserver/request_limiter.go

   1 package httpserver
   2
   3 import (
   4         "net/http"
   5 )
   6
   7 // RequestCounter is an http.Handler that tracks the number of
   8 // requests in progress.
   9 type RequestCounter interface {
  10         http.Handler
  11
  12         // Current() returns the number of requests in progress.
  13         Current() int
  14
  15         // Max() returns the maximum number of concurrent requests
  16         // that will be accepted.
  17         Max() int
  18 }
  19
  20 type limiterHandler struct {
  21         requests chan struct{}
  22         handler  http.Handler
  23 }
  24
  25 // NewRequestLimiter returns a RequestCounter that delegates up to
  26 // maxRequests at a time to the given handler, and responds 503 to all
  27 // incoming requests beyond that limit.
  28 func NewRequestLimiter(maxRequests int, handler http.Handler) RequestCounter {
  29         return &limiterHandler{
  30                 requests: make(chan struct{}, maxRequests),
  31                 handler:  handler,
  32         }
  33 }
  34
  35 func (h *limiterHandler) Current() int {
  36         return len(h.requests)
  37 }
  38
  39 func (h *limiterHandler) Max() int {
  40         return cap(h.requests)
  41 }
  42
  43 func (h *limiterHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  44         select {
  45         case h.requests <- struct{}{}:
  46         default:
  47                 // reached max requests
  48                 resp.WriteHeader(http.StatusServiceUnavailable)
  49                 return
  50         }
  51         h.handler.ServeHTTP(resp, req)
  52         <-h.requests
  53 }