github.com/hellobchain/third_party@v0.0.0-20230331131523-deb0478a2e52/go-chi/chi/middleware/throttle.go (about)

     1  package middleware
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/hellobchain/newcryptosm/http"
     7  )
     8  
     9  const (
    10  	errCapacityExceeded = "Server capacity exceeded."
    11  	errTimedOut         = "Timed out while waiting for a pending request to complete."
    12  	errContextCanceled  = "Context was canceled."
    13  )
    14  
    15  var (
    16  	defaultBacklogTimeout = time.Second * 60
    17  )
    18  
    19  // Throttle is a middleware that limits number of currently processed requests
    20  // at a time.
    21  func Throttle(limit int) func(http.Handler) http.Handler {
    22  	return ThrottleBacklog(limit, 0, defaultBacklogTimeout)
    23  }
    24  
    25  // ThrottleBacklog is a middleware that limits number of currently processed
    26  // requests at a time and provides a backlog for holding a finite number of
    27  // pending requests.
    28  func ThrottleBacklog(limit int, backlogLimit int, backlogTimeout time.Duration) func(http.Handler) http.Handler {
    29  	if limit < 1 {
    30  		panic("chi/middleware: Throttle expects limit > 0")
    31  	}
    32  
    33  	if backlogLimit < 0 {
    34  		panic("chi/middleware: Throttle expects backlogLimit to be positive")
    35  	}
    36  
    37  	t := throttler{
    38  		tokens:         make(chan token, limit),
    39  		backlogTokens:  make(chan token, limit+backlogLimit),
    40  		backlogTimeout: backlogTimeout,
    41  	}
    42  
    43  	// Filling tokens.
    44  	for i := 0; i < limit+backlogLimit; i++ {
    45  		if i < limit {
    46  			t.tokens <- token{}
    47  		}
    48  		t.backlogTokens <- token{}
    49  	}
    50  
    51  	fn := func(h http.Handler) http.Handler {
    52  		t.h = h
    53  		return &t
    54  	}
    55  
    56  	return fn
    57  }
    58  
    59  // token represents a request that is being processed.
    60  type token struct{}
    61  
    62  // throttler limits number of currently processed requests at a time.
    63  type throttler struct {
    64  	h              http.Handler
    65  	tokens         chan token
    66  	backlogTokens  chan token
    67  	backlogTimeout time.Duration
    68  }
    69  
    70  // ServeHTTP is the primary throttler request handler
    71  func (t *throttler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    72  	ctx := r.Context()
    73  	select {
    74  	case <-ctx.Done():
    75  		http.Error(w, errContextCanceled, http.StatusServiceUnavailable)
    76  		return
    77  	case btok := <-t.backlogTokens:
    78  		timer := time.NewTimer(t.backlogTimeout)
    79  
    80  		defer func() {
    81  			t.backlogTokens <- btok
    82  		}()
    83  
    84  		select {
    85  		case <-timer.C:
    86  			http.Error(w, errTimedOut, http.StatusServiceUnavailable)
    87  			return
    88  		case <-ctx.Done():
    89  			http.Error(w, errContextCanceled, http.StatusServiceUnavailable)
    90  			return
    91  		case tok := <-t.tokens:
    92  			defer func() {
    93  				t.tokens <- tok
    94  			}()
    95  			t.h.ServeHTTP(w, r)
    96  		}
    97  		return
    98  	default:
    99  		http.Error(w, errCapacityExceeded, http.StatusServiceUnavailable)
   100  		return
   101  	}
   102  }