github.com/lmb/consul@v1.4.1/lib/semaphore/semaphore.go (about)

     1  // Package semaphore implements a simple semaphore that is based on
     2  // golang.org/x/sync/semaphore but doesn't support weights. It's advantage over
     3  // a simple buffered chan is that the capacity of the semaphore (i.e. the number
     4  // of slots available) can be changed dynamically at runtime without waiting for
     5  // all existing work to stop. This makes it easier to implement e.g. concurrency
     6  // limits on certain operations that can be reconfigured at runtime.
     7  package semaphore
     8  
     9  import (
    10  	"container/list"
    11  	"context"
    12  	"sync"
    13  )
    14  
    15  // Dynamic implements a semaphore whose capacity can be changed dynamically at
    16  // run time.
    17  type Dynamic struct {
    18  	size    int64
    19  	cur     int64
    20  	waiters list.List
    21  	mu      sync.Mutex
    22  }
    23  
    24  // NewDynamic returns a dynamic semaphore with the given initial capacity. Note
    25  // that this is for convenience and to match golang.org/x/sync/semaphore however
    26  // it's possible to use a zero-value semaphore provided SetSize is called before
    27  // use.
    28  func NewDynamic(n int64) *Dynamic {
    29  	return &Dynamic{
    30  		size: n,
    31  	}
    32  }
    33  
    34  // SetSize dynamically updates the number of available slots. If there are more
    35  // than n slots currently acquired, no further acquires will succeed until
    36  // sufficient have been released to take the total outstanding below n again.
    37  func (s *Dynamic) SetSize(n int64) error {
    38  	s.mu.Lock()
    39  	defer s.mu.Unlock()
    40  	s.size = n
    41  	return nil
    42  }
    43  
    44  // Acquire attempts to acquire one "slot" in the semaphore, blocking only until
    45  // ctx is Done. On success, returns nil. On failure, returns ctx.Err() and leaves
    46  // the semaphore unchanged.
    47  //
    48  // If ctx is already done, Acquire may still succeed without blocking.
    49  func (s *Dynamic) Acquire(ctx context.Context) error {
    50  	s.mu.Lock()
    51  	if s.cur < s.size {
    52  		s.cur++
    53  		s.mu.Unlock()
    54  		return nil
    55  	}
    56  
    57  	// Need to wait, add to waiter list
    58  	ready := make(chan struct{})
    59  	elem := s.waiters.PushBack(ready)
    60  	s.mu.Unlock()
    61  
    62  	select {
    63  	case <-ctx.Done():
    64  		err := ctx.Err()
    65  		s.mu.Lock()
    66  		select {
    67  		case <-ready:
    68  			// Acquired the semaphore after we were canceled.  Rather than trying to
    69  			// fix up the queue, just pretend we didn't notice the cancellation.
    70  			err = nil
    71  		default:
    72  			s.waiters.Remove(elem)
    73  		}
    74  		s.mu.Unlock()
    75  		return err
    76  
    77  	case <-ready:
    78  		return nil
    79  	}
    80  }
    81  
    82  // Release releases the semaphore. It will panic if release is called on an
    83  // empty semphore.
    84  func (s *Dynamic) Release() {
    85  	s.mu.Lock()
    86  	defer s.mu.Unlock()
    87  
    88  	if s.cur < 1 {
    89  		panic("semaphore: bad release")
    90  	}
    91  
    92  	next := s.waiters.Front()
    93  
    94  	// If there are no waiters, just decrement and we're done
    95  	if next == nil {
    96  		s.cur--
    97  		return
    98  	}
    99  
   100  	// Need to yield our slot to the next waiter.
   101  	// Remove them from the list
   102  	s.waiters.Remove(next)
   103  	// And trigger it's chan before we release the lock
   104  	close(next.Value.(chan struct{}))
   105  	// Note we _don't_ decrement inflight since the slot was yielded directly.
   106  }