sdk/go/httpserver/request_limiter.go

   1 package httpserver
   2
   3 import (
   4         "net/http"
   5 )
   6
   7 type limiterHandler struct {
   8         requests chan struct{}
   9         handler  http.Handler
  10 }
  11
  12 func NewRequestLimiter(maxRequests int, handler http.Handler) http.Handler {
  13         return &limiterHandler{
  14                 requests: make(chan struct{}, maxRequests),
  15                 handler:  handler,
  16         }
  17 }
  18
  19 func (h *limiterHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  20         select {
  21         case h.requests <- struct{}{}:
  22         default:
  23                 // reached max requests
  24                 resp.WriteHeader(http.StatusServiceUnavailable)
  25                 return
  26         }
  27         h.handler.ServeHTTP(resp, req)
  28         <-h.requests
  29 }