github.com/decred/dcrlnd@v0.7.6/multimutex/multimutex.go (about)

     1  package multimutex
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  )
     7  
     8  // cntMutex is a struct that wraps a counter and a mutex, and is used
     9  // to keep track of the number of goroutines waiting for access to the
    10  // mutex, such that we can forget about it when the counter is zero.
    11  type cntMutex struct {
    12  	cnt int
    13  	sync.Mutex
    14  }
    15  
    16  // Mutex is a struct that keeps track of a set of mutexes with
    17  // a given ID. It can be used for making sure only one goroutine
    18  // gets given the mutex per ID.
    19  type Mutex struct {
    20  	// mutexes is a map of IDs to a cntMutex. The cntMutex for
    21  	// a given ID will hold the mutex to be used by all
    22  	// callers requesting access for the ID, in addition to
    23  	// the count of callers.
    24  	mutexes map[uint64]*cntMutex
    25  
    26  	// mapMtx is used to give synchronize concurrent access
    27  	// to the mutexes map.
    28  	mapMtx sync.Mutex
    29  }
    30  
    31  // NewMutex creates a new Mutex.
    32  func NewMutex() *Mutex {
    33  	return &Mutex{
    34  		mutexes: make(map[uint64]*cntMutex),
    35  	}
    36  }
    37  
    38  // Lock locks the mutex by the given ID. If the mutex is already
    39  // locked by this ID, Lock blocks until the mutex is available.
    40  func (c *Mutex) Lock(id uint64) {
    41  	c.mapMtx.Lock()
    42  	mtx, ok := c.mutexes[id]
    43  	if ok {
    44  		// If the mutex already existed in the map, we
    45  		// increment its counter, to indicate that there
    46  		// now is one more goroutine waiting for it.
    47  		mtx.cnt++
    48  	} else {
    49  		// If it was not in the map, it means no other
    50  		// goroutine has locked the mutex for this ID,
    51  		// and we can create a new mutex with count 1
    52  		// and add it to the map.
    53  		mtx = &cntMutex{
    54  			cnt: 1,
    55  		}
    56  		c.mutexes[id] = mtx
    57  	}
    58  	c.mapMtx.Unlock()
    59  
    60  	// Acquire the mutex for this ID.
    61  	mtx.Lock()
    62  }
    63  
    64  // Unlock unlocks the mutex by the given ID. It is a run-time
    65  // error if the mutex is not locked by the ID on entry to Unlock.
    66  func (c *Mutex) Unlock(id uint64) {
    67  	// Since we are done with all the work for this
    68  	// update, we update the map to reflect that.
    69  	c.mapMtx.Lock()
    70  
    71  	mtx, ok := c.mutexes[id]
    72  	if !ok {
    73  		// The mutex not existing in the map means
    74  		// an unlock for an ID not currently locked
    75  		// was attempted.
    76  		panic(fmt.Sprintf("double unlock for id %v",
    77  			id))
    78  	}
    79  
    80  	// Decrement the counter. If the count goes to
    81  	// zero, it means this caller was the last one
    82  	// to wait for the mutex, and we can delete it
    83  	// from the map. We can do this safely since we
    84  	// are under the mapMtx, meaning that all other
    85  	// goroutines waiting for the mutex already
    86  	// have incremented it, or will create a new
    87  	// mutex when they get the mapMtx.
    88  	mtx.cnt--
    89  	if mtx.cnt == 0 {
    90  		delete(c.mutexes, id)
    91  	}
    92  	c.mapMtx.Unlock()
    93  
    94  	// Unlock the mutex for this ID.
    95  	mtx.Unlock()
    96  }