github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/mutex/tmutex.go (about)

     1  package mutex
     2  
     3  import "sync/atomic"
     4  
     5  // TMutex is a mutual exclusion primitive that implements TryLock in addition
     6  // to Lock and Unlock.
     7  type TMutex struct {
     8  	v  int32
     9  	ch chan struct{}
    10  }
    11  
    12  // New is ctor for TMutex
    13  func New() *TMutex {
    14  	return &TMutex{v: 1, ch: make(chan struct{}, 1)}
    15  }
    16  
    17  // Lock has the same semantics as normal Mutex.Lock
    18  func (m *TMutex) Lock() {
    19  	// Uncontended case.
    20  	if atomic.AddInt32(&m.v, -1) == 0 {
    21  		return
    22  	}
    23  
    24  	for {
    25  
    26  		// if v < 0, someone is already contending, just wait for lock release
    27  		// otherwise, SwapInt32 can only return one of -1, 0, 1
    28  		// 1 means no contention
    29  		// -1 means someone is already contending
    30  		// 0 means no one else is contending but the lock hasn't been released
    31  		// so for -1 or 0, just wait for lock release
    32  		if v := atomic.LoadInt32(&m.v); v >= 0 && atomic.SwapInt32(&m.v, -1) == 1 {
    33  			return
    34  		}
    35  
    36  		// Wait for the mutex to be released before trying again.
    37  		<-m.ch
    38  	}
    39  }
    40  
    41  // TryLock returns true on success
    42  func (m *TMutex) TryLock() bool {
    43  	v := atomic.LoadInt32(&m.v)
    44  	if v <= 0 {
    45  		return false
    46  	}
    47  	return atomic.CompareAndSwapInt32(&m.v, 1, 0)
    48  }
    49  
    50  // Unlock has the same semantics as normal Mutex.Unlock
    51  func (m *TMutex) Unlock() {
    52  	if atomic.SwapInt32(&m.v, 1) == 0 {
    53  		// There were no pending waiters.
    54  		return
    55  	}
    56  
    57  	// Wake some waiter up.
    58  	select {
    59  	case m.ch <- struct{}{}:
    60  	default:
    61  	}
    62  }