wp-go/cache/locks.go

package cache

import (
	"context"
	"github.com/fthvgb1/wp-go/safety"
	"sync"
	"sync/atomic"
)

type Locks[K comparable] struct {
	numFn   func() int
	locks   []*sync.Mutex
	m       *safety.Map[K, *sync.Mutex]
	counter *int64
}

func (l *Locks[K]) Flush(_ context.Context) {
	l.m.Flush()
	atomic.StoreInt64(l.counter, 0)
}

func (l *Locks[K]) SetNumFn(numFn func() int) {
	l.numFn = numFn
}

func NewLocks[K comparable](num func() int) *Locks[K] {
	var i int64
	return &Locks[K]{numFn: num, m: safety.NewMap[K, *sync.Mutex](), counter: &i}
}

func (l *Locks[K]) SetLockNum(num int) {
	if num > 0 {
		l.locks = make([]*sync.Mutex, num)
		for i := 0; i < num; i++ {
			l.locks[i] = &sync.Mutex{}
		}
	}
}

func (l *Locks[K]) GetLock(ctx context.Context, gMut *sync.Mutex, keys ...K) *sync.Mutex {
	k := keys[0]
	lo, ok := l.m.Load(k)
	if ok {
		return lo
	}
	num := l.numFn()
	if num == 1 {
		return gMut
	}
	gMut.Lock()
	defer gMut.Unlock()
	lo, ok = l.m.Load(k)
	if ok {
		return lo
	}
	if num <= 0 {
		lo = &sync.Mutex{}
		l.m.Store(k, lo)
		return lo
	}
	if len(l.locks) == 0 {
		l.SetLockNum(num)
	}
	counter := int(atomic.LoadInt64(l.counter))
	if counter > len(l.locks)-1 {
		atomic.StoreInt64(l.counter, 0)
		counter = 0
	}
	lo = l.locks[counter]
	l.m.Store(k, lo)
	atomic.AddInt64(l.counter, 1)
	if len(l.locks) < num {
		for i := 0; i < num-len(l.locks); i++ {
			l.locks = append(l.locks, &sync.Mutex{})
		}
	}
	return lo
}