package ops import ( "context" "fmt" "io" "gitlink.org.cn/cloudream/common/pkgs/future" "gitlink.org.cn/cloudream/common/pkgs/ioswitch/dag" "gitlink.org.cn/cloudream/common/pkgs/ioswitch/exec" "gitlink.org.cn/cloudream/common/pkgs/logger" "gitlink.org.cn/cloudream/common/utils/io2" ) func init() { exec.UseOp[*SendStream]() exec.UseOp[*GetStream]() exec.UseOp[*SendVar]() exec.UseOp[*GetVar]() } type SendStream struct { Input *exec.StreamVar `json:"input"` Send *exec.StreamVar `json:"send"` Worker exec.WorkerInfo `json:"worker"` } func (o *SendStream) Execute(ctx context.Context, e *exec.Executor) error { err := e.BindVars(ctx, o.Input) if err != nil { return err } defer o.Input.Stream.Close() cli, err := o.Worker.NewClient() if err != nil { return fmt.Errorf("new worker %v client: %w", o.Worker, err) } defer cli.Close() logger.Debugf("sending stream %v as %v to worker %v", o.Input.ID, o.Send.ID, o.Worker) // 发送后流的ID不同 err = cli.SendStream(ctx, e.Plan().ID, o.Send, o.Input.Stream) if err != nil { return fmt.Errorf("sending stream: %w", err) } return nil } type GetStream struct { Signal *exec.SignalVar `json:"signal"` Target *exec.StreamVar `json:"target"` Output *exec.StreamVar `json:"output"` Worker exec.WorkerInfo `json:"worker"` } func (o *GetStream) Execute(ctx context.Context, e *exec.Executor) error { cli, err := o.Worker.NewClient() if err != nil { return fmt.Errorf("new worker %v client: %w", o.Worker, err) } defer cli.Close() logger.Debugf("getting stream %v as %v from worker %v", o.Target.ID, o.Output.ID, o.Worker) str, err := cli.GetStream(ctx, e.Plan().ID, o.Target, o.Signal) if err != nil { return fmt.Errorf("getting stream: %w", err) } fut := future.NewSetVoid() // 获取后送到本地的流ID是不同的 o.Output.Stream = io2.AfterReadClosedOnce(str, func(closer io.ReadCloser) { fut.SetVoid() }) e.PutVars(o.Output) return fut.Wait(ctx) } type SendVar struct { Input exec.Var `json:"input"` Send exec.Var `json:"send"` Worker exec.WorkerInfo `json:"worker"` } func (o *SendVar) Execute(ctx context.Context, e *exec.Executor) error { err := e.BindVars(ctx, o.Input) if err != nil { return err } cli, err := o.Worker.NewClient() if err != nil { return fmt.Errorf("new worker %v client: %w", o.Worker, err) } defer cli.Close() logger.Debugf("sending var %v as %v to worker %v", o.Input.GetID(), o.Send.GetID(), o.Worker) exec.AssignVar(o.Input, o.Send) err = cli.SendVar(ctx, e.Plan().ID, o.Send) if err != nil { return fmt.Errorf("sending var: %w", err) } return nil } type GetVar struct { Signal *exec.SignalVar `json:"signal"` Target exec.Var `json:"target"` Output exec.Var `json:"output"` Worker exec.WorkerInfo `json:"worker"` } func (o *GetVar) Execute(ctx context.Context, e *exec.Executor) error { cli, err := o.Worker.NewClient() if err != nil { return fmt.Errorf("new worker %v client: %w", o.Worker, err) } defer cli.Close() logger.Debugf("getting var %v as %v from worker %v", o.Target.GetID(), o.Output.GetID(), o.Worker) err = cli.GetVar(ctx, e.Plan().ID, o.Target, o.Signal) if err != nil { return fmt.Errorf("getting var: %w", err) } exec.AssignVar(o.Target, o.Output) e.PutVars(o.Output) return nil } type SendStreamType struct { ToWorker exec.WorkerInfo } func (t *SendStreamType) Send(n *dag.Node, v *dag.StreamVar) *dag.StreamVar { v.To(n, 0) return n.OutputStreams[0] } func (t *SendStreamType) InitNode(node *dag.Node) { dag.NodeDeclareInputStream(node, 1) dag.NodeNewOutputStream(node, nil) } func (t *SendStreamType) GenerateOp(op *dag.Node) (exec.Op, error) { return &SendStream{ Input: op.InputStreams[0].Var, Send: op.OutputStreams[0].Var, Worker: t.ToWorker, }, nil } func (t *SendStreamType) String(node *dag.Node) string { return fmt.Sprintf("SendStream[]%v%v", formatStreamIO(node), formatValueIO(node)) } type SendVarType struct { ToWorker exec.WorkerInfo } func (t *SendVarType) Send(n *dag.Node, v *dag.ValueVar) *dag.ValueVar { v.To(n, 0) n.OutputValues[0].Type = v.Type return n.OutputValues[0] } func (t *SendVarType) InitNode(node *dag.Node) { dag.NodeDeclareInputValue(node, 1) dag.NodeNewOutputValue(node, 0, nil) } func (t *SendVarType) GenerateOp(op *dag.Node) (exec.Op, error) { return &SendVar{ Input: op.InputValues[0].Var, Send: op.OutputValues[0].Var, Worker: t.ToWorker, }, nil } func (t *SendVarType) String(node *dag.Node) string { return fmt.Sprintf("SendVar[]%v%v", formatStreamIO(node), formatValueIO(node)) } type GetStreamType struct { FromWorker exec.WorkerInfo } func (t *GetStreamType) Get(n *dag.Node, v *dag.StreamVar) *dag.StreamVar { v.To(n, 0) return n.OutputStreams[0] } func (t *GetStreamType) SignalVar(n *dag.Node) *dag.ValueVar { return n.OutputValues[0] } func (t *GetStreamType) InitNode(node *dag.Node) { dag.NodeDeclareInputStream(node, 1) dag.NodeNewOutputValue(node, dag.SignalValueVar, nil) dag.NodeNewOutputStream(node, nil) } func (t *GetStreamType) GenerateOp(op *dag.Node) (exec.Op, error) { return &GetStream{ Signal: op.OutputValues[0].Var.(*exec.SignalVar), Output: op.OutputStreams[0].Var, Target: op.InputStreams[0].Var, Worker: t.FromWorker, }, nil } func (t *GetStreamType) String(node *dag.Node) string { return fmt.Sprintf("GetStream[]%v%v", formatStreamIO(node), formatValueIO(node)) } type GetVaType struct { FromWorker exec.WorkerInfo } func (t *GetVaType) Get(n *dag.Node, v *dag.ValueVar) *dag.ValueVar { v.To(n, 0) n.OutputValues[1].Type = v.Type return n.OutputValues[1] } func (t *GetVaType) SignalVar(n *dag.Node) *dag.ValueVar { return n.OutputValues[0] } func (t *GetVaType) InitNode(node *dag.Node) { dag.NodeDeclareInputValue(node, 1) dag.NodeNewOutputValue(node, dag.SignalValueVar, nil) dag.NodeNewOutputValue(node, 0, nil) } func (t *GetVaType) GenerateOp(op *dag.Node) (exec.Op, error) { return &GetVar{ Signal: op.OutputValues[0].Var.(*exec.SignalVar), Output: op.OutputValues[1].Var, Target: op.InputValues[0].Var, Worker: t.FromWorker, }, nil } func (t *GetVaType) String(node *dag.Node) string { return fmt.Sprintf("GetVar[]%v%v", formatStreamIO(node), formatValueIO(node)) }