github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/kv/store_limiter.go (about)

     1  package kv
     2  
     3  import (
     4  	"context"
     5  
     6  	"go.uber.org/ratelimit"
     7  )
     8  
     9  type StoreLimiter struct {
    10  	Store   Store
    11  	Limiter ratelimit.Limiter
    12  }
    13  
    14  func NewStoreLimiter(s Store, l ratelimit.Limiter) *StoreLimiter {
    15  	return &StoreLimiter{
    16  		Store:   s,
    17  		Limiter: l,
    18  	}
    19  }
    20  
    21  func (s *StoreLimiter) Get(ctx context.Context, partitionKey, key []byte) (*ValueWithPredicate, error) {
    22  	_ = s.Limiter.Take()
    23  	return s.Store.Get(ctx, partitionKey, key)
    24  }
    25  
    26  func (s *StoreLimiter) Set(ctx context.Context, partitionKey, key, value []byte) error {
    27  	_ = s.Limiter.Take()
    28  	return s.Store.Set(ctx, partitionKey, key, value)
    29  }
    30  
    31  func (s *StoreLimiter) SetIf(ctx context.Context, partitionKey, key, value []byte, valuePredicate Predicate) error {
    32  	_ = s.Limiter.Take()
    33  	return s.Store.SetIf(ctx, partitionKey, key, value, valuePredicate)
    34  }
    35  
    36  func (s *StoreLimiter) Delete(ctx context.Context, partitionKey, key []byte) error {
    37  	_ = s.Limiter.Take()
    38  	return s.Store.Delete(ctx, partitionKey, key)
    39  }
    40  
    41  func (s *StoreLimiter) Scan(ctx context.Context, partitionKey []byte, options ScanOptions) (EntriesIterator, error) {
    42  	_ = s.Limiter.Take()
    43  	return s.Store.Scan(ctx, partitionKey, options)
    44  }
    45  
    46  func (s *StoreLimiter) Close() {
    47  	s.Store.Close()
    48  }