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         "sync/atomic"
  10
  11         "github.com/prometheus/client_golang/prometheus"
  12 )
  13
  14 // RequestCounter is an http.Handler that tracks the number of
  15 // requests in progress.
  16 type RequestCounter interface {
  17         http.Handler
  18
  19         // Current() returns the number of requests in progress.
  20         Current() int
  21
  22         // Max() returns the maximum number of concurrent requests
  23         // that will be accepted.
  24         Max() int
  25 }
  26
  27 type limiterHandler struct {
  28         requests chan struct{}
  29         handler  http.Handler
  30         count    int64 // only used if cap(requests)==0
  31 }
  32
  33 // NewRequestLimiter returns a RequestCounter that delegates up to
  34 // maxRequests at a time to the given handler, and responds 503 to all
  35 // incoming requests beyond that limit.
  36 //
  37 // "concurrent_requests" and "max_concurrent_requests" metrics are
  38 // registered with the given reg, if reg is not nil.
  39 func NewRequestLimiter(maxRequests int, handler http.Handler, reg *prometheus.Registry) RequestCounter {
  40         h := &limiterHandler{
  41                 requests: make(chan struct{}, maxRequests),
  42                 handler:  handler,
  43         }
  44         if reg != nil {
  45                 reg.MustRegister(prometheus.NewGaugeFunc(
  46                         prometheus.GaugeOpts{
  47                                 Namespace: "arvados",
  48                                 Name:      "concurrent_requests",
  49                                 Help:      "Number of requests in progress",
  50                         },
  51                         func() float64 { return float64(h.Current()) },
  52                 ))
  53                 reg.MustRegister(prometheus.NewGaugeFunc(
  54                         prometheus.GaugeOpts{
  55                                 Namespace: "arvados",
  56                                 Name:      "max_concurrent_requests",
  57                                 Help:      "Maximum number of concurrent requests",
  58                         },
  59                         func() float64 { return float64(h.Max()) },
  60                 ))
  61         }
  62         return h
  63 }
  64
  65 func (h *limiterHandler) Current() int {
  66         if cap(h.requests) == 0 {
  67                 return int(atomic.LoadInt64(&h.count))
  68         }
  69         return len(h.requests)
  70 }
  71
  72 func (h *limiterHandler) Max() int {
  73         return cap(h.requests)
  74 }
  75
  76 func (h *limiterHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  77         if cap(h.requests) == 0 {
  78                 atomic.AddInt64(&h.count, 1)
  79                 defer atomic.AddInt64(&h.count, -1)
  80                 h.handler.ServeHTTP(resp, req)
  81                 return
  82         }
  83         select {
  84         case h.requests <- struct{}{}:
  85         default:
  86                 // reached max requests
  87                 resp.WriteHeader(http.StatusServiceUnavailable)
  88                 return
  89         }
  90         h.handler.ServeHTTP(resp, req)
  91         <-h.requests
  92 }