115 lines
2.0 KiB
Go
115 lines
2.0 KiB
Go
package utils
|
|
|
|
import (
|
|
"fmt"
|
|
"hash/fnv"
|
|
"sort"
|
|
"sync"
|
|
)
|
|
|
|
const shardCount = 256
|
|
|
|
type LockManager struct {
|
|
shards [shardCount]sync.RWMutex
|
|
}
|
|
|
|
type LockType int
|
|
|
|
const (
|
|
ReadLock LockType = iota
|
|
WriteLock
|
|
)
|
|
|
|
type Lock struct {
|
|
Path string
|
|
shard uint32
|
|
lockType LockType
|
|
}
|
|
|
|
func NewLockManager() *LockManager {
|
|
return &LockManager{}
|
|
}
|
|
|
|
func hashPath(path string) uint32 {
|
|
h := fnv.New32a()
|
|
h.Write([]byte(path))
|
|
return h.Sum32() % shardCount
|
|
}
|
|
|
|
func (lm *LockManager) AcquireRLock(path string) *Lock {
|
|
shard := hashPath(path)
|
|
lm.shards[shard].RLock()
|
|
|
|
return &Lock{
|
|
Path: path,
|
|
shard: shard,
|
|
lockType: ReadLock,
|
|
}
|
|
}
|
|
|
|
func (lm *LockManager) AcquireLock(path string) *Lock {
|
|
shard := hashPath(path)
|
|
lm.shards[shard].Lock()
|
|
|
|
return &Lock{
|
|
Path: path,
|
|
shard: shard,
|
|
lockType: WriteLock,
|
|
}
|
|
}
|
|
|
|
func (lm *LockManager) ReleaseLock(lock *Lock) error {
|
|
if lock == nil {
|
|
return fmt.Errorf("lock is nil")
|
|
}
|
|
if lock.shard >= shardCount {
|
|
return fmt.Errorf("invalid shard index")
|
|
}
|
|
|
|
switch lock.lockType {
|
|
case ReadLock:
|
|
lm.shards[lock.shard].RUnlock()
|
|
case WriteLock:
|
|
lm.shards[lock.shard].Unlock()
|
|
default:
|
|
return fmt.Errorf("unknown lock type")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (lm *LockManager) AcquireOrderedLock(paths ...string) []*Lock {
|
|
sortedPaths := make([]string, len(paths))
|
|
copy(sortedPaths, paths)
|
|
sort.Strings(sortedPaths)
|
|
|
|
locks := make([]*Lock, 0, len(sortedPaths))
|
|
for _, path := range sortedPaths {
|
|
lock := lm.AcquireLock(path)
|
|
locks = append(locks, lock)
|
|
}
|
|
return locks
|
|
}
|
|
|
|
func (lm *LockManager) AcquireOrderedRLock(paths ...string) []*Lock {
|
|
sortedPaths := make([]string, len(paths))
|
|
copy(sortedPaths, paths)
|
|
sort.Strings(sortedPaths)
|
|
|
|
locks := make([]*Lock, 0, len(sortedPaths))
|
|
for _, path := range sortedPaths {
|
|
lock := lm.AcquireRLock(path)
|
|
locks = append(locks, lock)
|
|
}
|
|
return locks
|
|
}
|
|
|
|
func (lm *LockManager) ReleaseLocks(locks []*Lock) error {
|
|
for i := len(locks) - 1; i >= 0; i-- {
|
|
if err := lm.ReleaseLock(locks[i]); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|