github.com/divyam234/rclone@v1.64.1/fs/accounting/tpslimit.go (about)

     1  package accounting
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/divyam234/rclone/fs"
     7  	"golang.org/x/time/rate"
     8  )
     9  
    10  var (
    11  	tpsBucket *rate.Limiter // for limiting number of http transactions per second
    12  )
    13  
    14  // StartLimitTPS starts the token bucket for transactions per second
    15  // limiting if necessary
    16  func StartLimitTPS(ctx context.Context) {
    17  	ci := fs.GetConfig(ctx)
    18  	if ci.TPSLimit > 0 {
    19  		tpsBurst := ci.TPSLimitBurst
    20  		if tpsBurst < 1 {
    21  			tpsBurst = 1
    22  		}
    23  		tpsBucket = rate.NewLimiter(rate.Limit(ci.TPSLimit), tpsBurst)
    24  		fs.Infof(nil, "Starting transaction limiter: max %g transactions/s with burst %d", ci.TPSLimit, tpsBurst)
    25  	}
    26  }
    27  
    28  // LimitTPS limits the number of transactions per second if enabled.
    29  // It should be called once per transaction.
    30  func LimitTPS(ctx context.Context) {
    31  	if tpsBucket != nil {
    32  		tbErr := tpsBucket.Wait(ctx)
    33  		if tbErr != nil && tbErr != context.Canceled {
    34  			fs.Errorf(nil, "HTTP token bucket error: %v", tbErr)
    35  		}
    36  	}
    37  }