github.com/abiosoft/semaphore@v0.0.0-20240818083615-bc6b5b10c137/semaphore.go (about)

     1  package semaphore
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  // Semaphore is an implementation of semaphore.
    10  type Semaphore struct {
    11  	permits int
    12  	avail   int
    13  	channel chan struct{}
    14  	aMutex  sync.Mutex   // acquire
    15  	rMutex  sync.Mutex   // release
    16  	pMutex  sync.RWMutex // number of permits
    17  }
    18  
    19  // New creates a new Semaphore with specified number of permits.
    20  func New(permits int) *Semaphore {
    21  	if permits < 1 {
    22  		panic("Invalid number of permits. Less than 1")
    23  	}
    24  
    25  	// fill channel buffer
    26  	channel := make(chan struct{}, permits)
    27  	for i := 0; i < permits; i++ {
    28  		channel <- struct{}{}
    29  	}
    30  
    31  	return &Semaphore{
    32  		permits: permits,
    33  		avail:   permits,
    34  		channel: channel,
    35  	}
    36  }
    37  
    38  // Acquire acquires one permit. If it is not available, the goroutine will block until it is available.
    39  func (s *Semaphore) Acquire() {
    40  	s.aMutex.Lock()
    41  	defer s.aMutex.Unlock()
    42  
    43  	s.pMutex.Lock()
    44  	s.avail--
    45  	s.pMutex.Unlock()
    46  
    47  	<-s.channel
    48  }
    49  
    50  // AcquireMany is similar to Acquire() but for many permits.
    51  //
    52  // The number of permits acquired is at most the number of permits in the semaphore.
    53  // i.e. if n = 5 and s was created with New(2), at most 2 permits will be acquired.
    54  func (s *Semaphore) AcquireMany(n int) {
    55  	if n > s.permits {
    56  		n = s.permits
    57  	}
    58  
    59  	for ; n > 0; n-- {
    60  		s.Acquire()
    61  	}
    62  }
    63  
    64  // AcquireContext is similar to AcquireMany() but takes a context. Returns true if successful
    65  // or false if the context is done first.
    66  func (s *Semaphore) AcquireContext(ctx context.Context, n int) bool {
    67  	acquired := make(chan struct{}, 1)
    68  	reverse := make(chan bool, 1)
    69  	go func() {
    70  		s.AcquireMany(n)
    71  		acquired <- struct{}{}
    72  		if <-reverse {
    73  			s.ReleaseMany(n)
    74  		}
    75  		close(acquired)
    76  		close(reverse)
    77  	}()
    78  
    79  	select {
    80  	case <-ctx.Done():
    81  		reverse <- true
    82  		return false
    83  	case <-acquired:
    84  		reverse <- false
    85  		return true
    86  	}
    87  }
    88  
    89  // AcquireWithin is similar to AcquireMany() but cancels if duration elapses before getting the permits.
    90  // Returns true if successful and false if timeout occurs.
    91  func (s *Semaphore) AcquireWithin(n int, d time.Duration) bool {
    92  	ctx, cancel := context.WithTimeout(context.Background(), d)
    93  	defer cancel()
    94  
    95  	return s.AcquireContext(ctx, n)
    96  }
    97  
    98  // Release releases one permit.
    99  func (s *Semaphore) Release() {
   100  	s.rMutex.Lock()
   101  	defer s.rMutex.Unlock()
   102  
   103  	s.channel <- struct{}{}
   104  
   105  	s.pMutex.Lock()
   106  	s.avail++
   107  	s.pMutex.Unlock()
   108  }
   109  
   110  // ReleaseMany releases n permits.
   111  //
   112  // The number of permits released is at most the number of permits in the semaphore.
   113  // i.e. if n = 5 and s was created with New(2), at most 2 permits will be released.
   114  func (s *Semaphore) ReleaseMany(n int) {
   115  	if n > s.permits {
   116  		n = s.permits
   117  	}
   118  
   119  	for ; n > 0; n-- {
   120  		s.Release()
   121  	}
   122  }
   123  
   124  // AvailablePermits gives number of available unacquired permits.
   125  func (s *Semaphore) AvailablePermits() int {
   126  	s.pMutex.RLock()
   127  	defer s.pMutex.RUnlock()
   128  
   129  	if s.avail < 0 {
   130  		return 0
   131  	}
   132  	return s.avail
   133  }
   134  
   135  // DrainPermits acquires all available permits and return the number of permits acquired.
   136  func (s *Semaphore) DrainPermits() int {
   137  	n := s.AvailablePermits()
   138  	if n > 0 {
   139  		s.AcquireMany(n)
   140  	}
   141  	return n
   142  }