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

package publock

import (
	"context"
	"fmt"
	"sync"
	"time"

	"gitlink.org.cn/cloudream/common/pkgs/future"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/publock/types"
	clirpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/client"
)

type AcquireOptionFn func(opt *types.AcquireOption)

func WithTimeout(timeout time.Duration) AcquireOptionFn {
	return func(opt *types.AcquireOption) {
		opt.Timeout = timeout
	}
}

func WithReason(reason string) AcquireOptionFn {
	return func(opt *types.AcquireOption) {
		opt.Reason = reason
	}
}

type PubLock struct {
	core       *Core
	cliCli     *clirpc.Client
	pubChan    clirpc.PubLockMessageChan
	lock       sync.Mutex
	acquirings map[string]*acquireInfo
	releasing  map[string]*releaseInfo
	nextCtxID  int64
}

func NewMaster() *PubLock {
	core := NewCore()
	return &PubLock{
		core: core,
	}
}

func NewSlave(cli *clirpc.Client) *PubLock {
	return &PubLock{
		cliCli:     cli,
		acquirings: make(map[string]*acquireInfo),
		releasing:  make(map[string]*releaseInfo),
	}
}

func (s *PubLock) BeginReentrant() *Reentrant {
	r := &Reentrant{
		p: s,
	}
	r.T = r
	return r
}

func (p *PubLock) BeginMutex() *MutexBuilder {
	m := &MutexBuilder{
		pub: p,
	}
	m.T = m
	return m
}

func (p *PubLock) Start() {
	if p.core != nil {
		p.core.Start()
		return
	}
}

func (p *PubLock) Stop() {
	p.lock.Lock()
	defer p.lock.Unlock()

	if p.core != nil {
		p.core.Stop()
		return
	}

	if p.pubChan != nil {
		p.pubChan.Close()
		p.pubChan = nil
	}

	p.cliCli.Release()
	p.cliCli = nil
}

func (p *PubLock) Acquire(req types.LockRequest, opt types.AcquireOption) (LockedRequest, error) {
	p.lock.Lock()

	if p.core != nil {
		p.lock.Unlock()
		return p.core.Acquire(req, opt)
	}

	if p.pubChan == nil {
		p.pubChan = p.cliCli.PubLockChannel(context.Background())
		go p.receivingChan()
	}

	acqID := fmt.Sprintf("%v", p.nextCtxID)
	p.nextCtxID++

	cerr := p.pubChan.Send(&types.AcquireMsg{ContextID: acqID, Request: req, Option: opt})
	if cerr != nil {
		p.lock.Unlock()
		return LockedRequest{}, cerr.ToError()
	}

	callback := future.NewSetValue[types.RequestID]()
	info := &acquireInfo{
		Request:  req,
		Callback: callback,
	}
	p.acquirings[acqID] = info
	p.lock.Unlock()

	reqID, err := callback.Wait(context.Background())
	if err != nil {
		return LockedRequest{}, err
	}

	return LockedRequest{
		Req:   req,
		ReqID: reqID,
	}, nil
}

func (p *PubLock) Release(reqID types.RequestID) {
	log := logger.WithField("Mod", "PubLock")

	p.lock.Lock()

	if p.core != nil {
		p.lock.Unlock()
		p.core.release(reqID)
		return
	}

	if p.pubChan == nil {
		p.pubChan = p.cliCli.PubLockChannel(context.Background())
		go p.receivingChan()
	}

	relID := fmt.Sprintf("%v", p.nextCtxID)
	p.nextCtxID++

	cerr := p.pubChan.Send(&types.ReleaseMsg{ContextID: relID, RequestID: reqID})
	if cerr != nil {
		p.lock.Unlock()
		log.Warnf("unlock %v: %v", reqID, cerr.ToError())
		return
	}

	callback := future.NewSetVoid()
	info := &releaseInfo{
		RequestID: reqID,
		Callback:  callback,
	}
	p.releasing[relID] = info
	p.lock.Unlock()

	err := callback.Wait(context.Background())
	if err != nil {
		log.Warnf("unlock %v: %v", reqID, err)
		return
	}

	log.Tracef("unlock %v", reqID)
}

func (p *PubLock) receivingChan() {
	log := logger.WithField("Mod", "PubLock")
	for {
		msg, cerr := p.pubChan.Receive()
		if cerr != nil {
			p.lock.Lock()
			for _, info := range p.acquirings {
				info.Callback.SetError(cerr.ToError())
			}
			p.acquirings = make(map[string]*acquireInfo)

			for _, info := range p.releasing {
				info.Callback.SetError(cerr.ToError())
			}
			p.releasing = make(map[string]*releaseInfo)
			p.lock.Unlock()

			log.Warnf("receive channel error: %v", cerr.ToError())
			return
		}

		p.lock.Lock()

		switch msg := msg.(type) {
		case *types.AcquireResultMsg:
			info, ok := p.acquirings[msg.ContextID]
			if !ok {
				continue
			}

			if msg.Success {
				info.Callback.SetValue(msg.RequestID)
			} else {
				info.Callback.SetError(fmt.Errorf(msg.Error))
			}
			delete(p.acquirings, msg.ContextID)

		case *types.ReleaseResultMsg:
			info, ok := p.releasing[msg.ContextID]
			if !ok {
				continue
			}

			info.Callback.SetVoid()
			delete(p.releasing, msg.ContextID)
		}

		p.lock.Unlock()
	}
}

type releaseInfo struct {
	RequestID types.RequestID
	Callback  *future.SetVoidFuture
}