github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/handy/lock.go (about)

     1  package handy
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"unsafe"
     7  )
     8  
     9  const mutexLocked = 1 << iota // mutex is locked
    10  
    11  type Lock struct {
    12  	sync.Mutex
    13  }
    14  
    15  // WithLock run code within protection of lock.
    16  func (l *Lock) WithLock(f func()) {
    17  	defer l.LockDeferUnlock()()
    18  	f()
    19  }
    20  
    21  func (l *Lock) TryLock() bool {
    22  	state := (*int32)(unsafe.Pointer(&l.Mutex))
    23  	return atomic.CompareAndSwapInt32(state, 0, mutexLocked)
    24  }
    25  
    26  func (l *Lock) LockDeferUnlock() func() {
    27  	l.Lock()
    28  	return l.Unlock
    29  }
    30  
    31  // LockUnlock lock on mutex and return unlock.
    32  // e.g. defer LockUnlock(mutex)
    33  func LockUnlock(m *sync.Mutex) func() {
    34  	m.Lock()
    35  	return m.Unlock
    36  }
    37  
    38  type RWLock struct {
    39  	sync.RWMutex
    40  }
    41  
    42  // WithLock run code within protection of lock.
    43  func (l *RWLock) WithLock(f func()) {
    44  	defer l.LockDeferUnlock()()
    45  	f()
    46  }
    47  
    48  // WithRLock run code within protection of lock.
    49  func (l *RWLock) WithRLock(f func()) {
    50  	defer l.RLockDeferRUnlock()()
    51  	f()
    52  }
    53  
    54  func (l *RWLock) TryLock() bool {
    55  	state := (*int32)(unsafe.Pointer(&l.RWMutex))
    56  	return atomic.CompareAndSwapInt32(state, 0, mutexLocked)
    57  }
    58  
    59  func (l *RWLock) LockDeferUnlock() func() {
    60  	l.Lock()
    61  	return l.Unlock
    62  }
    63  
    64  func (l *RWLock) RLockDeferRUnlock() func() {
    65  	l.RLock()
    66  	return l.RUnlock
    67  }
    68  
    69  // RWLockUnlock lock on mutex and return unlock.
    70  // e.g. defer LockUnlock(mutex)
    71  func RWLockUnlock(m *sync.RWMutex) func() {
    72  	m.Lock()
    73  	return m.Unlock
    74  }
    75  
    76  // RWRLockRUnlock lock on mutex and return unlock.
    77  // e.g. defer LockUnlock(mutex)
    78  func RWRLockRUnlock(m *sync.RWMutex) func() {
    79  	m.RLock()
    80  	return m.RUnlock
    81  }