github.com/rudderlabs/rudder-go-kit@v0.30.0/sync/plocker.go (about)

     1  package sync
     2  
     3  import "sync"
     4  
     5  // PartitionLocker is a lock that can be used to lock different partitions at the same time.
     6  type PartitionLocker struct {
     7  	l sync.Mutex // protects s
     8  	s map[string]*lockInfo
     9  }
    10  
    11  // NewPartitionLocker returns a new PartitionLocker.
    12  func NewPartitionLocker() *PartitionLocker {
    13  	return &PartitionLocker{
    14  		s: make(map[string]*lockInfo),
    15  	}
    16  }
    17  
    18  // Lock locks the lock. If the lock is locked, it waits until the lock is unlocked.
    19  func (p *PartitionLocker) Lock(id string) {
    20  	p.l.Lock()
    21  	li := p.lockInfo(id)
    22  	li.refs++
    23  	p.l.Unlock() // unlock before locking mu to avoid unnecessary blocking
    24  	li.mu.Lock()
    25  }
    26  
    27  // Unlock unlocks the lock. If the lock is not locked, it panics.
    28  func (p *PartitionLocker) Unlock(id string) {
    29  	p.l.Lock()
    30  	defer p.l.Unlock()
    31  	li := p.lockInfo(id)
    32  	li.mu.Unlock()
    33  	li.refs--
    34  	if li.refs == 0 {
    35  		delete(p.s, id)
    36  	}
    37  }
    38  
    39  // lockInfo returns the lockInfo for the given id. If the lockInfo does not exist, it is created.
    40  func (l *PartitionLocker) lockInfo(key string) *lockInfo {
    41  	mu, ok := l.s[key]
    42  	if !ok {
    43  		mu = &lockInfo{}
    44  		l.s[key] = mu
    45  	}
    46  	return mu
    47  }
    48  
    49  type lockInfo struct {
    50  	mu   sync.Mutex // the partition lock
    51  	refs int        // number of references to this lock
    52  }