github.com/grailbio/base@v0.0.11/limiter/limiter.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package limiter implements a concurrency limiter with support for
     6  // contexts.
     7  package limiter
     8  
     9  import "context"
    10  
    11  // A Limiter enforces concurrency limits among a set of goroutines.
    12  // It maintains a bucket of tokens; a number of tokens (e.g.,
    13  // representing the cost of an operation) must be acquired by a
    14  // goroutine before proceeding. A limiter is not fair: tokens are not
    15  // granted in FIFO order; rather, waiters are picked randomly to be
    16  // granted new tokens.
    17  //
    18  // A nil limiter issues an infinite number of tokens.
    19  type Limiter struct {
    20  	c      chan int
    21  	waiter chan struct{}
    22  }
    23  
    24  // New creates a new limiter with 0 tokens.
    25  func New() *Limiter {
    26  	l := &Limiter{make(chan int, 1), make(chan struct{}, 1)}
    27  	l.waiter <- struct{}{}
    28  	return l
    29  }
    30  
    31  // Acquire blocks until the goroutine is granted the desired number
    32  // of tokens, or until the context is done.
    33  func (l *Limiter) Acquire(ctx context.Context, need int) error {
    34  	if l == nil {
    35  		return ctx.Err()
    36  	}
    37  	select {
    38  	case <-l.waiter:
    39  	case <-ctx.Done():
    40  		return ctx.Err()
    41  	}
    42  	defer func() {
    43  		l.waiter <- struct{}{}
    44  	}()
    45  
    46  	var have int
    47  	for {
    48  		select {
    49  		case n := <-l.c:
    50  			have += n
    51  			if m := have - need; m >= 0 {
    52  				l.Release(m)
    53  				return nil
    54  			}
    55  		case <-ctx.Done():
    56  			l.Release(have)
    57  			return ctx.Err()
    58  		}
    59  	}
    60  }
    61  
    62  // Release adds a number of tokens back into the limiter.
    63  func (l *Limiter) Release(n int) {
    64  	if l == nil {
    65  		return
    66  	}
    67  	if n == 0 {
    68  		return
    69  	}
    70  	for {
    71  		select {
    72  		case l.c <- n:
    73  			return
    74  		case have := <-l.c:
    75  			n += have
    76  		}
    77  	}
    78  }
    79  
    80  type LimiterIfc interface {
    81  	Release(n int)
    82  	Acquire(ctx context.Context, need int) error
    83  }