github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/request/tpslimiter.go (about)

     1  package request
     2  
     3  import (
     4  	"context"
     5  	"golang.org/x/time/rate"
     6  	"sync"
     7  )
     8  
     9  var globalTPSLimiter = NewTPSLimiter()
    10  
    11  type TPSLimiter interface {
    12  	Limit(ctx context.Context, token string, tps float64, burst int)
    13  }
    14  
    15  func NewTPSLimiter() TPSLimiter {
    16  	return &multipleBucketLimiter{
    17  		buckets: make(map[string]*rate.Limiter),
    18  	}
    19  }
    20  
    21  // multipleBucketLimiter implements TPSLimiter with multiple bucket support.
    22  type multipleBucketLimiter struct {
    23  	mu      sync.Mutex
    24  	buckets map[string]*rate.Limiter
    25  }
    26  
    27  // Limit finds the given bucket, if bucket not exist or limit is changed,
    28  // a new bucket will be generated.
    29  func (m *multipleBucketLimiter) Limit(ctx context.Context, token string, tps float64, burst int) {
    30  	m.mu.Lock()
    31  	bucket, ok := m.buckets[token]
    32  	if !ok || float64(bucket.Limit()) != tps || bucket.Burst() != burst {
    33  		bucket = rate.NewLimiter(rate.Limit(tps), burst)
    34  		m.buckets[token] = bucket
    35  	}
    36  	m.mu.Unlock()
    37  
    38  	bucket.Wait(ctx)
    39  }