github.com/go-graphite/carbonapi@v0.17.0/limiter/limiter.go (about)

     1  package limiter
     2  
     3  import (
     4  	"context"
     5  )
     6  
     7  // ServerLimiter provides interface to limit amount of requests
     8  type RealLimiter struct {
     9  	m   map[string]chan struct{}
    10  	cap int
    11  }
    12  
    13  // NewServerLimiter creates a limiter for specific servers list.
    14  func NewServerLimiter(servers []string, l int) ServerLimiter {
    15  	if l <= 0 {
    16  		return &NoopLimiter{}
    17  	}
    18  
    19  	sl := make(map[string]chan struct{})
    20  
    21  	for _, s := range servers {
    22  		sl[s] = make(chan struct{}, l)
    23  	}
    24  
    25  	limiter := &RealLimiter{
    26  		m:   sl,
    27  		cap: l,
    28  	}
    29  	return limiter
    30  }
    31  
    32  func (sl RealLimiter) Capacity() int {
    33  	return sl.cap
    34  }
    35  
    36  // Enter claims one of free slots or blocks until there is one.
    37  func (sl RealLimiter) Enter(ctx context.Context, s string) error {
    38  	if sl.m == nil {
    39  		return nil
    40  	}
    41  
    42  	select {
    43  	case sl.m[s] <- struct{}{}:
    44  		return nil
    45  	case <-ctx.Done():
    46  		return ErrTimeout
    47  	}
    48  }
    49  
    50  // Frees a slot in limiter
    51  func (sl RealLimiter) Leave(ctx context.Context, s string) {
    52  	if sl.m == nil {
    53  		return
    54  	}
    55  
    56  	<-sl.m[s]
    57  }