65 lines
1.7 KiB
Go
65 lines
1.7 KiB
Go
|
// Copyright (C) 2019 Storj Labs, Inc.
|
||
|
|
// See LICENSE for copying information
|
||
|
|
|
||
|
|
package overlay
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"storj.io/storj/pkg/storj"
|
||
|
|
)
|
||
|
|
|
||
|
|
// KeyLock provides per-key RW locking. Locking is key-specific, meaning lock
|
||
|
|
// contention only exists for locking calls using the same key parameter. As
|
||
|
|
// with all locks, do not call Unlock() or RUnlock() before a corresponding
|
||
|
|
// Lock() or RLock() call.
|
||
|
|
//
|
||
|
|
// Note on memory usage: internally KeyLock lazily and atomically creates a
|
||
|
|
// separate sync.RWMutex for each key. To maintain synchronization guarantees,
|
||
|
|
// these interal mutexes are not freed until the entire KeyLock instance is
|
||
|
|
// freed.
|
||
|
|
type KeyLock struct {
|
||
|
|
locksMu sync.Mutex
|
||
|
|
locks map[storj.NodeID]*sync.RWMutex
|
||
|
|
}
|
||
|
|
|
||
|
|
// UnlockFunc is the function to unlock the associated successful lock
|
||
|
|
type UnlockFunc func()
|
||
|
|
|
||
|
|
// NewKeyLock create a new KeyLock
|
||
|
|
func NewKeyLock() *KeyLock {
|
||
|
|
return &KeyLock{
|
||
|
|
locks: make(map[storj.NodeID]*sync.RWMutex),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Lock the provided key. Returns the unlock function.
|
||
|
|
func (l *KeyLock) Lock(nodeID storj.NodeID) UnlockFunc {
|
||
|
|
lock := l.getLock(nodeID)
|
||
|
|
lock.Lock()
|
||
|
|
return lock.Unlock
|
||
|
|
}
|
||
|
|
|
||
|
|
// RLock the provided key. Returns the unlock function.
|
||
|
|
func (l *KeyLock) RLock(nodeID storj.NodeID) UnlockFunc {
|
||
|
|
lock := l.getLock(nodeID)
|
||
|
|
lock.RLock()
|
||
|
|
return lock.RUnlock
|
||
|
|
}
|
||
|
|
|
||
|
|
// getLock will atomically load the RWMutex for this key. If one does not yet
|
||
|
|
// exist, it will be lazily and atomically created. The resulting RWMutex is
|
||
|
|
// returned.
|
||
|
|
func (l *KeyLock) getLock(nodeID storj.NodeID) *sync.RWMutex {
|
||
|
|
l.locksMu.Lock()
|
||
|
|
defer l.locksMu.Unlock()
|
||
|
|
|
||
|
|
lo, ok := l.locks[nodeID]
|
||
|
|
if ok {
|
||
|
|
return lo
|
||
|
|
}
|
||
|
|
mu := &sync.RWMutex{}
|
||
|
|
l.locks[nodeID] = mu
|
||
|
|
return mu
|
||
|
|
}
|