JCS-pub/client/internal/publock/core.go

package publock

import (
	"context"
	"fmt"
	"sync"

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

type Core struct {
	lock         *sync.Mutex
	provdersTrie *trie.Trie[types.LockProvider]
	acquirings   []*acquireInfo
	acquired     map[types.RequestID]types.LockRequest
	nextReqID    int64
}

func NewCore() *Core {
	svc := &Core{
		lock:         &sync.Mutex{},
		provdersTrie: trie.NewTrie[types.LockProvider](),
		acquired:     make(map[types.RequestID]types.LockRequest),
	}

	svc.provdersTrie.Create([]any{lockprovider.UserSpaceLockPathPrefix, trie.WORD_ANY}).Value = lockprovider.NewUserSpaceLock()
	svc.provdersTrie.Create([]any{lockprovider.PackageLockPathPrefix, trie.WORD_ANY}).Value = lockprovider.NewPackageLock()
	return svc
}

func (s *Core) Start() {

}

func (s *Core) Stop() {

}

type acquireInfo struct {
	Request  types.LockRequest
	Callback *future.SetValueFuture[types.RequestID]
	LastErr  error
}

func (svc *Core) Acquire(req types.LockRequest, opt types.AcquireOption) (LockedRequest, error) {
	ctx := context.Background()
	if opt.Timeout != 0 {
		var cancel func()
		ctx, cancel = context.WithTimeout(ctx, opt.Timeout)
		defer cancel()
	}

	// 就地检测锁是否可用
	svc.lock.Lock()
	defer svc.lock.Unlock()

	reqID, err := svc.tryAcquireOne(req)
	if err != nil {
		return LockedRequest{}, err
	}

	if reqID != "" {
		svc.acquired[reqID] = req
		return LockedRequest{
			Req:   req,
			ReqID: reqID,
		}, nil
	}

	// 就地检测失败，那么就需要异步等待锁可用
	info := &acquireInfo{
		Request:  req,
		Callback: future.NewSetValue[types.RequestID](),
	}
	svc.acquirings = append(svc.acquirings, info)

	// 等待的时候不加锁
	svc.lock.Unlock()
	reqID, err = info.Callback.Wait(ctx)
	svc.lock.Lock()

	if err == nil {
		svc.acquired[reqID] = req
		return LockedRequest{
			Req:   req,
			ReqID: reqID,
		}, nil
	}

	if err != future.ErrCanceled {
		lo2.Remove(svc.acquirings, info)
		return LockedRequest{}, err
	}

	// 如果第一次等待是超时错误，那么在锁里再尝试获取一次结果
	reqID, err = info.Callback.TryGetValue()
	if err == nil {
		svc.acquired[reqID] = req
		return LockedRequest{
			Req:   req,
			ReqID: reqID,
		}, nil
	}

	lo2.Remove(svc.acquirings, info)
	return LockedRequest{}, err
}

func (s *Core) release(reqID types.RequestID) {
	s.lock.Lock()
	defer s.lock.Unlock()

	req, ok := s.acquired[reqID]
	if !ok {
		return
	}

	s.releaseRequest(reqID, req)
	s.tryAcquirings()
}

func (a *Core) tryAcquirings() {
	for i := 0; i < len(a.acquirings); i++ {
		req := a.acquirings[i]

		reqID, err := a.tryAcquireOne(req.Request)
		if err != nil {
			req.LastErr = err
			continue
		}

		req.Callback.SetValue(reqID)
		a.acquirings[i] = nil
	}

	a.acquirings = lo2.RemoveAllDefault(a.acquirings)
}

func (s *Core) tryAcquireOne(req types.LockRequest) (types.RequestID, error) {
	err := s.testOneRequest(req)
	if err != nil {
		return "", err
	}

	reqID := types.RequestID(fmt.Sprintf("%d", s.nextReqID))
	s.nextReqID++

	s.applyRequest(reqID, req)
	return reqID, nil
}

func (s *Core) testOneRequest(req types.LockRequest) error {
	for _, lock := range req.Locks {
		n, ok := s.provdersTrie.WalkEnd(lock.Path)
		if !ok || n.Value == nil {
			return fmt.Errorf("lock provider not found for path %v", lock.Path)
		}

		err := n.Value.CanLock(lock)
		if err != nil {
			return err
		}
	}

	return nil
}

func (s *Core) applyRequest(reqID types.RequestID, req types.LockRequest) {
	for _, lock := range req.Locks {
		p, _ := s.provdersTrie.WalkEnd(lock.Path)
		p.Value.Lock(reqID, lock)
	}
}

func (s *Core) releaseRequest(reqID types.RequestID, req types.LockRequest) {
	for _, lock := range req.Locks {
		p, _ := s.provdersTrie.WalkEnd(lock.Path)
		p.Value.Unlock(reqID, lock)
	}
}

type LockedRequest struct {
	Req   types.LockRequest
	ReqID types.RequestID
}