sdk/go/httpserver/request_limiter.go

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