github.com/viant/toolbox@v0.34.5/batch_limiter.go (about)

     1  package toolbox
     2  
     3  import "sync"
     4  
     5  //BatchLimiter represents a batch limiter
     6  type BatchLimiter struct {
     7  	queue chan uint8
     8  	group *sync.WaitGroup
     9  	Mutex *sync.RWMutex
    10  }
    11  
    12  //Acquire takes token form a channel, or wait if  no more elements in a a channel
    13  func (r *BatchLimiter) Acquire() {
    14  	<-r.queue
    15  }
    16  
    17  //Add adds element to wait group
    18  func (r *BatchLimiter) Add(delta int) {
    19  	r.group.Add(delta)
    20  }
    21  
    22  //Done flags wait group as done, returns back a token to a channel
    23  func (r *BatchLimiter) Done() {
    24  	r.group.Done()
    25  	r.queue <- uint8(1)
    26  }
    27  
    28  //Wait wait on wait group
    29  func (r *BatchLimiter) Wait() {
    30  	r.group.Wait()
    31  }
    32  
    33  //NewBatchLimiter creates a new batch limiter with batch size and total number of elements
    34  func NewBatchLimiter(batchSize, total int) *BatchLimiter {
    35  	var queue = make(chan uint8, batchSize)
    36  	for i := 0; i < batchSize; i++ {
    37  		queue <- uint8(1)
    38  	}
    39  	result := &BatchLimiter{
    40  		queue: queue,
    41  		group: &sync.WaitGroup{},
    42  		Mutex: &sync.RWMutex{},
    43  	}
    44  
    45  	result.group.Add(total)
    46  	return result
    47  }