JCS-pub/client/internal/publock/lockprovider/package.go

package lockprovider

import (
	"fmt"

	"gitlink.org.cn/cloudream/common/utils/lo2"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/publock/types"
)

const (
	PackageLockPathPrefix     = "Package"
	PackageStorageIDPathIndex = 1
	PackageBuzyLock           = "Buzy"
	PackagePinLock            = "Pin"
)

type PackageLock struct {
	stgLocks  map[string]*PackageStorageLock
	dummyLock *PackageStorageLock
}

func NewPackageLock() *PackageLock {
	return &PackageLock{
		stgLocks:  make(map[string]*PackageStorageLock),
		dummyLock: NewPackageStorageLock(),
	}
}

// CanLock 判断这个锁能否锁定成功
func (l *PackageLock) CanLock(lock types.Lock) error {
	nodeLock, ok := l.stgLocks[lock.Path[PackageStorageIDPathIndex]]
	if !ok {
		// 不能直接返回nil，因为如果锁数据的格式不对，也不能获取锁。
		// 这里使用一个空Provider来进行检查。
		return l.dummyLock.CanLock(lock)
	}

	return nodeLock.CanLock(lock)
}

// 锁定。在内部可以不用判断能否加锁，外部需要保证调用此函数前调用了CanLock进行检查
func (l *PackageLock) Lock(reqID types.RequestID, lock types.Lock) error {
	stgID := lock.Path[PackageStorageIDPathIndex]

	nodeLock, ok := l.stgLocks[stgID]
	if !ok {
		nodeLock = NewPackageStorageLock()
		l.stgLocks[stgID] = nodeLock
	}

	return nodeLock.Lock(reqID, lock)
}

// 解锁
func (l *PackageLock) Unlock(reqID types.RequestID, lock types.Lock) error {
	stgID := lock.Path[PackageStorageIDPathIndex]

	nodeLock, ok := l.stgLocks[stgID]
	if !ok {
		return nil
	}

	return nodeLock.Unlock(reqID, lock)
}

// Clear 清除内部所有状态
func (l *PackageLock) Clear() {
	l.stgLocks = make(map[string]*PackageStorageLock)
}

type PackageStorageLock struct {
	buzyReqIDs []types.RequestID
	pinReqIDs  []types.RequestID

	lockCompatibilityTable *LockCompatibilityTable
}

func NewPackageStorageLock() *PackageStorageLock {
	compTable := &LockCompatibilityTable{}

	sdLock := PackageStorageLock{
		lockCompatibilityTable: compTable,
	}

	compTable.
		Column(PackageBuzyLock, func() bool { return len(sdLock.buzyReqIDs) > 0 }).
		Column(PackagePinLock, func() bool { return len(sdLock.pinReqIDs) > 0 })

	comp := LockCompatible()
	uncp := LockUncompatible()

	compTable.MustRow(comp, uncp)
	compTable.MustRow(uncp, comp)

	return &sdLock
}

// CanLock 判断这个锁能否锁定成功
func (l *PackageStorageLock) CanLock(lock types.Lock) error {
	return l.lockCompatibilityTable.Test(lock)
}

// 锁定
func (l *PackageStorageLock) Lock(reqID types.RequestID, lock types.Lock) error {
	switch lock.Name {
	case PackageBuzyLock:
		l.buzyReqIDs = append(l.buzyReqIDs, reqID)
	case PackagePinLock:
		l.pinReqIDs = append(l.pinReqIDs, reqID)
	default:
		return fmt.Errorf("unknow lock name: %s", lock.Name)
	}

	return nil
}

// 解锁
func (l *PackageStorageLock) Unlock(reqID types.RequestID, lock types.Lock) error {
	switch lock.Name {
	case PackageBuzyLock:
		l.buzyReqIDs = lo2.Remove(l.buzyReqIDs, reqID)
	case PackagePinLock:
		l.pinReqIDs = lo2.Remove(l.pinReqIDs, reqID)
	default:
		return fmt.Errorf("unknow lock name: %s", lock.Name)
	}

	return nil
}