github.com/nebulouslabs/sia@v1.3.7/sync/limiter.go (about)

     1  package sync
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // A Limiter restricts access to a resource.
     8  //
     9  // Units of the resource are reserved via Request and returned via Release.
    10  // Conventionally, a caller who reserves n units is responsible for ensuring
    11  // that all n are eventually returned. Once the number of reserved units
    12  // exceeds the Limiters limit, further calls to Request will block until
    13  // sufficient units are returned via Release.
    14  //
    15  // This Limiter differs from others in that it allows requesting more than the
    16  // limit. This request is only fulfilled once all other units have been
    17  // returned. Once the request is fulfilled, calls to Request will block until
    18  // enough units have been returned to bring the total outlay below the limit.
    19  // This design choice prevents any call to Request from blocking forever,
    20  // striking a balance between precise resource management and flexibility.
    21  type Limiter struct {
    22  	limit   int
    23  	current int
    24  	mu      chan struct{} // can't select on sync.Mutex
    25  	cond    *sync.Cond
    26  }
    27  
    28  // Request blocks until n units are available. If n is greater than m's limit,
    29  // Request blocks until all of m's units have been released.
    30  //
    31  // Request is unbiased with respect to n: calls with small n do not starve
    32  // calls with large n.
    33  //
    34  // Request returns true if the request was canceled, and false otherwise.
    35  func (l *Limiter) Request(n int, cancel <-chan struct{}) bool {
    36  	// acquire mutex
    37  	select {
    38  	case <-cancel:
    39  		return true
    40  	case lock := <-l.mu:
    41  		// unlock
    42  		defer func() { l.mu <- lock }()
    43  	}
    44  
    45  	// spawn goroutine to handle cancellation
    46  	var cancelled bool
    47  	done := make(chan struct{})
    48  	defer close(done)
    49  	go func() {
    50  		select {
    51  		case <-cancel:
    52  			l.cond.L.Lock()
    53  			cancelled = true
    54  			l.cond.L.Unlock()
    55  			l.cond.Signal()
    56  		case <-done:
    57  		}
    58  	}()
    59  
    60  	// wait until request can be satisfied
    61  	l.cond.L.Lock()
    62  	for l.current+n > l.limit && l.current != 0 && !cancelled {
    63  		l.cond.Wait()
    64  	}
    65  	defer l.cond.L.Unlock()
    66  	if !cancelled {
    67  		l.current += n
    68  	}
    69  	return cancelled
    70  }
    71  
    72  // Release returns n units to l, making them available to future callers. It
    73  // is legal for n to be larger than l's limit, as long as n was previously
    74  // passed to Request.
    75  func (l *Limiter) Release(n int) {
    76  	l.cond.L.Lock()
    77  	l.current -= n
    78  	if l.current < 0 {
    79  		l.cond.L.Unlock()
    80  		panic("units released exceeds units requested")
    81  	}
    82  	l.cond.L.Unlock()
    83  	l.cond.Signal()
    84  }
    85  
    86  // SetLimit sets the limit of l. It is legal to interpose calls to SetLimit
    87  // between Request/Release pairs.
    88  func (l *Limiter) SetLimit(limit int) {
    89  	l.cond.L.Lock()
    90  	l.limit = limit
    91  	l.cond.L.Unlock()
    92  	l.cond.Signal()
    93  }
    94  
    95  // NewLimiter returns a Limiter with the supplied limit.
    96  func NewLimiter(limit int) *Limiter {
    97  	l := &Limiter{
    98  		limit: limit,
    99  		mu:    make(chan struct{}, 1),
   100  		cond:  sync.NewCond(new(sync.Mutex)),
   101  	}
   102  	l.mu <- struct{}{}
   103  	return l
   104  }