go-mixed/locker/source_locker.go

package locker

import (
	"context"

	redis_store "github.com/charlienet/go-mixed/locker/redis"
	"github.com/charlienet/go-mixed/redis"
)

type countLocker struct {
	rw    rwLocker
	Count int32
}

type SourceLocker struct {
	m                 RWLocker
	distributedLocker DistributedLocker
	locks             map[string]*countLocker
	err               error
}

func NewSourceLocker() *SourceLocker {
	l := &SourceLocker{
		locks: make(map[string]*countLocker),
	}

	l.m = NewRWLocker()
	return l
}

func (s *SourceLocker) WithRedis(key string, clients ...redis.Client) *SourceLocker {
	redisStore := redis_store.NewRedisStore(key, clients...)
	return s.WithDistributedLocker(redisStore)
}

func (s *SourceLocker) WithDistributedLocker(distributed DistributedLocker) *SourceLocker {
	s.distributedLocker = distributed
	return s
}

func (s *SourceLocker) Lock(key string) {
	s.m.Lock()
	l, ok := s.locks[key]
	if ok {
		l.Count++
	} else {
		n := NewRWLocker()
		l = &countLocker{rw: n, Count: 1}
		s.locks[key] = l
	}
	s.m.Unlock()

	l.rw.Lock()
}

func (s *SourceLocker) Unlock(key string) {
	s.m.Lock()

	if l, ok := s.locks[key]; ok {
		l.Count--
		l.rw.Unlock()

		if s.distributedLocker != nil {
			s.distributedLocker.Unlock(context.Background(), key)
		}

		if l.Count == 0 {
			delete(s.locks, key)
		}
	}
	s.m.Unlock()
}

func (s *SourceLocker) TryLock(key string) bool {
	s.m.Lock()
	l, ok := s.locks[key]
	if ok {
		l.Count++
	} else {
		n := NewRWLocker()
		l = &countLocker{rw: n, Count: 1}
		s.locks[key] = l
	}
	s.m.Unlock()

	if !l.rw.TryLock() {
		s.m.Lock()
		l.Count--
		if l.Count == 0 {
			delete(s.locks, key)
		}
		s.m.Unlock()
		return false
	}
	return true
}