github.com/fluhus/gostuff@v0.4.1-0.20240331134726-be71864f2b5d/nlock/nlock.go (about)

     1  // Package nlock provides a lock that can be locked n times simultaneously.
     2  package nlock
     3  
     4  import (
     5  	"fmt"
     6  	"sync"
     7  )
     8  
     9  // An NLock is a lock that allows n holders to hold the lock simultaneously.
    10  type NLock struct {
    11  	c *sync.Cond // Sync object.
    12  	n int        // Max number of holders.
    13  	i int        // Current number of holders.
    14  }
    15  
    16  // New creates a new lock for n maximum holders.
    17  func New(n int) *NLock {
    18  	if n <= 0 {
    19  		panic(fmt.Sprintf("Bad n: %v, needs to be at least 1.", n))
    20  	}
    21  	return &NLock{sync.NewCond(&sync.Mutex{}), n, 0}
    22  }
    23  
    24  // Lock locks the lock. Will block if n calls to lock were made, that were not
    25  // unlocked.
    26  func (n *NLock) Lock() {
    27  	n.c.L.Lock()
    28  	defer n.c.L.Unlock()
    29  	for n.i >= n.n {
    30  		n.c.Wait()
    31  	}
    32  	n.i++
    33  }
    34  
    35  // Unlock releases one holder of the lock. Panics if lock has 0 holders.
    36  func (n *NLock) Unlock() {
    37  	n.c.L.Lock()
    38  	defer n.c.L.Unlock()
    39  	if n.i <= 0 {
    40  		panic("Attempt to unlock a not locked lock.")
    41  	}
    42  	n.i--
    43  	n.c.Signal()
    44  }
    45  
    46  // TryLock attempts to obtain lock without waiting. Returns true if succeeded,
    47  // or false if not.
    48  func (n *NLock) TryLock() bool {
    49  	n.c.L.Lock()
    50  	defer n.c.L.Unlock()
    51  	if n.i >= n.n {
    52  		return false
    53  	}
    54  	n.i++
    55  	return true
    56  }