github.com/flowerwrong/netstack@v0.0.0-20191009141956-e5848263af28/tmutex/tmutex.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package tmutex provides the implementation of a mutex that implements an
    16  // efficient TryLock function in addition to Lock and Unlock.
    17  package tmutex
    18  
    19  import (
    20  	"sync/atomic"
    21  )
    22  
    23  // Mutex is a mutual exclusion primitive that implements TryLock in addition
    24  // to Lock and Unlock.
    25  type Mutex struct {
    26  	v  int32
    27  	ch chan struct{}
    28  }
    29  
    30  // Init initializes the mutex.
    31  func (m *Mutex) Init() {
    32  	m.v = 1
    33  	m.ch = make(chan struct{}, 1)
    34  }
    35  
    36  // Lock acquires the mutex. If it is currently held by another goroutine, Lock
    37  // will wait until it has a chance to acquire it.
    38  func (m *Mutex) Lock() {
    39  	// Uncontended case.
    40  	if atomic.AddInt32(&m.v, -1) == 0 {
    41  		return
    42  	}
    43  
    44  	for {
    45  		// Try to acquire the mutex again, at the same time making sure
    46  		// that m.v is negative, which indicates to the owner of the
    47  		// lock that it is contended, which will force it to try to wake
    48  		// someone up when it releases the mutex.
    49  		if v := atomic.LoadInt32(&m.v); v >= 0 && atomic.SwapInt32(&m.v, -1) == 1 {
    50  			return
    51  		}
    52  
    53  		// Wait for the mutex to be released before trying again.
    54  		<-m.ch
    55  	}
    56  }
    57  
    58  // TryLock attempts to acquire the mutex without blocking. If the mutex is
    59  // currently held by another goroutine, it fails to acquire it and returns
    60  // false.
    61  func (m *Mutex) TryLock() bool {
    62  	v := atomic.LoadInt32(&m.v)
    63  	if v <= 0 {
    64  		return false
    65  	}
    66  	return atomic.CompareAndSwapInt32(&m.v, 1, 0)
    67  }
    68  
    69  // Unlock releases the mutex.
    70  func (m *Mutex) Unlock() {
    71  	if atomic.SwapInt32(&m.v, 1) == 0 {
    72  		// There were no pending waiters.
    73  		return
    74  	}
    75  
    76  	// Wake some waiter up.
    77  	select {
    78  	case m.ch <- struct{}{}:
    79  	default:
    80  	}
    81  }