github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/p2p/middleware.go (about)

     1  package p2p
     2  
     3  import (
     4  	"sync/atomic"
     5  
     6  	logging "github.com/ipfs/go-log/v2"
     7  	"github.com/libp2p/go-libp2p/core/network"
     8  )
     9  
    10  var log = logging.Logger("shrex/middleware")
    11  
    12  type Middleware struct {
    13  	// concurrencyLimit is the maximum number of requests that can be processed at once.
    14  	concurrencyLimit int64
    15  	// parallelRequests is the number of requests currently being processed.
    16  	parallelRequests atomic.Int64
    17  	// numRateLimited is the number of requests that were rate limited.
    18  	numRateLimited atomic.Int64
    19  }
    20  
    21  func NewMiddleware(concurrencyLimit int) *Middleware {
    22  	return &Middleware{
    23  		concurrencyLimit: int64(concurrencyLimit),
    24  	}
    25  }
    26  
    27  // DrainCounter returns the current value of the rate limit counter and resets it to 0.
    28  func (m *Middleware) DrainCounter() int64 {
    29  	return m.numRateLimited.Swap(0)
    30  }
    31  
    32  func (m *Middleware) RateLimitHandler(handler network.StreamHandler) network.StreamHandler {
    33  	return func(stream network.Stream) {
    34  		current := m.parallelRequests.Add(1)
    35  		defer m.parallelRequests.Add(-1)
    36  
    37  		if current > m.concurrencyLimit {
    38  			m.numRateLimited.Add(1)
    39  			log.Debug("concurrency limit reached")
    40  			err := stream.Close()
    41  			if err != nil {
    42  				log.Debugw("server: closing stream", "err", err)
    43  			}
    44  			return
    45  		}
    46  		handler(stream)
    47  	}
    48  }