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

     1  package sync
     2  
     3  import "sync"
     4  
     5  // PartitionRWLocker is a read-write lock that can be used to lock different partitions at the same time.
     6  type PartitionRWLocker struct {
     7  	l sync.Mutex // protects s
     8  	s map[string]*rwLockInfo
     9  }
    10  
    11  // NewPartitionRWLocker returns a new PartitionRWLocker.
    12  func NewPartitionRWLocker() *PartitionRWLocker {
    13  	return &PartitionRWLocker{
    14  		s: make(map[string]*rwLockInfo),
    15  	}
    16  }
    17  
    18  // Lock locks the lock for writing. If the lock is locked for reading or writing, it waits until the lock is unlocked.
    19  func (p *PartitionRWLocker) 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  // RLock locks the lock for reading. If the lock is locked for writing, it waits until the lock is unlocked.
    28  func (p *PartitionRWLocker) RLock(id string) {
    29  	p.l.Lock()
    30  	li := p.lockInfo(id)
    31  	li.refs++
    32  	p.l.Unlock() // unlock before locking mu to avoid unnecessary blocking
    33  	li.mu.RLock()
    34  }
    35  
    36  // Unlock unlocks the lock for writing. If the lock is locked for reading or not locked for writing, it panics.
    37  func (p *PartitionRWLocker) Unlock(id string) {
    38  	p.l.Lock()
    39  	defer p.l.Unlock()
    40  	li := p.lockInfo(id)
    41  	li.mu.Unlock()
    42  	li.refs--
    43  	if li.refs == 0 {
    44  		delete(p.s, id)
    45  	}
    46  }
    47  
    48  // RUnlock unlocks the lock for reading. If the lock is locked for writing or not locked for reading, it panics.
    49  func (p *PartitionRWLocker) RUnlock(id string) {
    50  	p.l.Lock()
    51  	defer p.l.Unlock()
    52  	li := p.lockInfo(id)
    53  	li.mu.RUnlock()
    54  	li.refs--
    55  	if li.refs == 0 {
    56  		delete(p.s, id)
    57  	}
    58  }
    59  
    60  // RWMutexFor returns a new RWMutex scoped to the given id.
    61  func (p *PartitionRWLocker) RWMutexFor(id string) *RWMutex {
    62  	return &RWMutex{
    63  		plock: p,
    64  		id:    id,
    65  	}
    66  }
    67  
    68  // lockInfo returns the lockInfo for the given id. If the lockInfo does not exist, it is created.
    69  func (p *PartitionRWLocker) lockInfo(id string) *rwLockInfo {
    70  	li, ok := p.s[id]
    71  	if !ok {
    72  		li = &rwLockInfo{}
    73  		p.s[id] = li
    74  	}
    75  	return li
    76  }
    77  
    78  type rwLockInfo struct {
    79  	mu   sync.RWMutex // the partition lock
    80  	refs int          // number of references to this lock
    81  }
    82  
    83  // RWMutex is a read-write lock
    84  type RWMutex struct {
    85  	plock *PartitionRWLocker
    86  	id    string
    87  }
    88  
    89  // Lock locks the lock for writing. If the lock is locked for reading or writing, it waits until the lock is unlocked.
    90  func (m *RWMutex) Lock() {
    91  	m.plock.Lock(m.id)
    92  }
    93  
    94  // Unlock unlocks the lock for writing. If the lock is locked for reading or not locked for writing, it panics.
    95  func (m *RWMutex) Unlock() {
    96  	m.plock.Unlock(m.id)
    97  }
    98  
    99  // RLock locks the lock for reading. If the lock is locked for writing, it waits until the lock is unlocked.
   100  func (m *RWMutex) RLock() {
   101  	m.plock.RLock(m.id)
   102  }
   103  
   104  // RUnlock unlocks the lock for reading. If the lock is locked for writing or not locked for reading, it panics.
   105  
   106  func (m *RWMutex) RUnlock() {
   107  	m.plock.RUnlock(m.id)
   108  }