github.com/chenbh/concourse/v6@v6.4.2/atc/wrappa/concurrent_request_limits_wrappa.go (about)

     1  package wrappa
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"code.cloudfoundry.org/lager"
     7  	"github.com/chenbh/concourse/v6/atc/metric"
     8  	"github.com/tedsuo/rata"
     9  )
    10  
    11  type ConcurrentRequestLimitsWrappa struct {
    12  	logger                  lager.Logger
    13  	concurrentRequestPolicy ConcurrentRequestPolicy
    14  }
    15  
    16  func NewConcurrentRequestLimitsWrappa(
    17  	logger lager.Logger,
    18  	concurrentRequestPolicy ConcurrentRequestPolicy,
    19  ) Wrappa {
    20  	return ConcurrentRequestLimitsWrappa{
    21  		logger:                  logger,
    22  		concurrentRequestPolicy: concurrentRequestPolicy,
    23  	}
    24  }
    25  
    26  func (wrappa ConcurrentRequestLimitsWrappa) Wrap(
    27  	handlers rata.Handlers,
    28  ) rata.Handlers {
    29  	wrapped := rata.Handlers{}
    30  
    31  	for action, handler := range handlers {
    32  		pool, found := wrappa.concurrentRequestPolicy.HandlerPool(action)
    33  		if found {
    34  			inflight := &metric.Gauge{}
    35  			limitHit := &metric.Counter{}
    36  
    37  			metric.ConcurrentRequests[action] = inflight
    38  			metric.ConcurrentRequestsLimitHit[action] = limitHit
    39  
    40  			wrapped[action] = wrappa.wrap(
    41  				pool,
    42  				handler,
    43  				inflight, limitHit,
    44  			)
    45  		} else {
    46  			wrapped[action] = handler
    47  		}
    48  	}
    49  
    50  	return wrapped
    51  }
    52  
    53  func (wrappa ConcurrentRequestLimitsWrappa) wrap(
    54  	pool Pool,
    55  	handler http.Handler,
    56  	inflight *metric.Gauge, limitHit *metric.Counter,
    57  ) http.Handler {
    58  	if pool.Size() == 0 {
    59  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    60  			wrappa.logger.Debug("endpoint-disabled")
    61  			limitHit.Inc()
    62  
    63  			w.WriteHeader(http.StatusNotImplemented)
    64  		})
    65  	}
    66  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    67  		if !pool.TryAcquire() {
    68  			wrappa.logger.Info("concurrent-request-limit-reached")
    69  			limitHit.Inc()
    70  
    71  			w.WriteHeader(http.StatusServiceUnavailable)
    72  			return
    73  		}
    74  
    75  		defer inflight.Dec()
    76  		defer pool.Release()
    77  
    78  		inflight.Inc()
    79  		handler.ServeHTTP(w, r)
    80  	})
    81  }