JCS-pub/common/pkgs/rpc/channel.go

package rpc

import (
	"sync"

	"gitlink.org.cn/cloudream/common/utils/serder"
	"gitlink.org.cn/cloudream/jcs-pub/common/ecode"
)

var ErrChannelClosed = Failed(ecode.ChannelClosed, "channel closed")

type ChanSender[T any] interface {
	Send(val T) *CodeError
	Close()
	// 关闭连接，并发送错误码。注意：客户端部分的Channel调用此函数时设置的err不会被送到服务端，因为GRPC没有提供这样的机制。
	CloseWithError(err *CodeError)
}

type ChanReceiver[T any] interface {
	Receive() (T, *CodeError)
	Close()
	// 关闭连接，并发送错误码。注意：客户端部分的Channel调用此函数时设置的err不会被送到服务端，因为GRPC没有提供这样的机制。
	CloseWithError(err *CodeError)
}

type BidChan[Recv, Send any] interface {
	ChanSender[Send]
	ChanReceiver[Recv]
}

type fusedChannel[Recv, Send any] struct {
	err *CodeError
}

func (f *fusedChannel[Recv, Send]) Receive() (Recv, *CodeError) {
	var val Recv
	return val, f.err
}
func (f *fusedChannel[Recv, Send]) Send(val Send) *CodeError {
	return f.err
}
func (f *fusedChannel[Recv, Send]) Close() {
}
func (f *fusedChannel[Recv, Send]) CloseWithError(err *CodeError) {
}

func NewFusedChan[Recv, Send any](err *CodeError) BidChan[Recv, Send] {
	return &fusedChannel[Recv, Send]{err: err}
}

type bidChanClient[Recv, Send any] struct {
	cli      BidChannelAPIClient
	cancelFn func()
	lastErr  *CodeError
	lock     sync.Mutex
}

func NewBidChanClient[Recv, Send any](cli BidChannelAPIClient, cancelFn func()) BidChan[Recv, Send] {
	return &bidChanClient[Recv, Send]{cli: cli, cancelFn: cancelFn}
}

func (c *bidChanClient[Recv, Send]) Send(val Send) *CodeError {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.lastErr != nil {
		err := c.lastErr
		return err
	}

	data, err := serder.ObjectToJSONEx(val)
	if err != nil {
		c.cancelFn()
		c.lastErr = Failed(ecode.OperationFailed, err.Error())
		return Failed(ecode.OperationFailed, err.Error())
	}
	c.lock.Unlock()

	err = c.cli.Send(&Request{Payload: data})

	c.lock.Lock()

	if err != nil {
		c.cancelFn()
		if c.lastErr == nil {
			c.lastErr = ExtractCodeError(err)
		}
		return c.lastErr
	}

	return nil
}

func (c *bidChanClient[Recv, Send]) Receive() (Recv, *CodeError) {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.lastErr != nil {
		var def Recv
		return def, c.lastErr
	}

	c.lock.Unlock()
	resp, err := c.cli.Recv()
	c.lock.Lock()

	if err != nil {
		c.cancelFn()

		cerr := ExtractCodeError(err)
		if c.lastErr == nil {
			c.lastErr = cerr
		}

		var def Recv
		return def, cerr
	}

	resp2, err := serder.JSONToObjectEx[Recv](resp.Payload)
	if err != nil {
		c.cancelFn()

		cerr := Failed(ecode.OperationFailed, err.Error())
		if c.lastErr == nil {
			c.lastErr = cerr
		}

		var def Recv
		return def, cerr
	}

	return resp2, nil
}

func (c *bidChanClient[Recv, Send]) Close() {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.lastErr != nil {
		return
	}

	c.cli.CloseSend()
	c.lastErr = ErrChannelClosed
}

func (c *bidChanClient[Recv, Send]) CloseWithError(err *CodeError) {
	c.lock.Lock()
	defer c.lock.Unlock()

	if c.lastErr != nil {
		return
	}

	c.cancelFn()
	c.lastErr = err
}

type bidChanServer[Recv, Send any] struct {
	svr     BidChannelAPIServer
	errChan chan *CodeError
	lastErr *CodeError
	lock    sync.Mutex
}

func NewBidChanServer[Recv, Send any](svr BidChannelAPIServer, errChan chan *CodeError) BidChan[Recv, Send] {
	return &bidChanServer[Recv, Send]{svr: svr, errChan: errChan}
}

func (s *bidChanServer[Recv, Send]) Send(val Send) *CodeError {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.lastErr != nil {
		return s.lastErr
	}

	data, err := serder.ObjectToJSONEx(val)
	if err != nil {
		s.lastErr = Failed(ecode.OperationFailed, err.Error())
		s.errChan <- s.lastErr
		return s.lastErr
	}

	s.lock.Unlock()
	err = s.svr.Send(&Response{Payload: data})
	s.lock.Lock()

	if err != nil {
		cerr := ExtractCodeError(err)
		if s.lastErr == nil {
			s.lastErr = cerr
			s.errChan <- cerr
		}
		return cerr
	}

	return nil
}

func (s *bidChanServer[Recv, Send]) Receive() (Recv, *CodeError) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.lastErr != nil {
		var def Recv
		return def, s.lastErr
	}

	s.lock.Unlock()
	req, err := s.svr.Recv()
	s.lock.Lock()

	if err != nil {
		cerr := ExtractCodeError(err)
		if s.lastErr == nil {
			s.lastErr = cerr
			s.errChan <- cerr
		}

		var def Recv
		return def, cerr
	}

	req2, err := serder.JSONToObjectEx[Recv](req.Payload)
	if err != nil {
		cerr := Failed(ecode.OperationFailed, err.Error())
		if s.lastErr == nil {
			s.lastErr = cerr
			s.errChan <- cerr
		}

		var def Recv
		return def, cerr
	}

	return req2, nil
}
func (s *bidChanServer[Recv, Send]) Close() {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.lastErr != nil {
		return
	}

	s.lastErr = ErrChannelClosed
	s.errChan <- nil
}

func (s *bidChanServer[Recv, Send]) CloseWithError(err *CodeError) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.lastErr != nil {
		return
	}

	s.lastErr = err
	s.errChan <- err
}