diff --git a/libs/utils/striplock.go b/libs/utils/striplock.go new file mode 100644 index 0000000000..bf5db5195c --- /dev/null +++ b/libs/utils/striplock.go @@ -0,0 +1,64 @@ +package utils + +import ( + "sync" + + "github.com/celestiaorg/celestia-node/share" +) + +// StripLock provides a way to lock operations by height and data hash. +// It uses a fixed number of mutexes to avoid excessive memory allocation +// while still providing good concurrency characteristics. +type StripLock struct { + heights []*sync.RWMutex + datahashes []*sync.RWMutex +} + +// MultiLock represents multiple locks that need to be acquired together. +type MultiLock struct { + mu []*sync.RWMutex +} + +// NewStripLock creates a new StripLock with the specified size of mutex pools. +func NewStripLock(size int) *StripLock { + heights := make([]*sync.RWMutex, size) + datahashes := make([]*sync.RWMutex, size) + for i := 0; i < size; i++ { + heights[i] = &sync.RWMutex{} + datahashes[i] = &sync.RWMutex{} + } + return &StripLock{heights, datahashes} +} + +// ByHeight returns a mutex for the given height. +func (l *StripLock) ByHeight(height uint64) *sync.RWMutex { + lkIdx := height % uint64(len(l.heights)) + return l.heights[lkIdx] +} + +// ByHash returns a mutex for the given data hash. +func (l *StripLock) ByHash(datahash share.DataHash) *sync.RWMutex { + // Use the last 2 bytes of the hash as key to distribute the locks + last := uint16(datahash[len(datahash)-1]) | uint16(datahash[len(datahash)-2])<<8 + lkIdx := last % uint16(len(l.datahashes)) + return l.datahashes[lkIdx] +} + +// ByHashAndHeight returns a MultiLock for both the hash and height. +func (l *StripLock) ByHashAndHeight(datahash share.DataHash, height uint64) *MultiLock { + return &MultiLock{[]*sync.RWMutex{l.ByHash(datahash), l.ByHeight(height)}} +} + +// Lock acquires all contained locks. +func (m *MultiLock) Lock() { + for _, lk := range m.mu { + lk.Lock() + } +} + +// Unlock releases all contained locks. +func (m *MultiLock) Unlock() { + for _, lk := range m.mu { + lk.Unlock() + } +}