JCS-pub/hub/internal/rpc/ioswitch.go

package rpc

import (
	"context"
	"time"

	"gitlink.org.cn/cloudream/common/consts/errorcode"
	"gitlink.org.cn/cloudream/common/pkgs/future"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/common/utils/io2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc"
	hubrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/hub"
)

func (s *Service) ExecuteIOPlan(ctx context.Context, req *hubrpc.ExecuteIOPlan) (*hubrpc.ExecuteIOPlanResp, *rpc.CodeError) {
	log := logger.WithField("PlanID", req.Plan.ID)
	log.Infof("begin execute io plan")

	sw := exec.NewExecutor(req.Plan, exec.Location{
		IsDriver:   false,
		WorkerName: req.WorkerName,
	})

	s.swWorker.Add(sw)
	defer s.swWorker.Remove(sw)

	execCtx := exec.NewWithContext(ctx)
	exec.SetValueByType(execCtx, s.stgPool)
	ret, err := sw.Run(execCtx)
	if err != nil {
		log.Warnf("running io plan: %v", err)
		return nil, rpc.Failed(errorcode.OperationFailed, "%v", err)
	}

	log.Infof("plan finished")
	return &hubrpc.ExecuteIOPlanResp{
		Result: ret,
	}, nil
}

func (s *Service) SendIOStream(ctx context.Context, req *hubrpc.SendIOStream) (*hubrpc.SendIOStreamResp, *rpc.CodeError) {
	logger.
		WithField("PlanID", req.PlanID).
		WithField("VarID", req.VarID).
		Debugf("stream input")

	// 同一批Plan中每个节点的Plan的启动时间有先后，但最多不应该超过30秒
	ctx2, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()

	sw := s.swWorker.FindByIDContexted(ctx2, exec.PlanID(req.PlanID))
	if sw == nil {
		return nil, rpc.Failed(errorcode.DataNotFound, "plan not found")
	}

	fut := future.NewSetVoid()

	varID := exec.VarID(req.VarID)
	sw.PutVar(varID, &exec.StreamValue{Stream: io2.DelegateReadCloser(req.Stream, func() error {
		fut.SetVoid()
		return nil
	})})

	err := fut.Wait(ctx)
	if err != nil {
		return nil, rpc.Failed(errorcode.OperationFailed, "%v", err)
	}

	return &hubrpc.SendIOStreamResp{}, nil
}

func (s *Service) GetIOStream(ctx context.Context, req *hubrpc.GetIOStream) (*hubrpc.GetIOStreamResp, *rpc.CodeError) {
	logger.
		WithField("PlanID", req.PlanID).
		WithField("VarID", req.VarID).
		Debugf("stream output")

	// 同上
	ctx2, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()

	sw := s.swWorker.FindByIDContexted(ctx2, exec.PlanID(req.PlanID))
	if sw == nil {
		return nil, rpc.Failed(errorcode.DataNotFound, "plan not found")
	}

	sw.PutVar(req.SignalID, req.Signal)

	strVar, err := exec.BindVar[*exec.StreamValue](sw, ctx, exec.VarID(req.VarID))
	if err != nil {
		return nil, rpc.Failed(errorcode.OperationFailed, "bind var: %v", err)
	}

	return &hubrpc.GetIOStreamResp{
		Stream: strVar.Stream,
	}, nil
}

func (s *Service) SendIOVar(ctx context.Context, req *hubrpc.SendIOVar) (*hubrpc.SendIOVarResp, *rpc.CodeError) {
	ctx2, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()

	sw := s.swWorker.FindByIDContexted(ctx2, exec.PlanID(req.PlanID))
	if sw == nil {
		return nil, rpc.Failed(errorcode.DataNotFound, "plan not found")
	}

	sw.PutVar(req.VarID, req.Value)
	return &hubrpc.SendIOVarResp{}, nil
}

func (s *Service) GetIOVar(ctx context.Context, req *hubrpc.GetIOVar) (*hubrpc.GetIOVarResp, *rpc.CodeError) {
	ctx2, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()

	sw := s.swWorker.FindByIDContexted(ctx2, exec.PlanID(req.PlanID))
	if sw == nil {
		return nil, rpc.Failed(errorcode.DataNotFound, "plan not found")
	}

	sw.PutVar(req.SignalID, req.Signal)

	v, err := sw.BindVar(ctx, exec.VarID(req.VarID))
	if err != nil {
		return nil, rpc.Failed(errorcode.OperationFailed, "bind var: %v", err)
	}

	return &hubrpc.GetIOVarResp{
		Value: v,
	}, nil
}