github.com/go4org/go4@v0.0.0-20200104003542-c7e774b10ea0/syncutil/sem.go (about)

     1  package syncutil
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"sync"
     7  )
     8  
     9  type debugT bool
    10  
    11  var debug = debugT(false)
    12  
    13  func (d debugT) Printf(format string, args ...interface{}) {
    14  	if bool(d) {
    15  		log.Printf(format, args...)
    16  	}
    17  }
    18  
    19  // Sem implements a semaphore that can have multiple units acquired/released
    20  // at a time.
    21  type Sem struct {
    22  	c         *sync.Cond // Protects size
    23  	max, free int64
    24  }
    25  
    26  // NewSem creates a semaphore with max units available for acquisition.
    27  func NewSem(max int64) *Sem {
    28  	return &Sem{
    29  		c:    sync.NewCond(new(sync.Mutex)),
    30  		free: max,
    31  		max:  max,
    32  	}
    33  }
    34  
    35  // Acquire will deduct n units from the semaphore.  If the deduction would
    36  // result in the available units falling below zero, the call will block until
    37  // another go routine returns units via a call to Release.  If more units are
    38  // requested than the semaphore is configured to hold, error will be non-nil.
    39  func (s *Sem) Acquire(n int64) error {
    40  	if n > s.max {
    41  		return fmt.Errorf("sem: attempt to acquire more units than semaphore size %d > %d", n, s.max)
    42  	}
    43  	s.c.L.Lock()
    44  	defer s.c.L.Unlock()
    45  	for {
    46  		debug.Printf("Acquire check max %d free %d, n %d", s.max, s.free, n)
    47  		if s.free >= n {
    48  			s.free -= n
    49  			return nil
    50  		}
    51  		debug.Printf("Acquire Wait max %d free %d, n %d", s.max, s.free, n)
    52  		s.c.Wait()
    53  	}
    54  }
    55  
    56  // Release will return n units to the semaphore and notify any currently
    57  // blocking Acquire calls.
    58  func (s *Sem) Release(n int64) {
    59  	s.c.L.Lock()
    60  	defer s.c.L.Unlock()
    61  	debug.Printf("Release max %d free %d, n %d", s.max, s.free, n)
    62  	s.free += n
    63  	s.c.Broadcast()
    64  }